MODEL EXPLORER / MODERN LLMS
Eight queries, fewer key/value heads
Which attention heads can share a key/value cache?
RECORDED NUMERICAL EXAMPLE
Keep queries, share keys and values.
Head connections from the recorded configuration; cache bytes calculated as 2 × batch × tokens × KV heads × head width × bytes per value. This CPU reference does not measure decoding speed.
All eight query heads remain distinct. Lines show which K/V head each query consults. K/V weights are group means of the same MHA weights, so outputs need not match. Cache size is an analytical estimate for one layer, batch 1, head width 2 and float32, excluding temporary expansion and framework overhead.
Connection weights
| Query heads | Shared key/value heads | Weight |
|---|---|---|
| Q 0 | KV 0 | 1 |
| Q 1 | KV 1 | 1 |
| Q 2 | KV 2 | 1 |
| Q 3 | KV 3 | 1 |
| Q 4 | KV 4 | 1 |
| Q 5 | KV 5 | 1 |
| Q 6 | KV 6 | 1 |
| Q 7 | KV 7 | 1 |
Bars share one scale within this example. Values are rounded to four significant digits.
- Analytical K + V cache · bytes
- 16384
- KV heads
- 8
- Query heads per KV head
- 1
Verified by the local recipe
- MHA with duplicated K/V weights matches grouped attention for 8, 2 and 1 KV heads
- Future token changes cannot affect earlier causal outputs
- Compact per-group K/V calculation reproduces the final query output
- Analytical KV bytes scale with the number of KV heads
Download the source and run.py to reproduce these checks. Choosing an example here replays recorded values.
THE SAME MODEL, ON YOUR MACHINE
Pick up where
the graph leaves off.
Compare MHA, GQA and MQA while keeping eight query heads. Follow the head mapping and calculate compact KV-cache storage.
Use in VS Code ↗1. Save these files in one folder
2. Reproduce the example
Use Python 3.13 in a dedicated environment, matching the tested recipe. These commands are for macOS / Linux shells.
python3.13 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
python run.pyExpected output: [3,1,4,16]
Input 1: [1,4,16] · float32 randn · CPU / eval · seed 0
3. Open the source in TensorViz
Install the preview, open this folder in local VS Code, and choose Open as Graph above HeadSharingComparison in model.py. Select your Python environment. Save before choosing Validate.
Installation and supported setup ↗
The recipe was run with Python 3.13.13 and Torch 2.7.1. This is an untrained example; a successful shape check does not measure accuracy.
ABOUT THIS EXAMPLE
Compare MHA, GQA and MQA while keeping eight query heads. Follow the head mapping and calculate compact KV-cache storage.
Component · 2023
MHA / MQA / GQA
Eight query heads of width two with 8, 2 or 1 KV heads. Transparent repeat_interleave CPU implementation; no optimized attention kernel.
Source, capture & limitations +
Untrained width-16 comparison. Averaging K/V weights is an illustrative conversion without the paper's uptraining. Cache estimates exclude temporary repeated tensors, attention activations and allocator overhead; no quality or latency claim.
Original TensorViz teaching example. PyTorch provides the underlying operators.
No separate redistribution license has been declared for these project examples.
- Content revision
803295b72f9f4cbd- Source SHA-256
2d7dbb8bead241d951d07146b0a15f05cd19c7ae0d0bb5a90bf1ad1dd974aa16- Captured
- 2026-09-17 · Python 3.13.13 / Torch 2.7.1
Reduced causal MHA/GQA/MQA CPU forwards, independent per-head cache calculation, and duplicated-weight equivalence. Analytical cache estimates are not runtime measurements.
Read the model manifest ↗