[GG] kv: restore nvfp4_ds_mla writer after consolidation - #115
[GG] kv: restore nvfp4_ds_mla writer after consolidation#115voipmonitor wants to merge 4 commits into
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: 29 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 |
voipmonitor
changed the title
|
Superseded by the canonical dev/gilded-gnosis consolidation. The corresponding implementation is present in commit(s): c266677,505acd14a4,6428dee087,67d5d7cc9f,bb88242e51. Closing the old-base PR so future work targets the canonical GG branch. |
2 participants
Summary
Forward-port the write side and public dtype from the already merged #82 onto
dev/gilded-gnosis. The original authors and commits are preserved.GG already contains the newer B12X read/plan path for the 432-byte
nvfp4_ds_mlarecord, including GLM-specific shape handling, but the consolidation lost:CacheDTypeCLI value;This left a split implementation: the reader referenced
nvfp4_ds_mla, whilevllm serve --kv-cache-dtype nvfp4_ds_mlafailed argument parsing before startup. The generic upstreamnvfp4cache dtype is not an alias: it has a different non-MLA record contract.Integration details
The three conflicts with GG were resolved in favor of GG because GG already has newer versions of those reader-side changes:
Only the missing writer-side pieces are introduced. The v17 runtime SM100+ guard and independent public-dispatch test reset are retained.
Validation
git diff --checkpasses.CacheDTypecontainsnvfp4_ds_mla.torch.uint8.(num_blocks, block_size, 432).vllm serve --helpexposesnvfp4_ds_mla.Closes the FF-to-GG consolidation gap for merged PR #82.