Multi-Head Latent Attention KV Cache

The multi-head latent attention KV cache can erase most of a model's apparent state footprint while leaving a serving team with the wrong memory estimate. The fix is to account for what persists per token and layer, then separately price every tensor that is reconstructed, sharded, read, or quantized during decode.

This guide builds that ledger from tensor shapes instead of repeating a headline compression ratio. It separates MLA cache bytes, DeepSeek-V2 attention, latent KV compression, and the MLA decode bottleneck into claims that can be measured independently.

Multi-head latent attention KV cache accounting starts here

KV-cache estimates fail when architecture names substitute for shapes. For every attention layer, write the batch dimension, cached token count, number of KV heads, head dimension, latent rank, positional slice, element width, and replication factor. Mark each tensor as persistent, ephemeral, or a weight. Only persistent token-dependent tensors belong in the cache capacity total. Ephemeral reconstruction still matters, but it belongs in peak workspace and bandwidth measurements rather than in bytes retained for the next decode step.

For conventional multi-head attention, a useful first estimate is layers × batch × tokens × 2 × kvHeads × headDim × bytesPerElement. The factor of two represents keys and values. Grouped-query attention reduces kvHeads; cache quantization reduces element width. KV-cache optimization is the adjacent ledger for those architectures. MLA changes the stored representation itself: a joint low-rank latent carries content needed to reconstruct keys and values, while a separate positional key slice supports decoupled rotary position encoding.

Do not paste dimensions from one DeepSeek checkpoint into another. Read the model configuration and the serving implementation that actually allocates cache pages. Record local KV dimensions after tensor or sequence parallel partitioning, not only global model dimensions. Then reconcile calculated bytes against an allocation trace for one token, one layer, one sequence, and one batch. That tiny case exposes duplicated buffers, alignment, page metadata, and precision conversions before a 128K-context test makes the discrepancy expensive. The multi-head latent attention KV cache ledger anchors that reconciliation.

Name exactly what MLA keeps between tokens

The DeepSeek-V2 paper introduces multi-head latent attention as joint low-rank compression of keys and values. During generation, the reusable content state is the compressed cKV vector rather than fully expanded K and V for every head. A decoupled RoPE component carries positional information that cannot simply be absorbed into a fixed projection. A first-order persistent estimate is therefore layers × batch × tokens × (kvRank + ropeDim) × bytesPerElement, adjusted for the implementation's exact layout and replication.

That formula is a boundary, not a complete performance model. Query latents are computed for the current token and normally do not persist. Projection weights consume model memory, not token-growing cache. Expanded head-specific keys and values may appear as intermediates or may be algebraically absorbed into query and output projections. If a profiler shows those expanded tensors surviving across decode steps, the implementation is not realizing the architectural cache promise, even if its class is named MLA.

The DeepSeek-V2 reference code is useful for connecting paper notation to configuration names, but its repository also warns that open implementation performance differs from its optimized internal path. Treat it as an architectural reference, not a throughput guarantee. Your ledger should include allocator granularity, page size, reserved capacity, data type, alignment, and whether the positional slice is co-packed with the latent. Compare requested bytes, committed bytes, and bytes read per generated token. Those three numbers answer different operational questions and should never share one unlabeled chart. The multi-head latent attention KV cache needs all three views.

Calculate capacity with a falsifiable fixture

A cache calculator earns trust by predicting small allocations exactly before it projects fleet capacity. Start with BF16 or FP16 at two bytes per element, one sequence, a short token count, and two layers. Calculate conventional K and V bytes and MLA latent-plus-RoPE bytes by hand. Run the same configuration through a deterministic helper, double tokens, then triple batch. Both totals must scale linearly. Change the latent rank by one and confirm the byte delta equals layers × batch × tokens × bytesPerElement.

Next add real serving details as explicit terms rather than mystery multipliers. Page allocation rounds token capacity up to a block boundary. Tensor parallelism may shard expanded heads yet replicate the shared latent on every rank. Prefix sharing changes physical ownership across sequences. Cache quantization may attach scales per group. Speculative decoding temporarily reserves slots for tokens that are rejected. PagedAttention fragmentation measurement shows why logical cache bytes and committed page bytes need separate lines.

Keep the basic equation intact beside the enriched allocator model. When observed memory differs, attribute the residual to named terms: page slack, metadata, workspace, graph capture, allocator reserve, or replication. A ratio without the baseline dimensions cannot be reviewed. The runnable artifact below deliberately models only persistent logical tensors, validates every dimension, exposes the cached shape, and returns both bytes and compression ratio. Its repeatability check prevents environment or clock state from entering a function that should be pure arithmetic. Multi-head latent attention KV cache estimates should remain pure arithmetic too.

Runnable artifact: The multi-head latent attention KV cache calculator covers exact bytes, token and batch scaling, invalid dimensions, compression, shape reporting, and deterministic repeat execution. Replace the illustrative dimensions with the deployed checkpoint's configuration before comparing its result with an allocator trace.

Save this as mla-kv-cache-ledger.mjs and run node mla-kv-cache-ledger.mjs. Expected final line: PASS: 8 MLA cache ledger assertions.

import assert from "node:assert/strict";

const positiveInteger = (value, name) => {
  if (!Number.isInteger(value) || value <= 0) throw new RangeError(name + "_invalid");
  return value;
};

export function cacheLedger(config) {
  const layers = positiveInteger(config.layers, "layers");
  const batch = positiveInteger(config.batch, "batch");
  const tokens = positiveInteger(config.tokens, "tokens");
  const bytesPerElement = positiveInteger(config.bytesPerElement, "bytesPerElement");
  const kvHeads = positiveInteger(config.kvHeads, "kvHeads");
  const headDim = positiveInteger(config.headDim, "headDim");
  const kvRank = positiveInteger(config.kvRank, "kvRank");
  const ropeDim = positiveInteger(config.ropeDim, "ropeDim");
  const tokenLayers = layers * batch * tokens;
  const standardBytes = tokenLayers * 2 * kvHeads * headDim * bytesPerElement;
  const latentBytes = tokenLayers * (kvRank + ropeDim) * bytesPerElement;
  return Object.freeze({
    standardBytes,
    latentBytes,
    savedBytes: standardBytes - latentBytes,
    compressionRatio: latentBytes / standardBytes,
    cachedShape: Object.freeze([layers, batch, tokens, kvRank + ropeDim]),
  });
}

const base = Object.freeze({ layers: 2, batch: 1, tokens: 4, bytesPerElement: 2,
  kvHeads: 8, headDim: 128, kvRank: 512, ropeDim: 64 });
let assertions = 0;
const check = (fn) => { fn(); assertions += 1; };

check(() => assert.equal(cacheLedger(base).standardBytes, 32_768));
check(() => assert.equal(cacheLedger(base).latentBytes, 9_216));
check(() => assert.deepEqual(cacheLedger(base).cachedShape, [2, 1, 4, 576]));
check(() => assert.equal(cacheLedger({ ...base, tokens: 8 }).latentBytes, 18_432));
check(() => assert.equal(cacheLedger({ ...base, batch: 3 }).standardBytes, 98_304));
check(() => assert.ok(cacheLedger(base).compressionRatio < 0.3));
check(() => assert.throws(() => cacheLedger({ ...base, ropeDim: 0 }), /ropeDim_invalid/));
check(() => {
  const first = cacheLedger(base);
  const second = cacheLedger(base);
  assert.deepEqual(first, second);
  assert.equal(Object.isFrozen(first.cachedShape), true);
});
assert.equal(assertions, 8);
console.log("PASS: 8 MLA cache ledger assertions");

Price reconstruction where decode actually runs

MLA exchanges persistent memory traffic for projection work. One execution path reconstructs head-specific representations from the latent at decode time. Another algebraically absorbs projection matrices so attention operates without materializing the same expanded tensors. Those paths can have identical cache capacity and very different kernel shapes, arithmetic intensity, launch count, and temporary memory. Benchmark the path your serving engine selects for the exact batch, context, precision, and accelerator.

Decode is often bandwidth-sensitive because each new token reads historical state, but shrinking the state can move the bottleneck toward matrix-vector or small matrix multiplication. Record HBM bytes read, achieved bandwidth, tensor-core utilization, kernel duration, and tokens per second together. A higher compute percentage is not automatically a regression if wall time falls; a spectacular cache reduction is not automatically a win if reconstruction becomes launch-bound at small batch. FlashAttention IO-aware attention provides the same discipline: reason about movement through the memory hierarchy, not FLOPs in isolation.

Warm the kernels, pin clocks where the environment permits, and separate prefill from decode. Sweep context length along one axis and concurrent sequences along another. At short context, weight reads or launch overhead may dominate. At long context, latent-cache traffic becomes visible. At high batch, reconstruction GEMMs may use the hardware better. Report the crossover points rather than one average. An operator can then route workloads deliberately instead of assuming MLA has one universal speedup. The multi-head latent attention KV cache must earn that operating point.

Audit the latent bottleneck as learned state

Compression is not just a systems trick; it changes what the network must preserve in a shared representation. Through the Bottleneck studies a 114M-parameter MLA transformer and reports, within that experimental scope, a content-heavy latent, positional separation through RoPE, localized circuit behavior, and unused effective capacity. Those results are an informative mechanistic probe, not evidence that every large production checkpoint has identical rank needs or internal organization.

Use the study to formulate tests. Sweep latent rank during a controlled training or conversion experiment and plot quality against persistent bytes. Probe entity, position, retrieval, and induction behaviors separately rather than relying only on aggregate perplexity. If quantizing the latent, compare disruptions by layer and context length. A dimension that appears statistically underused may still carry rare behavior, multilingual structure, or long-context retrieval that the probe set misses. Capacity reduction needs behavioral controls as well as singular values.

Keep systems and model conclusions connected but distinct. A globally over-provisioned latent suggests an opportunity to train a smaller cache representation; it does not authorize truncating a trained checkpoint at inference. A content-position separation can explain why the positional slice remains outside cKV; it does not prove an implementation cached the correct slice. The release packet should pair shape and allocation evidence with task evaluation. This avoids two symmetrical mistakes: declaring compression safe because bytes fell, or declaring it useless because one unoptimized reconstruction kernel ran slowly. The multi-head latent attention KV cache needs both kinds of evidence.

Test sharding before multiplying by GPU count

A single shared latent may resist the same head-wise partition used by conventional attention. If every tensor-parallel rank loads or stores the complete latent cache, per-GPU capacity can still be low while aggregate memory traffic is replicated. Write the ledger twice: logical bytes for the request and physical bytes summed across ranks. Add collectives, duplicated positional state, and any all-gathered reconstruction intermediates. Then vary tensor-parallel width while holding model, batch, and context fixed.

The useful plot has per-rank cache bytes, aggregate HBM reads per token, collective time, and end-to-end latency. Ideal weight sharding can hide a latent replication penalty at one width and expose it at another. Compare against a one-rank or lower-TP control where possible. The field note on tensor parallelism for LLM inference explains why a smaller local tensor does not guarantee scaling once communication and replicated work enter the trace.

Also inspect sequence and data parallel alternatives. Replicating a compact latent may be acceptable if it removes a costly collective; sharding it may help capacity but create coordination on every token. The correct decision depends on topology, bandwidth, concurrency, and whether weights or cache dominate each rank. State the policy as a measured crossover: for example, choose absorbed projections below a batch threshold and a different kernel above it, or cap TP width for long-context MLA workloads. Architecture alone cannot select that operating point for the multi-head latent attention KV cache.

Publish a cache receipt, not a compression slogan

The qualification record should name checkpoint revision, layer count, head geometry, latent and RoPE dimensions, element widths, page size, parallel layout, kernel path, accelerator, software commit, batch, context, and sampling window. Include the one-token allocation control, calculator output, allocator reconciliation, profiler trace, and rejected configurations. If the observed ratio differs from the paper, the receipt should make the reason reconstructable without implying either result is wrong.

Gate on outcomes that matter to the deployment. Set capacity limits for logical and committed cache bytes, latency limits across representative context and concurrency buckets, and quality limits for retrieval and generation controls. Include failure probes for rank mismatch, an accidentally cached expanded tensor, latent replication, quantization metadata, and an engine fallback to a generic attention path. Re-run after checkpoint conversion, serving-engine upgrades, precision changes, and parallel-topology changes.

The multi-head latent attention KV cache is successful when the stored state, reconstructed work, and learned bottleneck all have explicit evidence. The cache calculator proves the capacity claim; the tensor ledger prevents category errors; the profiler identifies where cost moved; and model evaluation checks what compression preserved. Together they replace a borrowed percentage with a deployable boundary: exactly how many bytes each token adds, exactly which work each decode step repeats, and exactly where this MLA implementation wins on this hardware.