{"slug": "batching-by-length-instead-of-looping-item-by-item-for-slm-optimization", "title": "Batching by Length Instead of Looping Item by Item for SLM Optimization", "summary": "Sorting support tickets by token length before batching, rather than looping item by item, is the third and final optimization in a KDnuggets series on small language model (SLM) narrow automation, with benchmarks run on Qwen2.5-0.5B-Instruct in float16 via Hugging Face Transformers on an M2 MacBook Air with 24GB RAM. The article argues that batch size 1 is memory-bandwidth bound, since the hardware streams every weight out of memory to serve a single sequence, and that padding each batch to its own local maximum instead of the dataset's global maximum avoids wasted compute on padding tokens.", "body_md": "# Batching by Length Instead of Looping Item by Item for SLM Optimization\n\nWe finish off our short series on SLM optimization with the third entry, focused on batching by length instead of looping item by item.\n\nPrevious articles in this series discussed [constraining output space](https://www.kdnuggets.com/constraining-output-space-small-language-model-narrow-automation-optimization) as well as [reusing the prompt prefix with a key-value cache](https://www.kdnuggets.com/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization), both framed as approaches to small language model (SLM) narrow automation optimization. Let's finish this series off with the third entry, focused on batching by length instead of looping item by item.\n\nAs in our previous articles, all benchmarks below use [**Qwen2.5-0.5B-Instruct**](https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct) in float16 through **[Hugging Face Transformers](https://huggingface.co/docs/transformers/index)**, running on an M2 Macbook Air with 24GB RAM and a 16-core Neural Engine.\n\nDon't forget to set up a Python environment and install your requirements:\n\n```\npip install torch transformers accelerate\n```\n\nWe will continue to use the support ticket framing [from our first article](https://www.kdnuggets.com/constraining-output-space-small-language-model-narrow-automation-optimization).\n\n## Why Batch by Length Instead of Looping Item by Item\n\nProcessing one ticket per forward pass is the single largest source of waste in the whole pipeline. At batch size 1, a small model is memory-bandwidth bound rather than compute-bound: the hardware streams every weight out of memory in order to serve one sequence, then does it again for the next one, and the arithmetic units sit mostly idle in between. This is true on a GPU and it is true on the CPU we have been using throughout this series, which is where a 0.5B model most often actually runs.\n\nBatching amortizes that weight read across many sequences. But the obvious implementation introduces its own waste, since sequences in a batch must be padded to a common length. Real-world text has a long tail: if the longest item in your dataset is a few hundred tokens and the median is well under a hundred, padding every batch to the global maximum means most of what you compute is padding.\n\nThe answer is to **sort by token length** before forming batches, so each batch contains similarly sized items and pads to its own local maximum.\n\n## Looping Item by Item\n\nHere is the per-item baseline on a realistic length distribution, with the constrained scoring from the first article in this series carried forward so that each item costs exactly one forward pass:\n\n``` python\nimport os\nimport time\nimport inspect\nimport torch\nimport numpy as np\nfrom transformers import AutoTokenizer, AutoModelForCausalLM\n\nMODEL_ID = \"Qwen/Qwen2.5-0.5B-Instruct\"\n\ntorch.set_num_threads(os.cpu_count() or 1)\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\nif tokenizer.pad_token is None:\n    tokenizer.pad_token = tokenizer.eos_token\ntokenizer.padding_side = \"left\"  # keeps the last real token at index -1\n\nmodel = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)\nmodel.eval()\n\nLABELS = [\"billing\", \"technical\", \"account\"]\n\n# constrained scoring, carried over from the first article in this series\nlabel_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS]\nassert len(set(label_first_ids)) == len(LABELS), (\n    \"Labels share a first token; score full label sequences instead.\"\n)\nlabel_first_ids = torch.tensor(label_first_ids, device=model.device)\n\n# a causal LM returns a logit vector for every position by default; at batch 32 by\n# 400 tokens that is a multi-gigabyte tensor we would immediately throw away, so\n# ask for the last position only where the installed version supports it\n_forward_params = inspect.signature(model.forward).parameters\nif \"logits_to_keep\" in _forward_params:\n    LAST_LOGIT_ONLY = {\"logits_to_keep\": 1}\nelif \"num_logits_to_keep\" in _forward_params:\n    LAST_LOGIT_ONLY = {\"num_logits_to_keep\": 1}\nelse:\n    LAST_LOGIT_ONLY = {}\n\ndef build_prompt(ticket):\n    messages = [\n        {\n            \"role\": \"system\",\n            \"content\": \"You classify support tickets. Answer with exactly one of: billing, technical, account.\",\n        },\n        {\"role\": \"user\", \"content\": f\"Ticket: {ticket}\\nCategory:\"},\n    ]\n    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n\n# simulate a long-tailed ticket length distribution: most items are short, a few\n# are very long. Each ticket keeps a real, classifiable sentence at the front and\n# is extended with filler, so length varies without the label signal disappearing.\nBASE_TICKETS = [\n    \"My card was charged twice for the same invoice.\",\n    \"The mobile app crashes whenever I open the settings page.\",\n    \"I need to change the email address on my profile.\",\n]\nFILLER = (\n    \"I have been waiting for a response for several days now and would really \"\n    \"appreciate an update on this whenever someone gets a chance to look at it.\"\n).split()\n\nrng = np.random.default_rng(0)\ntarget_words = np.clip(rng.lognormal(np.log(60), 0.9, size=600), 12, 400).astype(int)\n\ndef make_ticket(base, n_words):\n    words = base.split()\n    while len(words) < n_words:\n        words += FILLER\n    return \" \".join(words[:n_words])\n\ntickets_var = [make_ticket(BASE_TICKETS[i % 3], int(n)) for i, n in enumerate(target_words)]\n\nprompts = [build_prompt(t) for t in tickets_var]\ntoken_lengths = [len(tokenizer(p, add_special_tokens=False)[\"input_ids\"]) for p in prompts]\n\nprint(\n    f\"Prompt lengths: min {min(token_lengths)}, \"\n    f\"median {int(np.median(token_lengths))}, \"\n    f\"max {max(token_lengths)} tokens\"\n)\nprint(f\"Padding every item to the global maximum would process \"\n      f\"{max(token_lengths) * len(prompts) / sum(token_lengths):.1f}x the necessary tokens\")\n\n# time inference\nbaseline_predictions = []\nstart = time.time()\n\nfor n, prompt in enumerate(prompts, start=1):\n\n    # this loop runs for minutes on CPU, so report progress rather than sitting silent\n    if n % 100 == 0:\n        print(f\"  {n}/{len(prompts)} tickets ({(time.time() - start) / n:.2f}s each)\", flush=True)\n\n    inputs = tokenizer(prompt, add_special_tokens=False, return_tensors=\"pt\").to(model.device)\n    with torch.no_grad():\n        logits = model(**inputs, **LAST_LOGIT_ONLY).logits[0, -1, :]\n    baseline_predictions.append(LABELS[int(logits[label_first_ids].argmax())])\n\nduration_loop = time.time() - start\n\n# output task metrics\nprint(f\"One at a time: {duration_loop:.2f} seconds ({len(prompts) / duration_loop:.1f} items/sec)\")\n```\n\nOutput:\n\n```\nPrompt lengths: min 48, median 94, max 449 tokens\nPadding every item to the global maximum would process 3.7x the necessary tokens\n  100/600 tickets (0.24s each)\n  200/600 tickets (0.23s each)\n  300/600 tickets (0.23s each)\n  400/600 tickets (0.23s each)\n  500/600 tickets (0.24s each)\n  600/600 tickets (0.24s each)\nOne at a time: 144.35 seconds (4.2 items/sec)\n```\n\n## Batching by Length\n\nNow the batched version. It runs the same data twice: once in arbitrary order, to isolate what batching alone is worth, and once sorted by length, to show what the sorting adds on top. Both runs track how much of the processed token budget went to padding.\n\n``` python\nimport os\nimport time\nimport inspect\nimport torch\nimport numpy as np\nfrom transformers import AutoTokenizer, AutoModelForCausalLM\n\nMODEL_ID = \"Qwen/Qwen2.5-0.5B-Instruct\"\nBATCH_SIZE = 32\n\ntorch.set_num_threads(os.cpu_count() or 1)\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\nif tokenizer.pad_token is None:\n    tokenizer.pad_token = tokenizer.eos_token\n\n# keep the last real token at index -1\ntokenizer.padding_side = \"left\"\n\nmodel = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)\nmodel.eval()\n\nLABELS = [\"billing\", \"technical\", \"account\"]\nlabel_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS]\nassert len(set(label_first_ids)) == len(LABELS), (\n    \"Labels share a first token; score full label sequences instead.\"\n)\nlabel_first_ids = torch.tensor(label_first_ids, device=model.device)\n\n_forward_params = inspect.signature(model.forward).parameters\nif \"logits_to_keep\" in _forward_params:\n    LAST_LOGIT_ONLY = {\"logits_to_keep\": 1}\nelif \"num_logits_to_keep\" in _forward_params:\n    LAST_LOGIT_ONLY = {\"num_logits_to_keep\": 1}\nelse:\n    LAST_LOGIT_ONLY = {}\n\ndef build_prompt(ticket):\n    messages = [\n        {\n            \"role\": \"system\",\n            \"content\": \"You classify support tickets. Answer with exactly one of: billing, technical, account.\",\n        },\n        {\"role\": \"user\", \"content\": f\"Ticket: {ticket}\\nCategory:\"},\n    ]\n    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n\nBASE_TICKETS = [\n    \"My card was charged twice for the same invoice.\",\n    \"The mobile app crashes whenever I open the settings page.\",\n    \"I need to change the email address on my profile.\",\n]\nFILLER = (\n    \"I have been waiting for a response for several days now and would really \"\n    \"appreciate an update on this whenever someone gets a chance to look at it.\"\n).split()\n\nrng = np.random.default_rng(0)\ntarget_words = np.clip(rng.lognormal(np.log(60), 0.9, size=600), 12, 400).astype(int)\n\ndef make_ticket(base, n_words):\n    words = base.split()\n    while len(words) < n_words:\n        words += FILLER\n    return \" \".join(words[:n_words])\n\ntickets_var = [make_ticket(BASE_TICKETS[i % 3], int(n)) for i, n in enumerate(target_words)]\nprompts = [build_prompt(t) for t in tickets_var]\ntoken_lengths = [len(tokenizer(p, add_special_tokens=False)[\"input_ids\"]) for p in prompts]\n\nprint(\n    f\"Prompt lengths: min {min(token_lengths)}, \"\n    f\"median {int(np.median(token_lengths))}, \"\n    f\"max {max(token_lengths)} tokens\"\n)\n\ndef run_batched(order, batch_size):\n    predictions = [None] * len(prompts)\n    processed_tokens = real_tokens = 0\n    start = time.time()\n    for i in range(0, len(order), batch_size):\n        idx = order[i:i + batch_size]\n        batch = tokenizer(\n            [prompts[j] for j in idx],\n            add_special_tokens=False,\n            padding=True,\n            return_tensors=\"pt\",\n        ).to(model.device)\n        processed_tokens += batch[\"input_ids\"].numel()\n        real_tokens += int(batch[\"attention_mask\"].sum())\n        with torch.no_grad():\n            logits = model(**batch, **LAST_LOGIT_ONLY).logits[:, -1, :]\n        best = logits[:, label_first_ids].argmax(dim=-1)\n        for slot, choice in zip(idx, best.tolist(), strict=True):\n            predictions[slot] = LABELS[choice]\n    return predictions, time.time() - start, processed_tokens, real_tokens\n\ndef classify_one(prompt):\n    \"\"\"Reference path: a single unpadded sequence. Used only to verify.\"\"\"\n    inputs = tokenizer(prompt, add_special_tokens=False, return_tensors=\"pt\").to(model.device)\n    with torch.no_grad():\n        logits = model(**inputs, **LAST_LOGIT_ONLY).logits[0, -1, :]\n    return LABELS[int(logits[label_first_ids].argmax())]\n\norder = sorted(range(len(prompts)), key=lambda i: token_lengths[i])\npredictions, duration_batched, processed, real = run_batched(order, BATCH_SIZE)\n\nprint(f\"Length-bucketed batching: {duration_batched:.2f} seconds ({len(prompts) / duration_batched:.1f} items/sec)\")\nprint(f\"Padding overhead: {100 * (1 - real / processed):.1f}% of processed tokens were padding\")\n\n# correctness: padded rows must score the same as unpadded ones. check a spread of\n# lengths rather than all 600, since the point is to catch a padding-side or\n# position bug, and such a bug shows up on the very first padded row.\nprobe = order[::60]\nmismatches = [i for i in probe if classify_one(prompts[i]) != predictions[i]]\nprint(f\"Batched vs unbatched agreement on {len(probe)} probes: {len(probe) - len(mismatches)}/{len(probe)}\")\nassert not mismatches, f\"Batched path disagrees at indices {mismatches}\"\n\n# estimate the per-item cost on the same probe set, then scale\nstart = time.time()\nfor i in probe:\n    classify_one(prompts[i])\n```\n\nOutput:\n\n```\nLoading weights: 100%|████████████████████████████████████████████████████████████████████| 290/290 [00:01<00:00, 252.11it/s]\nPrompt lengths: min 48, median 94, max 449 tokens\nLength-bucketed batching: 79.60 seconds (7.5 items/sec)\nPadding overhead: 7.6% of processed tokens were padding\nBatched vs unbatched agreement on 10 probes: 10/10\n```\n\nA large throughput increase on identical hardware and an identical model, purely from how the work was scheduled. Note that the two batched runs do the same arithmetic per real token and differ only in how much padding they carry, so the gap between them is a direct measurement of what the sort buys you.\n\n- Setting `padding_side = \"left\"` is required here, not a choice. With right padding,`logits[:, -1, :]` would land on a pad token for every row shorter than the batch maximum, deceptively producing garbage predictions. Left padding guarantees index`-1` is the true final token of every sequence.\n- Left padding does shift each row's absolute token positions, because a plain `model(**batch)` call numbers positions from zero across the padded width rather than deriving them from the attention mask. For a rotary-embedding model like Qwen2.5 this is harmless, since attention depends only on the*relative* distance between tokens and every real token in a row shifts by the same amount. For a model with learned absolute position embeddings it would not be harmless, and you would need to pass`position_ids` built from the mask. Either way, the agreement check against the per-item loop is what tells you which situation you are in.\n- Ask for the last position's logits only. By default a causal LM returns a logit vector for every input position, and the vocabulary here is around 150k entries: at batch 32 by 400 tokens in float32 that is a multi-gigabyte tensor allocated and discarded on every single batch. `logits_to_keep=1` (named`num_logits_to_keep` in older versions of Transformers) suppresses it. This is negligible at batch size 1 with a short prompt, which is why it never came up in the earlier articles, and it dominates everything else once you start batching long inputs.\n- Sorting before chunking holds padding overhead to a few percent. The same data through fixed-size batches in arbitrary order pushes it far higher, which is the difference the script measures directly rather than asserting: every padding token is a token the hardware processed for no reason.\n- Sorting reorders the data, so keep the original indices around and write results back to their proper slots, as the `order` list does above. Losing the alignment between inputs and predictions is an easy and very expensive bug, and unlike a crash it produces plausible-looking output.\n- Pick `BATCH_SIZE` by measuring, not by intuition, which is what the sweep at the end of the script is for. Throughput climbs steeply and then plateaus once you saturate compute; past that point you are only increasing the odds of an out-of-memory error on your longest bucket. The best value depends on your hardware and on your length distribution, so it is worth re-running the sweep when either changes.\n\nOne caveat: prefix caching and batching need care to combine. The cache we built in the previous article has a batch dimension of 1, so reusing it across a batch means expanding every key and value tensor along that dimension to match, and cropping it back correctly afterwards. It is worth doing when your prefix is long, but do it deliberately and verify the predictions against the unbatched path rather than assuming the two optimizations compose for free.\n\n## Wrapping Up\n\nThis has been our third and final attempt at optimizing SLMs for our narrow automation series, and our target technique this time was length-bucketed batching. This technique replaces a one-item-at-a-time loop with sorted batches that pad to their own local maximum, which gets the hardware out of the memory-bandwidth-bound regime without paying for the padding that naive batching would introduce. By implementing it, we process the same 600 tickets in a fraction of the wall-clock time, with the same predictions coming out the other end.\n\nNone of these techniques makes the model \"smarter.\" Each one was verified by checking that its output was identical to the slower path it replaced, and that check is the whole reason to trust the speedup numbers at all. An optimization that changes your predictions is not an optimization, it is a regression with a stopwatch attached.\n\nTo recap the series:\n\n- **[Constrained scoring](https://www.kdnuggets.com/constraining-output-space-small-language-model-narrow-automation-optimization)** replaces free-form generation and string parsing with a single forward pass restricted to the valid label set, making malformed output structurally impossible while handing you a confidence score for routing edge cases to humans\n- **[Prefix key-value caching](https://www.kdnuggets.com/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization)** computes the static instruction block once instead of once per item, and pays off in proportion to how much of your prompt never changes\n- **[Length-bucketed batching](https://www.kdnuggets.com/batching-by-length-instead-of-looping-item-by-item-for-slm-optimization)** gets the hardware out of the memory-bandwidth-bound constraint and keeps padding overhead near zero by grouping similarly sized inputs together\n\n \n\n \n\n[**\\[Matthew Mayo\\](https://www.kdnuggets.com/wp-content/uploads/./profile-pic.jpg)**](https://www.linkedin.com/in/mattmayo13/) ([**@mattmayo13**](https://twitter.com/mattmayo13)) holds a master's degree in computer science and a graduate diploma in data mining. As managing editor of [KDnuggets](https://www.kdnuggets.com/) & [Statology](https://www.statology.org/), and contributing editor at [Machine Learning Mastery](https://machinelearningmastery.com/), Matthew aims to make complex data science concepts accessible. His professional interests include natural language processing, language models, machine learning algorithms, and exploring emerging AI. He is driven by a mission to democratize knowledge in the data science community. Matthew has been coding since he was 6 years old.", "url": "https://wpnews.pro/news/batching-by-length-instead-of-looping-item-by-item-for-slm-optimization", "canonical_source": "https://www.kdnuggets.com/batching-by-length-instead-of-looping-item-by-item-for-slm-optimization", "published_at": "2026-09-25 14:00:07+00:00", "updated_at": "2026-09-25 14:33:17.482128+00:00", "lang": "en", "topics": ["large-language-models", "ai-infrastructure", "machine-learning", "natural-language-processing"], "entities": ["Qwen2.5-0.5B-Instruct", "Hugging Face Transformers", "KDnuggets", "M2 MacBook Air"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/batching-by-length-instead-of-looping-item-by-item-for-slm-optimization", "markdown": "https://wpnews.pro/news/batching-by-length-instead-of-looping-item-by-item-for-slm-optimization.md", "text": "https://wpnews.pro/news/batching-by-length-instead-of-looping-item-by-item-for-slm-optimization.txt", "jsonld": "https://wpnews.pro/news/batching-by-length-instead-of-looping-item-by-item-for-slm-optimization.jsonld"}}