Repository navigation
feat(dsa): implement fp8 KV cache (e4m3) with a bf16 indexer cache - #1776
HadarIngonyama wants to merge 2 commits into
Conversation
`--kv-cache-dtype` already advertises `fp8_e4m3` and `fp8_e5m2` as choices, but `_init_kv_cache_dtype` raises for both. This implements them for the MLA latent cache and converts at the two kernel boundaries that assume a bf16 cache: the qblock prefill upcasts the KV block to the query dtype before the score dot, and `sparse_mla_page_level` casts new KV to the cache dtype before the v2 kernel stores it. No per-tensor scale is needed. fp8 precision is relative, so rescaling by a power of two buys nothing, and the latent cache holds post-`kv_a_layernorm` values and post-rope `k_pe`, both O(1) -- well inside e4m3's 448 range and above its 0.0156 subnormal floor. e4m3 is preferred over e5m2 for the extra mantissa bit (~6% vs ~12.5% worst-case relative error); measured kernel error is exactly 2x worse on e5m2. The indexer key cache deliberately stays bf16. The two caches fail differently under quantization: the latent cache feeds a softmax, where per-element error averages out over hundreds of keys, while the indexer cache feeds a top-k that *selects* which pages are read at all, so error there does not average -- it flips a page in or out, and a page the answer lives on is either fetched or lost. At 128 of 768 dims over 21 of 78 layers the indexer is 5.1% of the cache, so respecting that asymmetry is cheap: 1.90x capacity instead of 2.00x. Measured on a 16-device TPU v7x slice with GLM-5.2-FP8: 1.90x the token capacity of bf16, prefill unchanged (1.839 vs 1.847 s at 8k), and needle recall 39/40 against bf16's 38/40 over an 80-sample sweep at 500,593 and 975,485 tokens. Adds `test_dsa_fp8_kv_tpu.py`, a single-device gate that the 4-packed fp8 paged cache compiles through both readers of the cache -- the qblock prefill and the page-level decode -- and that the error against an identical bf16 cache stays within bounds set at ~1.4x the measured values (e4m3 3.47% rel_l2 on prefill and 3.06% on decode; e5m2 6.86% and 6.32%). It also asserts e5m2 is the worse of the two, which is what the extra mantissa bit should buy and a check that the dtypes are plumbed through rather than aliased. Registered in `unit-test-tpu-v6e-1`; a Mosaic refusal to lower the 4-packed layout reports a skip rather than a failure, since that is a platform statement and not a regression.
|
@cjx0709 would you have time to review this? We can't request reviewers from a fork. This PR implements The three design points we'd most like your eyes on are in the description: no per-tensor scale, e4m3 over e5m2, and why the indexer cache stays bf16. The scheduler's capacity math accounts for the two dtypes separately. We merged current main; CI is rerunning now. The earlier red run was a 40-minute timeout in One overlap to be aware of: zorrofox's #1811 also touches cc: @zorrofox |
|
Heads-up before review: we found a gap in this PR. The new KV is cast to the cache dtype only in the decode path ( The fix is to cast once at the backend's entry instead, plus a test on the dense path. We'll push it as soon as we can; the design points in the description are unaffected, so feel free to review those in the meantime. |
|
Hi @HadarIngonyama @tsvikas, I opened #1819 for FP8 KV on MHA/GQA and standard MLA. It leaves DSA to this PR: FP8 with DSA errors out and points here. +1 on casting at the backend entry; I did the same for standard MLA. One thing I hit on v7x: Data point for your no-scale choice: on Kimi-K2.5 (standard MLA, v7x), unscaled e4m3 gave 2x KV capacity, GSM8K 98.4 vs 98.4 and MMLU 88.4 vs 89.6 (FP8 vs bf16, N=250). |
Motivation
--kv-cache-dtypealready advertisesfp8_e4m3andfp8_e5m2as choices, but_init_kv_cache_dtyperaisesValueErrorfor both — onlyautoandbf16areactually implemented. This implements the fp8 options for the MLA latent KV cache.
Long-context DSA serving is KV-capacity bound: how much context a device can hold,
and how many bytes the attend has to pull from HBM, are both set by the latent
cache. Halving its element size addresses both — it roughly doubles the context
that fits, and halves the HBM→VMEM traffic of the attend, which is one of the
largest ops in a long-context prefill.
Modifications
model_runner_kv_cache_mixin.py: resolvefp8_e4m3/fp8_e5m2through adtype table instead of raising.
memory_pool.py: addMLATokenToKVPool.indexer_dtype_for(). The DSA indexerkey cache stays bf16 when the latent cache goes fp8;
allocate_buffersnowtakes a per-buffer dtype, and
profiled_bytes_per_token/_indexer_buffer_bytesaccount for the two dtypes separately so the scheduler'scapacity math stays correct.
sparse_mla_prefill_qblock.py: the qblock prefill upcasts the KV block to thequery dtype once, before the score dot. Mixing fp8 with bf16 in that dot, or
down-casting
pto the cache dtype, would put fp8's ~6% relative error on thesoftmax probabilities — the one tensor whose small values carry the attention
tail.
sparse_mla.py:sparse_mla_page_levelcasts new KV to the cache dtype beforethe v2 kernel stores it, since the kernel writes straight into the packed cache
and would otherwise reinterpret rather than convert.
python/sgl_jax/test/test_dsa_fp8_kv_tpu.py: new, and registered in theunit-test-tpu-v6e-1suite intest/srt/run_suite.py.Three design points worth flagging for review:
No per-tensor scale. Unlike int8, fp8 precision is relative, so rescaling by
a power of two buys nothing. The latent cache holds post-
kv_a_layernormvalues(unit RMS) and post-rope
k_pe, both O(1) — well inside e4m3's 448 range andabove its 0.0156 subnormal floor.
e4m3 over e5m2. The range headroom above is already ample, so the extra
mantissa bit is the better trade (~6% vs ~12.5% worst-case relative error).
Measured kernel error is exactly 2x worse on e5m2.
The indexer key cache deliberately stays bf16. The two caches fail differently
under quantization. The latent cache feeds a softmax, where per-element error
averages out over hundreds of keys. The indexer cache feeds a top-k that selects
which pages are read at all, so error there does not average — it flips a page in
or out, and a page the answer lives on is either fetched or lost. At 128 of 768
dims over 21 of 78 layers the indexer is 5.1% of the cache, so respecting the
asymmetry is cheap: 1.90x capacity instead of 2.00x.
Accuracy Tests
Unit (new, in this PR). Single-device gate that the 4-packed fp8 paged cache
compiles through both readers of the cache — the qblock prefill and the
sparse_mla_page_leveldecode path, which reach it by different routes — and thatthe output error against an otherwise identical bf16 cache stays within bounds.
fp8 makes
kv_packing4 instead of 2, so the paged cache becomes[P, ps//4, 4, Dk_pad]; whether Mosaic will DMA and reshape a 4-packedsub-32-bit array inside these kernels is a question only the compiler answers,
which is why this gate exists.
Measured on this commit, on a TPU v7x, error against the bf16 cache:
The e5m2/e4m3 ratio of 1.87–2.10x is the extra mantissa bit, and the test asserts
that ordering as a check that the two dtypes are genuinely plumbed through rather
than aliased. Bounds are set at roughly 1.4x the observed error. The test is
registered in
unit-test-tpu-v6e-1; note it has only been run on v7x, and ifMosaic declines to lower the 4-packed layout on another TPU generation it reports
a skip rather than a failure — happy to make that a hard failure if you'd prefer
to find out loudly.
End-to-end retrieval. An 80-sample needle sweep, both dtypes, same session:
5 needle variants × 5 depths at 500,593 tokens and 5 × 3 at 975,485 tokens.
fp8 recalls 39/40 against bf16's 38/40 — no measurable fp8 penalty. A
teacher-forced uplift probe with 7 distractors gives 5.32 nats for fp8 vs 5.49
for bf16, with the needle at p≈1.0 in nearly every cell.
Two scope limits, stated plainly. This is long-context retrieval accuracy, not a
general-capability benchmark — no GSM8K or MMLU on the fp8 cache yet. And the
end-to-end numbers here and below come from our serving deployment, which adds
16-way context parallelism over the KV cache; that change is not in this PR,
and holding a 500k–1M token context on one slice is not possible without it. This
PR is the dtype plumbing those runs used, and the unit numbers above are the part
measured on this commit alone.
Benchmarking and Profiling
test reports the same page buffer as 0.6 MiB in e4m3 against 1.2 MiB in bf16.
In a full model the bf16 indexer cache dilutes that to 1.90x tokens per byte
of HBM.
out to be a function of the larger pool fp8 enables, not of fp8 itself: capped
to bf16's pool, fp8 decoded faster. Profiling attributed it to a per-layer
full-pool copy that scales with pool size, not to unpack/convert — the heaviest
genuine conversion op measured 0.0 ms/core. Worth knowing if you enable fp8 to
grow the pool rather than to save bytes at a fixed pool.
The prefill and decode figures carry the same caveat as the retrieval numbers: they
are from the context-parallel deployment, not from this patch in isolation.
Checklist