MLA in full detail

Here is a comment vLLM leaves in the middle of its MLA metadata builder, explaining why the chunked-prefill workspace is capped:

        chunked_prefill_workspace_size = min(
            # Try for 8 full length request or at least 4 pages per-request
            max(
                8 * model_config.max_model_len,
                4 * scheduler_config.max_num_seqs * cache_config.block_size,
            ),
            # For long-context models try not to over-allocate limiting
            # kv-cache space, limiting it to 64k tokens,
            # which would result in the workspace being:
            #   2*(576)*(64*1024) = 144mb
            # (assuming 576 MLA head dim, and fp16)
            # which would result in up-projected context being
            #   2*(192*128)*(64*1024) = 3gb
            # (assuming 192 QK head dim, 128 heads, and fp16)
            64 * 1024,
        )

Work the comment's own arithmetic: a 64k-row latent workspace is $2 \times 576 \times 65{,}536 = 75{,}497{,}472$ bytes, and up-projecting it gives $2 \times 192 \times 128 \times 65{,}536 = 3{,}221{,}225{,}472$ bytes. So 72 MiB of cached latent becomes 3 GB of keys the moment you decompress it — a factor of $192 \times 128 / 576 = 42.7$, and that only counts the query-key half; the value half is separate. (The comment's own "144mb" label is twice what its formula on the line above computes; the 3 GB figure checks out exactly.) This is the central tension of Multi-head Latent Attention. The cache is small because you refuse to store per-head K and V. Attention, however, is defined over per-head K and V. Something has to bridge the gap, and the naive bridge — up-project the latent back into heads, then run ordinary attention — undoes the entire saving in a transient buffer, once per layer, once per step.

§3.5 established what MLA buys: DeepSeek-V3 spends 68.6 KiB per token against Llama-3-70B's 320 KiB with GQA-8, and its decode arithmetic intensity is 242 FLOP/byte at TP=1 — 82 % of the H100's 295 FLOP/byte ridge, against 8 for GQA-8. It also established the catch: the latent cache is replicated across tensor-parallel ranks, so the cluster-level advantage is gone by TP ≈ 4.7 and the per-rank intensity decays as $1/\text{TP}$ — 242 at TP=1, 30 at TP=8. This chapter owns the mechanism underneath all of that: what exactly is cached, why the RoPE slice has to travel separately, how the absorbed form works, why prefill refuses to use it, and what the whole arrangement costs in kernels and complexity.

Ordinary attention caches, per token per layer, one K vector and one V vector for every KV head. MLA caches one vector, full stop — a low-rank latent $c_t$ of width $r_{kv}$ that every head shares, plus a small position-carrying slice of width $d_{\text{rope}}$ that every head also shares. Per-head keys and values do not exist in HBM. They exist only as a function of $c_t$ and two weight matrices, $W_{UK}$ and $W_{UV}$, that are the same for every token and every step.

Associativity moves key expansion onto the query side and value expansion after attention. That does not require pre-multiplying all query/output weights once at loading. The cited implementation reshapes weights at load time, then applies query and output batch matrix multiplications at runtime; it avoids materializing a potentially huge combined output matrix. The attention kernel consumes the latent cache, with two cost-oriented formulations:

MLA has two possible ways of computing, a data-movement friendly approach and a
compute friendly approach. We generally want to use the compute friendly
approach for "prefill" (i.e. the ratio Sq / Skv is relatively large, often near
1) and the data-movement friendly approach for "decode" (i.e. the ratio
Sq / Skv is small, often near 0).

NOTE what we deem small and large is currently determined by if it is labelled
prefill or decode by the scheduler, but this is something we should probably
tune.

Read the two CACHED nodes together: $512 + 64 = 576$ values per token per layer, and in bf16 that is 1,152 bytes — the number §2.1's $c_{\text{MLA}} = L\,(r_{kv} + d_{\text{rope}})\,b$ multiplies by 61 layers to get 68.6 KiB per token. The two NEVER CACHED nodes are $128 \times (128 + 128) = 32{,}768$ values per token per layer — 65,536 bytes, 57× larger than the thing you cached, and that is before the 64-wide RoPE slice is broadcast onto every head to form the full 192-wide key.

Every number in this chapter comes from a config read in one of the two repos, not from memory. vLLM ships a metrics test that constructs a full DeepseekV3Config, which is the cleanest in-tree statement of the shape:

def test_mla_config_parser():
    """Test MLAConfigParser extracts MLA-specific fields from DeepseekV3Config."""
    hf_config = DeepseekV3Config(
        hidden_size=7168,
        num_attention_heads=128,
        num_hidden_layers=61,
        kv_lora_rank=512,
        qk_nope_head_dim=128,
        qk_rope_head_dim=64,
        v_head_dim=128,
        q_lora_rank=1536,
    )

vLLM's MLA file header restates the same five attention dims with an explicit "in DSV3" annotation, and SGLang's MHA forward module repeats four of them in a config comment (python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py:L83-L88). Three independent in-tree statements, all agreeing:

H           hidden size
N           number of attention heads
Lq          latent dimension for Q              1536 in DSV3
Lkv         latent dimension for K/V            512 in DSV3
P           nope dimension, no rope.            128 in DSV3
R           rope dimension, goes through rope.  64 in DSV3
V           V head dim.                         128 in DSV3

## Vector/Matrix Definitions

h_t         hidden states (input to attention)  shape [Sq, H]
q_c         latent/compressed Q                 shape [Sq, Lq]
q_nope      uncompressed Q (no-rope)            shape [Sq, N, P]
q_pe        uncompressed Q (rope)               shape [Sq, N, R]
kv_c        latent/compressed KV                shape [Skv, Lkv]
k_pe        decoupled k position embeddings     shape [Skv, R]
new_kv_c    new kv_c from current iter          shape [Sq, Lkv]
new_k_pe    new k_pe from current iter          shape [Sq, R]
cache_kv_c  cached k_c from previous iters      shape [C, Lkv]
cache_k_pe  cached k_pe from previous iters     shape [C, R]
W_DQ        project h_t to q_c                  shape [H, Lq]
W_UQ        project q_c to q_nope               shape [Lq, N * P]
W_QR        project q_c to q_pe                 shape [Lq, N * R]
W_DKV       project h_t to kv_c                 shape [H, Lkv]
W_UK        project kv_c to k_nope              shape [Lkv, N, P]
W_KR        project h_t to k_pe                 shape [H, R]
W_UV        project kv_c to v                   shape [Lkv, N, V]
W_O         project v to h_t                    shape [N * V, H]

The 64-wide k_pe slice sitting next to the latent looks like an inelegance, and it is the single thing that makes MLA complicated. It exists because rotary position embedding and weight absorption are mutually exclusive, and it is worth deriving that rather than asserting it.

Write the un-absorbed score between query token $i$ and key token $j$, ignoring the softmax scale. With no RoPE anywhere:

Here column-vector notation defines $W_{UK}\in\mathbb{R}^{d_{\mathrm{nope}}\times r}$. Code may store its transpose for row-vector multiplication; keep that convention explicit.

The right-hand form never builds $k_j$. Now suppose RoPE were applied to the reconstructed key, as it is in an ordinary transformer. RoPE is a block-diagonal rotation $R_p$ that depends on the position $p$, and it is applied to the query at its position and the key at its position:

The relative-position property that makes RoPE attractive is exactly what kills absorption. To fold $W_{UK}$ into the query you need a single matrix that multiplies $q_i$ once and then meets every $c_j$. Here the matrix you would have to fold is $R_{j-i} W_{UK}$, and it is different for every distance $j-i$. On a 128k-context model that is 128k distinct $512 \times 128$ matrices per head. Absorption is not slower in that world; it does not exist.

DeepSeek's answer, which the code calls "decoupled", is to split the head into two halves that are handled by different rules. The $d_{\text{nope}} = 128$ half never sees RoPE, so $W_{UK}$ absorbs into it cleanly. The $d_{\text{rope}} = 64$ half carries all the position information and is not produced from the latent at all — it is projected straight from the hidden state by $W_{KR}$, rotated, and cached verbatim, with one copy shared by all 128 query heads. The score becomes a sum of two terms:

You pay 64 extra cached values per token to keep the other 512 absorbable. That is the trade, and it is why the cache cell is 576 rather than 512. In the code the two halves are literally adjacent slices of one tensor — vLLM splits kv_lora into kv_c and k_pe and gives k_pe a head dim of 1 (vllm/model_executor/layers/mla.py:L189-L192), and both engines store them in a single 576-wide cache row so the decode kernel reads one contiguous span.

Now the part that is genuinely gnarly. We will do the query side, then the value side, then count the cost.

Step 1 — what the naive path does

Per decode step, per layer, with $S$ tokens of context and one query token, the naive path is:

1. Read $S$ latents from cache: $S \times 576$ values.
2. Up-project: k_nope = c @ W_UK gives $[S, 128, 128]$; v = c @ W_UV gives $[S, 128, 128]$. Concatenate k_pe broadcast over heads to get $K$ of shape $[S, 128, 192]$.
3. Run ordinary MHA with QK head dim 192, V head dim 128.

Step 2 costs $2 \cdot r_{kv} \cdot h \cdot (d_{\text{nope}} + d_v) = 2 \cdot 512 \cdot 128 \cdot 256 = 33{,}554{,}432$ FLOPs per context token, and materialises $128 \times (192 + 128) \times 2 = 81{,}920$ bytes per context token against the 1,152 bytes you read. It is 71× the traffic you saved and it happens every step. This is the 72 MiB → 3 GB workspace comment from §1, in units of one token.

Step 2 — fold W_UK into the query

Take the absorbed score from the previous section and read it as an instruction rather than an identity. $\left( W_{UK}^{\top} q^{\text{nope}}_i \right)$ depends only on the query. Compute it once, per query token, per head:

Batched over heads that is a bmm of $[N, B, P] \times [N, P, L_{kv}] \to [N, B, L_{kv}]$ — which is precisely what vLLM writes:

                # Pads the head_dim if necessary (for the underlying kernel)
                N, B, P = mqa_q_nope.shape
                W_UK_T = self.W_UK_T_dcp_qrep if qrep_decode else self.W_UK_T
                assert W_UK_T is not None
                _, _, L = W_UK_T.shape
# ...
                # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
                torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope)

                # Convert from (N, B, L) to (B, N, L)
                mqa_ql_nope = mqa_ql_nope.transpose(0, 1)

The transposed $W_{UK}$ was built once at load time, out of the checkpoint's kv_b_proj, which stores $[W_{UK}; W_{UV}]$ concatenated per head:

        kv_b_proj_weight = kv_b_proj_weight.view(
            self.kv_lora_rank,
            self.num_heads,
            self.qk_nope_head_dim + self.v_head_dim,
        )

        W_UK, W_UV = kv_b_proj_weight.split(
            [self.qk_nope_head_dim, self.v_head_dim], dim=-1
        )
# ...
            # Convert from (L, N, V) to (N, L, V)
            replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)
            # Convert from (L, N, P) to (N, P, L)
            replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True)

SGLang does exactly the same split at weight-load time, into fields it calls w_kc and w_vc:

            w_kc, w_vc = w.unflatten(
                0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
            ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)

Step 3 — attend against the latent directly

Concatenate $\tilde{q}$ with the un-absorbed rotary half to get a query of width $r_{kv} + d_{\text{rope}} = 576$, and hand the cache to the kernel as both K and V — K being the full 576-wide row, V being its leading 512-wide latent slice. vLLM's Triton MLA backend shows the shapes with no ceremony:

        # Add a head dim of 1
        kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)
        kv_c_cache = kv_c_and_k_pe_cache[..., : self.kv_lora_rank]
        PAGE_SIZE = kv_c_and_k_pe_cache.size(1)
# ...
        decode_attention_fwd(
            q,
            kv_c_and_k_pe_cache,
            kv_c_cache,
            o,
            lse,
            block_table,
            seq_lens,
            attn_logits,
            num_kv_splits,
            self.scale,
            PAGE_SIZE,
            k_scale=layer._k_scale,
            v_scale=layer._k_scale,
            is_mla=True,
        )

SGLang's call is the same idea with the same tensor passed twice:

                attn_output = self.attn_mqa(
                    q_nope_out,
                    k_nope,
                    k_nope,
                    forward_batch,
                    q_rope=q_pe,
                    k_rope=k_pe,

Note the head count: the query still has $h = 128$ heads, but there is exactly one K/V "head" — the shared latent. Attention has become MQA in its data-movement shape while remaining MHA in its expressive power. That is what vLLM's header means by "the attention simulates a multi-head attention, while the compute is similar to multi-query attention".

Step 4 — fold W_UV into the output

The kernel now returns $[B, h, r_{kv}]$ — attention outputs living in latent space, not value space. The softmax weights $\alpha_{ij}$ are unaffected by the change of basis, so:

The sum in brackets is exactly what the kernel produced. Apply $W_{UV}$ afterwards, once per query token rather than once per context token, and then $W_O$ as usual. In principle $W_{UV}$ and $W_O$ could be pre-multiplied into a single $[h \cdot r_{kv}, d]$ matrix; neither engine does that, because it would inflate the output-projection weight from $128 \cdot 128 \cdot 7168$ to $128 \cdot 512 \cdot 7168$ parameters — 4× the weight bytes to save one small bmm. Both keep them separate:

    def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
        # Convert from (B, N, L) to (N, B, L)
        x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
        out = out.view(-1, self.num_heads, self.v_head_dim)
# ...
            # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
            torch.bmm(x, self.W_UV, out=out.transpose(0, 1))

What it costs

Now count, per layer, one query token, $S$ context tokens, TP=1. All arithmetic below is derived from the shapes table; nothing here was measured.

Three things fall out of that table, and they are the honest description of MLA.

One: the two folds cost exactly what the up-projection they replace costs. $2h\,d_{\text{nope}} r_{kv} + 2h\,r_{kv} d_v = 2 h\, r_{kv} (d_{\text{nope}} + d_v)$, which is the same product as kv_b_proj applied to one token. Absorption is not free, but it is not extra either — it moves a fixed cost from the context axis to the query axis. That is the whole trick in one line.

Two: the attention core gets 3.4× more expensive. $278{,}528 / 81{,}920 = 3.4$. The effective per-head key width went from 192 to 576 and the value width from 128 to 512, because every head now attends over the full latent rank. This is the price of the memory saving, and it is why §3.5's intensity number is 242 rather than something GQA-like: the numerator grew while the denominator shrank.

Three: the FLOP-only comparison favors absorption from context length 2 for one query under these dimensions. This is not a proof of faster execution for every decode backend: launches, occupancy, available kernels and existing expanded-cache reuse can change latency.

To put the 3.4× in wall-clock terms: at 4,096 tokens of context, TP=1, batch 1, the absorbed attention core over all 61 layers is 69.6 GFLOP, which is 70 µs at the H100's 989.4 TFLOP/s, against 287.8 MB of cache reads, which is 86 µs at 3.35 TB/s (both derived). Bandwidth still wins, but by 1.2× — the ridge/intensity ratio $295/242$ — where a GQA-8 model at $I=8$ is bandwidth-bound by $295/8 \approx 37\times$. Push to TP=8 and the FLOPs divide by eight while the bytes do not — 8.8 µs of compute against the same 86 µs of reads. That is the $1/\text{TP}$ intensity decay from §3.5, seen from the kernel's side.

Everything above assumed one query token. Change that and the accounting inverts, because the naive path's up-projection is paid per context token and then shared by every query in the block, while the absorbed path's folds are paid per query token. Let $Q$ be the query-block length and $S$ the context. For a rectangular block attending all S keys, the FLOP comparison is:

Both bold terms are the same constant, $33{,}554{,}432$. Divide through by $S$ and the condition becomes $33{,}554{,}432 \cdot (Q/S) + 196{,}608\,Q < 33{,}554{,}432$. Two consequences:

That is the derivation behind vLLM's "compute friendly for prefill, data-movement friendly for decode", and it is why both engines carry two complete attention implementations for one layer. SGLang makes the duality structural — one MLA layer owns two RadixAttention objects with different head geometry:

        self.attn_mqa = RadixAttention(
            self.num_local_heads,
            self.kv_lora_rank + self.qk_rope_head_dim,
            self.scaling,
            num_kv_heads=1,
            layer_id=layer_id,
            v_head_dim=self.kv_lora_rank,
            quant_config=quant_config,
            prefix=add_prefix("attn_mqa", prefix),
        )
# ...
        self.attn_mha = RadixAttention(
            self.num_local_heads,
            self.qk_nope_head_dim + self.qk_rope_head_dim,
            self.scaling,
            num_kv_heads=self.num_local_heads,
            layer_id=layer_id,
            v_head_dim=self.v_head_dim,
            quant_config=quant_config,
            prefix=add_prefix("attn_mha", prefix),
        )

Read the four numbers that differ: attn_mqa is $(576, 1, 512)$ for (QK dim, KV heads, V dim); attn_mha is $(192, 128, 128)$. Same layer, same weights, two completely different kernel shapes.

The dispatch, and its actual threshold

SGLang chooses per forward batch, per backend, through a registry of handler functions. The shared handler is short enough to read in full:

def _handle_attention_backend(attn, forward_batch, backend_name):
    # Captured prefill (tc_piecewise or breakable) must keep a single attention
    # path: pin the absorbed MLA method — MHA one-shot/chunked shapes vary with
    # kv-len and cannot be captured.
    if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph():
        return AttnForwardMethod.MLA
# ...
    sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)
    disable_ragged = (
        backend_name in ["flashinfer", "flashmla"]
    ) and attn.flashinfer_mla_disable_ragged

    if (
        not disable_ragged
        and forward_batch.forward_mode.is_extend_without_speculative()
        and (
            (
                sum_extend_prefix_lens >= attn.chunked_prefix_cache_threshold
                and not attn.disable_chunked_prefix_cache
            )
            or sum_extend_prefix_lens == 0
        )
    ):
        if _support_mha_one_shot(attn, forward_batch, backend_name):
            return AttnForwardMethod.MHA_ONE_SHOT
        return AttnForwardMethod.MHA_CHUNKED_KV
    else:
        return _dispatch_mla_subtype(attn, forward_batch)

The two MHA exits are taken when the prefix is zero (a cold prefill — the degenerate case where absorption is always wrong) or when the prefix is large enough to amortise the up-projection. chunked_prefix_cache_threshold defaults to 8,192 tokens (python/sglang/srt/environ.py:L590, SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD = EnvInt(8192)); everything between — a short extend against a small existing prefix — falls through to absorbed MLA. The module's own comment states the rule: "For batches with smaller sum_prefix_length > 0, MLA kernel with absorption will be used instead" (python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py:L96-L101).

vLLM does not dispatch per batch; it dispatches per row, splitting the batch into a decode prefix and a prefill suffix and running both kernels in the same layer call:

        num_mqa_tokens = attn_metadata.num_decode_tokens
        num_mha_tokens = q.size(0) - num_mqa_tokens
# ...
            self.impl.forward_mha(  # type: ignore[attr-defined]
                q[num_mqa_tokens:],
                k_c_normed[num_mqa_tokens:],
                k_pe[num_mqa_tokens:],
                kv_cache,
                attn_metadata,
                self._k_scale,
                output=mha_output[num_mqa_tokens:num_actual_toks],
                output_scale=mha_output_scale,
            )

What counts as "decode" is reorder_batch_threshold, and it is a per-backend constant, not a global:

    query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.VARLEN
    reorder_batch_threshold: int = 512  # process small prefills with decode pathway

    # ... FlashMLAMetadataBuilder:
    query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
    reorder_batch_threshold: int = 128  # process small prefills with decode pathway
    # ^ TODO(matt): tune this

Compare those to the rectangular-model derivation. FlashMLA's 128 sits under its threshold of 171; FlashAttnMLA's 512 sits well above it, meaning vLLM routes query blocks through the absorbed kernel that a full-pair FLOP model favors MHA for. Actual causal pair counts and prefix reuse differ. That is not obviously wrong — the simple model also ignores the 71× traffic the naive path materialises, the chunked-context workspace it needs, and the merge-and-rescale pass that stitches chunks back together. It does mean the thresholds are empirical, which the TODO(matt): tune this and the header's "this is something we should probably tune" both concede. Only a measurement settles where the real crossover is, and neither repo publishes one at this SHA.

The layer-level structures differ in a way worth naming. vLLM has one MLAAttention module that owns both paths and both folded weights; the backend supplies only forward_mqa, while forward_mha lives in the shared MLACommonBaseImpl and delegates to a separately selected prefill backend. SGLang has one layer that owns two RadixAttention children and a registry of per-backend dispatch functions. Same decomposition, opposite direction: vLLM makes the decode kernel pluggable and the prefill path shared; SGLang makes both kernels concrete objects and the decision pluggable.

vLLM's MLA backend inventory and its separate ladder

§3.4 catalogued vLLM's non-MLA backends and handed the MLA ones here. Reading vllm/v1/attention/backends/registry.py:L44-L130 at this SHA: sixteen enum entries resolve into vllm/v1/attention/backends/mla/, and three more are MLA backends registered from model directories rather than the shared tree (FLASHMLA_SPARSE_DSV4, FLASHINFER_MLA_SPARSE_DSV4, ROCM_FLASHMLA_SPARSE_DSV4, all pointing into vllm/models/deepseek_v4/). Nineteen MLA members of the enum, then, against eighteen non-MLA ones out of thirty-seven in total — 19/37, a shade over half. One model family accounts for more than half of vLLM's attention backend surface.

They get their own priority ladder, entirely disjoint from the dense one, keyed on compute capability and on whether the cache is quantised:

    if use_mla:
        if device_capability.major == 10:
            # Sparse MLA backend priorities
            # See https://github.com/vllm-project/vllm/issues/35807 for
            # benchmark results
            if kv_cache_dtype is not None and is_quantized_kv_cache(kv_cache_dtype):
# ...
            return [
                AttentionBackendEnum.FLASHINFER_MLA,
                # R1 dims + FP8 KV only; rejected by supports_combination
                # otherwise. Behind FLASHINFER_MLA: wins past bs≈8, regresses
                # at bs≤2.
                AttentionBackendEnum.TOKENSPEED_MLA,
                AttentionBackendEnum.CUTLASS_MLA,
                AttentionBackendEnum.FLASH_ATTN_MLA,
                AttentionBackendEnum.FLASHMLA,
                AttentionBackendEnum.TRITON_MLA,
                *sparse_backends,
            ]
        elif device_capability.major == 12:
            return [
                AttentionBackendEnum.TRITON_MLA,
                AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
            ]
        else:
            return [
                AttentionBackendEnum.FLASH_ATTN_MLA,
                AttentionBackendEnum.FLASHMLA,
                AttentionBackendEnum.FLASHINFER_MLA,
                AttentionBackendEnum.TRITON_MLA,
                AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
                AttentionBackendEnum.FLASHMLA_SPARSE,
            ]

Three branches, three orderings, and on SM120 a list of only two — Blackwell consumer parts get Triton or nothing. The use_mla flag is checked before anything else, so an MLA model and a dense model on the same GPU take entirely different code paths through selection. SGLang reaches the same place with a string: its "flashinfer" factory branches on runner.use_mla_backend to build either FlashInferAttnBackend or FlashInferMLAAttnBackend (python/sglang/srt/layers/attention/attention_registry.py:L42-L66), and its MLA classes form an inheritance chain rather than a flat list — CutlassMLABackend, FlashMLABackend and TRTLLMMLABackend all subclass FlashInferMLAAttnBackend, and CuteDslMLABackend and TokenspeedMLABackend subclass TRTLLMMLABackend.

One token, DeepSeek-V3, TP=1, FLASH_ATTN_MLA, 4,096 tokens of context. Layer 30.

1. MultiHeadLatentAttentionWrapper.forward (vllm/model_executor/layers/mla.py:L170-L192) runs fused_qkv_a_proj on the hidden state, producing $1536 + 576 = 2112$ values, and splits them into q_c [1536] and kv_lora [576]. kv_lora splits again into kv_c [512] and k_pe [64]; k_pe gets a head dim of 1.
2. Same function, L194-L203: q_b_proj maps q_c to $128 \times 192 = 24{,}576$ values, viewed as [128, 192], and rotary_emb rotates the trailing 64 columns of $q$ and the whole of k_pe in place.
3. MLAAttention.forward writes the 576-wide row into the paged cache and calls forward_impl. num_mqa_tokens = attn_metadata.num_decode_tokens = 1, num_mha_tokens = 0, so the MHA branch is skipped entirely.
4. forward_impl splits $q$ into mqa_q_nope [1, 128, 128] and mqa_q_pe [1, 128, 64], transposes the former to [128, 1, 128], and runs torch.bmm(mqa_q_nope, W_UK_T) against the [128, 128, 512] folded weight, yielding [128, 1, 512] → [1, 128, 512]. This is the absorption.
5. The tuple (mqa_ql_nope, mqa_q_pe) goes to impl.forward_mqa, which concatenates it to [1, 128, 576] and launches the kernel against the paged latent cache — K is the 576-wide row, V is its 512-wide prefix. Output: [1, 128, 512].
6. self._v_up_proj(attn_out, out=mqa_output_slice) runs torch.bmm(x, self.W_UV) with $x$ as [128, 1, 512] and W_UV as [128, 512, 128], writing [1, 128, 128] = 16,384 values.
7. Back in vllm/model_executor/layers/mla.py:L226, o_proj maps those 16,384 values to 7,168.

The two bmms in steps 4 and 6 are the entire absorbed-weight machinery: 33.55 MFLOP together, under 3 % of the 1.14 GFLOP attention core they enable at this context length (derived).

§2.1 flagged that DeepSeek-V3.2 adds a per-token cost the $c_{\text{MLA}}$ formula does not cover, and deferred it here. V3.2 replaces dense attention over the latent with sparse attention: a lightweight "indexer" scores every cached position against the current query, takes the top-$k$, and the MLA kernel attends only to those rows. The indexer needs its own key per token, and that key lives in its own cache.

Both engines size it identically, and both arrive at 132 bytes per token per layer. SGLang computes it explicitly:

        index_head_dim = get_dsa_index_head_dim(kvc.model_config.hf_config)
        indexer_size_per_token = (
            index_head_dim + index_head_dim // DSATokenToKVPool.quant_block_size * 4
        )
# ...
        return int(
            indexer_size_per_token * num_indexer_layers * element_size * indexer_ratio
        )

With index_head_dim = 128 (asserted at python/sglang/srt/mem_cache/memory_pool.py:L4412: assert index_head_dim == 128), quant_block_size = 128 and index_k_with_scale_buffer_dtype = torch.uint8 (python/sglang/srt/mem_cache/memory_pool.py:L4366-L4367), that is $128 + 1 \times 4 = 132$ bytes. vLLM's indexer states the same layout in a comment — "Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes" (vllm/v1/attention/backends/mla/indexer.py:L481) — and registers the indexer's storage as a second MLAAttentionSpec with head_size=132 and dtype=torch.uint8:

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
        return MLAAttentionSpec(
            block_size=self.cache_config.block_size,
            num_kv_heads=1,
            head_size=self.head_dim,
            dtype=self.dtype,
        )  # Only has one vector instead of K + V
# ...
        self.k_cache = DeepseekV32IndexerCache(
            head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
            dtype=torch.uint8,
            prefix=f"{prefix}.k_cache",
            cache_config=cache_config,
        )

Cost, derived: 132 bytes × 61 layers = 8,052 bytes = 7.86 KiB per token, on top of the 68.6 KiB latent — an 11.5 % increase in the per-token cell if every layer indexes. It is usually less: both engines honour a skip pattern in which a layer reuses the previous layer's top-$k$ indices and therefore needs no indexer key of its own. SGLang counts only the non-skipping layers (python/sglang/srt/model_executor/pool_configurator.py:L377-L409, via dsa_layer_skips_topk), and vLLM computes the same predicate from index_topk_freq / index_topk_pattern / index_skip_topk_offset when deciding whether to build an Indexer at all (vllm/model_executor/models/deepseek_v2.py:L1119-L1145).

The interaction with the latent cache is that the indexer selects rows of it. SGLang passes the result straight into the same absorbed decode call as a keyword — self.attn_mqa(..., topk_indices=topk_indices) (python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py:L722-L735) — so the sparse and dense paths share one kernel signature and one cache. Two caches, one indexed by the other; the absorbed-weight machinery is unchanged.

Do not apply an equal-width GQA cache formula to this MLA geometry. The expanded per-token cache is L*h*(d_nope+d_rope+d_v)*b = 61*128*320*2 bytes, or 4.765625 MiB. The latent cache is 61*576*2 = 70,272 bytes (68.625 KiB). Their ratio is 71.11. The asymmetric 192-wide keys and 128-wide values must both be counted; 2*L*h*d_h*b with one head width cannot directly produce that comparison.

This is a single-head, row-vector reference with a separate already-rotated positional slice. Expanded and absorbed scores, attention probabilities and output must agree in float64. Keep the original unabsorbed QK scale: replacing it with the latent width's square root changes the function. Multi-head execution repeats the fold per head and applies the output projection afterward; finite-precision paths need tolerance-based validation.

import numpy as np

rng = np.random.default_rng(42)
queries, context, rank, key_dim, value_dim, rope_dim = 3, 5, 4, 2, 3, 2
c = rng.normal(size=(context, rank))
q = rng.normal(size=(queries, key_dim))
q_rope = rng.normal(size=(queries, rope_dim))
k_rope = rng.normal(size=(context, rope_dim))
uk = rng.normal(size=(rank, key_dim))
uv = rng.normal(size=(rank, value_dim))
scale = (key_dim + rope_dim) ** -0.5
expanded_scores = (q @ (c @ uk).T + q_rope @ k_rope.T) * scale
absorbed_scores = ((q @ uk.T) @ c.T + q_rope @ k_rope.T) * scale
np.testing.assert_allclose(absorbed_scores, expanded_scores, atol=1e-12)
def softmax(a):
    e = np.exp(a - a.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)
a = softmax(expanded_scores)
np.testing.assert_allclose(a @ (c @ uv), (a @ c) @ uv, atol=1e-12)
assert 3*5 + 3*4//2 == 21  # prefix P=5, Q=3 causal pairs
print("MLA score/output parity and causal-pair accounting pass.")

Read the two formulations against each other, in the file that defines both:

V=~/Documents/other_git_repos/vllm
sed -n '66,118p' $V/vllm/model_executor/layers/attention/mla_attention.py   # both paths, side by side
grep -n "reorder_batch_threshold" $V/vllm/v1/attention/backends/mla/*.py    # every backend's decode window
grep -rn "use_mla" $V/vllm/platforms/cuda.py                                # the separate ladder

Then inspect the dispatch predicate before flipping a threshold. The cited prefix condition is prefix_len >= (threshold or 0), among other eligibility checks. Setting the threshold to zero makes that condition easier to satisfy for the chunked/MHA path; it does not force absorption. Profile the resolved branch under both settings:

SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD=0 python -m sglang.launch_server \
  --model-path deepseek-ai/DeepSeek-V3 --tp 8 --attention-backend flashmla
# Compare the default 8192 only after checking the actual eligible dispatch branch.
# Record kernel shapes; threshold zero does not force absorbed MQA.

On vLLM, the equivalent knob is the backend choice, because the threshold is a backend constant: --attention-backend FLASHMLA gives a 128-token decode window, --attention-backend FLASH_ATTN_MLA gives 512, and TRITON_MLA gives 1. Same model, same cache, three different definitions of "decode".

• DeepSeek-V2 (arXiv 2405.04434) — introduces MLA, the decoupled RoPE strategy and the absorption identity; the source both engines' file headers name as their main reference. Its 93.3 % KV-reduction and 5.76× throughput claims are the model authors' own, measured on their hardware.
• DeepSeek-V3 (arXiv 2412.19437) — the 61-layer, 128-head configuration used throughout this chapter.
• FlashInfer PR #551 — the other reference vLLM's header cites; the first widely-read implementation of the absorbed decode kernel.
• vllm#12601 — DeepSeek-V3 MLA with FP8 KV cache, and vllm#10927, the first version to store the latent rather than per-head K/V.
• vllm#35807 — the benchmark thread _get_backend_priorities cites by URL for the SM100 sparse-MLA ordering. The closest thing to a published measurement of MLA backend ranking that exists in-tree.
• deepseek-ai/FlashMLA — the kernel behind FLASHMLA and FLASHMLA_SPARSE, written by the model authors for exactly this attention shape.

Next: §7.3 on models that replace the growing cache with a constant-size recurrent state — the other way to attack the same bottleneck. §5.3 owns DP attention, the direct engineering answer to MLA's replicated cache.