feat(mla): add compact NVFP4 FP8-RoPE cache writer - #37
Conversation
WalkthroughAdds a CUDA writer for NVFP4 MLA KV-cache records with FP8 RoPE, exposes it through a validated Torch/Python API, and tests record encoding plus compatibility with multi-split decode and multi-tile prefill readers. ChangesNVFP4 MLA FP8 RoPE writer
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant concat_and_cache_nvfp4_mla_fp8_rope
participant TorchOperator
participant ConcatAndCacheNvfp4MlaFp8RopeKernel
participant kv_cache
Caller->>concat_and_cache_nvfp4_mla_fp8_rope: provide kv_c, k_pe, kv_cache, slot_mapping
concat_and_cache_nvfp4_mla_fp8_rope->>TorchOperator: dispatch validated tensors
TorchOperator->>ConcatAndCacheNvfp4MlaFp8RopeKernel: launch one CTA per token
ConcatAndCacheNvfp4MlaFp8RopeKernel->>kv_cache: write packed NoPE and FP8 RoPE records
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 13143d9a-936e-4169-970e-6eacadf12509
📒 Files selected for processing (2)
b12x/attention/mla/kv_cache.pytests/test_attention_mla_kv_cache.py
| slot = Int64(slot_mapping[token]) | ||
| if slot >= Int64(0): | ||
| # Slot capacity is host-asserted < 2^31, so the block/offset | ||
| # split runs in Int32 (the byte offset below is Int64). | ||
| slot32 = slot.to(Int32) | ||
| block_idx = slot32 // Int32(self.block_size) | ||
| block_off = slot32 % Int32(self.block_size) | ||
| dst = ( | ||
| get_ptr_as_int64(kv_cache, 0) | ||
| + block_idx.to(Int64) * block_stride | ||
| + (block_off * entry_stride).to(Int64) | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Bound cache geometry and every nonnegative slot before pointer arithmetic.
block_size == 0 reaches division by zero, while any slot at or above cache capacity reaches an out-of-bounds global write because Line 196 only checks the lower bound. Require a positive block size and reject or kernel-guard oversized slots before computing dst; add zero-block and high-slot guard tests.
Also applies to: 436-440, 471-472
| for i in cutlass.range_constexpr(_PE_DIM // 2): | ||
| pair = _ld_global_u32( | ||
| get_ptr_as_int64(k_pe, token * k_pe_stride + Int32(2 * i)) | ||
| ) | ||
| if cutlass.const_expr(self.is_bf16): | ||
| f0, f1 = _bf16x2_to_f32x2(pair) | ||
| else: | ||
| f0, f1 = f16x2_to_f32x2(pair) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -C4 \
'_ld_global_u32|assumed_align=|data_ptr\(\) %|stride\(0\) %' \
b12x/attention/mla/kv_cache.py \
tests/test_attention_mla_kv_cache.pyRepository: lukealonso/b12x
Length of output: 5026
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '300,470p' b12x/attention/mla/kv_cache.pyRepository: lukealonso/b12x
Length of output: 6250
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C3 \
'k_pe.*data_ptr\(\)|k_pe.*stride\(0\)|assumed_align=16|assumed_align=2|from_dlpack\(' \
b12x/attention/mla/kv_cache.py tests -g '!**/.git/**'Repository: lukealonso/b12x
Length of output: 14172
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '180,290p' b12x/attention/mla/kv_cache.pyRepository: lukealonso/b12x
Length of output: 5167
Match alignment assumptions to the actual access sizes. b12x/attention/mla/kv_cache.py:253-260, 336-339, 460-467
k_peis read withld.global.b32, but the launch only assumes 2-byte alignment and never checks for a 4-byte-aligned row start; an odd offset or stride can misalign every load.kv_cacheis declared withassumed_align=16, but the host checks only enforce 8-byte alignment. Lower the assumption to 8 or enforce 16-byte alignment consistently.
2 participants
Summary
b12x.attention.mla.kv_cache.concat_and_cache_nvfp4_mla_fp8_ropenvfp4_ds_mlarecord consumed by the existing SM120 sparse-MLA readers:[0, 256): packed E2M1 NoPE values[256, 288): 32 E4M3 group scales[288, 292): fp32 RoPE scale[292, 304): zero padding[304, 368): 64 E4M3 RoPE valuesThe API intentionally lives on the explicit
b12x.attention.mla.kv_cachesubmodule. It is not imported fromb12x.attention.mla.__init__, so users that do not select this cache format do not pay import or JIT-registration side effects.Motivation
The compact GLM MLA cache format is already understood by b12x's
ScaleFormat.NVFP4_E4M3decode and prefill readers, but the package did not expose the matching writer. Downstream vLLM integration consequently depended on a separately mounted private shared object. Moving the writer to b12x gives the record ABI one owner and lets downstream callers use a normal package API.This PR makes no serving-throughput claim; it is an API/correctness extraction. The paired vLLM PR only replaces the unavailable external loader with this package API.
Correctness
The focused GPU suite covers:
Development validation also compared complete records against the established 368-byte writer over 15 remap cases: BF16/FP16, block sizes 16/64, token counts 1/37/256, multiple seeds, and skipped slots.
Test plan
pytest tests/test_attention_mla_kv_cache.py -vv— 18 passed on RTX PRO 6000 Blackwell / SM120ruff check b12x/attention/mla/kv_cache.py tests/test_attention_mla_kv_cache.py— passedruff format --check b12x/attention/mla/kv_cache.py tests/test_attention_mla_kv_cache.py— passedSummary by CodeRabbit
New Features
Bug Fixes
Tests