---
title: Dynamic Rebatching for Efficient Inference
url: https://www.emergentmind.com/topics/drex
type: topic
---

# Dynamic Rebatching for Efficient Inference

Dynamic rebatching is an execution algorithm for early-exit (EE) large language model inference that independently separates requests at each early-exit point according to their confidence decisions. Requests satisfying the exit criterion generate tokens immediately, while requests requiring deeper computation are logically buffered, regrouped, and forwarded through subsequent layers. In “Dynamic Rebatching for Efficient Early-Exit Inference with DREX” [2512.15705], Dynamic Rebatching is distinguished from DREX: the former is the central execution mechanism, whereas DREX is the complete inference-serving system that adds copy-free buffering, virtual-memory handling of missing KV-cache state, Adaptive Rebatching Threshold (ART) selection, and SLA-aware scheduling.

## 1. Terminological identification and scope

The supplied material contains no literal placeholder named `[PERSON_NAME]`. In the context of the early-exit inference paper, the intended term is **Dynamic Rebatching**. It should not be treated as synonymous with DREX, whose name expands to “Dynamic Rebatching for Efficient Early-Exit Inference.”

Dynamic Rebatching addresses the mismatch between independent EE decisions and conventional batch execution. At an EE ramp, each request or token is evaluated separately. A request whose confidence exceeds the configured threshold may exit at the ramp, whereas a request with insufficient confidence continues through deeper layers. Conventional batching frameworks generally require a single action for the entire batch, despite these heterogeneous decisions.

A separate supplied paper, “DexFLEX: Contact-Aware Foundation Controller for Command-Guided Dexterity” [2610.14372], also contains corrupted occurrences of `[PERSON_NAME]`. In that paper, the placeholder alternately replaces a method name, author names, task names, and possibly a teleoperation system. The available text therefore does not support assigning `[PERSON_NAME]` to a unique individual across both papers. For the early-exit topic, however, the identifiable technical term is Dynamic Rebatching.

## 2. Motivation: split decisions in early-exit inference

An EE language model contains intermediate EE ramps. After the layers preceding a ramp have executed, each request is evaluated through the model’s `getIndividualDecision(Tensor)` interface. The result is an `EEMask` containing one binary decision per request: \(1\) denotes early exit and \(0\) denotes continuation.

A batch may therefore contain both exiting and continuing requests. For example, requests \(s_2\) and \(s_4\) may satisfy the exit criterion while \(s_1\) and \(s_3\) continue to deeper layers. Traditional grouped policies resolve this split by imposing one batch-wide action:

- **Consensus** requires all requests to agree before the batch exits.
- **Greedy** exits the entire batch when at least one request wants to exit.
- **Majority** follows the majority decision.

These policies produce two distinct forms of deviation from the underlying EE model. An **involuntary exit** occurs when a request is forced to exit despite failing its own confidence criterion. An **involuntary stay** occurs when a request that is ready to exit is forced to continue. For batch sizes 4 and 8, the reported rates are:

| Policy | Involuntary exit, BS=4 | Involuntary exit, BS=8 |
|---|---:|---:|
| Consensus | 0 | 0 |
| Majority | 22.3% | 34.2% |
| Greedy | 25.9% | 35.1% |

Consensus avoids involuntary exits but produces involuntary stays of 29.7% and 32.1% for batch sizes 4 and 8, respectively. Majority produces both involuntary exits and involuntary stays, while Greedy produces involuntary exits but no involuntary stays. Dynamic Rebatching removes the need for a batch-wide decision by allowing every request to follow its own EE result.

## 3. Execution algorithm

Dynamic Rebatching proceeds through shallow execution, independent decision-making, logical buffering, conditional flushing, and deep execution.

First, the complete batch executes the layers preceding an EE ramp. The EE ramp evaluates every request independently and produces the corresponding `EEMask`.

Second, requests satisfying the confidence threshold immediately generate their output tokens. They are removed from the current deep-layer computation and returned to the scheduler for a subsequent decoding iteration unless they have reached EOS.

Third, requests failing the exit criterion enter a per-model rebatching buffer. In the example above, \(s_1\) and \(s_3\) are placed in the buffer. The buffer stores logical membership, request identifiers, and execution ordering rather than constructing a new tensor containing copied hidden states or KV caches.

Fourth, the buffer manager and scheduler decide whether to hold or flush the buffered requests. Continuing requests may be held until a sufficiently large deep batch forms, or flushed earlier when the batch is large enough, waiting time is excessive, or SLA deadlines are approaching. The basic batch-efficiency condition is

$$
b_{\mathrm{buffer}} \geq b_{\mathrm{scheduler}},
$$

where \(b_{\mathrm{buffer}}\) denotes the number of waiting requests and \(b_{\mathrm{scheduler}}\) denotes the size of the next fresh scheduler batch.

Fifth, when the buffer is flushed, only the selected continuing requests execute the deeper layers. Requests may have originated from different earlier batches. Attention kernels receive an index mapping identifying the physical KV-cache entries corresponding to logical requests. DREX specifically uses FlashAttention’s `cache_batch_idx` mechanism to support this reordering and selection without copying the underlying hidden states or KV caches.

The execution model distinguishes three iteration types: a **full iteration**, in which all requests execute all layers and require time \(t_f\); a **shallow iteration**, in which requests execute only through the EE ramp and require time \(t_s\); and a **deep iteration**, in which buffered continuing requests execute the remaining layers and require time \(t_d\). For a continuing request, the rebatching path is \(t_s+t_d\), and its overhead is

$$
c=t_s+t_d-t_f.
$$

The savings from skipping the deep layers are

$$
\mathrm{savings}=t_f-t_s=t_d-c.
$$

The reported examples give \(c=5.35\) ms and \(t_d=11.10\) ms for Llama-EE-13B with batch size 8, and \(c=7.92\) ms and \(t_d=33.30\) ms for Llama-EE-70B with batch size 8. For Llama-EE-70B, the rebatching overhead is reported as less than 6% of a standard iteration time.

## 4. Copy-free buffering and missing KV-cache state

A naïve rebatching implementation would gather hidden states and KV-cache entries into new contiguous tensors whenever a batch splits. DREX instead implements a copy-free rebatching buffer. The buffer maintains logical membership and ordering, while the underlying tensor storage remains in place. Kernels receive virtual indices identifying active sequences, so the deep batch is represented as a logical view rather than a physically copied batch.

This design relies on attention-kernel mappings, particularly `cache_batch_idx`. The paper reports that the resulting overhead is largely independent of model size and sequence length because hidden states and KV caches are not physically moved.

Early exits also create missing KV-cache state. If a token exits at layer \(i\), its KV entries for layers \(i+1,\ldots,L\) are not computed. A later token in the same sequence may nevertheless need those entries when it traverses the deeper layers. Two alternatives are discussed: KV recomputation, which incurs additional computation and heterogeneous attention-mask management, and physical state-copying, which avoids recomputation but creates memory redundancy.

DREX uses memory-efficient state-copying with virtual GPU memory mappings. When a token exits at layer \(i\), its last computed KV entry is identified, and the virtual KV-cache tensors for subsequent layers are mapped to the same physical memory block. The mappings are read-only for subsequent attention operations, so no physical duplication is performed. If the sequence later enters deeper layers, those layers consume the mapped entries, while the cache manager passes references and lazily fills or materializes missing state as necessary.

The implementation uses NVIDIA low-level GPU virtual memory support, the `vAttention` `vMemMap` interface, and vTensor-style virtual tensor management. The reported evaluation shows up to 18.3% lower CUDA memory-operation volume under the Greedy policy and a 5.7% average reduction. For a 4,000-token generation with Llama-EE-13B, traditional EE-LLM implementation is reported to incur up to 10.9% redundant KV cache, whereas virtual sharing avoids that redundancy.

## 5. Adaptive profitability and SLA-aware scheduling

Dynamic Rebatching is not necessarily profitable for every split. If only a small number of requests exit, the buffering, synchronization, scheduling, and post-processing costs may exceed the computation saved by skipping deeper layers.

Let \(b\) be the total batch size, \(b'\) the number of requests that want to exit, \(c\) the rebatching overhead, and \(t_d-c\) the saving per early-exiting request. Rebatching is profitable when

$$
b'(t_d-c)>(b-b')c.
$$

Rearrangement yields

$$
b'>\frac{c}{t_d}b.
$$

This threshold is called the **Adaptive Rebatching Threshold (ART)**. DREX updates \(c\) and \(t_d\) every 100 steps using profiled and averaged iteration latencies. For Llama-EE-13B with batch size 8, the reported ART is

$$
\frac{5.35}{11.10}\cdot 8=3.86,
$$

so at least four requests must want to exit. For Llama-EE-70B, the ART is

$$
\frac{7.92}{33.30}\cdot 8=1.90,
$$

allowing smaller exit subsets because the deeper layers are more expensive. For multiple EE ramps, the condition generalizes to

$$
ART(i)=\frac{c}{t_d^i}b.
$$

When ART rejects a split, the consequence may be an involuntary stay, because requests continue through deeper layers despite being eligible for exit. It does not create an involuntary exit.

Buffering can improve throughput while increasing request completion time (RCT). DREX therefore uses an SLA-aware scheduler that considers the next fresh scheduler batch size, the number of buffered requests, and the age of the oldest buffered request. The expected remaining completion time is defined as

$$
r_{\mathrm{expected}}=age+L-l,
$$

where \(age\) is the current age in iterations, \(L\) is maximum output length, and \(l\) is current generated length. The SLA completion requirement is converted from time to iterations using the profiled standard iteration time.

The scheduler can flush a small buffer before the ideal batch size is reached when requests approach their SLA deadlines. It also modifies vLLM-style preemption by preferentially evicting sequences in the rebatching buffer. Their KV caches may be offloaded to CPU memory or discarded for later recomputation. The reported results describe a throughput/RCT trade-off: with low SLA pressure, Rebatching improves throughput by 11.4% over Consensus, while average RCT is approximately 1.4 times higher and tail RCT 3 times higher. Increasing SLA pressure sacrifices throughput for lower completion time, and extreme pressure approaches conservative Consensus behavior.

## 6. Evaluation and system significance

DREX is implemented on Sarathi-Serve with approximately 1,500 lines of Python, 10 lines of CUDA for vAttention integration, and approximately 500 lines of Python for EE utilities. Evaluated models include Llama-EE-13B and Llama-EE-70B, using an A100 node with 80 GB of VRAM for Llama-EE-13B and an H200 node with 141 GB of VRAM for Llama-EE-70B. The workload is HELM-style CNN/Daily Mail text summarization, filtered to fewer than 2,048 input tokens and comprising 2,160 examples.

The evaluation measures throughput, RCT, EE-token confidence, BERTScore, EE proportion, involuntary exit, and involuntary stay. Baselines include Latency-only, Consensus, Majority, Greedy, Non-EE, and Dynamic Rebatching/DREX.

DREX improves throughput by 2–10.3% over the baseline EE policies, excluding Greedy, and by as much as 12% over non-EE Llama-EE-70B. It maintains equal or better P95 confidence than the alternatives. Greedy can achieve higher throughput but exhibits severely degraded confidence: its average P95 confidence score is reported as 96% worse, with a P95 confidence score as low as 0.03.

For multiple exits in Llama-EE-70B, with exits at layer 40 at confidence 0.7 and layer 60 at confidence 0.9, DREX improves throughput by 4.5–8.6% across batch sizes relative to the other policies. Its BERTScore is approximately comparable, within about 1%, and 11% higher than Greedy despite Greedy’s higher throughput.

The results also distinguish confidence preservation from semantic quality. BERTScore may decline as more tokens exit even when confidence is high, because current EE confidence does not perfectly correlate with end-to-end summary quality. This decline is attributed to the relationship between confidence and summary quality rather than to involuntary exits.

Dynamic Rebatching’s central guarantee is therefore structural: the system does not force a request to exit merely because another request in the batch exits. Independent EE decisions, conditional buffering, virtual cache mappings, ART-based profitability analysis, and SLA-aware flushing jointly preserve the underlying model’s intended per-token exit policy while exploiting available computation skipping.

Source: https://www.emergentmind.com/topics/drex