@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7702ee87db

[Bugfix][DSA] Write nvfp4_ds_mla from the fused norm+rope kernel (#55538)

Signed-off-by: Shiyang Chen <shiychen@nvidia.com>

sychen52 committed Sep 14, 2026 at 13:25 UTC 7702ee87dba0d8eed7f201e77a0a6613aac738a8

4 files changed +186 -20

csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu

+2 -12

@@ -109,9 +109,6 @@ __device__ __forceinline__ void nvfp4_unpack16_e2m1(uint64_t raw,

109 #endif

110 }

111

112 -// Rounded UP to the next e4m3 value so amax / sf never exceeds the payload

113 -// range: round-to-nearest can land up to 33% low near e4m3's 2^-9 floor, which

114 -// would saturate the tile's largest values.

112 __device__ __forceinline__ float nvfp4_tile_scale(const float* vals,

113 float max_val,

114 uint8_t* sf_out) {

@@ -121,15 +118,8 @@ __device__ __forceinline__ float nvfp4_tile_scale(const float* vals,

118 amax = fmaxf(amax, fabsf(vals[i]));

119 }

120 const float sf_f = fmaxf(amax / max_val, 0.001953125f); // 2^-9

124 - __nv_fp8_e4m3 sf8(sf_f);

125 - uint8_t sf_bits = *reinterpret_cast<const uint8_t*>(&sf8);

126 - if (float(sf8) < sf_f && sf_bits < 0x7E) {

127 - // Bump to the next e4m3 value (bit patterns of positive e4m3 values are

128 - // monotonic; 0x7E = 448 is the max finite value).

129 - sf_bits += 1;

130 - sf8 = *reinterpret_cast<const __nv_fp8_e4m3*>(&sf_bits);

131 - }

132 - *sf_out = sf_bits;

121 + const __nv_fp8_e4m3 sf8(sf_f);

122 + *sf_out = *reinterpret_cast<const uint8_t*>(&sf8);

123 return float(sf8);

124 }

125

tests/kernels/test_fused_deepseek_v32_norm_rope.py

+120

@@ -573,6 +573,126 @@ def test_fused_norm_rope_ds_mla(num_tokens: int):

573 assert (topk == 7).all(), "topk buffer should be untouched (no indexer)"

574

575

576 +E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]

577 +

578 +

579 +def quantize_to_e2m1(x: torch.Tensor) -> torch.Tensor:

580 + """Round to nearest e2m1, saturating to +-6 like cvt.rn.satfinite.e2m1x2.f32."""

581 + mags = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=x.device)

582 + a = x.float().abs().clamp_max(6.0)

583 + mids = (mags[:-1] + mags[1:]) / 2

584 + code = torch.bucketize(a, mids, right=True)

585 + on_tie = (a.unsqueeze(-1) == mids).any(dim=-1)

586 + tie_code = torch.bucketize(a, mids, right=False)

587 + code = torch.where(on_tie, tie_code + (tie_code & 1), code)

588 + return (torch.signbit(x).to(torch.uint8) << 3) | code.to(torch.uint8)

589 +

590 +

591 +def nvfp4_sf_byte(s: torch.Tensor) -> torch.Tensor:

592 + """Scale-factor byte permutation shared with FlashMLA: 8*(s&3) + (s>>2)."""

593 + return 8 * (s % 4) + (s // 4)

594 +

595 +

596 +@pytest.mark.skipif(

597 + not current_platform.is_device_capability(100),

598 + reason="nvfp4_ds_mla requires SM100 (Blackwell)",

599 +)

600 +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])

601 +def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int):

602 + """nvfp4_ds_mla MLA cache layout (FlashMLA sparse, SM100 only).

603 +

604 + Per-token 352-byte entry: 256 B of 512 e2m1 NoPE packed 2/byte (low nibble

605 + = even element) | 64 B unscaled e4m3 RoPE | 32 B byte-permuted e4m3 tile

606 + scales, one per 16 NoPE elements.

607 + """

608 + torch.manual_seed(5)

609 + dev = "cuda"

610 + max_pos = 8192

611 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos

612 +

613 + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)

614 + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)

615 + # Tile 0 is deliberately low-magnitude: amax/6 lands between two e4m3

616 + # subnormals (spaced a flat 2^-9 there), so round-to-nearest picks a scale

617 + # BELOW amax/6 and the tile's peak saturates at +-6. Asserted below, so a

618 + # future change to the scale rule cannot silently stop covering this.

619 + # amax must land in (6, 9) * 2^-9 after rms_norm; 7.5 sits mid-band, with

620 + # margin for the per-token RMS to vary.

621 + kv_c[:, :16] = (

622 + 7.5 * (2.0**-9) * torch.linspace(0.2, 1.0, 16, device=dev, dtype=torch.bfloat16)

623 + )

624 + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)

625 + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)

626 + kvw = torch.ones(KV_LORA, device=dev, dtype=torch.bfloat16)

627 + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)

628 +

629 + bs = max_pos

630 + mla_cache = torch.zeros(1, bs, 352, device=dev, dtype=torch.uint8)

631 + slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)

632 + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)

633 +

634 + q_out = K.fused_norm_rope(

635 + pos,

636 + q_c,

637 + qw,

638 + EPS,

639 + kv_c,

640 + kvw,

641 + EPS,

642 + k_pe,

643 + mla_cos_sin,

644 + None,

645 + None,

646 + None,

647 + EPS,

648 + None,

649 + topk,

650 + slot_mapping=slot,

651 + indexer_k_cache=None,

652 + mla_kv_cache=mla_cache,

653 + mla_kv_cache_dtype="nvfp4_ds_mla",

654 + mla_k_scale=None,

655 + has_indexer=False,

656 + index_rope_interleave=False,

657 + )

658 +

659 + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (nvfp4_ds_mla)")

660 +

661 + kv_ref = rms_norm(kv_c, kvw) # [N, 512] fp32

662 + kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) # [N, 64]

663 + tiles = kv_ref.view(num_tokens, 32, 16)

664 + amax = tiles.abs().amax(dim=-1)

665 + scale_target = torch.clamp_min(amax / 6.0, 2.0**-9)

666 + ref_scale = scale_target.to(FP8) # round-to-nearest

667 +

668 + # The crafted tile must actually exercise the saturating path.

669 + assert (ref_scale[:, 0].float() < scale_target[:, 0]).all(), (

670 + "tile 0 should round its scale DOWN; adjust the crafted magnitude"

671 + )

672 + assert (amax[:, 0] / ref_scale[:, 0].float() > 6.0).all(), (

673 + "tile 0 should saturate e2m1; adjust the crafted magnitude"

674 + )

675 +

676 + ref_codes = quantize_to_e2m1(tiles / ref_scale.float().unsqueeze(-1))

677 + ref_codes = ref_codes.reshape(num_tokens, KV_LORA)

678 +

679 + cache = mla_cache[0, :num_tokens] # [N, 352] uint8

680 + packed = cache[:, : KV_LORA // 2]

681 + got_codes = torch.stack([packed & 0xF, packed >> 4], dim=-1)

682 + got_codes = got_codes.reshape(num_tokens, KV_LORA) # low nibble = even elem

683 + torch.testing.assert_close(got_codes, ref_codes, rtol=0, atol=0)

684 +

685 + got_rope = cache[:, KV_LORA // 2 : KV_LORA // 2 + ROPE_DIM].view(FP8)

686 + assert_fp8(got_rope, kpe_ref.to(FP8), "nvfp4_ds_mla RoPE e4m3")

687 +

688 + perm = nvfp4_sf_byte(torch.arange(32, device=dev))

689 + got_scale = cache[:, KV_LORA // 2 + ROPE_DIM :].view(FP8)[:, perm]

690 + torch.testing.assert_close(got_scale.float(), ref_scale.float(), rtol=0, atol=0)

691 +

692 + # No indexer on this call: top-k buffer must be untouched.

693 + assert (topk == 7).all(), "topk buffer should be untouched (no indexer)"

694 +

695 +

696 def test_fused_norm_rope_supports_large_token_count():

697 """Keep the token count off CUDA grid-y at its 65,536-block boundary."""

698 num_tokens = 65536

vllm/models/deepseek_v32/attention.py

+3 -1

@@ -237,7 +237,9 @@ class DeepseekV32Attention(MLAAttention):

237

238 fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)

239 self._fp8_query = fp8_attention and self.impl.supports_quant_query_input

240 - self._fp8_kv_needs_view = fp8_attention and self.kv_cache_dtype != "fp8_ds_mla"

240 + self._fp8_kv_needs_view = fp8_attention and not self.kv_cache_dtype.endswith(

241 + "_ds_mla"

242 + )

243

244 self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False)

245

vllm/models/deepseek_v32/common/kernels.py

+61 -7

@@ -150,6 +150,7 @@ def _fused_norm_rope_kernel(

150 mla_cache_ds_scale_ptr,

151 mla_cache_ds_rope_ptr,

152 MLA_CACHE_DS_MLA: tl.constexpr,

153 + MLA_CACHE_NVFP4: tl.constexpr,

154 MLA_NUM_TILES: tl.constexpr,

155 MLA_TILE_DIM: tl.constexpr,

156 # Top k indices

@@ -261,6 +262,56 @@ def _fused_norm_rope_kernel(

262 mla_block_idx = slot_idx // mla_block_size

263 mla_block_off = slot_idx % mla_block_size

264

265 + if MLA_CACHE_NVFP4:

266 + # nvfp4_ds_mla layout (352 B/token, KV_DIM == 512):

267 + # [0, 256) 512 e2m1 NoPE, packed 2/byte (low nibble = even)

268 + # [256, 320) 64 e4m3 RoPE, unscaled

269 + # [320, 352) 32 e4m3 NoPE scales, byte-permuted

270 + # Keep in lockstep with concat_and_cache_nvfp4_ds_mla (CUDA) and

271 + # nvfp4_sf_byte() in FlashMLA: byte(s) = 8*(s&3) + (s>>2).

272 + byte_base = (

273 + mla_block_idx * mla_cache_block_stride

274 + + mla_block_off * mla_cache_entry_stride

275 + )

276 + kv_2d = tl.reshape(kv_c.to(tl.float32), (MLA_NUM_TILES, MLA_TILE_DIM))

277 + amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)

278 + # sf = e4m3(max(amax/6, 2^-9)), round-to-nearest like

279 + # cvt_warp_fp16_to_fp4.

280 + sf8 = tl.maximum(amax * (1.0 / 6.0), 0.001953125).to(tl.float8e4nv)

281 + q = kv_2d * (1.0 / sf8.to(tl.float32))

282 + qp = tl.reshape(q, (KV_DIM // 2, 2))

283 + sel = tl.arange(0, 2)[None, :]

284 + even = tl.sum(tl.where(sel == 0, qp, 0.0), axis=1)

285 + odd = tl.sum(tl.where(sel == 1, qp, 0.0), axis=1)

286 + # Same instruction the CUDA writer uses, so the two are bit-exact

287 + # by construction: dest byte = {high: odd, low: even}.

288 + packed = tl.inline_asm_elementwise(

289 + "{ .reg .b8 t; .reg .b32 w;"

290 + " cvt.rn.satfinite.e2m1x2.f32 t, $2, $1;"

291 + " mov.b32 w, {t, t, t, t}; and.b32 $0, w, 255; }",

292 + "=r,f,f",

293 + [even, odd],

294 + dtype=tl.int32,

295 + is_pure=True,

296 + pack=1,

297 + ).to(tl.uint8)

298 + tl.store(

299 + mla_cache_ptr + byte_base + tl.arange(0, KV_DIM // 2),

300 + packed.to(tl.float8e4nv, bitcast=True),

301 + )

302 + # RoPE: unscaled e4m3, interleaved like the fp8_ds_mla path

303 + rope_base = mla_cache_ptr + byte_base + KV_DIM // 2

304 + tl.store(rope_base + dim_off * 2, r1.to(tl.float8e4nv))

305 + tl.store(rope_base + dim_off * 2 + 1, r2.to(tl.float8e4nv))

306 + # scales, permuted: byte(s) = 8*(s&3) + (s>>2)

307 + sf_tile = tl.arange(0, MLA_NUM_TILES)

308 + sf_perm = 8 * (sf_tile % 4) + (sf_tile // 4)

309 + tl.store(

310 + mla_cache_ptr + byte_base + (KV_DIM // 2 + 64) + sf_perm,

311 + tl.reshape(sf8, (MLA_NUM_TILES,)),

312 + )

313 + return

314 +

315 if MLA_CACHE_DS_MLA:

316 # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element

317 # tile of the NoPE is dynamically quantized to fp8 with its own

@@ -491,7 +542,8 @@ def fused_norm_rope(

542 idx_cache_stride = 0

543

544 # --- MLA KV cache setup ---

494 - mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla"

545 + mla_cache_nvfp4 = mla_kv_cache_dtype == "nvfp4_ds_mla"

546 + mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla" or mla_cache_nvfp4

547 mla_cache_fp8 = is_quantized_kv_cache(mla_kv_cache_dtype) and not mla_cache_ds_mla

548 mla_num_tiles = 1

549 mla_ds_scale_view = torch.empty(0, dtype=torch.float32, device=device)

@@ -499,11 +551,12 @@ def fused_norm_rope(

551 if mla_kv_cache is not None:

552 mla_block_size = mla_kv_cache.shape[1]

553 if mla_cache_ds_mla:

502 - # 656-byte custom layout addressed in bytes; mla_cache_ptr is the

503 - # 1-byte fp8 view, so block/entry strides are byte offsets and the

504 - # fp32/bf16 views share the same buffer.

505 - assert kv_dim == 512, "fp8_ds_mla requires kv_lora_rank == 512"

506 - mla_num_tiles = kv_dim // 128

554 + # Custom byte-addressed layout: 656 B/token for fp8_ds_mla, 352 B

555 + # for nvfp4_ds_mla. mla_cache_ptr is the 1-byte fp8 view, so

556 + # block/entry strides are byte offsets and the fp32/bf16 views

557 + # (fp8_ds_mla only) share the same buffer.

558 + assert kv_dim == 512, "ds_mla layouts require kv_lora_rank == 512"

559 + mla_num_tiles = kv_dim // 16 if mla_cache_nvfp4 else kv_dim // 128

560 u8_cache = mla_kv_cache.view(torch.uint8)

561 mla_block_stride = u8_cache.stride(0)

562 mla_entry_stride = u8_cache.stride(1)

@@ -596,7 +649,8 @@ def fused_norm_rope(

649 mla_k_scale,

650 mla_ds_scale_view,

651 mla_ds_rope_view,

599 - mla_cache_ds_mla,

652 + mla_cache_ds_mla and not mla_cache_nvfp4,

653 + mla_cache_nvfp4,

654 mla_num_tiles,

655 kv_dim // mla_num_tiles if mla_cache_ds_mla else 1,

656 # Top k indices buffer