mla: add inline per-token outer scales to 368-byte NVFP4 records - #86
mla: add inline per-token outer scales to 368-byte NVFP4 records#86yatesdr wants to merge 3 commits into
Conversation
📝 WalkthroughWalkthroughThe change adds optional per-token latent scaling for NVFP4 MLA cache records and wires it through public APIs, unified decode, MG prefill, shared-memory staging, device dequantization, compilation keys, validation, and regression tests. ChangesNVFP4 per-token latent scaling
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant MLAAPI
participant CacheWriter
participant UnifiedKernel
participant NVFP4Math
Caller->>MLAAPI: enable latent_scale_per_token
MLAAPI->>UnifiedKernel: validate traits and launch mode
CacheWriter->>UnifiedKernel: provide NVFP4 records with token scales
UnifiedKernel->>NVFP4Math: stage and load candidate-row scales
NVFP4Math-->>Caller: compute scaled decode or prefill output
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
The nvfp4_ds_mla writer quantizes at an implicit outer scale of 1.0,
which parks small-magnitude tokens' E4M3 group scales in subnormals
(GLM-5.2 kv_c spans ~240x across layers; shallow layers land at
max_abs ~0.02). The static per-layer calibration file corrects the
positioning but is checkpoint-specific and goes stale. This mode
makes the record self-describing instead:
writer (per_token_scale=True):
s_t = token_amax / (6*448) # warp-bfly reduce over group amaxes
stored fp32 at record bytes [292, 296) (existing zero pad; the
record stays 368 bytes) and every group scale byte is encoded
relative to s_t, so the largest group's E4M3 scale sits at the
top of range by construction -- no calibration artifact, for any
checkpoint, ever.
readers (latent_scale_per_token=True):
value = e2m1 * e4m3(group scale) * s_t, with s_t scalar-gathered
per candidate into the kv_sc smem buffer (decode io + MG io,
same idiom as the DSV4 footer gather; 8-aligned nc.v2 load at
[288, 296), second word). QK hoists per candidate row; the
launch-scalar latent_scale is dead in this mode.
Threading: traits field + central fail-closed checks (NVFP4_E4M3 +
368-byte fp8-rope record only), decode/MG smem kv_sc allocation
(BI x 4 per buffer), api/kernel entry kwargs (default off), compile
spec bumps (writer 2->3, decode 18->19, MG prefill 4->5) so stale
cubins can never run against the new record semantics. Mode off
keeps every specialization byte-identical.
Tests: tests/attention/test_mla_kv_cache_per_token_scale.py -- record
ABI (bit-exact s_t, pad, mode-independent rope lane), the positioning
invariant (max group scale >= 256 per token; the static writer's
shallow-magnitude subnormal defect reproduced as a unit test),
accuracy dominance at shallow magnitudes, zero-token edge.
DRAFT: authored without CUDA; not yet compiled. Validation runbook:
glm52-opt design/nvfp4-dynamic-second-level-scale-phaseA-addendum.md.
Co-authored-by: Fable <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01YHuYXXq9krPdvcgpzY4EfE
…cords Both production-head tests (multisplit decode, multitile MG prefill) now run against writer records in per_token_scale mode with the reference dequant applying the in-record fp32 outer scale. Validated on CN4 SM120. Co-authored-by: Sol <noreply@openai.com> Claude-Session: https://claude.ai/code/session_01YHuYXXq9krPdvcgpzY4EfE
2b72c29 to
7f3d270
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/attention/test_attention_mla_kv_cache.py (1)
217-244: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winStale
_PAD_OFFSETname now points at the latent-scale field.
_PAD_OFFSET = 292is reused to read the new fp32 latent scale (records[:, _PAD_OFFSET:_PAD_OFFSET+4]), but the sibling test file (test_mla_kv_cache_per_token_scale.py) names the same byte range_LATENT_SCALE_OFFSET = 292and reserves_PAD_OFFSET = 296for the actual trailing pad. The offsets are numerically correct here, but the name no longer describes what it points to, which will confuse future readers/editors of this file.♻️ Suggested rename for clarity
-_PAD_OFFSET = 292 +_LATENT_SCALE_OFFSET = 292 +_PAD_OFFSET = 296and update the two usages accordingly (
records[:, _ROPE_SCALE_OFFSET:_LATENT_SCALE_OFFSET]for rope scales,records[:, _LATENT_SCALE_OFFSET:_LATENT_SCALE_OFFSET+4]for the new latent scale).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/attention/test_attention_mla_kv_cache.py` around lines 217 - 244, Rename the offset symbol used for the fp32 latent-scale field from _PAD_OFFSET to _LATENT_SCALE_OFFSET, matching the sibling test’s layout naming. Update both references in _dequantize_records: the rope-scale slice should end at _LATENT_SCALE_OFFSET, and the latent-scale slice should start there; preserve the existing numeric offsets and dequantization behavior.tests/attention/test_mla_kv_cache_per_token_scale.py (1)
69-110: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTriplicated dequant reference logic across two test files.
_dequantize_two_leveland_dequantize_staticreimplement the same e2m1/group-scale decode already present as_dequantize_recordsintests/attention/test_attention_mla_kv_cache.py. Three near-identical copies of the record layout logic (already showing naming drift between the two files, see the_PAD_OFFSETcomment there) means any future record-layout change requires synced edits in three places.Consider hoisting a single parametrized dequant helper (e.g., into
tests/_reference/helpers.py, which already hostsE2M1_TO_FLOAT32) and having both test files call it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/attention/test_mla_kv_cache_per_token_scale.py` around lines 69 - 110, Extract the shared e2m1 and group-scale decoding logic from _dequantize_two_level, _dequantize_static, and _dequantize_records into one parametrized helper in the existing reference helpers module. Support the optional latent-scale multiplication needed by the two-level path, then update both test files to call the helper and remove their duplicated record-layout logic.sparkinfer/attention/_shared/mla/kv_cache.py (1)
33-50: 🗄️ Data Integrity & Integration | 🔵 TrivialPer-token scale mode has no runtime cross-check between writer and reader.
The 368-byte record is byte-identical between
per_token_scale=True/False; only the semantic content at[292, 296)differs (a valid fp32 scale vs. zero). If a decode/prefill caller'slatent_scale_per_tokenflag ever drifts out of sync with how the record was actually written (e.g. mixed writer/reader configs during a rollout), a legacy record read in per-token mode yieldslatent_scale=0for every token — a silent all-zero/degenerate output rather than an exception, since nothing in the record itself signals which mode was used to write it.This is presumably an accepted tradeoff given the "server-static" per-instance design noted here, but worth calling out explicitly for anyone changing this config at runtime or per-request in the future.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sparkinfer/attention/_shared/mla/kv_cache.py` around lines 33 - 50, Ensure the per-token scale mode remains server-static and cannot drift between record writers and readers: validate that the shared `per_token_scale`/`latent_scale_per_token` configuration is consistent for the instance, rather than attempting to infer the mode from the 368-byte record. Preserve the existing record layout and legacy-record behavior, and document the configuration invariant near the relevant MLA cache setup.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@sparkinfer/attention/_shared/mla/kv_cache.py`:
- Around line 33-50: Ensure the per-token scale mode remains server-static and
cannot drift between record writers and readers: validate that the shared
`per_token_scale`/`latent_scale_per_token` configuration is consistent for the
instance, rather than attempting to infer the mode from the 368-byte record.
Preserve the existing record layout and legacy-record behavior, and document the
configuration invariant near the relevant MLA cache setup.
In `@tests/attention/test_attention_mla_kv_cache.py`:
- Around line 217-244: Rename the offset symbol used for the fp32 latent-scale
field from _PAD_OFFSET to _LATENT_SCALE_OFFSET, matching the sibling test’s
layout naming. Update both references in _dequantize_records: the rope-scale
slice should end at _LATENT_SCALE_OFFSET, and the latent-scale slice should
start there; preserve the existing numeric offsets and dequantization behavior.
In `@tests/attention/test_mla_kv_cache_per_token_scale.py`:
- Around line 69-110: Extract the shared e2m1 and group-scale decoding logic
from _dequantize_two_level, _dequantize_static, and _dequantize_records into one
parametrized helper in the existing reference helpers module. Support the
optional latent-scale multiplication needed by the two-level path, then update
both test files to call the helper and remove their duplicated record-layout
logic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 400788aa-b60c-4584-be1f-0dfc8359eeb4
📒 Files selected for processing (13)
sparkinfer/attention/_shared/mla/api.pysparkinfer/attention/_shared/mla/decode_math.pysparkinfer/attention/_shared/mla/io.pysparkinfer/attention/_shared/mla/io_mg.pysparkinfer/attention/_shared/mla/kernel.pysparkinfer/attention/_shared/mla/kv_cache.pysparkinfer/attention/_shared/mla/prefill.pysparkinfer/attention/_shared/mla/prefill_mg.pysparkinfer/attention/_shared/mla/smem.pysparkinfer/attention/_shared/mla/smem_mg.pysparkinfer/attention/_shared/mla/traits.pytests/attention/test_attention_mla_kv_cache.pytests/attention/test_mla_kv_cache_per_token_scale.py
Name the inline scale accurately, share the host reference decoder, and pin the largest-group positioning invariant without claiming every group avoids E4M3 subnormals. Assisted-by: OpenAI Codex
yatesdr
changed the title
1 participant
Summary
Add an optional per-token outer scale to the existing 368-byte
nvfp4_ds_mlaKV-cache record used withKV_FP8_ROPE=1.The writer derives the scale from each token's 512-value compressed latent,
stores it in four bytes of existing padding, and writes the 32 E4M3 group
scales relative to it. Decode and MG-prefill readers apply the stored scale
during dequantization. Record size, page geometry, RoPE encoding, and
transport size do not change.
The mode is off by default and selected once per server process. Companion
vLLM PR local-inference-lab/vllm#189 makes writer/reader mode immutable,
rejects incompatible SparkInfer APIs, and separates persistent DRAM/NVMe
cache namespaces by record ABI.
Why
The existing NVFP4 NoPE lane encodes 512 E2M1 values with one E4M3 scale per
16-value group and an implicit outer scale of 1.0. GLM-5.2's cached latent
magnitude varies substantially by layer. In low-magnitude layers, many group
scales therefore fall into E4M3's subnormal range and lose mantissa
precision.
This affects the hidden-state trajectory consumed by the sparse indexer. In
the tested GLM-5.2 NF3/NVFP4 configuration, a frozen 343,727-token retrieval
prompt omitted the required needle tokens from the exact top-2,048 candidate
set until layers 62–74, while a 245,497-token control retrieved correctly.
Changing only KV scale positioning restored the frozen failures and the
randomized depth ladder.
PR #145 addresses the same scale-positioning issue with a checkpoint-specific
per-layer calibration file. This change derives the scale from each token at
write time, so no calibration artifact is required. Static calibration and
dynamic per-token scaling remain mutually exclusive.
Record format
For token
t:where 6 is the E2M1 maximum and 448 is the E4M3 maximum.
Readers reconstruct
e2m1 * e4m3(group_scale) * s_t.This positions the largest group scale in each token near the top of the
E4M3 range. Smaller groups in the same token may still be subnormal; the
implementation does not claim otherwise.
The bytes are not independently self-identifying. Safe persistent-cache use
therefore requires the immutable mode and external ABI namespace implemented
by vLLM #189.
Implementation
kv_scshared-memorybuffer and apply it in QK and P·V dequantization.
latent_scale_per_token; writer, decode, andprefill compile-spec versions are incremented.
NVFP4 layout.
Validation
Hardware: 4× RTX PRO 6000 (SM120), TP4/DCP4. Model: GLM-5.2 NF3 hybrid.
Long-context gates used MTP3, CUDA graphs,
max_model_len=480000,max_num_batched_tokens=3072, exact top-k, cold requests, and a fresh cachenamespace.
Correctness and retrieval
The GPU suite covers static and dynamic records, zero-token handling,
multisplit production decode, and multitile MG prefill. Ruff check and format
also pass for all four cleanup-touched files.
KL divergence
KL(BF16 reference || candidate), n=3, same image/tokens/runner/runtime,
2,047 scored positions; lower is better:
The matched mean improves by 4.92%. This KLD gate exercises the cache
record/scale path at 2,048-token context; it does not exercise sparse top-k.
Performance and capacity
The MTP0 comparison uses identical image bits, model, TP/DCP posture, fixed
prompts, two uncounted warmups per arm, no speculative drafts, and recorded
GPU clocks/temperatures. It isolates reader-kernel cost from MTP acceptance
changes and passes the pre-committed 1% decode-overhead limit.
Compatibility
vLLM #189. Static calibration namespaces include the scale-file SHA-256.
[292,296)without modification.Reproduction
Companion vLLM PR: local-inference-lab/vllm#189.
Behavioral evaluation image:
ghcr.io/yatesdr/glm52-serve@sha256:db82fdcb5756d4a547853ba1330538bdd8a3dc0c6443c29bc49ba77b69b51cd1.The later PR-head commits contain documentation, tests, and cache-contract
hardening; they do not change the measured writer/reader arithmetic.
Limitations
Implementation and text were drafted with AI assistance. All changed lines
were independently reviewed, and measurements were executed and validated by
the submitting operators. This is not a duplicate of #145: #145 supplies
checkpoint-specific static calibration; this PR implements calibration-free
per-token scaling.