[GLM5] Support FlashMLA FP8 KV cache (Hopper & Blackwell)

tests/kernels/test_fused_deepseek_v32_norm_rope.py

@@ -292,6 +292,76 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
     assert (topk == 7).all(), "topk buffer should be untouched on shared layer"
 
 
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
+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.
+    """
+    torch.manual_seed(5)
+    dev = "cuda"
+    max_pos = 8192
+    pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
+
+    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
+    mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+
+    bs = max_pos
+    mla_cache = torch.zeros(1, bs, 656, device=dev, dtype=torch.uint8)
+    slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
+    topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
+
+    q_out = K.fused_norm_rope(
+        pos,
+        q_c,
+        qw,
+        EPS,
+        kv_c,
+        kvw,
+        EPS,
+        k_pe,
+        mla_cos_sin,
+        None,
+        None,
+        None,
+        EPS,
+        None,
+        topk,
+        slot_mapping=slot,
+        indexer_k_cache=None,
+        mla_kv_cache=mla_cache,
+        mla_kv_cache_dtype="fp8_ds_mla",
+        mla_k_scale=None,
+        has_indexer=False,
+        index_rope_interleave=False,
+    )
+
+    assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (ds_mla)")
+
+    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_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, KV_LORA).to(FP8)
+
+    cache = mla_cache[0, :num_tokens]  # [N, 656] uint8
+    nope = cache[:, :KV_LORA].view(FP8)
+    scales = cache.view(torch.float32)[:, KV_LORA // 4 : KV_LORA // 4 + 4]
+    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)
+    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.
+    assert (topk == 7).all(), "topk buffer should be untouched (no indexer)"
+
+
 # ── fused_q ──────────────────────────────────────────────────────────────────
 
 
@@ -400,6 +470,75 @@ def test_fused_q_no_indexer(num_tokens: int):
     assert_fp8(mqa[:, :, KV_LORA:], (qpe_ref / s).to(FP8), "mqa q_pe")
 
 
+@pytest.mark.parametrize("num_tokens", [1, 17, 512])
+@pytest.mark.parametrize("has_indexer", [True, False])
+def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
+    """bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is
+    produced (bf16, unquantized); ql_nope is consumed directly by the caller."""
+    torch.manual_seed(6)
+    dev = "cuda"
+    max_pos = 8192
+    pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
+
+    q_pe = torch.randn(
+        num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
+    )
+    ql_nope = torch.randn(
+        num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
+    )
+    q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
+    q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+
+    index_q = index_w = idx_cos_sin = None
+    if has_indexer:
+        index_q = torch.randn(
+            num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
+        )
+        index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
+        idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+
+    iq_fp8, iw_out, q_pe_out = K.fused_q(
+        pos,
+        q_pe,
+        q_cos_sin,
+        index_q,
+        idx_cos_sin,
+        ql_nope,
+        q_scale,
+        index_w,
+        INDEX_HEAD_DIM**-0.5,
+        INDEX_HEADS**-0.5,
+        has_indexer=has_indexer,
+        index_rope_interleave=False,
+        quantize_mqa=False,
+    )
+
+    # MQA query: only the RoPE'd q_pe, bf16, unquantized.
+    assert q_pe_out.dtype == torch.bfloat16
+    assert q_pe_out.shape == (num_tokens, NUM_HEADS, ROPE_DIM)
+    qpe_ref = rope(
+        q_pe.float(),
+        pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
+        q_cos_sin,
+        interleave=True,
+    )
+    assert_bf16(q_pe_out, qpe_ref, "bf16 q_pe RoPE")
+
+    # Indexer-Q is unchanged on this path (still UE8M0 fp8 + folded weights).
+    if has_indexer:
+        assert index_q is not None
+        iq_ref = rope(
+            index_q.float(),
+            pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
+            idx_cos_sin,
+            interleave=False,
+        )
+        q_ref, scale_ref = ue8m0_quant(iq_ref)
+        assert_fp8(iq_fp8, q_ref, "indexer-Q fp8 (bf16-query path)")
+        iw_ref = index_w * scale_ref * (INDEX_HEAD_DIM**-0.5) * (INDEX_HEADS**-0.5)
+        torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
+
+
 # ── fused_eh_norm (MTP) ──────────────────────────────────────────────────────
 
 

vllm/models/deepseek_v32/nvidia/attention.py

@@ -269,20 +269,30 @@ def __init__(
         # Runtime toggle for index_share_for_mtp_iteration: MTP draft step 0
         # computes the top-k, steps 1+ set this True to reuse it.
         self.skip_topk = False
-        # Single fused fp8 path: Triton fused norm/rope/cache + fused-q write a
-        # single fp8 MQA query and the contiguous [kv_c; k_pe] MLA cache layout.
-        # This requires an fp8 KV cache and a sparse MLA backend that accepts a
-        # quantized query (FlashInfer sparse on SM100).
-        assert (
-            is_quantized_kv_cache(self.kv_cache_dtype)
-            and self.impl.supports_quant_query_input
-        ), (
-            "deepseek_v32 (nvidia) requires an fp8 KV cache served by the "
-            "FlashInfer sparse MLA backend (which accepts a quantized query). "
-            "Launch with --kv-cache-dtype fp8."
+        # Fused fp8 paths: Triton fused norm/rope/cache + fused-q. Two layouts,
+        # picked by the sparse MLA backend's query support:
+        #   * supports_quant_query_input (FlashInfer sparse, SM100): per-tensor
+        #     fp8 cache + a single packed fp8 MQA query.
+        #   * not supported (FlashMLA sparse, SM90/SM100): fp8_ds_mla cache
+        #     (per-128 block-scaled fp8 NoPE + unquantized bf16 RoPE) + a bf16
+        #     (ql_nope, q_pe) query tuple. FA3 cannot mix a bf16 query with an
+        #     fp8 KV cache, so FlashMLA (which dequantizes internally) is used.
+        #     FlashMLA sparse runs on both Hopper and Blackwell, so this is the
+        #     only DSA path on SM90 and an opt-in alternative on SM100.
+        assert is_quantized_kv_cache(self.kv_cache_dtype), (
+            "deepseek_v32 (nvidia) requires an fp8 KV cache served by a sparse "
+            "MLA backend. Launch with --kv-cache-dtype fp8 (FlashInfer sparse) "
+            "or --kv-cache-dtype fp8_ds_mla (FlashMLA sparse)."
         )
+        self._fp8_query = self.impl.supports_quant_query_input
+        if not self._fp8_query:
+            assert self.kv_cache_dtype == "fp8_ds_mla", (
+                "deepseek_v32 (nvidia) on a bf16-query sparse MLA backend "
+                "(FlashMLA sparse) requires the fp8_ds_mla KV cache layout. "
+                "Launch with --kv-cache-dtype fp8_ds_mla."
+            )
         # The paged KV cache is stored as uint8 and viewed as fp8 for the decode
-        # (per-tensor fp8; never the fp8_ds_mla layout on this path).
+        # (per-tensor fp8). The fp8_ds_mla layout is consumed as raw bytes.
         self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla"
         # GLM-5.2 uses interleaved indexer RoPE; DeepSeek-V3.2 uses NeoX.
         self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False)
@@ -465,6 +475,7 @@ def _fused_attention(
             indexer_n_head_scale,
             has_indexer=has_indexer,
             index_rope_interleave=self._index_rope_interleave,
+            quantize_mqa=self._fp8_query,
         )
 
         if attn_metadata is None:
@@ -496,8 +507,17 @@ def _fused_attention(
         kv_cache = self.kv_cache
         if self._fp8_kv_needs_view:
             kv_cache = kv_cache.view(torch.float8_e4m3fn)
+        if self._fp8_query:
+            # FlashInfer sparse: single packed fp8 query.
+            mqa_q_arg: torch.Tensor | tuple[torch.Tensor, torch.Tensor] = mqa_q[
+                :num_actual
+            ]
+        else:
+            # FlashMLA sparse: bf16 (ql_nope, q_pe) tuple. mqa_q is the RoPE'd
+            # q_pe; ql_nope is consumed directly.
+            mqa_q_arg = (ql_nope[:num_actual], mqa_q[:num_actual])
         attn_out, _ = self.impl.forward_mqa(  # type: ignore[attr-defined]
-            mqa_q[:num_actual], kv_cache, attn_metadata, self
+            mqa_q_arg, kv_cache, attn_metadata, self
         )
         x = attn_out.view(
             num_actual, self.num_local_heads, self.kv_lora_rank