# Optimizing Small Language Models Through Intelligent Batching Strategies
## Introduction
Small language models (SLMs) have emerged as a powerful tool for focused automation tasks, but their real-world speed often lags behind what their compact size would suggest. One of the primary culprits behind underperforming inference pipelines is a deceptively simple loop: processing one input item per forward pass. On CPUs and GPUs alike, serving a single sequence at a time forces the hardware to repeatedly stream all model weights out of memory, leaving computational units largely idle between passes.
This article explores an approach that radically improves throughput: sorting inputs by their token length before grouping them into batches. By ensuring each batch contains similarly sized sequences, the amount of wasted computation on padding tokens drops dramatically. We’ll walk through the concepts, the implementation, and the measurable improvements this technique yields.
## The Problem with Item-by-Item Processing
When an inference pipeline processes one input at a time, the hardware is clipping along at a fraction of its potential. At batch size 1, the system becomes memory-bandwidth bound rather than compute-bound: the processor streams every weight out of memory to handle a single short sequence, then does it again for the next one. The arithmetic units spend most of their time waiting for data to arrive from memory.
This inefficiency scales with the number of the samples in a dataset and becomes especially costly when dealing with a production volume of inputs. The arithmetic logic units sit practically dormant while memory transfers dominate execution time.
naive batching approach — grouping items together without regard to their length — worsens the problem in a different way. When sequences of vastly different lengths are grouped together, every item in the batch gets padded to match the longest sequence. In datasets with heavy-tailed length distributions (common in real-world text), the longest item can be several times longer than the median. Padding everything to that maximum means most of the computation targets meaningless padding tokens instead of actual content.
## The Sorting Solution
The answer lies in preprocessing the dataset before batching. By sorting all inputs by their token count and then forming contiguous chunks, each batch contains items of similar length. The padding required within any single batch stays minimal, which means the hardware spends almost all of its compute cycles on meaningful work.
This technique requires only two steps:
1. **Measure** the token length of every input in the dataset.
2. **Sort** by length and divide into fixed-size groups.
The overhead of sorting is trivial compared to the inference savings. In practical benchmarks, this approach keeps padding overhead to roughly a single digit percentage of total processed tokens, compared to values closer to three or four times the necessary computation on unordered data.
## A Practical Demonstration
To illustrate the technique, consider a classification task where support tickets are sorted into categories such as billing, technical, or account issues. The setup uses a Qwen2.5-0.5B-Instruct model running in float16 precision. The hardware platform is a laptop-class chip with integrated memory and a neural processing engine, representative of the environments where 0.5B parameter models are most commonly deployed in practice.
### Dataset Characteristics
A synthetic dataset of 600 support tickets was generated with a realistic length distribution. Most tickets are under 100 tokens, while a long tail extends to approximately 450 tokens. The median sits around 94 tokens, while the maximum reaches close to 450 tokens. If every item were padded to the global maximum, the total token budget would expand by nearly four times what is actually necessary.
### Baseline: Sequential Processing
Processing each ticket individually takes approximately 144 seconds for the full dataset, yielding around 4 to 5 items per second. Every pass reads the model weights from memory, processes a single short sequence, and discards the result. The arithmetic units barely get warmed up.
### Optimized: Sorted Batching
The same 600 tickets, when sorted by length and processed in batches of 32, complete in roughly 80 seconds — approximately 7 to 8 items per second. That represents better than a 1.8x improvement in throughput using identical hardware and an identical model. The padding overhead stays under 8% of total processed tokens, a small price for a substantial speedup.
## Key Implementation Considerations
Several details matter when putting this approach into practice.
**Left Padding Is Essential** When padding sequences within a batch, the padding tokens must be added to the left side of each input. This guarantees that the last non-padding token always occupies the final position index regardless of where padding begins. If right padding were used instead, the final token index would land on a pad for every sequence shorter than the batch maximum, producing meaningless predictions.
**Fetch Only the Final Logits** A causal language model produces a logit vector for every position in the input by default. When batching long sequences at moderate batch sizes, the resulting logit tensor can consume multiple gigabytes per batch — and it is thrown away immediately except for the slot corresponding to the final position. Modern transformer libraries offer parameters to suppress all but the final set of logits during the forward pass (`logits_to_keep` or `num_logits_to_keep` depending on version). This optimization is negligible for single short inputs but dominant when batching longer sequences.
**Preserve Original Indexing** Sorting reorders the dataset, so the mapping between sorted positions and original input slots must be preserved. Writing predictions back into an output array using the original indices prevents a subtle but expensive bug where correct-looking outputs end up attached to the wrong inputs.
**Choose Batch Size Empirically** The ideal batch size depends on both the hardware capabilities and the length distribution of the dataset. Throughput typically rises steeply as batch size increases, then plateaus once compute capacity is saturated. Beyond that point, increasing the batch size only raises the risk of memory exhaustion on the longest items in a bucket. Running a quick sweep to identify the sweet spot is recommended whenever the dataset or hardware changes.
**Composing with Prefix Caching** The sorted batching approach can be combined with prompt prefix caching from prior optimization work. However, the cache structure needs careful handling when the batch dimension grows beyond 1. Key and value tensors must be expanded along the batch axis to match each group, and the cache entries need to correspond correctly to the sequences consuming them. Always verify predictions against the unbatched path when combining optimizations to confirm they compose correctly rather than interfering with each other.
## FAQ
### What hardware is needed to benefit from length-bucketed batching?
The technique benefits any environment where small language models are deployed — CPUs with limited cores, integrated GPU setups, and dedicated GPU hardware alike. The greatest wins are seen on memory-bandwidth-constrained devices (like laptops or edge servers), since those are where the overhead of per-item processing is most painful.
### Does sorting change the model’s predictions?
No. Sorting is purely a scheduling optimization that changes *when* each input is processed, not *how* it is processed. Verification against the unbatched sequential path should always confirm identical outputs.
### How much speedup can I expect?
Results vary by model size, hardware, and dataset length distribution, but in the demonstrated scenario a well-tuned sorted batching approach achieved roughly an 1.8x reduction in total processing time. Datasets with greater length variance tend to benefit even more, since naive padding waste is proportionally larger.
### What batch size should I use?
Batch size 32 is a common starting point, but the optimal value should be measured for your specific model, hardware, and data. If you see diminishing returns or out-of-memory errors on your longest sequences, reduce the batch size. If the GPU or CPU is well below utilization, increase it.
### Does this work with any language model architecture?
Length-bucketed batching is agnostic to architecture, but left padding is safe for models using rotary position embeddings (like the Qwen family). For models relying on learned absolute position embeddings, you will need to supply custom position IDs derived from the attention mask to preserve correctness.
### Can I combine this with other SLM optimizations?
Yes. Length-bucketed batching pairs well with constrained scoring (restricting outputs to a known label set) and prefix key-value caching (reusing computation on static prompt portions). Each optimization addresses a different bottleneck, and when verified independently, they tend to compose additively.
## Conclusion
Length-bucketed batching is one of the most impactful yet straightforward optimizations available for small language model inference pipelines. By sorting inputs before forming batches, practitioners can achieve substantial throughput gains — often nearly double the speed of sequential processing — without changing the model, the hardware, or the predictions themselves.
The core insight is counterintuitively simple: the biggest source of waste in a per-item loop is the repeated weight loading for each individual pass, and naive grouping amplifies this by forcing every item to pay for the longest sequence in its batch. Sorting breaks the cycle by grouping similarly sized inputs together so that padding stays minimal and the hardware stays busy doing useful computation.
This approach sits within a broader optimization mindset: instead of trying to make the model itself faster, restructure the work so that the hardware spends less time waiting and more time computing. Whether deployed on a laptop CPU, a consumer GPU, or a data center accelerator, the same principle applies — measure your data, sort your inputs, and batch intelligently.
The elegance of this technique is that it requires no model retraining, no specialized libraries beyond a standard transformer toolkit, and no approximation of any kind. The speedup is real, the outputs are identical, and the implementation can be added to an existing pipeline in a single pass over the dataset.
Thank you for reading
