Variable-Length Batching

The primary goal of variable-length batching is to reduce the computational waste caused by padding tokens. In static batching, all sequences in a batch are padded to the length of the longest sequence, forcing the GPU to perform useless calculations on dummy tokens. Variable-length techniques group sequences by similar lengths or use specialized data structures to process only the real tokens, dramatically improving FLOP efficiency and reducing iteration time.

• Example: A batch with sequences of lengths 10, 50, and 100 tokens would require 300 tokens of computation with padding (100x3). Variable-length batching can process these lengths natively, performing only 160 token computations.

Efficient implementation relies on data structures and compute kernels designed for non-uniform data. A ragged tensor (or jagged array) is a core data structure that stores a concatenated list of variable-length sequences alongside metadata (offsets) to track boundaries. Specialized GPU kernels are then required to operate directly on these ragged tensors, performing operations like attention and matrix multiplication without explicit padding.

Frameworks like PyTorch support ragged tensors via custom extensions or third-party libraries. The performance gain comes from eliminating the memory reads, writes, and arithmetic operations on padding tokens, which is critical in memory-bound phases like decoding.

Variable-length batching is a foundational enabler for continuous batching (also known as iteration-level scheduling). In a continuous batching system, new requests can join the active batch, and completed sequences can exit, at each decoding iteration. This dynamic batch composition inherently contains sequences at different stages of generation, resulting in highly variable lengths.

Without variable-length support, the scheduler would be forced to pad all active sequences to the length of the longest one at every step, nullifying the benefits of continuous execution. Therefore, variable-length processing is essential for achieving high GPU utilization in modern, interactive inference servers.

This technique has a direct and measurable impact on key inference metrics:

• GPU Memory: Reduces the peak memory footprint by avoiding storage of large padding tensors. This allows for larger effective batch sizes within fixed memory constraints.
• Throughput: Increases the number of tokens processed per second (Tokens/sec) by eliminating wasted compute on padding. The gain is most significant when sequence length variance within a batch is high.
• Latency: Can reduce tail latency by preventing short requests from being delayed behind long ones due to excessive padding (mitigating head-of-line blocking).

The trade-off is increased complexity in memory management and kernel implementation compared to uniform, padded tensors.

The scheduler's algorithm for grouping requests is critical. Naive grouping can lead to fragmentation. Effective heuristics include:

• Length-Aware Binning: Grouping incoming requests into bins based on prompt length or current generated length.
• Dynamic Re-packing: Periodically re-organizing the active batch in memory to consolidate free space left by completed sequences.
• Cost-Aware Scheduling: Estimating the computational cost of a sequence based on its length and prioritizing grouping to balance the load across streaming multiprocessors (SMs) on the GPU.

These heuristics work alongside the batching policy (e.g., using a batch window or batch timeout) to decide when to form a new batch for the prefilling phase.

It's important to distinguish variable-length batching from related techniques:

• Static Batching: Uses a fixed batch composition where all sequences are padded to a uniform length. Simple but inefficient for variable workloads.
• Dynamic Batching: Groups requests based on a time window, but often still uses internal padding. It's a scheduling policy that can be combined with variable-length execution.
• Variable-Length Batching: Specifically refers to the execution mechanism that avoids padding. It is the computational core that makes efficient dynamic and continuous batching possible.

In practice, a production inference server uses a dynamic batching policy that schedules requests into a request queue, and a variable-length execution engine to process them.

The efficient implementation of variable-length batching in any framework relies on low-level computational primitives.

• Ragged Tensors: A data structure that represents a nested variable-length list as a single tensor object. It stores values in a flat buffer and uses a separate structure to track offsets. This allows frameworks to pass batches of variable-length sequences to optimized kernels without manual padding.
• Custom CUDA Kernels: Specialized kernels are written to directly read from ragged tensor formats or to use sequence length arrays to skip computations on padding tokens. This eliminates the idle cycles spent processing meaningless padding.
• Compiler Support: Compilers like TorchInductor (PyTorch) and XLA (JAX/TensorFlow) are increasingly generating code that can efficiently handle dynamic shapes, reducing the need for framework-specific kernels. These enablers are what transform the high-level concept of variable-length batching into realized GPU efficiency.

An inference optimization technique that dynamically groups incoming requests into batches at each model iteration. This allows:

• New requests to join an active batch.
• Completed sequences to exit the batch immediately.
• Maximized GPU utilization by eliminating idle time between static batches. It is the overarching scheduling paradigm that enables variable-length batching within each iteration.

A request scheduling strategy that groups multiple inference queries based on a time window or queue size. It contrasts with static batching by forming batches adaptively as requests arrive. While dynamic batching decides when to form a batch, variable-length batching solves the problem of how to efficiently execute sequences of different lengths within that batch once formed.

The fine-grained mechanism within continuous batching where the scheduler makes grouping decisions for the active requests at each decoding step. This is where variable-length batching is applied. The scheduler must decide, for every iteration:

• Which sequences are still active.
• How to pack their varying token counts for the next forward pass.
• This enables real-time adaptation to the changing composition of the batch.

A specialized data structure that efficiently stores and operates on sequences of varying lengths without padding. It is a fundamental enabler for variable-length batching.

• Structure: Contains a flat values array and a nested row-splits array defining sequence boundaries.
• Benefit: Eliminates the memory and compute waste of padding tokens.
• Use: Specialized GPU kernels can process ragged tensors directly, performing the variable-length batching logic at the hardware level.

The process of adding dummy tokens to shorter sequences so all sequences in a batch have a uniform length. This is the problem variable-length batching aims to solve.

• Overhead: Padding tokens consume memory bandwidth and compute cycles without contributing to the useful output.
• Impact: Inefficiency grows with greater variance in sequence lengths. Variable-length batching via ragged tensors or specialized kernels directly minimizes this overhead.

Two fundamental performance regimes in inference that variable-length batching must navigate.

• Memory-Bound: Operation limited by memory bandwidth (e.g., autoregressive decoding). Here, variable-length batching aims to reduce wasted memory reads/writes for padding.
• Compute-Bound: Operation limited by processor speed (e.g., prefill phase). Here, variable-length batching must ensure efficient, aligned compute operations despite irregular shapes. Optimized kernels for variable-length batching are designed to perform well in both regimes.