# Batching by Length Instead of Looping Item by Item for SLM Optimization

> Source: <https://www.kdnuggets.com/batching-by-length-instead-of-looping-item-by-item-for-slm-optimization>
> Published: 2026-09-25 14:00:07+00:00

# Batching by Length Instead of Looping Item by Item for SLM Optimization

We finish off our short series on SLM optimization with the third entry, focused on batching by length instead of looping item by item.

Previous 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.

As 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.

Don't forget to set up a Python environment and install your requirements:

```
pip install torch transformers accelerate
```

We 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).

## Why Batch by Length Instead of Looping Item by Item

Processing 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.

Batching 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.

The answer is to **sort by token length** before forming batches, so each batch contains similarly sized items and pads to its own local maximum.

## Looping Item by Item

Here 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:

``` python
import os
import time
import inspect
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"

torch.set_num_threads(os.cpu_count() or 1)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"  # keeps the last real token at index -1

model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)
model.eval()

LABELS = ["billing", "technical", "account"]

# constrained scoring, carried over from the first article in this series
label_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS]
assert len(set(label_first_ids)) == len(LABELS), (
    "Labels share a first token; score full label sequences instead."
)
label_first_ids = torch.tensor(label_first_ids, device=model.device)

# a causal LM returns a logit vector for every position by default; at batch 32 by
# 400 tokens that is a multi-gigabyte tensor we would immediately throw away, so
# ask for the last position only where the installed version supports it
_forward_params = inspect.signature(model.forward).parameters
if "logits_to_keep" in _forward_params:
    LAST_LOGIT_ONLY = {"logits_to_keep": 1}
elif "num_logits_to_keep" in _forward_params:
    LAST_LOGIT_ONLY = {"num_logits_to_keep": 1}
else:
    LAST_LOGIT_ONLY = {}

def build_prompt(ticket):
    messages = [
        {
            "role": "system",
            "content": "You classify support tickets. Answer with exactly one of: billing, technical, account.",
        },
        {"role": "user", "content": f"Ticket: {ticket}\nCategory:"},
    ]
    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

# simulate a long-tailed ticket length distribution: most items are short, a few
# are very long. Each ticket keeps a real, classifiable sentence at the front and
# is extended with filler, so length varies without the label signal disappearing.
BASE_TICKETS = [
    "My card was charged twice for the same invoice.",
    "The mobile app crashes whenever I open the settings page.",
    "I need to change the email address on my profile.",
]
FILLER = (
    "I have been waiting for a response for several days now and would really "
    "appreciate an update on this whenever someone gets a chance to look at it."
).split()

rng = np.random.default_rng(0)
target_words = np.clip(rng.lognormal(np.log(60), 0.9, size=600), 12, 400).astype(int)

def make_ticket(base, n_words):
    words = base.split()
    while len(words) < n_words:
        words += FILLER
    return " ".join(words[:n_words])

tickets_var = [make_ticket(BASE_TICKETS[i % 3], int(n)) for i, n in enumerate(target_words)]

prompts = [build_prompt(t) for t in tickets_var]
token_lengths = [len(tokenizer(p, add_special_tokens=False)["input_ids"]) for p in prompts]

print(
    f"Prompt lengths: min {min(token_lengths)}, "
    f"median {int(np.median(token_lengths))}, "
    f"max {max(token_lengths)} tokens"
)
print(f"Padding every item to the global maximum would process "
      f"{max(token_lengths) * len(prompts) / sum(token_lengths):.1f}x the necessary tokens")

# time inference
baseline_predictions = []
start = time.time()

for n, prompt in enumerate(prompts, start=1):

    # this loop runs for minutes on CPU, so report progress rather than sitting silent
    if n % 100 == 0:
        print(f"  {n}/{len(prompts)} tickets ({(time.time() - start) / n:.2f}s each)", flush=True)

    inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
    with torch.no_grad():
        logits = model(**inputs, **LAST_LOGIT_ONLY).logits[0, -1, :]
    baseline_predictions.append(LABELS[int(logits[label_first_ids].argmax())])

duration_loop = time.time() - start

# output task metrics
print(f"One at a time: {duration_loop:.2f} seconds ({len(prompts) / duration_loop:.1f} items/sec)")
```

Output:

```
Prompt lengths: min 48, median 94, max 449 tokens
Padding every item to the global maximum would process 3.7x the necessary tokens
  100/600 tickets (0.24s each)
  200/600 tickets (0.23s each)
  300/600 tickets (0.23s each)
  400/600 tickets (0.23s each)
  500/600 tickets (0.24s each)
  600/600 tickets (0.24s each)
One at a time: 144.35 seconds (4.2 items/sec)
```

## Batching by Length

Now 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.

``` python
import os
import time
import inspect
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
BATCH_SIZE = 32

torch.set_num_threads(os.cpu_count() or 1)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# keep the last real token at index -1
tokenizer.padding_side = "left"

model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)
model.eval()

LABELS = ["billing", "technical", "account"]
label_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS]
assert len(set(label_first_ids)) == len(LABELS), (
    "Labels share a first token; score full label sequences instead."
)
label_first_ids = torch.tensor(label_first_ids, device=model.device)

_forward_params = inspect.signature(model.forward).parameters
if "logits_to_keep" in _forward_params:
    LAST_LOGIT_ONLY = {"logits_to_keep": 1}
elif "num_logits_to_keep" in _forward_params:
    LAST_LOGIT_ONLY = {"num_logits_to_keep": 1}
else:
    LAST_LOGIT_ONLY = {}

def build_prompt(ticket):
    messages = [
        {
            "role": "system",
            "content": "You classify support tickets. Answer with exactly one of: billing, technical, account.",
        },
        {"role": "user", "content": f"Ticket: {ticket}\nCategory:"},
    ]
    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

BASE_TICKETS = [
    "My card was charged twice for the same invoice.",
    "The mobile app crashes whenever I open the settings page.",
    "I need to change the email address on my profile.",
]
FILLER = (
    "I have been waiting for a response for several days now and would really "
    "appreciate an update on this whenever someone gets a chance to look at it."
).split()

rng = np.random.default_rng(0)
target_words = np.clip(rng.lognormal(np.log(60), 0.9, size=600), 12, 400).astype(int)

def make_ticket(base, n_words):
    words = base.split()
    while len(words) < n_words:
        words += FILLER
    return " ".join(words[:n_words])

tickets_var = [make_ticket(BASE_TICKETS[i % 3], int(n)) for i, n in enumerate(target_words)]
prompts = [build_prompt(t) for t in tickets_var]
token_lengths = [len(tokenizer(p, add_special_tokens=False)["input_ids"]) for p in prompts]

print(
    f"Prompt lengths: min {min(token_lengths)}, "
    f"median {int(np.median(token_lengths))}, "
    f"max {max(token_lengths)} tokens"
)

def run_batched(order, batch_size):
    predictions = [None] * len(prompts)
    processed_tokens = real_tokens = 0
    start = time.time()
    for i in range(0, len(order), batch_size):
        idx = order[i:i + batch_size]
        batch = tokenizer(
            [prompts[j] for j in idx],
            add_special_tokens=False,
            padding=True,
            return_tensors="pt",
        ).to(model.device)
        processed_tokens += batch["input_ids"].numel()
        real_tokens += int(batch["attention_mask"].sum())
        with torch.no_grad():
            logits = model(**batch, **LAST_LOGIT_ONLY).logits[:, -1, :]
        best = logits[:, label_first_ids].argmax(dim=-1)
        for slot, choice in zip(idx, best.tolist(), strict=True):
            predictions[slot] = LABELS[choice]
    return predictions, time.time() - start, processed_tokens, real_tokens

def classify_one(prompt):
    """Reference path: a single unpadded sequence. Used only to verify."""
    inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
    with torch.no_grad():
        logits = model(**inputs, **LAST_LOGIT_ONLY).logits[0, -1, :]
    return LABELS[int(logits[label_first_ids].argmax())]

order = sorted(range(len(prompts)), key=lambda i: token_lengths[i])
predictions, duration_batched, processed, real = run_batched(order, BATCH_SIZE)

print(f"Length-bucketed batching: {duration_batched:.2f} seconds ({len(prompts) / duration_batched:.1f} items/sec)")
print(f"Padding overhead: {100 * (1 - real / processed):.1f}% of processed tokens were padding")

# correctness: padded rows must score the same as unpadded ones. check a spread of
# lengths rather than all 600, since the point is to catch a padding-side or
# position bug, and such a bug shows up on the very first padded row.
probe = order[::60]
mismatches = [i for i in probe if classify_one(prompts[i]) != predictions[i]]
print(f"Batched vs unbatched agreement on {len(probe)} probes: {len(probe) - len(mismatches)}/{len(probe)}")
assert not mismatches, f"Batched path disagrees at indices {mismatches}"

# estimate the per-item cost on the same probe set, then scale
start = time.time()
for i in probe:
    classify_one(prompts[i])
```

Output:

```
Loading weights: 100%|████████████████████████████████████████████████████████████████████| 290/290 [00:01<00:00, 252.11it/s]
Prompt lengths: min 48, median 94, max 449 tokens
Length-bucketed batching: 79.60 seconds (7.5 items/sec)
Padding overhead: 7.6% of processed tokens were padding
Batched vs unbatched agreement on 10 probes: 10/10
```

A 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.

- 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.
- 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.
- 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.
- 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.
- 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.
- 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.

One 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.

## Wrapping Up

This 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.

None 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.

To recap the series:

- **[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
- **[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
- **[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

 

 

[**\[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.