iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Ariel Bereslavsky
2026-09-19 17:15:07 -07:00
committed by GitHub
co-authored by Ariel Bereslavsky Codex Misha Goin Yongye Zhu mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 674b6d95d6
commit a7fda4c88b
8 changed files with 76 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/cache_kernels.cu
+3 -3
View File
@@ -19,7 +19,6 @@
#include <algorithm>
#include <cassert>
#include <cfloat>
#include <cstdlib>
#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)) {
csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu
+10 -9
View File
@@ -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 <cuda_runtime.h>
#include <cfloat>
#include <type_traits>
#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 <typename scalar_t, bool APPLY_ROPE = false>
@@ -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<float*>(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)
tests/kernels/attention/test_cache.py
+5 -5
View File
@@ -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)
tests/kernels/attention/test_flashmla_sparse.py
+31 -8
View File
@@ -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])
tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py
+16 -2
View File
@@ -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
)
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+5 -4
View File
@@ -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.
vllm/models/deepseek_v32/common/kernels.py
+3 -4
View File
@@ -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)
vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py
+3 -3
View File
@@ -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