Knowledge note

vLLM MHA Code Path: From QKV to Paged KV Cache

A code-reading trace of decoder multi-head attention in vLLM, from model-side QKV projection and RoPE through runtime metadata, paged KV writes, backend dispatch, and the output projection.

Updated September 16, 2026Confidence: Medium

vLLM MHA Code Path: From QKV to Paged KV Cache

Repository: vllm-project/vllm

Inspected commit: 2d24355eb87b716fc1169e66731dc0386ed1a3a2

Scope: decoder self-attention in V1, using Llama’s model module and the FlashAttention backend as a concrete path. The same generic layer also carries grouped-query attention (GQA) and multi-query attention (MQA); encoder, cross-attention, multi-head latent attention (MLA), sparse attention, and platform-specific alternatives are outside this trace.

Verification boundary: clean, commit-pinned static reading. No model was served and no GPU kernel or performance benchmark was run. A scoped freshness check on 2026-08-16 found relevant upstream changes, but the repository’s 14-day evidence interval deferred a new snapshot until 2026-08-27.

The Short Version

vLLM’s multi-head attention (MHA) is not one function. It is a four-part handoff:

  1. A model-specific module projects hidden states into Q/K/V, applies rotary position embedding (RoPE), and later applies the output projection.
  2. The reusable Attention layer reshapes flat tensors into heads and crosses a compilation-safe custom-op boundary.
  3. GPUModelRunner supplies the changing serving state: request lengths, block tables, and the physical slot assigned to every new K/V token.
  4. The selected hardware backend writes new K/V into the paged cache and runs the actual attention kernel over that cache.

Intuition: the model file describes the Transformer layer; the model runner describes this iteration’s batch; the backend describes this device’s execution. The generic attention layer is the seam joining all three.

For the attention equation itself, start with The Transformer. For the memory mapping underneath the kernel, see vLLM Block Table Management.

The Most Important Distinction

The model-side module and the reusable attention layer are different objects. In the Llama example, LlamaAttention.__init__ owns the learned QKV/output projections and constructs a generic Attention instance. Its LlamaAttention.forward() is correspondingly small: fused QKV projection, split, RoPE on Q/K, generic attention, then the row-parallel output projection.

The generic Attention object owns no QKV projection weights. It owns the head geometry, cache contract, backend implementation, and forward dispatch. That separation lets many model files reuse the same serving machinery while platforms replace the backend without rewriting the model.

When Each Piece Runs

Time Owner State produced or changed Why it exists
Model construction Model module Local Q/K/V head counts, fused projection modules, RoPE, generic Attention Express the checkpoint architecture and tensor-parallel layout.
Model construction Generic attention + selector One validated backend class and its implementation object Choose a kernel family compatible with dtype, cache format, mask, platform, and configuration.
KV-cache initialization Model runner Allocated cache tensors bound to layer names Give every attention layer its persistent paged storage.
Every engine iteration Model runner + metadata builder Packed query boundaries, sequence lengths, block tables, slot mappings Translate scheduled requests into the tensors kernels understand.
Every layer forward Model module + generic attention Q/K/V heads, new cache entries, attention output Perform the Transformer computation for the scheduled tokens.

The persistent-cache step is explicit in GPUModelRunner.initialize_kv_cache_tensors(): the runner allocates or reshapes cache storage, handles cross-layer sharing, then binds layer-name-specific views into the static forward context.

One Path Covers MHA, GQA, and MQA

The control path does not branch on three separate attention classes. It is parameterized by query-head count $H_q$ and key/value-head count $H_{kv}$. LlamaAttention.__init__ partitions query heads across tensor-parallel (TP) ranks; KV heads are either partitioned or replicated so every rank has at least one.

Variant Head relation Local Q shape Local K/V shape Cache consequence
MHA $H_q = H_{kv}$ [tokens, Hq/TP, D] [tokens, Hkv/TP, D] One K/V head per query head.
GQA $H_q > H_{kv} > 1$ [tokens, Hq/TP, D] Partitioned or replicated [tokens, Hkv_local, D] Several query heads share each K/V head.
MQA $H_{kv} = 1$ [tokens, Hq/TP, D] One replicated K/V head per TP rank Smallest K/V width; all query heads share it.

Here $D$ is head_dim. The generic layer enforces that num_heads is divisible by num_kv_heads, so the kernel can infer the number of query heads sharing each K/V head. This is why a Llama-family GQA model and an ordinary MHA model can use the same Python forward path.

How the Backend Is Chosen

The generic layer first records head and cache geometry, then requests a backend with dtype, KV-cache dtype, mask features, attention type, and optional quantization constraints in Attention.__init__. The selector adds configuration-level constraints—including explicit backend overrides, cache block size, KV transfer, prefix masks, and prefill context parallelism—in get_attn_backend().

The final choice is delegated to the active platform by _cached_get_attn_backend(). Only after that choice does Attention.__init__ instantiate the backend-specific implementation.

Important: FlashAttention below is a worked backend, not a promise that every NVIDIA run selects it. Hardware generation, dtype, head size, cache layout, configured overrides, sliding windows, sinks, and other features can select or reject a backend before the first forward.

Worked Trace: One Decode Token

Assume one request is decoding one new token in one decoder layer. Its old K/V already occupies non-contiguous physical pages.

  1. Map the new token to a physical cache slot. The scheduler has already allocated blocks. The worker’s GPUModelRunner._get_slot_mappings() exposes a token-to-physical-slot vector for each KV-cache group and maps it back to every layer in that group. Padding slots become -1.

  2. Build the batch’s attention metadata. In GPUModelRunner.execute_model(), the runner combines slot mappings with query lengths, sequence lengths, and block-table tensors, then lets each backend’s metadata builder produce the per-layer structure it needs.

  3. Inject metadata without changing every model signature. The runner enters set_forward_context(...) around the model call. Model modules therefore pass only Q/K/V; the generic attention op retrieves dynamic serving state by layer name.

  4. Project and position the token. LlamaAttention.forward() computes one fused QKV projection, splits the local Q/K/V widths, applies RoPE to Q and K, and calls self.attn(q, k, v). No block table appears in this model-level method.

  5. Turn flat widths into head tensors. The generic Attention.forward() preallocates output and views Q as [tokens, Hq_local, D] while viewing K/V as [tokens, Hkv_local, D].

  6. Cross the compilation boundary. The remainder of Attention.forward() chooses a direct call or opaque torch.ops.vllm call. When cache update is a separate operation, a dummy tensor creates a data dependency so torch.compile cannot reorder attention before the cache write.

  7. Resolve this layer’s serving state. The custom op uses get_attention_context() to retrieve per-layer metadata, the bound KV-cache tensor, and this layer’s slot mapping from the forward context.

  8. Scatter the new K/V into its page. FlashAttention declares the cache write to be separate from its forward in FlashAttentionBackend. The generic unified_kv_cache_update() delegates to the selected implementation, and FlashAttentionImpl.do_kv_cache_update() performs the reshape_and_cache_flash scatter using slot_mapping.

  9. Read the whole logical history through the block table. The generic unified_attention_with_output() calls self.impl.forward(...). FlashAttention first views the paged tensor as separate K and V caches and extracts query boundaries, sequence lengths, and the block table in FlashAttentionImpl.forward().

  10. Launch attention and return to the model. The normal non-cascade path calls flash_attn_varlen_func(...) with the paged K/V views, query/sequence lengths, causal/window rules, and block table. The result returns through the preallocated output tensor; Llama flattens local heads and applies o_proj.

The state change happens at step 8. Steps 9-10 read the newly updated cache, so the current token can attend to itself while causal masking excludes future positions.

Why vLLM Is Structured This Way

Model code stays close to checkpoint semantics

The Llama module still reads like ordinary inference code: projection, split, RoPE, attention, output projection. Cache paging and batch packing do not leak into every model implementation.

Prefill and decode share an interface

Both phases pass packed token tensors through Attention.forward(). Query boundaries and sequence lengths in metadata tell the backend whether a request contributes a long prompt chunk, a short extension, or one decode token. This keeps the model path stable while the batch composition changes every iteration.

Cache write and cache read use different maps

  • slot_mapping answers: where should each new K/V token be written?
  • Block table answers: which physical pages make up each request’s logical history?

Conflating them is a common source-reading mistake. The first is token-to-slot scatter metadata; the second is request-to-page gather metadata consumed by PagedAttention-style kernels.

Backend replacement is a load-time decision

The hot model path does not repeatedly ask which kernel family to use. Backend selection and implementation construction happen before execution; per-step work is limited to metadata creation, cache update, and kernel dispatch.

Failure and Debugging Surfaces

Symptom First boundary to inspect Typical mismatch
Wrong output only with TP > 1 Model head layout and fused QKV projection Query heads partitioned but KV replication/partition count is wrong.
Corruption after preemption or prefix reuse Worker slot_mapping and block table New K/V scattered to one page while attention reads a different logical page list.
Backend rejected at startup Selector constraints Unsupported head size, dtype, cache dtype/layout, mask feature, or block size.
Failure only under torch.compile or CUDA graphs Generic custom-op boundary Cache side effect or output buffer dependency is not preserved.
Correct prefill but wrong decode Per-step metadata and cache write Query/sequence lengths or current-token slot differs between phases.
Different behavior across GPUs/platforms Platform-selected backend The model path is shared, but kernel support and cache-update ownership differ.
  1. vllm/model_executor/models/llama.py — recognize familiar Transformer math.
  2. vllm/model_executor/layers/attention/attention.py — understand the serving adapter and compile-safe dispatch.
  3. vllm/v1/worker/gpu_model_runner.py — see where batch and cache metadata enter the forward.
  4. vllm/v1/attention/selector.py — understand why a backend was selected.
  5. One file in vllm/v1/attention/backends/ — follow cache layout, metadata, and the kernel call for the actual device.

One Thing to Remember

vLLM’s MHA path is a stateful serving pipeline around a familiar Transformer core: the model creates Q/K/V, the worker describes where the current batch lives, the generic attention layer preserves the handoff, and the backend writes and reads the paged KV cache.

Go Deeper