From 6329d2ffd95e7804846e92fafc7ac9d73cb2a732 Mon Sep 17 00:00:00 2001 From: haosdent Date: Sat, 17 Oct 2026 12:11:13 +0800 Subject: [PATCH] [Perf] Emit the DSA indexer weights in BF16 on SM100 On SM100 the MQA logits kernels take BF16 weights, and mqa_logits_mx_inputs converted the FP32 weights fused_q wrote with a copy kernel on every indexer call. Let fused_q (Triton and CuTeDSL) write the weights in the dtype the logits kernels take, so the conversion is a no-op. The weights are still computed in FP32 and rounded once, so the logits are unchanged. Signed-off-by: haosdent --- .../test_fused_deepseek_v32_norm_rope.py | 15 +++++++++++++-- vllm/models/deepseek_v32/common/kernels.py | 18 +++++++++++------- .../deepseek_v32/nvidia/ops/fused_q_cutedsl.py | 13 ++++++++++--- 3 files changed, 34 insertions(+), 12 deletions(-) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index ec61b94abd5e..457d11a244b2 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -67,6 +67,17 @@ class ModelConfig(NamedTuple): # ── reference helpers ──────────────────────────────────────────────────────── +def assert_index_weights(out: torch.Tensor, ref: torch.Tensor) -> None: + """Scaled indexer weights come out in the MQA logits kernels' dtype: BF16 + on SM100 (one BF16 step of tolerance), FP32 elsewhere.""" + if mqa_logits_use_mx_scales(): + assert out.dtype == torch.bfloat16 + torch.testing.assert_close(out.float(), ref, rtol=8e-3, atol=1e-3) + else: + assert out.dtype == torch.float32 + torch.testing.assert_close(out, ref, rtol=1e-3, atol=1e-3) + + def make_cos_sin(max_pos: int, rot_dim: int, device) -> torch.Tensor: """cos||sin cache: row[pos] = [cos(theta)(rot/2), sin(theta)(rot/2)].""" half = rot_dim // 2 @@ -846,7 +857,7 @@ def test_fused_q(num_tokens: int, index_interleave: bool, cfg: ModelConfig): q_ref, scale_ref = ue8m0_quant(iq_ref) assert_fp8(iq_fp8, q_ref, "indexer-Q fp8") iw_ref = index_w * scale_ref * softmax_scale * head_scale - torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3) + assert_index_weights(iw_out, iw_ref) @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) @@ -989,7 +1000,7 @@ def test_fused_q_bf16_query( iw_ref = ( index_w * scale_ref * (cfg.index_head_dim**-0.5) * (cfg.index_heads**-0.5) ) - torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3) + assert_index_weights(iw_out, iw_ref) def test_fused_q_triton_supports_large_token_count(): diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index f876120a6ef2..68ac098f2ae8 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -958,7 +958,7 @@ def _fused_q_kernel( index_weights *= index_weights_head_scale tl.store( index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx, - index_weights, + index_weights.to(index_weights_out_ptr.dtype.element_ty), ) @@ -980,11 +980,12 @@ def fused_q( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Fuse the MQA-query and indexer-query RoPE/quantization. - Returns ``(index_q_fp8, index_weights_out, mqa_q)``. When ``quantize_mqa`` - is True (FlashInfer sparse, fp8 query) ``mqa_q`` is a single fp8 tensor - packing ``[ql_nope; q_pe]``. When False (FlashMLA sparse, bf16 query) it is - the RoPE'd ``q_pe`` in bf16; the caller pairs it with ``ql_nope`` as the - ``(ql_nope, q_pe)`` tuple the backend expects. + Returns ``(index_q_fp8, index_weights_out, mqa_q)``. ``index_weights_out`` + is in the dtype the MQA logits kernels take: BF16 on SM100, else FP32. When + ``quantize_mqa`` is True (FlashInfer sparse, fp8 query) ``mqa_q`` is a + single fp8 tensor packing ``[ql_nope; q_pe]``. When False (FlashMLA sparse, + bf16 query) it is the RoPE'd ``q_pe`` in bf16; the caller pairs it with + ``ql_nope`` as the ``(ql_nope, q_pe)`` tuple the backend expects. """ assert positions.ndim == 1 assert positions.dtype == torch.int64 @@ -1046,7 +1047,10 @@ def fused_q( mqa_q = q_pe_out index_q_fp8 = torch.empty_like(index_q, dtype=_FP8_DTYPE) - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + index_weights_out = torch.empty_like( + index_weights, + dtype=torch.bfloat16 if mqa_logits_use_mx_scales() else torch.float32, + ) if cutedsl_kernel is not None: cutedsl_kernel( positions, diff --git a/vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py b/vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py index e0c58a7e787d..b3c9d3f8fa30 100644 --- a/vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py +++ b/vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py @@ -79,11 +79,12 @@ def fused_q_cutedsl( if has_indexer: idx_rope_type = _TORCH_TO_CUTE_DTYPE[idx_rope_cache.dtype] idx_weights_type = _TORCH_TO_CUTE_DTYPE[idx_weights.dtype] + idx_weights_out_type = _TORCH_TO_CUTE_DTYPE[idx_weights_out.dtype] else: idx_dim = num_idx_heads = 0 idx_q = idx_rope_cache = idx_q_fp8 = None idx_weights = idx_weights_out = None - idx_rope_type = idx_weights_type = None + idx_rope_type = idx_weights_type = idx_weights_out_type = None rope_type = _TORCH_TO_CUTE_DTYPE[rope_cache.dtype] compiled = FusedQKernel.compile( @@ -95,6 +96,7 @@ def fused_q_cutedsl( num_idx_heads, idx_rope_type, idx_weights_type, + idx_weights_out_type, index_rope_interleave, ) compiled( @@ -470,7 +472,9 @@ def indexer( # scale indexer weights if sublane_id == 0: w = idx_weights[token_id, head_id].to(Float32) - idx_weights_out[token_id, head_id] = w * scale * weight_scale + idx_weights_out[token_id, head_id] = (w * scale * weight_scale).to( + idx_weights_out.element_type + ) @staticmethod @cache @@ -483,6 +487,7 @@ def compile( num_idx_heads: int, idx_rope_type: type[cutlass.Numeric] | None, idx_weights_type: type[cutlass.Numeric] | None, + idx_weights_out_type: type[cutlass.Numeric] | None, index_rope_interleave: bool, ): num_tokens = cute.sym_int() @@ -521,7 +526,9 @@ def compile( Float8E4M3FN, (num_tokens, num_idx_heads, idx_dim), divisibility=16 ) index_weights_out = _make_fake_tensor( - Float32, (num_tokens, num_idx_heads), divisibility=4 + idx_weights_out_type, + (num_tokens, num_idx_heads), + divisibility=4 if idx_weights_out_type is Float32 else 8, ) else: index_q = index_rope_cache = index_q_fp8 = None