[gg-rebased] kv: restore nvfp4_ds_mla writer - #122
Conversation
Adds an opt-in "nvfp4_ds_mla" KV cache dtype for the B12X sparse-MLA backend: the 512-dim MLA latent is stored as packed NVFP4 (E2M1 data + per-16-group E4M3 scales) instead of fp8, shrinking the per-token record from 656 B to 432 B per layer (256 B FP4 NoPE + 32 B E4M3 scales + 16 B alignment pad + 128 B BF16 RoPE) for +39-48% KV pool at equal budget. Behavior is unchanged unless opted in: every change is gated on kv_cache_dtype == "nvfp4_ds_mla", and fp8_ds_mla serving takes byte-identical code paths — including the b12x call signatures. The scale_format / caps kwargs are forwarded to b12x ONLY for the FP4 record, so fp8 serving keeps working on a b12x tree without the nvfp4 read-path port. Write side: csrc concat_and_cache_nvfp4_mla, in-tree in libtorch_stable/cache_kernels.cu (+ ops.h decl, _C_cache_ops schema), guarded by ENABLE_NVFP4_SM100/SM120 with a clear error on pre-Blackwell builds. _custom_ops falls back to loading a companion vllm/_nvfp4_mla_cache_C.so iff the main build lacks the op, so the feature can also ship as an overlay on an existing image. Read side: requires the b12x ScaleFormat.NVFP4_E4M3 (== 2) decode/extend path (companion b12x PR to follow); until that lands, requesting nvfp4_ds_mla fails loudly at plan construction with an unexpected-kwarg error. B12X_MLA_SPARSE only; FLASHMLA_SPARSE still canonicalizes to fp8_ds_mla. Validated on GLM-5.2 753B @ TP4/DCP4 on 4x RTX PRO 6000 (SM120): KV pool 454,510 vs 307,547 tokens (+47.8%) at util 0.96; GPQA-Diamond 174/198 vs 175/198 for fp8 KV on the same checkpoint (statistically tied); NIAH 30/30 from 4k to 360k; needle retrieved at 460k depth; decode speed within noise of fp8 at matched context; zero OOMs. Signed-off-by: David Young <davidseanyoung@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JzPRoS8j7b78iivwSmFv4y
…group scales Quantizes a random MLA latent through the op and dequantizes the cache record with a torch reference (E2M1 nibble table x per-group E4M3 scales): asserts the stored scales match E4M3(group_amax/6) within half a mantissa step, bounds the per-element NoPE error by the E2M1 grid half-gap (1.25x group scale), and checks the 16-byte pad is zeroed, the 16-bit RoPE lane is copied verbatim, and unmapped slots stay untouched. Skips cleanly without CUDA, on ROCm, and below SM100. Signed-off-by: David Young <davidseanyoung@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JzPRoS8j7b78iivwSmFv4y
|
Warning Review limit reached
Next review available in: 57 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (7)
✨ 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 |
67d5d7c
into
local-inference-lab:dev/gilded-gnosis-rebase
3 participants
Summary
Forward-port the missing
nvfp4_ds_mlawrite side and public cache dtype from #115 ontodev/gilded-gnosis-rebaseat533b037f35.The rebased tree contains the B12X reader/plan contract but still lacks the public CLI dtype and SM100/SM120 writer. This PR restores:
CacheDTyperegistration;The original authorship from the already validated writer implementation is preserved.
Port validation
ruff check: passruff format --check: passgit diff --check: passSupersedes #115 for the upstream-rebased GG line.