Yash Datta

MRV2: how vLLM rewrote its model runner, in six architectural moves

6,700 lines became modules none over ~1,300, with 56% more output throughput on Qwen3-0.6B / 1×GB200 and 6.3% lower TPOT on GLM-4.7-FP8 / 4×GB200. The rewrite works by deleting six assumptions V1 hardcoded.

LLM inference 8 min read

TL;DR. vLLM rebuilt its execution core as Model Runner V2 (MRV2), announced March 2026 and shipping since v0.17 behind the VLLM_USE_V2_MODEL_RUNNER=1 flag (not the default yet). The gpu_model_runner.py monolith had grown past 6,700 lines; the largest file in MRV2 is now under 1,300. vLLM’s own benchmarks: 25K vs 16K output tok/s on Qwen3-0.6B with 1×GB200 (a 56% gain), and 6.3% lower TPOT on 4×GB200 with GLM-4.7-FP8 and MTP=1.

Six architectural decisions carry the speedup. Each one deletes an assumption V1 hardcoded.

1. Decouple persistent state from step inputs

V1’s “persistent batch” used the same tensors for state and model inputs. When a request joined or finished mid-batch, it forced tensor-wide reordering to maintain layout, plus a redundant CachedRequestState backup because rows could be overwritten while the request was still alive.

MRV2 splits them:

# Pre-allocate once. Max 1024 rows.
# Each live request owns a permanent row until it finishes.
# Step inputs are gathered from state by the attention backend's ordering.
state = PersistentState(max_num_reqs=1024)
inputs = gather_inputs(state, scheduling_order)

Preemption is treated as completion: on resume the request is re-added as fresh state, so CachedRequestState disappears. Gather is a parallel GPU op with negligible overhead.

Takeaway: when your hot-path state table doubles as an I/O format, every scheduling event becomes a layout migration. Separate the two and you get back simple row allocation.

2. Kill the async barrier without adding a lock

Async scheduling means CPU prepares step N+1 while GPU executes step N. The CPU can mutate a buffer while the GPU is mid-async-copy from it: a classic race.

V1’s fix: an explicit async barrier around critical sections. Bug-prone, kills overlap.

MRV2’s fix is one line:

# V1 (race condition):
self.states = torch.zeros(..., pin_memory=True)
self.states[req_idx] = new_req.data
states_gpu = self.states.to("cuda", non_blocking=True)  # CPU may still write here

# MRV2:
self.states = torch.zeros(..., pin_memory=False)   # persistent, unpinned
self.states[req_idx] = new_req.data
tmp = self.states.pin_memory()                     # snapshot
states_gpu = tmp.to("cuda", non_blocking=True)     # GPU reads tmp, CPU keeps writing self.states

The throwaway pinned copy is the snapshot boundary. Now CPU writes and GPU reads never touch the same bytes.

Takeaway: a copy can be cheaper than a barrier when the data is small and the barrier kills overlap.

3. StagedWriteTensor: ragged diffs, one kernel launch

Block tables are too big to copy CPU→GPU every step. Updates are ragged: different rows, different offsets, different lengths.

state = StagedWriteTensor(size=(1024, 1000), dtype=torch.int32, device="cuda")

state.stage_write(row=2, start=3, value=[3, 1, 2])     # ragged
state.stage_write(row=0, start=1, value=[-1, -2, -5])  # ragged

state.apply_write()  # packs diffs → one H2D copy → one kernel

Base tensor stays on GPU. Diffs stage on CPU. Pack into a contiguous buffer, one H2D copy, one kernel applies all of them.

Takeaway: ragged update patterns have a natural compile step: batch the edits, serialize once, launch once.

4. Input prep on the GPU (via Triton)

V1 built input_ids, positions, query_start_loc, seq_lens on the CPU in Python. That is slow for large batches, and it is a sync point: the GPU often knows values the CPU does not yet (e.g., accepted tokens from speculative decoding).

MRV2 runs input prep as Triton kernels on the GPU. Plus UVA (Universal Virtual Addressing) for large CPU-resident tensors like prefill_token_ids, so GPU kernels access them directly without duplicating to GPU memory.

Takeaway: when the GPU already knows the value, do not round-trip to the CPU to compute it. And when the CPU-resident data is huge, map it instead of copying it.

5. Gumbel-Max sampling: no softmax materialization

Standard LLM sampling: compute logits → softmax → sample from categorical → get token.

Gumbel-Max identity: argmax(logits + gumbel_noise) is an exact sample from softmax(logits). Never materialize the probability distribution.

MRV2’s Triton Gumbel kernel uses stateless in-kernel RNG (seed → noise, no RNG state threaded between steps). Pairs with two more tricks:

  • Top-k first, then logprobs: V1 materialized full-vocabulary logprobs then took top-k. MRV2 finds top-k indices from logits directly, then computes logprobs only for those indices.
  • Chunked prompt logprobs: chunk inside a single long prompt to avoid memory spikes.

Takeaway: if you only need a sample or the top-k, you never need the normalized distribution. Gumbel-Max draws an exact softmax sample without forming softmax; top-k indices come straight from logits, with logprobs computed only for those.

6. Explicit CUDA graph management

V1’s CUDA graph handling was implicit, interleaved with warmup, memory profiling, empty DP forward passes, torch.compile tracing, all shoved into dummy_run.

MRV2:

  • execute_model supports dummy runs without affecting state.
  • dummy_run delegates to execute_model for profiling/warmup.
  • CUDA graph capture is its own path, owned by a CUDAGraphManager using standard PyTorch APIs.

One concrete win: MRV2 can capture multiple draft-model forward passes into a single CUDA graph, which matters for speculative-decoding throughput.

Takeaway: when one function is overloaded as “do all setup-like things,” the divergence bugs compound. Separate paths are worth the LOC.

The architecture in one picture

MRV2 dataflow: the CPU schedules and stages ragged diffs; the GPU holds persistent state, preps inputs with Triton kernels, runs the model under a CUDA graph, and samples with Gumbel-Max, with no CPU sync in the hot path.
Figure 1. The hot path. CPU schedules and stages diffs that apply with one H2D copy plus one kernel. Inputs prep on-device. The model runs under a CUDA graph. The sampler never forms softmax.

What this teaches beyond vLLM

Five patterns that transfer to any throughput-bound system:

  1. Separate state from wire format. Anything that mutates frequently shouldn’t also be your serialization layout.
  2. Snapshot instead of synchronize when the cost of a copy is less than the cost of stalling the pipeline.
  3. Batch the ragged edits. Stage → pack → one launch. Universally cheaper than N small operations.
  4. Compute on the machine that owns the data. CPU-resident bookkeeping that the GPU already knows → move the compute to the GPU (or UVA the tensor in place).
  5. Mathematical equivalents are an optimization surface. Softmax → Gumbel-Max. Full-vocab logprobs → top-k-first-then-logprobs. The textbook form is rarely the cheapest form.

Status and further reading

MRV2 is still moving fast. As of writing, LoRA, DBO, logits processors, and speculative decoding beyond Eagle / Eagle3 / MTP remain V1-only; EPLB and piecewise CUDA graphs for pipeline parallelism have already landed in MRV2. Enable with VLLM_USE_V2_MODEL_RUNNER=1; it is not the default yet.

The primary sources, all worth a read on their own: