[Perf] Emit the DSA indexer weights in BF16 on SM100

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():

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,

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