MHA → MQA → GQA → MLA

Take Llama-3-70B's attention geometry — $L = 80$ layers, $h = 64$ query heads, head dimension $d_h = 128$ — and store keys and values the way the original Transformer did, one K and one V per query head. Feed the cell-size formula from §2.1:

One 8,192-token conversation would need 20 GiB of KV cache. On an 80 GB H100 running the model at TP=8 with roughly 50 GiB of pool per card (§2.1, derived), the whole eight-GPU node would hold twenty concurrent 8k requests — exactly one eighth of the 162 the shipped model manages, because MHA's per-rank cell is eight times GQA-8's. Twenty. On half a million dollars of hardware.

The shipped Llama-3-70B does not do that. It declares num_key_value_heads: 8, and the same arithmetic gives 320 KiB per token, 2.5 GiB at 8k, and 162 concurrent requests. The QK and PV attention FLOPs remain unchanged for fixed query heads and dimensions, but K/V projection FLOPs decrease; all 64 query heads still attend over the context. What changed is how many distinct K/V vectors have to exist in HBM, and therefore how many bytes the decode kernel drags across the memory bus on every single step.

That single config field is the whole subject of this chapter. Its four settings — $h_{kv} = h$, $h_{kv} = 1$, $h_{kv} = h/g$, and "there is no $h_{kv}$, cache a latent instead" — are MHA, MQA, GQA and MLA.

Attention has two head counts that people habitually conflate. The query head count $h$ decides how many independent attention patterns the layer can express — it is a capacity knob, and it costs FLOPs and query-projection parameters. The KV head count $h_{kv}$ decides how many distinct key/value vectors must be written to and re-read from the cache — it is a bandwidth knob, and it costs bytes on every decode step forever. The variants in this chapter all hold $h$ fixed and shrink $h_{kv}$. Query heads are broadcast against whatever KV heads exist; the broadcast is free in memory, which is exactly the point.

Read the strips, not the fans. Every doubling of the group size $g = h/h_{kv}$ halves the copper strip. That is the entire mechanism, and it is why the progression is a line and not a set of unrelated ideas.

Two results already derived elsewhere in this book meet here, and they move in the same direction. From §2.1, the per-token cell size:

with $L$ layers, $h$ query heads, $h_{kv}$ KV heads, head dimension $d_h$, $b$ bytes per stored element, and $g$ the group size — how many query heads share one KV head. From §0.4, the decode attention kernel's arithmetic intensity:

Context length $s$ cancelled. Head dimension cancelled. What survives is $g$ and the dtype. So one knob buys two things at once, and this is the sentence to remember:

On an H100 SXM the ridge point is $I^{*} = 295$ FLOP/byte (§0.4), so every one of these variants is still deep in the bandwidth-bound region — GQA-8 at $I = 8$ reaches about 2.7% of the card's bf16 peak. Nobody escapes the roofline by picking a group size. What they buy is a factor, and factors of 8 are worth having.

What each variant costs in weights, too

The KV projections shrink with $h_{kv}$ as well, which is often forgotten. For Llama-3-70B ($d = 8192$, $h \cdot d_h = 8192$), MHA would need $W_K$ and $W_V$ of $8192 \times 8192$ each, 134.2M parameters per layer; GQA-8 needs $8192 \times 1024$ each, 16.8M per layer. Over 80 layers that is 9.4B parameters saved — 17.5 GiB of bf16 weights, or roughly one whole H100's worth of the eight-card node. GQA is a weight optimisation that happens to be a bandwidth optimisation, which is part of why it was cheap for model builders to adopt.

The four variants, arithmetically

*MLA's row is DeepSeek-V3's own shape ($L = 61$, $r_{kv} = 512$, $d_{\text{rope}} = 64$, $h = 128$), not Llama-3-70B's — the architectures are not interchangeable. Its "$h_{kv} = 1$" is what the engines allocate, not what the config declares; see §5. Its intensity is derived below and is a genuinely different number from anything in the GQA family.

MQA: correct, and too aggressive

Shazeer's 2019 paper proposed exactly the extreme: one KV head, shared by every query head. The abstract's claim is the honest one — "We verify experimentally that the resulting models can indeed be much faster to decode, and incur only minor quality degradation from the baseline" (arXiv:1911.02150). Four years later the GQA paper opens by naming the two problems that stopped MQA from becoming the default: "MQA can lead to quality degradation, and moreover it may not be desirable to train a separate model just for faster inference" (arXiv:2305.13245). Its introduction adds training instability to the charge sheet. On the paper's T5-XXL comparison, MQA scores 46.6 average against MHA's 47.2 — small, but consistently down, and one KV head across 64 query heads is a very thin channel through which every attention pattern in the layer must be expressed.

The second objection is the practical one. A 70B checkpoint is not retrained because a serving team wants a smaller cache. The GQA paper's contribution is as much a conversion recipe as an architecture: mean-pool the existing KV heads within each group and uptrain with "5% of original pre-training compute". That is what made the switch tractable for everyone shipping models.

Why g = 4 or 8 and not 32

The GQA paper's own numbers show the shape of the curve: MHA-XXL takes 1.51 s per sample at 47.2 average, MQA-XXL takes 0.24 s at 46.6, and GQA-8-XXL takes 0.28 s at 47.1 (Table 1). GQA-8 recovers essentially all of MQA's speed and essentially all of MHA's quality. In the paper, GQA-8 means eight KV groups, not universally eight queries per group. Its measured speed approaches MQA on that workload because, past that point, KV traffic is no longer what the decode step is waiting for — the weights are. Halving KV bytes again when KV is 10% of the step's traffic buys 5%.

But there is a second, purely mechanical reason $h_{kv} = 8$ is the near-universal choice, and it is written in the serving code rather than the paper. Tensor parallelism shards KV heads across ranks by integer division, and once $h_{kv}$ drops below the TP size the heads get replicated instead of sharded — you pay the full cell size on every rank and gain nothing. Eight is the number of GPUs in a node. Setting $h_{kv} = 8$ makes TP=8 the exact point where each rank owns precisely one KV head, which is the largest useful $g$ for the standard deployment shape.

The sharpest form of that argument: at TP=8, a GQA-8 model and an MQA model have the same per-rank cache. Both give max(1, h_kv // 8) = 1 KV head per GPU, and both cost 40 KiB per token per rank on Llama-3-70B's geometry. MQA's entire remaining byte advantage over GQA-8 evaporates at this common eight-way tensor-parallel shape — while its quality cost does not. Choosing $h_{kv} = 8$ is choosing the smallest KV cache that TP=8 can still shard. §6 makes this quantitative.

MLA: not another point on the line

MLA does not choose an $h_{kv}$. It stores, per token per layer, a single compressed latent of width $r_{kv} + d_{\text{rope}}$ and reconstructs per-head keys and values inside the kernel. §2.1 derived the resulting cell size, which has neither a head count nor a factor of two:

The mechanism — the absorbed-weight trick, and why prefill and decode use different formulations — is §7.2's subject, and this chapter deliberately does not go there. What belongs here is the shape of the trade, and it is visible in the intensity. At decode, all $h$ query heads attend over one shared latent: each head computes scores against $r_{kv} + d_{\text{rope}} = 576$ dimensions and an output over $r_{kv} = 512$ dimensions, while the bytes read are one 576-wide latent per position:

Against the H100's ridge of 295, DeepSeek-V3's decode attention at TP=1 sits at 82% of the way to compute-bound — a place no GQA model gets near. That is the honest description of MLA: it did not find a cheaper way to store attention state, it found a way to pay in FLOPs instead of bytes. GQA moves you along the roofline's bandwidth ramp; MLA moves closer to the ridge but 242 FLOP/byte remains below 295; actual tile rereads can move it farther left. Which is better depends entirely on which resource you are short of, and — as §6 shows — on how many GPUs you have.

Neither engine has an "attention variant" enum with four members. MHA, MQA and GQA are all the same code path with a different integer; MLA is a separate layer class. SGLang states this in its type system with a two-member enum — there is no GQA member because GQA is not a different architecture, just a different head count:

class AttentionArch(IntEnum):
    MLA = auto()
    MHA = auto()

Where h_kv comes from

vLLM resolves the total KV head count by trying a list of config aliases and falling back to the query head count — that fallback is the MHA case, and it is why a pre-2023 checkpoint with no num_key_value_heads field simply behaves as MHA:

    def get_total_num_kv_heads(self) -> int:
        attributes = [
            # For Falcon:
            "n_head_kv",
            "num_kv_heads",
            # For LLaMA-2:
            "num_key_value_heads",
            # For ChatGLM:
            "multi_query_group_num",
            # For Step3p5:
            "num_attention_groups",
        ]
        # For non-grouped-query attention models, the number of KV heads is
        # equal to the number of attention heads.
        default_factory = self.get_total_num_attention_heads
        return getattr_iter(
            self.hf_text_config, attributes, default_factory=default_factory
        )

SGLang's equivalent carries a fossil of the MQA era: a boolean multi_query flag, used by Falcon and GPTBigCode, that short-circuits to one KV head — with a comment recording that this path never got TP support:

        if not new_decoder_arch_falcon and getattr(
            self.hf_text_config, "multi_query", False
        ):
            # Multi-query attention, only one KV head.
            # Currently, tensor parallelism is not supported in this case.
            return 1

Where TP divides it

vLLM performs the division on the config object, and the MLA early-return is the first appearance of this chapter's punchline — MLA returns 1 before the TP division exists:

    def get_num_kv_heads(
        self,
        parallel_config: ParallelConfig,
        arch_config: ModelArchitectureConfig | None = None,
    ) -> int:
        """Returns the number of KV heads per GPU.

        Pass ``arch_config`` (from ``model_arch_config[layer_idx]``) to size a
        single layer of a heterogeneous model rather than the model as a whole.
        """
        if self.use_mla:
            # When using MLA during decode it becomes MQA
            return 1

        arch_config = arch_config or self.model_arch_config
        total_num_kv_heads = arch_config.total_num_kv_heads
        # If tensor parallelism is used, we divide the number of KV heads by
        # the tensor parallel size. We will replicate the KV heads in the
        # case where the number of KV heads is smaller than the tensor
        # parallel size so each GPU has at least one KV head.
        return max(1, total_num_kv_heads // parallel_config.tensor_parallel_size)

The max(1, ...) is the replication clause, spelled out in the comment. Past $\text{TP} = h_{kv}$ the division saturates and extra ranks duplicate rather than shard.

A GQA model repeats the same division locally. Llama's attention module divides both head counts, asserts the two divisibility regimes explicitly, and — critically — resolves head_dim from the config with a fallback, not by assumption:

        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

        head_dim = getattr(config, "head_dim", None)
        self.head_dim = head_dim or self.hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim

Note that self.num_heads and self.num_kv_heads are divided by the same $\text{TP}$, so $g$ is invariant under tensor parallelism as long as $h_{kv} \geq \text{TP}$. Llama-3-70B at TP=8 has 8 query heads and 1 KV head per rank: still $g = 8$, still $I = 8$. Intensity is a per-rank property that TP does not change.

The per-rank counts go straight into the generic Attention layer, whose only statement about variants is a default and an assertion:

        if num_kv_heads is None:
            num_kv_heads = num_heads
        assert num_heads % num_kv_heads == 0, (
            f"num_heads ({num_heads}) is not divisible by num_kv_heads ({num_kv_heads})"
        )
        self.quant_config = quant_config
        self.layer_name = prefix

        self.num_heads = num_heads
        self.head_size = head_size
        self.head_size_v = self.head_size if head_size_v is None else head_size_v
        self.num_kv_heads = num_kv_heads

SGLang's layer names the same thing more honestly — the tp_ prefixes say out loud that these are already per-rank counts:

        super().__init__()
        self.tp_q_head_num = num_heads
        self.tp_k_head_num = num_kv_heads
        self.tp_v_head_num = num_kv_heads
        self.head_dim = head_dim
        self.qk_head_dim = head_dim
        self.v_head_dim = v_head_dim if v_head_dim != -1 else head_dim

Where the head count becomes bytes

In vLLM the layer publishes a cache spec, and the spec's page arithmetic is where $h_{kv}$ finally turns into memory. The factor of two is written as a sum of K and V head sizes rather than a multiplication, which is precisely the hook MLA needs:

    def __post_init__(self):
        if self.head_size_v is None:
            object.__setattr__(self, "head_size_v", self.head_size)

    @property
    def num_heads(self) -> int:
        if self.num_head_slots is not None:
            return self.num_head_slots
        return self.num_kv_heads

    @property
    def state_content_size_bytes(self) -> int:
        """Bytes per (head slot, stored state) cell of the page."""
        if self.state_content_bytes is not None:
            return self.state_content_bytes
        return (self.head_size + self.head_size_v) * get_dtype_size(self.dtype)

A GQA layer returns a FullAttentionSpec carrying its per-rank num_kv_heads:

            return FullAttentionSpec(
                block_size=block_size,
                num_kv_heads=self.num_kv_heads,
                head_size=self.head_size,
                head_size_v=self.head_size_v,
                dtype=self.kv_cache_torch_dtype,
                kv_quant_mode=quant_mode,
            )

An MLA layer is a different class entirely, and pins both fields as literals. Note self.num_kv_heads = 1 and head_size = kv_lora_rank + qk_rope_head_dim — no TP anywhere in either:

        self.num_heads = num_heads
        self.scale = scale
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank
        self.kv_b_proj = kv_b_proj
        self.dcp_q_replicate = dcp_q_replicate
        self.W_UK_T_dcp_qrep: torch.Tensor | None = None
        self.head_size = kv_lora_rank + qk_rope_head_dim
        self.layer_name = prefix
        self.indexer = indexer
        self.non_causal_multi_token_decode = non_causal_multi_token_decode
        self.sliding_window = sliding_window
        self.num_kv_heads = 1
        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
        kv_cache_dtype = kv_cache_dtype_str_to_dtype(
            self.kv_cache_dtype, vllm_config.model_config
        )
        common_kwargs = dict(
            block_size=vllm_config.cache_config.block_size,
            num_kv_heads=1,
            head_size=self.head_size,
            dtype=kv_cache_dtype,
            cache_dtype_str=self.kv_cache_dtype,
            kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
            # fp8_ds_mla: 656-byte custom layout (kv_lora_rank=512 +
            # qk_rope_head_dim=64, head_size=576). See flashmla_sparse.py.
            state_content_bytes=656 if self.kv_cache_dtype == "fp8_ds_mla" else None,
        )

MLAAttentionSpec then sets head_size_v = 0 (vllm/v1/kv_cache_interface.py:L392, quoted in §2.1), which deletes the factor of two from state_content_size_bytes above. Two literals — a 1 and a 0 — convert the general formula into $c_{\text{MLA}}$.

SGLang says the same thing at the model level, and its naming is the clearest statement in either codebase of what MLA is at decode time: DeepSeek's layer builds two RadixAttention objects, one literally called attn_mqa:

        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),
        )

The query head count is sharded — self.num_local_heads = num_heads // attn_tp_size (python/sglang/srt/models/deepseek_v2.py:L1773) — but num_kv_heads=1 is a literal, at every TP size. Which layer runs when, and why prefill wants a different one, is §7.2.

SGLang computes one scalar cell size per model and divides the pool by it. The function has two branches, and the difference between them is the whole argument of this section. First the GQA branch:

        else:
            n = model_config.get_num_kv_heads(tp_size, dcp_size)
            cell_size = (
                n
                * (model_config.head_dim + model_config.v_head_dim)
                * effective_num_layers
                * kv_size
            )

Now the MLA branch, from earlier in the same function. tp_size and dcp_size are bound at the top of the function and then never referenced here:

        kv_size = torch._utils._element_size(kv_cache_dtype)
        tp_size = get_parallel().attn_tp_size
        dcp_size = get_parallel().attn_dcp_size

        if kvc.use_mla_backend:
            from sglang.srt.mem_cache.kv_cache_configurator import (
                calculate_mla_kv_cache_dim,
            )

            cell_size = (
                calculate_mla_kv_cache_dim(
                    model_config=model_config,
                    kv_cache_dtype=kv_cache_dtype,
                    server_args=kvc.server_args,
                )
                * effective_num_layers
                * kv_size
            )

Confirmed at 7d89325: the MLA cell size is a pure function of the checkpoint's latent width, the layer count and the dtype. It is identical at TP=1 and TP=16. vLLM says the same thing differently, with the if self.use_mla: return 1 early return quoted in §5. MLA KV caches are replicated on every tensor-parallel rank.

The consequence is a capacity law. Let $P$ be the per-GPU KV pool in bytes. For a GQA model, each rank stores $c / \min(\text{TP}, h_{kv})$ bytes per token, so the cluster holds

Holding the per-rank byte pool P fixed, GQA logical-token capacity grows with min(TP,h_kv), while a replicated MLA latent has fixed capacity. In actual deployment P changes when weights are sharded, so adding ranks can improve either capacity even after KV-cell sharding saturates. The following comparison deliberately fixes P=40 GiB to isolate cache layout, not whole-deployment memory accounting.

MLA's per-token cell is 4.7× smaller and it still loses at TP=8. The crossover is where $\min(\text{TP}, h_{kv})$ passes $c / c_{\text{MLA}} = 320/68.6 = 4.66$, so the real-valued crossover is TP=4.66; a supported divisor such as TP=8 lies beyond it, while TP=5 is not valid for these head counts. GQA cell-size sharding stops once TP reaches the KV-head count; extra MLA ranks reduce per-rank weights and can thereby free more pool space, although latent bytes per token remain replicated.

The same asymmetry shows up in the intensity. GQA's $I = 2g/b$ is TP-invariant because both head counts shard together. MLA's is not: the query heads shard and the latent does not, so per-rank intensity is $2h_{\text{local}}(2r_{kv}+d_{\text{rope}})/(b(r_{kv}+d_{\text{rope}}))$ and falls linearly with TP — 242 at TP=1, 30 at TP=8, 15 at TP=16 (derived). MLA's spectacular roofline position is a TP=1 property that tensor parallelism spends.

The other escape hatch is decode context parallelism, which shards along the sequence axis rather than the head axis — vLLM implements it by dividing the per-request length, max_model_len = cdiv(max_model_len, dcp_world_size) (vllm/v1/kv_cache_interface.py:L300-L305), not by touching the cell size. That is why DCP works for MLA at all: there are no heads left to split.

Follow the number 8 from the checkpoint to a page of HBM, in vLLM, function by function.

1. config.json declares num_attention_heads: 64, num_key_value_heads: 8, hidden_size: 8192, num_hidden_layers: 80. No head_dim field.
2. ModelArchConfigConvertor.get_total_num_kv_heads (model_arch_config_convertor.py:L147-L164) scans the alias list, hits num_key_value_heads, returns 8.
3. get_head_size (model_arch_config_convertor.py:L99-L113) finds no head_dim, no hidden_size_per_head, and falls through to $8192 / 64 = $ 128.
4. LlamaAttention.__init__ (llama.py:L140-L158) divides: num_heads = 64 // 8 = 8, num_kv_heads = max(1, 8 // 8) = 1. The total_num_kv_heads >= tp_size assertion passes on equality.
5. Attention.__init__ (attention.py:L317-L328) asserts 8 % 1 == 0, stores num_kv_heads = 1, head_size = head_size_v = 128.
6. Attention.get_kv_cache_spec (attention.py:L648-L655) returns FullAttentionSpec(num_kv_heads=1, head_size=128, head_size_v=128, block_size=16).
7. AttentionSpec.state_content_size_bytes (kv_cache_interface.py:L245-L248) computes $(128 + 128) \times 2 = 512$ bytes per token per head slot; unpadded_page_size_bytes multiplies by num_heads = 1 and storage_block_size = 16 to give 8,192 bytes per page.
8. get_num_blocks (vllm/v1/core/kv_cache_utils.py:L1038-L1040) divides the measured pool by the page size and then by 80 layers. Per-rank cell size: $8192 / 16 \times 80 = 40{,}960$ bytes = 40 KiB.
9. The engine logs it (vllm/v1/core/kv_cache_utils.py:L1925-L1931): "GPU KV cache size: %s tokens, Maximum concurrency for %s tokens per request: %.2fx".

The same trace on DeepSeek-V3 diverges at step 2: ModelConfig.get_num_kv_heads returns 1 at the use_mla branch without ever reading parallel_config.tensor_parallel_size, and steps 4–8 run through MLAAttention and MLAAttentionSpec instead, with head_size = 576, head_size_v = 0, giving $576 \times 2 \times 61 = 70{,}272$ bytes per token on every rank.

        decode_context_parallel_size = parallel_config.decode_context_parallel_size
        if decode_context_parallel_size > 1 and not self.use_mla:
            total_num_kv_heads = self.get_total_num_kv_heads()
            if tensor_parallel_size <= total_num_kv_heads:
                raise ValueError(
                    "Decode context parallelism for GQA/MQA requires "
                    f"`--tensor-parallel-size` ({tensor_parallel_size}) to be "
                    "greater than the model's total number of KV heads "
                    f"({total_num_kv_heads}). Increase `--tensor-parallel-size` "
                    "or set `--decode-context-parallel-size 1`."
                )

A fifth failure mode is subtler and has no error message. Because $I = 2g/b$ is TP-invariant for GQA but MLA's per-rank intensity falls as $1/\text{TP}$, an MLA model that profiled beautifully on one card can look ordinary on eight — the attention kernel that was near the roofline ridge is now four to eight times further from it, and it is also reading a replicated cache. If you benchmark MLA at TP=1 and extrapolate, you will be wrong in a direction that flatters the architecture.

First, resolve the shapes the way the engines do, from the config rather than from arithmetic:

python3 - <<'PY'
import json, urllib.request
for repo in ["NousResearch/Meta-Llama-3-70B", "Qwen/Qwen3-32B"]:
    u = f"https://huggingface.co/{repo}/raw/main/config.json"
    c = json.load(urllib.request.urlopen(u))
    L  = c["num_hidden_layers"]
    h  = c["num_attention_heads"]
    hk = c.get("num_key_value_heads", h)          # MHA fallback, as in vLLM
    dh = c.get("head_dim") or c["hidden_size"] // h   # declared first, THEN divide
    cell = 2 * L * hk * dh * 2
    print(f"{repo:34s} L={L:3d} h={h:3d} h_kv={hk:3d} d_h={dh:3d} "
          f"g={h//hk:2d}  {cell/1024:7.1f} KiB/tok  I={h//hk} FLOP/byte")
PY

Then watch the TP division move the number. Launch the same model twice and read the KV log line; token capacity can grow by more than two from TP=1 to TP=2, since halving per-rank weights also increases the byte pool while the cell size halves:

# vLLM: grep the capacity line emitted by kv_cache_utils.py:L1925
vllm serve NousResearch/Meta-Llama-3-8B --tensor-parallel-size 1 --max-model-len 8192 2>&1 | grep "GPU KV cache size"
vllm serve NousResearch/Meta-Llama-3-8B --tensor-parallel-size 2 --max-model-len 8192 2>&1 | grep "GPU KV cache size"

# SGLang: the scheduler logs max_total_num_tokens (python/sglang/srt/managers/scheduler.py:L1096)
python3 -m sglang.launch_server --model-path NousResearch/Meta-Llama-3-8B --tp 1 2>&1 | grep max_total_num_tokens
python3 -m sglang.launch_server --model-path NousResearch/Meta-Llama-3-8B --tp 2 2>&1 | grep max_total_num_tokens

# Now do it at --tp 16 on a model with h_kv=8 and watch the token count NOT double.

The last line is the experiment worth doing. Llama-3-8B has $h_{kv} = 8$; at TP=16 the cell size stops shrinking, so the only remaining gain is the weight footprint, and the resident-token count will grow by far less than 2×.

• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (2019) — the MQA paper. Short, and the framing of incremental decoding as a memory-bandwidth problem is the origin of half this book.
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023) — Table 1 is the quality/speed evidence quoted above; §3 is the mean-pooling uptraining recipe that made conversion of existing checkpoints practical.
• DeepSeek-AI, DeepSeek-V2 (2024) — where MLA is introduced, with the 93.3% KV reduction and 5.76× throughput claims. The MHA-vs-MLA ablation is in its architecture section; the serving consequences are §7.2.
• vLLM PRs #10927 (DeepSeek-V2 MLA, the first version to store the latent rather than per-head K/V) and #12601 (DeepSeek-V3 MLA with fp8 compute) — the discussions are where the num_kv_heads = 1 / head_size_v = 0 encoding gets argued out.
• SGLang PR #4521 — "Reduce computation and communication in DP attention", part of the data-parallel-attention line of work that is the direct engineering response to the replication problem in §6.
• Pope et al., Efficiently Scaling Transformer Inference (2022) — the partitioning analysis that explains why the KV-head count and the accelerator count want to be the same number.