diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index c5e35c52b8..4b8f4c23a5 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -19,7 +19,6 @@ #include #include -#include #include #ifdef USE_ROCM @@ -589,8 +588,9 @@ __global__ void concat_and_cache_ds_mla_kernel( max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 16)); } - // Compute the scale for the tile - float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN); + // Both SM90 and SM100 readers preserve power-of-two fp32 scales exactly. + float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, 1e-4f); + tile_scale = exp2f(ceilf(log2f(tile_scale))); // The first lane of each half-warp writes the scale to kv_cache if ((lane_idx == 0) || (lane_idx == 16)) { diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu index 05c7bfc6d1..3892790c40 100644 --- a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu +++ b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -35,8 +35,9 @@ * * fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert): * - full key concat (bf16), and cache insert in DeepSeek's 656-byte - * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic - * scales, RoPE bf16), bit-compatible with concat_and_cache_ds_mla_kernel. + * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile + * power-of-two scales, RoPE bf16), bit-compatible with + * concat_and_cache_ds_mla_kernel. * * Both use Programmatic Dependent Launch (PDL) to overlap the tail of the * producing GEMMs on sm_90+, and are structured after @@ -68,7 +69,6 @@ #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" #endif #include -#include #include #ifdef USE_ROCM @@ -111,9 +111,9 @@ constexpr float kFp8Max = 224.0f; #else constexpr float kFp8Max = 448.0f; #endif -// Divisor for fp8_ds_mla per-tile dynamic scales (matches cache_kernels.cu). -// fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128), [512,528) 4 fp32 -// tile scales, [528,656) RoPE 64 bf16. +// Divisor for fp8_ds_mla per-tile power-of-two scales (matches +// cache_kernels.cu). fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128), +// [512,528) 4 fp32 tile scales, [528,656) RoPE 64 bf16. constexpr float kFp8ScaleDivisor = kFp8Max; // Copy 8 source elements (one uint4 of bf16/fp16) to `dst`. FP8=false stores a @@ -325,7 +325,7 @@ __device__ __forceinline__ void writeLatent576( } // Write [kv_c | k_pe] into the fp8_ds_mla 656B entry using one warp: NoPE 512 -// as fp8 in 4 tiles of 128 (per-tile dynamic absmax scale, 4 fp32 scales at +// as fp8 in 4 tiles of 128 (per-tile power-of-two scale, 4 fp32 scales at // [512,528)), RoPE 64 as bf16 at [528,656). Bit-compatible with // concat_and_cache_ds_mla_kernel. template @@ -349,7 +349,8 @@ __device__ __forceinline__ void writeDsMlaCache( for (int offset = 4; offset > 0; offset /= 2) { max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 8)); } - float const tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN); + float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, 1e-4f); + tile_scale = exp2f(ceilf(log2f(tile_scale))); if ((laneId & 7) == 0) { reinterpret_cast(row)[kKvLoraRank / 4 + tile] = tile_scale; } @@ -598,7 +599,7 @@ __global__ void fusedKimiK3MLAKVConcatPackKernel( // ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert // // Cache entry (656 bytes), matching concat_and_cache_ds_mla_kernel: -// [0, 512) NoPE 512 vals as fp8, 4 tiles of 128, each dynamically scaled +// [0, 512) NoPE 512 vals as fp8, 4 tiles of 128, power-of-two scaled // [512, 528) 4 fp32 per-tile scales // [528, 656) RoPE 64 vals as bf16 (unquantized) // The cache slot uses one warp: lane L quantizes NoPE elems [L*16, L*16+16) diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 0d5066d2f1..41f0a1138b 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -960,14 +960,14 @@ def test_concat_and_cache_ds_mla( tile_end = (tile_idx + 1) * 128 tile_data[:] = kv_c_data[tile_start:tile_end] - # tile_scale = tile_data.amax().to(torch.float32) / 448. - # NOTE: Using torch's amax() gives different results, - # so this must be manually computed. + # Using torch's amax() gives different results, so this must be + # manually computed. tile_data_float = tile_data.to(torch.float32) manual_max = abs(tile_data_float[0]) for j in range(1, 128): manual_max = max(manual_max, abs(tile_data_float[j])) - tile_scale = manual_max / 448.0 + raw_scale = torch.clamp(manual_max / 448.0, min=1e-4) + tile_scale = torch.exp2(torch.ceil(torch.log2(raw_scale))) ref_cache_32bit[kv_lora_rank // 4 + tile_idx] = tile_scale @@ -1008,7 +1008,7 @@ def test_concat_and_cache_ds_mla( ref_rope = ref_cache_slice.view(dtype)[kv_lora_rank // 2 + 8 :] torch.testing.assert_close(kv_nope, ref_nope, atol=0.001, rtol=0.1) - torch.testing.assert_close(kv_scales, ref_scales, atol=0.001, rtol=0.1) + torch.testing.assert_close(kv_scales, ref_scales, atol=0, rtol=0) torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1) diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index fb93557a32..bfcd6ae10a 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/tests/kernels/attention/test_flashmla_sparse.py @@ -109,8 +109,9 @@ def test_sparse_flashmla_metadata_smoke(): assert num_splits is None -def test_sparse_flashmla_decode_smoke(): +def test_sparse_flashmla_decode_matches_cache_writer_scales(): import vllm.v1.attention.ops.flashmla as fm + from vllm import _custom_ops as ops ok, reason = fm.is_flashmla_sparse_supported() if not ok: @@ -130,7 +131,7 @@ def test_sparse_flashmla_decode_smoke(): # Metadata q_seq_per_hk = seqlen_q * num_heads_q // num_heads_k # q_heads_per_hk = num_heads_q // num_heads_k - cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device) + cache_seqlens = torch.ones(batch_size, dtype=torch.int32, device=device) tile_md, num_splits = fm.get_mla_metadata( cache_seqlens, q_seq_per_hk, @@ -146,14 +147,33 @@ def test_sparse_flashmla_decode_smoke(): dtype=torch.bfloat16, device=device, ) - k_cache = torch.zeros( - (1, page_block_size, num_heads_k, bytes_per_token), - dtype=torch.uint8, - device=device, + cache_rows = torch.zeros( + (1, page_block_size, bytes_per_token), dtype=torch.uint8, device=device ) - indices = torch.zeros( - (batch_size, seqlen_q, topk), dtype=torch.int32, device=device + # 336 / 448 = 0.75 distinguishes arbitrary FP32 from power-of-two scales. + kv_c = torch.full((1, 512), 336.0, dtype=torch.bfloat16, device=device) + kv_c[:, -128:] = 0 + ops.concat_and_cache_mla( + kv_c, + torch.zeros((1, 64), dtype=torch.bfloat16, device=device), + cache_rows, + torch.zeros(1, dtype=torch.int64, device=device), + "fp8_ds_mla", + torch.ones(1, dtype=torch.float32, device=device), ) + scales = cache_rows[0, 0].view(torch.float32)[128:132] + expected_scales = torch.tensor([1.0, 1.0, 1.0, 2**-13], device=device) + torch.testing.assert_close(scales, expected_scales, rtol=0, atol=0) + cached_nope = cache_rows[0, 0, :512].view(torch.float8_e4m3fn).float() + expected = (cached_nope * scales.repeat_interleave(128)).to(torch.bfloat16) + + k_cache = cache_rows.unsqueeze(2) + indices = torch.full( + (batch_size, seqlen_q, topk), -1, dtype=torch.int32, device=device + ) + indices[..., 0] = 0 + # SM90 sparse decode only supports a fixed top-k width. Padding with -1 + # exercises the same cache data on SM90 and SM100 without topk_length. block_table = torch.zeros((batch_size, 128), dtype=torch.int32, device=device) out, lse = fm.flash_mla_with_kvcache( @@ -170,6 +190,9 @@ def test_sparse_flashmla_decode_smoke(): assert out.shape[0] == batch_size assert out.shape[-1] == head_dim_v assert lse.shape[0] == batch_size + torch.testing.assert_close( + out, expected.view(1, 1, 1, 512).expand_as(out), rtol=0, atol=0 + ) @pytest.mark.parametrize("h_q", [64, 128]) diff --git a/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py b/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py index e1582b7621..af41722605 100644 --- a/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py +++ b/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py @@ -222,7 +222,8 @@ def test_decode_concat_ignores_negative_slots_for_cache(cache_format: str) -> No _assert_cache_matches_reference(mixed_cache, reference_cache, initial_cache) -def test_ds_mla_cache_insert_bit_compatible_with_reference() -> None: +@pytest.mark.parametrize("decode", [False, True]) +def test_ds_mla_cache_insert_bit_compatible_with_reference(decode: bool) -> None: """Fused ds_mla insertion must match the reference bit-for-bit.""" torch.manual_seed(0) num_tokens, num_blocks = 33, 16 @@ -242,5 +243,18 @@ def test_ds_mla_cache_insert_bit_compatible_with_reference() -> None: got = ref.clone() scale = torch.ones(1, device="cuda", dtype=torch.float32) ops.concat_and_cache_mla(kv_c, k_pe.squeeze(1), ref, slots, "fp8_ds_mla", scale) - fused_mla_key_concat_ds_mla_insert(q, k_nope, k_pe, kv_c, got, slots) + if decode: + ql_nope = torch.randn( + num_tokens, NUM_HEADS, KV_LORA_RANK, device="cuda", dtype=dt + ) + fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, q[..., -ROPE_HEAD_DIM:], kv_c, k_pe, got, slots, ds_mla=True + ) + else: + fused_mla_key_concat_ds_mla_insert(q, k_nope, k_pe, kv_c, got, slots) assert torch.equal(ref, got) + rows = got.view(-1, DS_MLA_CACHE_ENTRY)[slots[slots >= 0]] + scales = rows[:, KV_LORA_RANK : KV_LORA_RANK + 16].view(torch.float32) + torch.testing.assert_close( + torch.log2(scales), torch.log2(scales).round(), rtol=0, atol=0 + ) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 7f74cde05d..61f977ddbd 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -507,8 +507,8 @@ def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool): def test_fused_norm_rope_ds_mla(num_tokens: int): """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100). - Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, dynamic float32 - scale) | 4 float32 scales | 64 bf16 (unquantized) RoPE. + Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, power-of-two scales + stored as float32) | 4 float32 scales | 64 bf16 (unquantized) RoPE. """ torch.manual_seed(5) dev = "cuda" @@ -557,7 +557,8 @@ def test_fused_norm_rope_ds_mla(num_tokens: int): kv_ref = rms_norm(kv_c, kvw) # [N, 512] fp32 kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) # [N, 64] tiles = kv_ref.view(num_tokens, 4, 128) - ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1.1754944e-38) + ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1e-4) + ref_scale = torch.exp2(torch.ceil(torch.log2(ref_scale))) ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, KV_LORA).to(FP8) cache = mla_cache[0, :num_tokens] # [N, 656] uint8 @@ -566,7 +567,7 @@ def test_fused_norm_rope_ds_mla(num_tokens: int): rope_off = KV_LORA // 2 + 8 rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + ROPE_DIM] - torch.testing.assert_close(scales, ref_scale, rtol=1e-2, atol=1e-6) + torch.testing.assert_close(scales, ref_scale, rtol=0, atol=0) assert_fp8(nope, ref_nope, "ds_mla NoPE fp8") assert_bf16(rope_vals, kpe_ref, "ds_mla RoPE bf16") # No indexer on this call: top-k buffer must be untouched. diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index fd02a51076..1808d12f64 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -315,7 +315,7 @@ def _fused_norm_rope_kernel( if MLA_CACHE_DS_MLA: # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element # tile of the NoPE is dynamically quantized to fp8 with its own - # float32 scale; the RoPE tail is stored unquantized in bf16. + # power-of-two scale stored as float32; the RoPE tail is bf16. # bytes [0, KV_DIM) : KV_DIM fp8 NoPE values # bytes [KV_DIM, KV_DIM + 16) : MLA_NUM_TILES float32 scales # bytes [KV_DIM + 16, ...) : 2 * KPE_HALF_ROT_DIM bf16 RoPE @@ -327,9 +327,8 @@ def _fused_norm_rope_kernel( ) kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM)) tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True) - # scale = amax / 448 (fp8 e4m3 max), matching the reference - # concat_and_cache_ds_mla kernel; floored to FLT_MIN. - tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1.1754944e-38) + tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1e-4) + tile_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(tile_scale))) kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,)) tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8) tile_off = tl.arange(0, MLA_NUM_TILES) diff --git a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py index 24ae1a1a82..79de3af0bc 100644 --- a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py +++ b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py @@ -76,9 +76,9 @@ def fused_mla_key_concat_ds_mla_insert( """Concat full K (bf16) and insert the latent in the fp8_ds_mla layout. The cache uses DeepSeek's 656-byte block-scaled layout (NoPE in 4 tiles of - 128 with per-tile dynamic fp8 scales, RoPE as bf16) -- self-scaling, so no - scale argument. Returns the bf16 full key; optionally rotates ``q`` and - writes ``kv_cache`` in place. + 128 with per-tile power-of-two scales stored as float32, RoPE as bf16) -- + self-scaling, so no scale argument. Returns the bf16 full key; optionally + rotates ``q`` and writes ``kv_cache`` in place. """ k_pe = k_pe.reshape(k_pe.shape[0], -1) tp, num_heads, qk_nope_head_dim = k_nope.shape