GQA and the live-token budget

Model:

layers = 32
query_heads = 32
head_dim = 128
kv_dtype = bf16 = 2 bytes
GPU = 48 GiB
weights + runtime = 22 GiB
fragmentation reserve = 4 GiB

Part A — MHA vs GQA vs MQA

bytes_per_token = layers × 2 × kv_heads × head_dim × 2

Compute bytes/token and KiB/token for:

1. MHA: kv_heads = 32
2. GQA: kv_heads = 8
3. MQA: kv_heads = 1

Then KV GiB for one request at 8K live tokens, each variant.

Part B — Three products, two head counts

Free KV pool = 48 - 22 - 4 = 22 GiB.

Workloads:

name	prompt	output	live tokens
Chat	2048	512	2560
Doc QA	16384	1024	17408
Agent	8192	4096	12288

For MHA and GQA-8, table:

workload	KV GiB / req	max concurrent (floor)

1. Which product dies first under MHA?
2. What does GQA-8 buy Doc QA in concurrent slots?
3. Why "the model fits" was never the serving question?

Part C — Admission

Write one admission rule the scheduler can implement, in live tokens or KV bytes, that:

• Protects Chat from a single 32K Doc QA
• Does not require a second GPU
• Names the metric that pages on-call before OOM

Part D — The lie in max_batch_size

A config says max_num_seqs = 16. Show two mixes of the table above where 16 is safe and where 16 OOMs under GQA-8. The moral: seqs are not the unit; live tokens are.

Acceptance

1. MHA: 32*2*32*128*2 = 524288 B/tok = 512 KiB/tok. GQA-8: 128 KiB/tok. MQA: 16 KiB/tok.
2. Doc QA MHA: 17408 × 512 KiB ≈ 8.5 GiB/req → ~2 concurrent. GQA-8: ~2.1 GiB/req → ~10.
3. Admission mentions live tokens or KV bytes, not only seq count.
4. A 16-seq mix of Doc QA can OOM while 16 chats cannot.

Stretch

MLA-style compressed KV (treat as 4 KiB/token, fictional). Recompute Doc QA concurrency. One sentence on why architecture changes are capacity changes.

Check

Free 22 GiB / 512 KiB ≈ 45K live tokens (MHA). / 128 KiB ≈ 180K live tokens (GQA). Chat 2560 tok: MHA ~17 seqs, GQA ~70 — then other limits. Doc QA 17408: MHA 2, GQA 10. Agent 12288: MHA 3, GQA 14. Alert on KV bytes or live tokens at 80% of the pool, not on GPU util.

Debrief

GQA is not a quality-only paper. It is how you turn the same 22 GiB into a product that can hold more than two long documents. If you cannot do this arithmetic in a design review, you are not yet serving.