[ROCm][Perf][MiniMax-M3.1] Emit the o_proj MXFP8 input from the NVFP4 sparse decode kernel

tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py

@@ -241,3 +241,104 @@ def test_decode_block_outer_pages_bitwise(num_kv_heads, decode_query_len):
     ]
     assert torch.isfinite(outs[0]).all()
     assert torch.equal(outs[0], outs[1])
+
+
+def _quantized_decode(q, cache, topk, table, seq_lens, qlen, extra_rows=0, **kw):
+    """The kernel's MXFP8 output (values, scales) of the rows plus ``extra_rows``
+    sentinel rows that must stay untouched."""
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        MXFP8_BLOCK_SIZE,
+    )
+
+    k_scale, v_scale = _scales()
+    rows, heads, dim = q.shape
+    out_q = torch.full(
+        (rows + extra_rows, heads * dim), 0x5A, dtype=torch.uint8, device="cuda"
+    ).view(torch.float8_e4m3fn)
+    out_s = torch.full(
+        (rows + extra_rows, heads * dim // MXFP8_BLOCK_SIZE),
+        0xA5,
+        dtype=torch.uint8,
+        device="cuda",
+    )
+    minimax_m3_nvfp4_sparse_decode_fly(
+        q, cache, topk, table, seq_lens, SM_SCALE, None, k_scale=k_scale,
+        v_scale=v_scale, decode_query_len=qlen, out_q=out_q, out_scale=out_s, **kw,
+    )  # fmt: skip
+    return out_q, out_s
+
+
+def _assert_mxfp8_equals_quantized_bf16(q, cache, topk, table, seq_lens, qlen, **kw):
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        mxfp8_e4m3_quantize,
+    )
+
+    rows = q.shape[0]
+    out_bf16 = _decode(q, cache, topk, table, seq_lens, qlen, **kw)
+    assert torch.isfinite(out_bf16).all()
+    ref_q, ref_s = mxfp8_e4m3_quantize(out_bf16.view(rows, -1))
+    got_q, got_s = _quantized_decode(
+        q, cache, topk, table, seq_lens, qlen, extra_rows=3, **kw
+    )
+    assert torch.equal(got_q[:rows].view(torch.uint8), ref_q.view(torch.uint8))
+    assert torch.equal(got_s[:rows], ref_s)
+    # Nothing past the rows is written.
+    assert (got_q[rows:].view(torch.uint8) == 0x5A).all()
+    assert (got_s[rows:] == 0xA5).all()
+
+
+def _assert_multi_request_mxfp8(num_reqs, qlen, ctx, **kw):
+    torch.manual_seed(num_reqs * 7 + qlen)
+    nb = (ctx + BLOCK - 1) // BLOCK
+    # Every request reads the same pages (read-only): a small cache at any context.
+    cache = torch.randint(
+        0, 256, (nb + 3, 2, BLOCK, 72), dtype=torch.uint8, device="cuda"
+    )
+    for side in range(2):
+        sf = cache.view(nb + 3, -1)[:, side * 9216 + 8192 : (side + 1) * 9216]
+        sf.masked_fill_((sf & 0x7F) == 0x7F, 0x38)  # no NaN scales
+    lens = [ctx - 97 * r for r in range(num_reqs)]
+    table = torch.arange(nb, dtype=torch.int32, device="cuda").repeat(num_reqs, 1)
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    topk = _selections(lens, qlen)
+    q = (torch.randn(num_reqs * qlen, GROUP, D, device="cuda") * 0.02).bfloat16()
+    _assert_mxfp8_equals_quantized_bf16(q, cache, topk, table, seq_lens, qlen, **kw)
+
+
+@pytest.mark.parametrize("ctx", [30_000, 150_000, 500_000])
+@pytest.mark.parametrize("qlen", [4, 8])
+@pytest.mark.parametrize("num_reqs", [1, 2, 3, 5, 8, 12, 16, 24])
+def test_decode_mxfp8_output_matches_quantized_bf16(num_reqs, qlen, ctx):
+    """The MXFP8 epilogue equals ``mxfp8_e4m3_quantize`` of the same kernel's
+    bf16 output, bit for bit, at the default split choice for each batch."""
+    _assert_multi_request_mxfp8(num_reqs, qlen, ctx)
+
+
+@pytest.mark.parametrize("splits", [2, 4])
+@pytest.mark.parametrize("num_reqs", [1, 5, 8, 9, 16, 17, 24])
+def test_decode_mxfp8_output_multi_request_xcd_splits(num_reqs, splits):
+    """DSpark 7 rows (8 per request) on the 8-XCD grid, explicit split counts
+    across the request ranges a split heuristic distinguishes."""
+    _assert_multi_request_mxfp8(num_reqs, 8, 150_000, splits=splits, xcd_count=8)
+
+
+@pytest.mark.parametrize("splits", [1, 2, 4, 8])
+@pytest.mark.parametrize("xcd_count", [0, 8])
+@pytest.mark.parametrize("num_kv_heads", [1, 4])
+@pytest.mark.parametrize("qlen", [1, 4, 8])
+def test_decode_mxfp8_output_splits_and_heads(qlen, num_kv_heads, xcd_count, splits):
+    """Every split count (the last workgroup quantizes the merged O), XCD grid
+    and KV-head count; padding requests (seq_len 0) give zero values and scale 0,
+    like the quantized bf16 zeros."""
+    torch.manual_seed(3)
+    lens = [700, 0, 3000, 20000, 0, 129]
+    table, num_blocks = _block_table([max(s, 1) for s in lens])
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    topk = _selections([max(s, 1) for s in lens], qlen, num_kv_heads)
+    cache = _nvfp4_cache(num_blocks, _scales(), num_kv_heads)
+    # Magnitudes over many binades, so block scales cover a wide E8M0 range.
+    mag = torch.pow(2.0, torch.randint(-6, 7, (topk.shape[1], GROUP * num_kv_heads, 1)))
+    q = (torch.randn(topk.shape[1], GROUP * num_kv_heads, D) * mag).cuda().bfloat16()
+    _assert_mxfp8_equals_quantized_bf16(
+        q, cache, topk, table, seq_lens, qlen, splits=splits, xcd_count=xcd_count
+    )

vllm/models/minimax_m3/amd/model.py

@@ -30,6 +30,7 @@
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.config import (
     CacheConfig,
+    CUDAGraphMode,
     VllmConfig,
     get_current_vllm_config,
 )
@@ -67,6 +68,11 @@
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    MXFP8_SCALE_DTYPE,
+    MXFP8_VALUE_DTYPE,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -790,6 +796,8 @@ def __init__(
         # One fused kernel inserts both caches, so NVFP4 K/V needs an NVFP4
         # index cache.
         self.use_nvfp4_kv = get_kv_quant_mode(self.kv_cache_dtype).is_nvfp4
+        # Whether the attention can emit MXFP8 for o_proj; resolved on first use.
+        self._o_proj_mxfp8: bool | None = None
         # Indexer and main attention are separate impls. This indexer is
         # ROCm-only, so the model selects it rather than the neutral
         # MiniMaxM3Indexer, which knows nothing about it and would pick its own.
@@ -1171,8 +1179,8 @@ def forward(
                 qkv, positions, fwd_slot_mapping
             )
             output = torch.empty_like(q)
-            attn_output = self._run_attention(q, index_q, output)
-            output, _ = self.o_proj(attn_output)
+            gfx942_attn = self._run_attention(q, index_q, output)
+            output, _ = self.o_proj(gfx942_attn)
             return output
 
         main_slot_mapping = fwd_slot_mapping[self.layer_name]
@@ -1183,10 +1191,24 @@ def forward(
                 main_slot_mapping,
                 fwd_slot_mapping[self.indexer.index_cache.prefix],
             )
-            output = torch.empty_like(q)
-            attn_output = self._run_attention(
-                q, None if self.skip_index_topk else index_q, output
-            )
+            index_arg = None if self.skip_index_topk else index_q
+            attn_output: torch.Tensor | QuantizedActivation
+            if self._emits_o_proj_mxfp8():
+                # Decode-only: the attention kernel writes MXFP8 for o_proj.
+                out_q = q.new_empty((q.shape[0], self.q_size), dtype=MXFP8_VALUE_DTYPE)
+                out_scale = q.new_empty(
+                    (q.shape[0], self.q_size // MXFP8_BLOCK_SIZE),
+                    dtype=MXFP8_SCALE_DTYPE,
+                )
+                self._run_attention(
+                    q, index_arg, None, out_q=out_q, out_scale=out_scale
+                )
+                attn_output = QuantizedActivation(
+                    out_q, out_scale, q.dtype, out_q.shape, kMxfp8Dynamic
+                )
+            else:
+                output = torch.empty_like(q)
+                attn_output = self._run_attention(q, index_arg, output)
             output, _ = self.o_proj(attn_output)
             return output
         q = qkv.new_empty((num_tokens, self.q_size))
@@ -1388,13 +1410,37 @@ def forward(
         output, _ = self.o_proj(attn_output)
         return output
 
+    def _emits_o_proj_mxfp8(self) -> bool:
+        """Whether this step's attention writes MXFP8 for ``o_proj`` itself:
+        the NVFP4 decode kernel can, ``o_proj`` takes MXFP8 input, and the
+        batch is decode-only (a captured piecewise graph or a mixed batch keeps
+        the bf16 output)."""
+        if self._o_proj_mxfp8 is None:
+            self._o_proj_mxfp8 = hasattr(
+                self.impl, "forward_decode_mxfp8"
+            ) and _can_fuse_input_norm_mxfp8(self.o_proj)
+        if not self._o_proj_mxfp8:
+            return False
+        ctx = get_forward_context()
+        if ctx.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE:
+            return False
+        all_md = ctx.attn_metadata
+        md = all_md.get(self.layer_name) if isinstance(all_md, dict) else None
+        return (
+            isinstance(md, MiniMaxM3SparseMetadata)
+            and md.num_prefills == 0
+            and md.num_decodes > 0
+        )
+
     @eager_break_during_capture
     def _run_attention(
         self,
         query: torch.Tensor,
         index_query: torch.Tensor | None,
-        output: torch.Tensor,
-    ) -> torch.Tensor:
+        output: torch.Tensor | None,
+        out_q: torch.Tensor | None = None,
+        out_scale: torch.Tensor | None = None,
+    ) -> torch.Tensor | None:
         # Single eager break around both: their split-K kernels read per-request
         # metadata and can't be captured into a cudagraph. The indexer writes its
         # top-k into the shared ``topk_indices_buffer``; the attend reads it back.
@@ -1453,6 +1499,13 @@ def _run_attention(
                         self.indexer(index_query)
                 else:
                     self.indexer(index_query)
+        if out_q is not None:
+            assert out_scale is not None
+            self.impl.forward_decode_mxfp8(  # type: ignore[attr-defined]
+                self, query, self.kv_cache, out_q, out_scale
+            )
+            return None
+        assert output is not None
         if self.use_aiter_sparse_pa:
             assert isinstance(self.impl, MiniMaxM3SparseAiterPAImpl)
             return self.impl.forward(

vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/host.py

@@ -30,8 +30,8 @@
 
 
 @functools.cache
-def _get_kernel(splits: int, topk: int, xcds: int):
-    return _kernel.compile_nvfp4_sparse_decode(splits, topk, xcds)
+def _get_kernel(splits: int, topk: int, xcds: int, fp8_out: bool = False):
+    return _kernel.compile_nvfp4_sparse_decode(splits, topk, xcds, fp8_out)
 
 
 def _query_xcd_count(device_index: int) -> int:
@@ -106,17 +106,23 @@ def minimax_m3_nvfp4_sparse_decode_fly(
     block_table: torch.Tensor,  # [reqs, max_blocks] int32
     seq_lens: torch.Tensor,  # [reqs] int32
     sm_scale: float,
-    output: torch.Tensor,  # [rows, H, 128] bf16
+    output: torch.Tensor | None,  # [rows, H, 128] bf16; None with out_q
     k_scale: torch.Tensor | None = None,
     v_scale: torch.Tensor | None = None,
     decode_query_len: int = 1,
     splits: int | None = None,
     xcd_count: int | None = None,
+    out_q: torch.Tensor | None = None,  # [rows, H * 128] float8_e4m3fn
+    out_scale: torch.Tensor | None = None,  # [rows, H * 4] uint8 (E8M0)
 ) -> None:
     """Sparse decode attention over the NVFP4 K/V pages.
 
     Row i is query row i % decode_query_len of request i // decode_query_len
     and attends to the tokens of its ``topk_idx`` blocks up to its position.
+
+    With ``out_q`` and ``out_scale`` the output is written as MXFP8 (E4M3 values
+    and row-major E8M0 scales per 32 dims, equal to ``mxfp8_e4m3_quantize`` of
+    the bf16 output) instead of bf16 into ``output``, which is then unused.
     """
     rows, num_heads, head_dim = q.shape
     if rows == 0:
@@ -125,8 +131,29 @@ def minimax_m3_nvfp4_sparse_decode_fly(
     topk = topk_idx.shape[-1]
     assert rows == seq_lens.shape[0] * decode_query_len
     assert head_dim == _kernel.D and num_heads == num_kv_heads * _kernel.HG
-    assert q.dtype == torch.bfloat16 and output.dtype == torch.bfloat16
-    assert q.is_contiguous() and output.is_contiguous()
+    fp8_out = out_q is not None
+    assert fp8_out == (out_scale is not None) and fp8_out == (output is None)
+    assert q.dtype == torch.bfloat16
+    assert q.is_contiguous()
+    if out_q is not None and out_scale is not None:
+        assert out_q.dtype == torch.float8_e4m3fn and out_scale.dtype == torch.uint8
+        assert out_q.is_contiguous() and out_scale.is_contiguous()
+        assert out_q.shape[0] >= rows and out_q.shape[1] == num_heads * head_dim
+        assert out_scale.shape[0] >= rows
+        assert out_scale.shape[1] == num_heads * head_dim // 32
+        out_ptr, out_q_ptr, out_scale_ptr = (
+            q.data_ptr(),
+            out_q.data_ptr(),
+            out_scale.data_ptr(),
+        )
+    else:
+        assert output is not None
+        assert output.dtype == torch.bfloat16 and output.is_contiguous()
+        out_ptr, out_q_ptr, out_scale_ptr = (
+            output.data_ptr(),
+            q.data_ptr(),
+            q.data_ptr(),
+        )
     assert kv_cache.dtype == torch.uint8 and kv_cache.stride(1) == _kernel.BLK * 72
     assert kv_cache.stride(3) == 1 and kv_cache.stride(0) < 2**31
     assert topk_idx.dtype == torch.int32 and topk_idx.shape[0] == num_kv_heads
@@ -156,13 +183,15 @@ def minimax_m3_nvfp4_sparse_decode_fly(
     vs = unit if v_scale is None else v_scale
     c = float(np.float32(sm_scale) * np.float32(LOG2E))
     run_compiled(
-        _get_kernel(splits, topk, xcds),
+        _get_kernel(splits, topk, xcds, fp8_out),
         q.data_ptr(),
         kv_cache.data_ptr(),
         topk_idx.data_ptr(),
         block_table.data_ptr(),
         seq_lens.data_ptr(),
-        output.data_ptr(),
+        out_ptr,
+        out_q_ptr,
+        out_scale_ptr,
         ks.data_ptr(),
         vs.data_ptr(),
         nvfp4_e4m3_lut(q.device).data_ptr(),