nvfp4 MLA KV cache + per-layer outer-scale + fp8-RoPE KV-cell (sm120) - #33
Conversation
3-way merged from fork-point 5af873a onto master; ~48 conflicts resolved toward the feature (see PR note). Companion: local-inference-lab/vllm#95. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughAdds NVFP4/E4M3 sparse MLA support with optional FP8 RoPE cache layouts, latent scaling, new trait metadata, cache geometry selection, decode kernel integration, and SM120 MG prefill routing. ChangesNVFP4 sparse MLA
Sequence Diagram(s)sequenceDiagram
participant Caller
participant sparse_mla_api
participant UnifiedMLATraits
participant UnifiedDecodeKernel
participant decode_math
Caller->>sparse_mla_api: NVFP4 parameters and latent_scale
sparse_mla_api->>UnifiedMLATraits: resolve scale format and fp8_rope
UnifiedMLATraits-->>UnifiedDecodeKernel: cache geometry and specialization
UnifiedDecodeKernel->>decode_math: stage, dequantize, and accumulate NVFP4 data
decode_math-->>Caller: sparse MLA output
🚥 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
b12x/attention/mla/traits.py (1)
15-24: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
KV_FP8_ROPEenv-gate caching acrosstraits.pyandapi.py. Both modules independently parse and cache the sameos.environ.get("KV_FP8_ROPE", "0") == "1"boolean at import time instead of sharing one source of truth; the shared root cause is thatapi.pynever importstraits.py's existing accessor.
b12x/attention/mla/traits.py#L15-L24: keep this as the single canonical definition ofKV_FP8_ROPE_ENV/KV_FP8_ROPE_ENABLED/kv_fp8_rope_enabled().b12x/attention/mla/api.py#L30-L31: drop the local_KV_FP8_ROPE_ENV/_KV_FP8_ROPE_ENABLEDconstants and have_resolve_kv_fp8_rope(Lines 141-148) calltraits.kv_fp8_rope_enabled()instead of its own cached copy.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: f45930c0-991e-40d7-bf10-dbd70020ffee
📒 Files selected for processing (8)
b12x/attention/mla/api.pyb12x/attention/mla/decode_math.pyb12x/attention/mla/io.pyb12x/attention/mla/io_mg.pyb12x/attention/mla/kernel.pyb12x/attention/mla/prefill.pyb12x/attention/mla/prefill_mg.pyb12x/attention/mla/traits.py
| _sm120_route = _use_sm120_sparse_mla(backend=backend, device=q_all.device) | ||
| # NVFP4 (scale_format=2) selection: explicit kwarg wins; otherwise the | ||
| # scratch/workspace planned kv_cache_dtype supplies it. None -> inferred | ||
| # from q_head_dim inside the kernel launchers (fp8 GLM default). | ||
| scale_format_for_call = ( | ||
| scale_format if scale_format is not None else getattr(workspace, "scale_format", None) | ||
| ) | ||
| if int(scale_format_for_call or -1) == 2 and fp8_rope is None: | ||
| # Normal vLLM calls arrive through b12x.integration.mla, whose stable | ||
| # public signature does not carry this new option. The allocated record | ||
| # is therefore the authoritative process-lifetime ABI at this boundary; | ||
| # derive the specialization from it instead of rereading a mutable env. | ||
| record_bytes = int(kv_cache.shape[-1]) | ||
| if record_bytes not in (368, 432): | ||
| raise ValueError( | ||
| "NVFP4 sparse MLA cache record must be 368 or 432 bytes, got " | ||
| f"{record_bytes}" | ||
| ) | ||
| fp8_rope_for_call = record_bytes == 368 | ||
| else: | ||
| fp8_rope_for_call = _resolve_kv_fp8_rope(fp8_rope) | ||
| if int(scale_format_for_call or -1) == 2: | ||
| expected_record_bytes = 368 if fp8_rope_for_call else 432 | ||
| if int(kv_cache.shape[-1]) != expected_record_bytes: | ||
| raise ValueError( | ||
| "NVFP4 sparse MLA cache record disagrees with KV_FP8_ROPE: " | ||
| f"got {int(kv_cache.shape[-1])} bytes, expected " | ||
| f"{expected_record_bytes}" | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm sparse_mla_reference has no NVFP4/scale_format awareness.
rg -n 'def sparse_mla_reference' -A 40 b12x/attention/mla/ | head -80
rg -n 'scale_format' b12x/attention/mla/reference.py b12x/attention/mla/compressed_reference.py 2>/dev/nullRepository: lukealonso/b12x
Length of output: 2935
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== api.py excerpt =="
sed -n '560,720p' b12x/attention/mla/api.py | cat -n
echo
echo "== reference.py excerpt =="
sed -n '1,220p' b12x/attention/mla/reference.py | cat -nRepository: lukealonso/b12x
Length of output: 17815
NVFP4 needs an SM120-only guard
scale_format=2 is validated here, but the non-_sm120_route path still falls through to sparse_mla_reference(), which only understands the packed NSA layout and has no NVFP4 handling. Add the same SM120 gate used for attn_sink so non-SM120 callers fail fast with a clear error instead of reaching the reference path with an unsupported cache format.
| if cutlass.const_expr(t.scale_format == ScaleFormat.NVFP4_E4M3): | ||
| s0_load_q_bf16_to_smem( | ||
| q_token, q_fp8_addr, q_rope_addr, | ||
| head_base, Int32(self.valid_hpb), tid, | ||
| d_nope=t.d_nope, d_rope=t.d_rope, hpb=t.hpb, | ||
| q_nope_bf16_stride=t.q_nope_stride, | ||
| q_rope_stride=L.q_rope_stride, | ||
| num_threads=self.math_threads, barrier_id=2, | ||
| ) | ||
| else: | ||
| s0_quantize_q_to_smem( | ||
| q_token, q_fp8_addr, q_sc_view, q_rope_addr, amax_view, | ||
| head_base, Int32(self.valid_hpb), tid, | ||
| d_nope=t.d_nope, d_rope=t.d_rope, | ||
| d_qk=t.d_nope + t.d_rope, | ||
| quant_tile=t.quant_tile, num_scales=t.num_scales, | ||
| hpb=(8 if self.native_h8 else t.hpb), | ||
| q_nope_stride=t.q_nope_stride, | ||
| q_rope_stride=L.q_rope_stride, | ||
| num_threads=self.math_threads, barrier_id=2, | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Locate the shared-storage class and layout builder to confirm q_fp8 sizing
# accounts for BF16 (2B/elem) vs FP8 (1B/elem) staging under scale_format==NVFP4_E4M3.
rg -n 'def get_unified_shared_storage_cls' -A 80 b12x/attention/mla/
rg -n 'def make_smem_layout\b' -A 60 b12x/attention/mla/
rg -n 'q_nope_bf16' b12x/attention/mla/
rg -n '\bq_fp8\b.*=' b12x/attention/mla/Repository: lukealonso/b12x
Length of output: 16803
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== kernel.py locations =="
sed -n '710,750p' b12x/attention/mla/kernel.py
echo
sed -n '1454,1498p' b12x/attention/mla/kernel.py
echo
echo "== smem layout =="
sed -n '189,240p' b12x/attention/mla/smem.py
echo
echo "== run_unified_decode and validation =="
sed -n '2230,2385p' b12x/attention/mla/kernel.py
echo
sed -n '560,610p' b12x/attention/mla/api.py
echo
echo "== callers of run_unified_decode =="
rg -n 'run_unified_decode\(' b12x/attention/mla/Repository: lukealonso/b12x
Length of output: 16642
Allocate a BF16 staging buffer for the NVFP4 path kernel.py:725-745, 1464-1491 writes BF16 Q-NoPE into q_fp8, but smem.py allocates q_fp8 as 1 byte/elem (hpb * q_nope_stride). That can overrun the packed shared-memory struct; mirror the dedicated q_nope_bf16 buffer used in prefill_mg.py, or widen this specialization.
There was a problem hiding this comment.
Actionable comments posted: 2
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: f45930c0-991e-40d7-bf10-dbd70020ffee
📒 Files selected for processing (8)
b12x/attention/mla/api.pyb12x/attention/mla/decode_math.pyb12x/attention/mla/io.pyb12x/attention/mla/io_mg.pyb12x/attention/mla/kernel.pyb12x/attention/mla/prefill.pyb12x/attention/mla/prefill_mg.pyb12x/attention/mla/traits.py
| model_type, compute_mode, scale_format = infer_model_type(q_head_dim, swa_k_cache.dtype) | ||
| traits = make_unified_traits(model_type, compute_mode, scale_format) | ||
| if scale_format_override is not None: | ||
| scale_format = int(scale_format_override) | ||
| traits = make_unified_traits( | ||
| model_type, | ||
| compute_mode, | ||
| scale_format, | ||
| fp8_rope=fp8_rope_override, | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Validate the NVFP4 cache ABI inside run_unified_decode.
Line 2360 derives 368/432-byte geometry from the override or environment, but this launcher does not verify swa_k_cache.shape[-1]. Direct callers can therefore select the wrong record stride and misread the flattened cache.
Proposed validation
if scale_format_override is not None:
scale_format = int(scale_format_override)
+ if scale_format == ScaleFormat.NVFP4_E4M3 and fp8_rope_override is None:
+ record_bytes = int(swa_k_cache.shape[-1])
+ if record_bytes not in (368, 432):
+ raise ValueError(
+ f"NVFP4 cache record must be 368 or 432 bytes, got {record_bytes}"
+ )
+ fp8_rope_override = record_bytes == 368
traits = make_unified_traits(
model_type,
compute_mode,
scale_format,
fp8_rope=fp8_rope_override,
)
+ if (
+ scale_format == ScaleFormat.NVFP4_E4M3
+ and int(swa_k_cache.shape[-1]) != int(traits.kv_gmem_stride)
+ ):
+ raise ValueError("NVFP4 cache record width disagrees with fp8_rope_override")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| model_type, compute_mode, scale_format = infer_model_type(q_head_dim, swa_k_cache.dtype) | |
| traits = make_unified_traits(model_type, compute_mode, scale_format) | |
| if scale_format_override is not None: | |
| scale_format = int(scale_format_override) | |
| traits = make_unified_traits( | |
| model_type, | |
| compute_mode, | |
| scale_format, | |
| fp8_rope=fp8_rope_override, | |
| ) | |
| model_type, compute_mode, scale_format = infer_model_type(q_head_dim, swa_k_cache.dtype) | |
| if scale_format_override is not None: | |
| scale_format = int(scale_format_override) | |
| if scale_format == ScaleFormat.NVFP4_E4M3 and fp8_rope_override is None: | |
| record_bytes = int(swa_k_cache.shape[-1]) | |
| if record_bytes not in (368, 432): | |
| raise ValueError( | |
| f"NVFP4 cache record must be 368 or 432 bytes, got {record_bytes}" | |
| ) | |
| fp8_rope_override = record_bytes == 368 | |
| traits = make_unified_traits( | |
| model_type, | |
| compute_mode, | |
| scale_format, | |
| fp8_rope=fp8_rope_override, | |
| ) | |
| if ( | |
| scale_format == ScaleFormat.NVFP4_E4M3 | |
| and int(swa_k_cache.shape[-1]) != int(traits.kv_gmem_stride) | |
| ): | |
| raise ValueError("NVFP4 cache record width disagrees with fp8_rope_override") |
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
|
Superseded by the clean current-master replacement #35. The replacement preserves the current Spark H8/H16 optimizations, fixes the BF16 shared-memory sizing and duplicate-definition issues found in this draft, adds direct ABI guards and focused tests, and is ready/non-draft and mergeable. Closing this conflicting draft so there is one canonical NVFP4 MLA PR. |
|
superseded by #35 |
3 participants
nvfp4 MLA KV cache + per-layer outer-scale + fp8-RoPE KV-cell (sm120)
Adds a 4-bit NVFP4 MLA KV cache path to the b12x MLA kernels, plus two quality/efficiency
refinements, all validated on GLM-5.2 (4x RTX PRO 6000 Blackwell, TP4/DCP4):
fp8 — ~1.5x the context pool vs fp8_ds_mla.
scale_format==2): in-kernels_lrestore on read; recovers latentquantization loss (teacher-forced KLD 0.184 -> 0.152).
(+15.8% pool, genuine 1M single-request). 64K needle retrieval 15/15 = 15/15, KL below noise.
Companion vLLM plumbing: local-inference-lab/vllm#95.
This is 3-way merged from the fork-point
5af873aonto currentmaster. The nvfp4/fp8-ropeadditions overlap
master's post-5af873aMLA changes in ~48 spots (decode_math/kernel/prefill/prefill_mg), which were resolved toward preserving the feature. Those resolutions need a careful
review against your recent MLA work — flagging so we can reconcile any infra changes I may have
shadowed. Opening as draft for exactly that reason.
Summary by CodeRabbit