---
title: Token-Balanced Batching
url: https://www.emergentmind.com/topics/token-balanced-batching
type: topic
---

# Token-Balanced Batching

Token-balanced batching is a family of batching regimes that aim to allocate computational work according to the number of tokens—rather than merely the number of requests or samples—in each batch. This approach is motivated by the need to maximize hardware utilization, stabilize learning dynamics, and control latency in both training and inference of large-scale models, particularly when serving requests or data with highly variable token counts. Token-balanced batching has found applications in distributed LLM inference (e.g., pipeline and tensor parallel settings), large-scale vision model training with variable input resolutions, and in queueing-theoretic analyses of inference workloads.

## 1. Core Principles and Motivations

Conventional batching strategies typically impose limits based on the number of requests per batch or the total token count. In settings where requests or samples have highly variable lengths—such as autoregressive LLM inference or training vision transformers on mixed-resolution images—this can result in suboptimal hardware utilization, large gradient imbalances, or unacceptably long tail latencies. Token-balanced batching explicitly targets the total token count per batch as the primary workload determinant. The central goals are:

- To maintain balanced computational loads across devices or pipeline stages.
- To smooth per-iteration or per-batch token counts, reducing "valleys" of under-utilization and eliminating peaks that risk memory exhaustion.
- To decouple batch sizing from the number of items, focusing instead on cumulative token workload.
- To guarantee fairness and optimization stability when items contribute vastly different token counts, e.g., mixing high- and low-resolution images or handling LLM requests with widely varying output lengths [2512.20157], [2407.05347], [2504.14775].

## 2. Algorithmic Approaches and Formulations

Algorithmic realizations of token-balanced batching vary by domain but share a computational bin-packing ethos: batches are constructed to greedily (or optimally) approach, but not exceed, a specified token budget, sometimes under constraints on the maximum number of items or specific resource limits.

### 2.1. LLM Inference: Memory-Centric and Elastic Batching

**Memory-centric batching** replaces the conventional cap on requests with a cap on the incremental key-value (KV) memory that additional prefill tokens would require, ensuring efficient use of GPU memory and maximizing the size of each token batch. The batch size in tokens $B_{\text{tok}}$ is thus

$$
\text{maximize } B_{\text{tok}} \quad \text{subject to } \sum_{i} \text{KVsize}(i) \leq M_\text{threshold},\quad \text{chunk size} \leq S_\text{chunk}
$$

where $M_\text{threshold}$ is set based on available KV-cache at each iteration [2412.03594].

**Elastic batching** (from a queueing-theoretic perspective) adapts batch formation so requests with shorter output lengths can exit earlier, rather than being forced to wait for the longest in the batch. This reduces mean waiting time and latency for the majority of requests [2407.05347]. The elastic-batch service time is modeled as:

$$
H^{elastic} = k_1 b + k_2 + k_3 b E[N] + k_4 n_b
$$

where $b$ is batch size, $E[N]$ average output length, and $n_b$ maximum output length in the batch.

### 2.2. Distributed Training: Input Packing for Variable-Length Data

In multi-resolution vision training, token-balanced batching relies on packing as many images as possible into a batch, provided the aggregate token count does not exceed a fixed maximum sequence length $C_{\text{max}}$ [2512.20157]. The bin-packing process per device is:

1. Order/shuffle candidate images.
2. Fill each sequence by greedily aggregating images until adding the next would exceed $C_{\text{max}}$.
3. Apply block-diagonal attention masks (e.g., FlexAttention) to maintain independence across images.
4. Normalize losses by per-image token count, stabilizing gradients across resolutions.

## 3. Scheduling, Grouping, and Runtime Coordination

In complex parallel or distributed systems, token-balanced batching is integrated with advanced scheduling and grouping policies.

- **Global Prefix Sharing**: BatchLLM constructs a radix prefix-tree over all prompts to identify prefix-sharing groups, maximizing KV-cache reuse and minimizing redundant computation [2412.03594].
- **Reordering by Decode Ratio**: Requests/groups with higher decoding-to-prefill ratios are prioritized so that decoding tokens can be interleaved with prefill workloads of later (longer) groups, which sustains GPU occupancy throughout batch iterations.
- **Independent Prefill and Decode Control Loops**: gLLM decouples prefill and decode token budgets using distinct "rate controllers" driven by live system state, such as pending token queues and observed KV-cache headroom. The budgets are:

  - **Prefill $\#P$**: Adjusted by tokens-waiting and KV-cache utilization controllers, throttled or suspended under memory pressure.
  - **Decode $\#D$**: Kept near-uniform across pipeline stages; assigned via
    $$
    \#D = \lceil \#RD / \text{PP}_\text{depth} \rceil
    $$

  where $\#RD$ is decode tokens and $\text{PP}_\text{depth}$ pipeline depth [2504.14775].

- **Non-blocking Scheduling and Preemptive Metadata Delivery**: High-performance runtimes (e.g., gLLM) use asynchronous message passing and overlap I/O with computation to eliminate pipeline bubbles due to imbalanced batch composition.

## 4. Empirical Results and Quantitative Impact

Empirical studies across LLM inference and vision model training have consistently demonstrated the advantages of token-balanced batching:

- **LLM Inference**:
  - gLLM achieves throughput increases from 11% up to 398% over strong baselines such as vLLM and SGLang, with marked reductions in end-to-end latency and improved deadline attainment for real-time workloads [2504.14775].
  - BatchLLM raises steady-state GPU occupancy from ~50–60% to >90%, dramatically increasing per-iteration token utilization and throughput by 1.3×–2.0× in practical settings [2412.03594].
  - Queuing-theoretic elastic batching yields mean waiting time reductions of 30–50% compared to static batch regimes; optimally chosen batch-size and token limits cut waiting times by an order of magnitude under realistic arrival rates [2407.05347].

- **Vision Model Training**:
  - Packing images up to a token budget (with loss normalization) prevents catastrophic forgetting of low-resolution data once high-resolution images are introduced, maintaining or improving representation accuracy for all resolutions [2512.20157].
  - Hardware throughput increases from ~7.5k tokens/s to ~20k tokens/s per GPU with token-balanced batching.

## 5. Analysis and Theoretical Underpinnings

Theoretical analyses of token-balanced batching draw from both queueing theory and optimization:

- **M/G/1 Bulk-Service Models**: Token-balanced batching is analyzed using M/G/1 (and M/Dⁿ/1) bulk-service queues, where batch processing time depends jointly on batch size and the longest token sequence. Mean waiting times admit closed-form upper bounds in terms of batch parameters, guiding the optimal choice of batch size and token limits for latency-sensitive applications [2407.05347].
- **Worst-Case Bounds**: Imposing a maximum token-length per request (clipping at the 95th percentile, for example) provides strict upper bounds on per-batch latency, shrinking tail waiting times even if a small proportion of requests are curtailed [2407.05347].

A summary of queueing regimes and their impact:

| Batching Policy    | Queue Model | Service Time Dependency             | Mean Waiting Time Improvement        |
|--------------------|-------------|-------------------------------------|--------------------------------------|
| Dynamic            | M/G/1 bulk  | $T(n, m) = f(n, m_{\max})$         | Baseline                             |
| Fixed-size         | M/Dⁿ/1      | $T(n, \mathbb{E}m)$                 | Cuts $E[W]$ at optimal $n$           |
| Elastic            | --          | Immediate completion per request    | Reduces $E[W}$ further by 30–50%     |

## 6. Implementation Considerations and Practicalities

Practical deployment of token-balanced batching requires attention to several factors:

- **Token Budget Selection**: Empirical tuning of $C_{\text{max}}$ (vision) or $M_\text{threshold}$ (LLM) is necessary to balance throughput, memory usage, and latency. In vision, $C_{\text{max}}$ is typically set to just fit a 768×768 image (≈2,500–4,096 tokens).
- **Per-item Loss Normalization**: When batch items contribute variable tokens, losses (e.g., MSE) are normalized by per-item token count ($1/N_q$) to prevent gradient bias towards longer samples [2512.20157].
- **Attention Masking**: Batching together distinct sequences or images in vision models requires attention masking (e.g., block-diagonal FlexAttention) to prevent cross-sample gradient contamination.
- **Workload Skew Management**: Aggressive batch packing can risk starvation or over-representation of certain item types (e.g., low-res images or short LLM completions); scheduling heuristics and group sorting by decode-to-prefill ratio alleviate this.

## 7. Limitations and Trade-offs

Token-balanced batching, while highly effective, introduces several trade-offs:

- **Latency vs. Throughput**: Maximizing token budget may delay short requests in low-load conditions; elastic/early-exit departures partially mitigate this.
- **Complex Scheduling Overhead**: Construction of global prefix groups, per-group reordering, and memory tracking impose modest, though generally sub-second, scheduling overhead for thousands of requests [2412.03594].
- **Fairness under Token Clipping**: Setting token-length caps can sacrifice the experience of the small fraction of requests seeking very long outputs; selection of percentile thresholds balances user utility against system efficiency.

A plausible implication is that in multi-tenant production environments, token-balanced batching must be configured in conjunction with application-level SLOs to realize full system benefit.

---

**References**:  
gLLM: [2504.14775]  
BatchLLM: [2412.03594]  
AMoE: [2512.20157]  
Queueing-Theoretic Analysis: [2407.05347]

Source: https://www.emergentmind.com/topics/token-balanced-batching