@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a7fda4c88b

[Bugfix] Fix SM100 fp8_ds_mla cache scales (#49435)

Signed-off-by: Ariel Bereslavsky <arielb@getimpala.ai> Signed-off-by: Ariel Bereslavsky <ariel@getimpala.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Ariel Bereslavsky <arielb@getimpala.ai> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Misha Goin <mgoin64@gmail.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Ariel Bereslavsky committed Sep 20, 2026 at 03:15 UTC a7fda4c88bfc421d31e33acc5e01e86ebe467ad8

8 files changed +76 -38

csrc/libtorch_stable/cache_kernels.cu

+3 -3

@@ -19,7 +19,6 @@

20 #include <algorithm>

21 #include <cassert>

22 -#include <cfloat>

22 #include <cstdlib>

24 #ifdef USE_ROCM

@@ -589,8 +588,9 @@ __global__ void concat_and_cache_ds_mla_kernel(

588 max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 16));

589 }

590

592 - // Compute the scale for the tile

593 - float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN);

591 + // Both SM90 and SM100 readers preserve power-of-two fp32 scales exactly.

592 + float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, 1e-4f);

593 + tile_scale = exp2f(ceilf(log2f(tile_scale)));

594

595 // The first lane of each half-warp writes the scale to kv_cache

596 if ((lane_idx == 0) || (lane_idx == 16)) {

csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu

+10 -9

@@ -35,8 +35,9 @@

35 *

36 * fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert):

37 * - full key concat (bf16), and cache insert in DeepSeek's 656-byte

38 - * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic

39 - * scales, RoPE bf16), bit-compatible with concat_and_cache_ds_mla_kernel.

38 + * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile

39 + * power-of-two scales, RoPE bf16), bit-compatible with

40 + * concat_and_cache_ds_mla_kernel.

41 *

42 * Both use Programmatic Dependent Launch (PDL) to overlap the tail of the

43 * producing GEMMs on sm_90+, and are structured after

@@ -68,7 +69,6 @@

69 #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"

70 #endif

71 #include <cuda_runtime.h>

71 -#include <cfloat>

72 #include <type_traits>

74 #ifdef USE_ROCM

@@ -111,9 +111,9 @@ constexpr float kFp8Max = 224.0f;

111 #else

112 constexpr float kFp8Max = 448.0f;

113 #endif

114 -// Divisor for fp8_ds_mla per-tile dynamic scales (matches cache_kernels.cu).

115 -// fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128), [512,528) 4 fp32

116 -// tile scales, [528,656) RoPE 64 bf16.

114 +// Divisor for fp8_ds_mla per-tile power-of-two scales (matches

115 +// cache_kernels.cu). fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128),

116 +// [512,528) 4 fp32 tile scales, [528,656) RoPE 64 bf16.

117 constexpr float kFp8ScaleDivisor = kFp8Max;

118

119 // Copy 8 source elements (one uint4 of bf16/fp16) to `dst`. FP8=false stores a

@@ -325,7 +325,7 @@ __device__ __forceinline__ void writeLatent576(

325 }

326

327 // Write [kv_c | k_pe] into the fp8_ds_mla 656B entry using one warp: NoPE 512

328 -// as fp8 in 4 tiles of 128 (per-tile dynamic absmax scale, 4 fp32 scales at

328 +// as fp8 in 4 tiles of 128 (per-tile power-of-two scale, 4 fp32 scales at

329 // [512,528)), RoPE 64 as bf16 at [528,656). Bit-compatible with

330 // concat_and_cache_ds_mla_kernel.

331 template <typename scalar_t, bool APPLY_ROPE = false>

@@ -349,7 +349,8 @@ __device__ __forceinline__ void writeDsMlaCache(

349 for (int offset = 4; offset > 0; offset /= 2) {

350 max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 8));

351 }

352 - float const tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN);

352 + float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, 1e-4f);

353 + tile_scale = exp2f(ceilf(log2f(tile_scale)));

354 if ((laneId & 7) == 0) {

355 reinterpret_cast<float*>(row)[kKvLoraRank / 4 + tile] = tile_scale;

356 }

@@ -598,7 +599,7 @@ __global__ void fusedKimiK3MLAKVConcatPackKernel(

599 // ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert

600 //

601 // Cache entry (656 bytes), matching concat_and_cache_ds_mla_kernel:

601 -// [0, 512) NoPE 512 vals as fp8, 4 tiles of 128, each dynamically scaled

602 +// [0, 512) NoPE 512 vals as fp8, 4 tiles of 128, power-of-two scaled

603 // [512, 528) 4 fp32 per-tile scales

604 // [528, 656) RoPE 64 vals as bf16 (unquantized)

605 // The cache slot uses one warp: lane L quantizes NoPE elems [L*16, L*16+16)

tests/kernels/attention/test_cache.py

+5 -5

@@ -960,14 +960,14 @@ def test_concat_and_cache_ds_mla(

960 tile_end = (tile_idx + 1) * 128

961 tile_data[:] = kv_c_data[tile_start:tile_end]

962

963 - # tile_scale = tile_data.amax().to(torch.float32) / 448.

964 - # NOTE: Using torch's amax() gives different results,

965 - # so this must be manually computed.

963 + # Using torch's amax() gives different results, so this must be

964 + # manually computed.

965 tile_data_float = tile_data.to(torch.float32)

966 manual_max = abs(tile_data_float[0])

967 for j in range(1, 128):

968 manual_max = max(manual_max, abs(tile_data_float[j]))

970 - tile_scale = manual_max / 448.0

969 + raw_scale = torch.clamp(manual_max / 448.0, min=1e-4)

970 + tile_scale = torch.exp2(torch.ceil(torch.log2(raw_scale)))

971

972 ref_cache_32bit[kv_lora_rank // 4 + tile_idx] = tile_scale

973

@@ -1008,7 +1008,7 @@ def test_concat_and_cache_ds_mla(

1008 ref_rope = ref_cache_slice.view(dtype)[kv_lora_rank // 2 + 8 :]

1009

1010 torch.testing.assert_close(kv_nope, ref_nope, atol=0.001, rtol=0.1)

1011 - torch.testing.assert_close(kv_scales, ref_scales, atol=0.001, rtol=0.1)

1011 + torch.testing.assert_close(kv_scales, ref_scales, atol=0, rtol=0)

1012 torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1)

1013

1014

tests/kernels/attention/test_flashmla_sparse.py

+31 -8

@@ -109,8 +109,9 @@ def test_sparse_flashmla_metadata_smoke():

109 assert num_splits is None

110

111

112 -def test_sparse_flashmla_decode_smoke():

112 +def test_sparse_flashmla_decode_matches_cache_writer_scales():

113 import vllm.v1.attention.ops.flashmla as fm

114 + from vllm import _custom_ops as ops

115

116 ok, reason = fm.is_flashmla_sparse_supported()

117 if not ok:

@@ -130,7 +131,7 @@ def test_sparse_flashmla_decode_smoke():

131 # Metadata

132 q_seq_per_hk = seqlen_q * num_heads_q // num_heads_k

133 # q_heads_per_hk = num_heads_q // num_heads_k

133 - cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device)

134 + cache_seqlens = torch.ones(batch_size, dtype=torch.int32, device=device)

135 tile_md, num_splits = fm.get_mla_metadata(

136 cache_seqlens,

137 q_seq_per_hk,

@@ -146,14 +147,33 @@ def test_sparse_flashmla_decode_smoke():

147 dtype=torch.bfloat16,

148 device=device,

149 )

149 - k_cache = torch.zeros(

150 - (1, page_block_size, num_heads_k, bytes_per_token),

151 - dtype=torch.uint8,

152 - device=device,

150 + cache_rows = torch.zeros(

151 + (1, page_block_size, bytes_per_token), dtype=torch.uint8, device=device

152 )

154 - indices = torch.zeros(

155 - (batch_size, seqlen_q, topk), dtype=torch.int32, device=device

153 + # 336 / 448 = 0.75 distinguishes arbitrary FP32 from power-of-two scales.

154 + kv_c = torch.full((1, 512), 336.0, dtype=torch.bfloat16, device=device)

155 + kv_c[:, -128:] = 0

156 + ops.concat_and_cache_mla(

157 + kv_c,

158 + torch.zeros((1, 64), dtype=torch.bfloat16, device=device),

159 + cache_rows,

160 + torch.zeros(1, dtype=torch.int64, device=device),

161 + "fp8_ds_mla",

162 + torch.ones(1, dtype=torch.float32, device=device),

163 )

164 + scales = cache_rows[0, 0].view(torch.float32)[128:132]

165 + expected_scales = torch.tensor([1.0, 1.0, 1.0, 2**-13], device=device)

166 + torch.testing.assert_close(scales, expected_scales, rtol=0, atol=0)

167 + cached_nope = cache_rows[0, 0, :512].view(torch.float8_e4m3fn).float()

168 + expected = (cached_nope * scales.repeat_interleave(128)).to(torch.bfloat16)

169 +

170 + k_cache = cache_rows.unsqueeze(2)

171 + indices = torch.full(

172 + (batch_size, seqlen_q, topk), -1, dtype=torch.int32, device=device

173 + )

174 + indices[..., 0] = 0

175 + # SM90 sparse decode only supports a fixed top-k width. Padding with -1

176 + # exercises the same cache data on SM90 and SM100 without topk_length.

177

178 block_table = torch.zeros((batch_size, 128), dtype=torch.int32, device=device)

179 out, lse = fm.flash_mla_with_kvcache(

@@ -170,6 +190,9 @@ def test_sparse_flashmla_decode_smoke():

190 assert out.shape[0] == batch_size

191 assert out.shape[-1] == head_dim_v

192 assert lse.shape[0] == batch_size

193 + torch.testing.assert_close(

194 + out, expected.view(1, 1, 1, 512).expand_as(out), rtol=0, atol=0

195 + )

196

197

198 @pytest.mark.parametrize("h_q", [64, 128])

tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py

+16 -2

@@ -222,7 +222,8 @@ def test_decode_concat_ignores_negative_slots_for_cache(cache_format: str) -> No

222 _assert_cache_matches_reference(mixed_cache, reference_cache, initial_cache)

223

224

225 -def test_ds_mla_cache_insert_bit_compatible_with_reference() -> None:

225 +@pytest.mark.parametrize("decode", [False, True])

226 +def test_ds_mla_cache_insert_bit_compatible_with_reference(decode: bool) -> None:

227 """Fused ds_mla insertion must match the reference bit-for-bit."""

228 torch.manual_seed(0)

229 num_tokens, num_blocks = 33, 16

@@ -242,5 +243,18 @@ def test_ds_mla_cache_insert_bit_compatible_with_reference() -> None:

243 got = ref.clone()

244 scale = torch.ones(1, device="cuda", dtype=torch.float32)

245 ops.concat_and_cache_mla(kv_c, k_pe.squeeze(1), ref, slots, "fp8_ds_mla", scale)

245 - fused_mla_key_concat_ds_mla_insert(q, k_nope, k_pe, kv_c, got, slots)

246 + if decode:

247 + ql_nope = torch.randn(

248 + num_tokens, NUM_HEADS, KV_LORA_RANK, device="cuda", dtype=dt

249 + )

250 + fused_mla_decode_q_concat_kv_cache_insert(

251 + ql_nope, q[..., -ROPE_HEAD_DIM:], kv_c, k_pe, got, slots, ds_mla=True

252 + )

253 + else:

254 + fused_mla_key_concat_ds_mla_insert(q, k_nope, k_pe, kv_c, got, slots)

255 assert torch.equal(ref, got)

256 + rows = got.view(-1, DS_MLA_CACHE_ENTRY)[slots[slots >= 0]]

257 + scales = rows[:, KV_LORA_RANK : KV_LORA_RANK + 16].view(torch.float32)

258 + torch.testing.assert_close(

259 + torch.log2(scales), torch.log2(scales).round(), rtol=0, atol=0

260 + )

tests/kernels/test_fused_deepseek_v32_norm_rope.py

+5 -4

@@ -507,8 +507,8 @@ def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool):

507 def test_fused_norm_rope_ds_mla(num_tokens: int):

508 """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100).

509

510 - Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, dynamic float32

511 - scale) | 4 float32 scales | 64 bf16 (unquantized) RoPE.

510 + Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, power-of-two scales

511 + stored as float32) | 4 float32 scales | 64 bf16 (unquantized) RoPE.

512 """

513 torch.manual_seed(5)

514 dev = "cuda"

@@ -557,7 +557,8 @@ def test_fused_norm_rope_ds_mla(num_tokens: int):

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

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

559 tiles = kv_ref.view(num_tokens, 4, 128)

560 - ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1.1754944e-38)

560 + ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1e-4)

561 + ref_scale = torch.exp2(torch.ceil(torch.log2(ref_scale)))

562 ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, KV_LORA).to(FP8)

563

564 cache = mla_cache[0, :num_tokens] # [N, 656] uint8

@@ -566,7 +567,7 @@ def test_fused_norm_rope_ds_mla(num_tokens: int):

567 rope_off = KV_LORA // 2 + 8

568 rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + ROPE_DIM]

569

569 - torch.testing.assert_close(scales, ref_scale, rtol=1e-2, atol=1e-6)

570 + torch.testing.assert_close(scales, ref_scale, rtol=0, atol=0)

571 assert_fp8(nope, ref_nope, "ds_mla NoPE fp8")

572 assert_bf16(rope_vals, kpe_ref, "ds_mla RoPE bf16")

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

vllm/models/deepseek_v32/common/kernels.py

+3 -4

@@ -315,7 +315,7 @@ def _fused_norm_rope_kernel(

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

318 - # float32 scale; the RoPE tail is stored unquantized in bf16.

318 + # power-of-two scale stored as float32; the RoPE tail is bf16.

319 # bytes [0, KV_DIM) : KV_DIM fp8 NoPE values

320 # bytes [KV_DIM, KV_DIM + 16) : MLA_NUM_TILES float32 scales

321 # bytes [KV_DIM + 16, ...) : 2 * KPE_HALF_ROT_DIM bf16 RoPE

@@ -327,9 +327,8 @@ def _fused_norm_rope_kernel(

327 )

328 kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM))

329 tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)

330 - # scale = amax / 448 (fp8 e4m3 max), matching the reference

331 - # concat_and_cache_ds_mla kernel; floored to FLT_MIN.

332 - tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1.1754944e-38)

330 + tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1e-4)

331 + tile_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(tile_scale)))

332 kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,))

333 tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8)

334 tile_off = tl.arange(0, MLA_NUM_TILES)

vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py

+3 -3

@@ -76,9 +76,9 @@ def fused_mla_key_concat_ds_mla_insert(

76 """Concat full K (bf16) and insert the latent in the fp8_ds_mla layout.

78 The cache uses DeepSeek's 656-byte block-scaled layout (NoPE in 4 tiles of

79 - 128 with per-tile dynamic fp8 scales, RoPE as bf16) -- self-scaling, so no

80 - scale argument. Returns the bf16 full key; optionally rotates ``q`` and

81 - writes ``kv_cache`` in place.

79 + 128 with per-tile power-of-two scales stored as float32, RoPE as bf16) --

80 + self-scaling, so no scale argument. Returns the bf16 full key; optionally

81 + rotates ``q`` and writes ``kv_cache`` in place.

82 """

83 k_pe = k_pe.reshape(k_pe.shape[0], -1)

84 tp, num_heads, qk_nope_head_dim = k_nope.shape