[ROCm][Perf][DSv4.1] TP1 MXFP8 GEMM tiles and a pre-quantized wo_b input for DEP8 on MI355X

.buildkite/test_areas/kernels.yaml

@@ -523,6 +523,7 @@ steps:
       - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
       - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
       - pytest -v -s tests/kernels/attention/test_rocm_triton_attn_dsv4.py::test_dsv41_adaptive_mla_swa_metadata_graph_replay
+      - pytest -v -s tests/kernels/attention/test_rocm_triton_attn_dsv4.py -k "wo_a_bmm or wo_b_takes_wo_a_mxfp8_gate"
 
 - label: ":nvidia: (L4) Quantization Kernels"
   key: kernels-quantization-test

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

@@ -1593,24 +1593,28 @@ def test_sparse_attn_decode_mxfp8_output() -> None:
 
 @requires_gfx950
 @pytest.mark.parametrize(
-    "num_tokens, n_groups",
+    "num_tokens, n_groups, o_lora_rank",
     # One case per tile tier of _mxfp8_wo_a_bmm_config, with partial M tiles.
     [
-        (1, 4),
-        (20, 4),
-        (48, 4),
-        (77, 4),
-        (130, 4),
-        (300, 4),
-        (700, 4),
-        (1000, 4),
-        (1100, 8),
+        (1, 4, 128),
+        (20, 4, 128),
+        (48, 4, 128),
+        (77, 4, 128),
+        (130, 4, 128),
+        (300, 4, 128),
+        (700, 4, 128),
+        (1000, 4, 128),
+        (300, 8, 128),
+        (700, 8, 128),
+        (1100, 8, 128),
+        (1600, 8, 128),
+        (1600, 8, 256),
     ],
 )
 @pytest.mark.parametrize("block_scales", [False, True])
 @torch.inference_mode()
 def test_rocm_mxfp8_wo_a_bmm(
-    num_tokens: int, n_groups: int, block_scales: bool
+    num_tokens: int, n_groups: int, o_lora_rank: int, block_scales: bool
 ) -> None:
     from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
         _as_block32_scale,
@@ -1622,7 +1626,7 @@ def test_rocm_mxfp8_wo_a_bmm(
 
     device = torch.device("cuda")
     torch.manual_seed(3)
-    o_lora_rank, group_dim = 128, 2048
+    group_dim = 2048
     a, a_scale = _mxfp8_e4m3_quantize_torch(
         torch.randn(num_tokens, n_groups * group_dim, device=device)
     )
@@ -1648,6 +1652,196 @@ def test_rocm_mxfp8_wo_a_bmm(
     torch.testing.assert_close(out.float(), expected.flatten(1), atol=5e-2, rtol=1e-2)
 
 
+@requires_gfx950
+@pytest.mark.parametrize(
+    "num_tokens, n_groups, o_lora_rank",
+    # A tile narrower than an MX block (separate quant), then each of the nine
+    # tiles with a fused epilogue: the work tiers at G = 4 and 8, the G = 8
+    # tiers past 256 tokens and the 256-wide G = 8 tile.
+    [
+        (1, 4, 128),
+        (77, 4, 128),
+        (600, 4, 128),
+        (900, 4, 128),
+        (24, 8, 128),
+        (100, 8, 128),
+        (224, 8, 128),
+        (300, 8, 128),
+        (800, 8, 128),
+        (1600, 8, 256),
+    ],
+)
+@pytest.mark.parametrize("block_scales", [False, True])
+@torch.inference_mode()
+def test_rocm_mxfp8_wo_a_bmm_mxfp8_matches_quantized_bmm(
+    num_tokens: int, n_groups: int, o_lora_rank: int, block_scales: bool
+) -> None:
+    """wo_b's MXFP8 input from the wo_a epilogue is bit-identical to quantizing
+    the bf16 wo_a output, so the fused o_proj path does not change results."""
+    from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+        _as_block32_scale,
+    )
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        _mxfp8_e4m3_quantize_torch,
+        mxfp8_e4m3_quantize,
+    )
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        rocm_mxfp8_wo_a_bmm,
+        rocm_mxfp8_wo_a_bmm_mxfp8,
+    )
+
+    device = torch.device("cuda")
+    torch.manual_seed(5)
+    group_dim = 2048
+    a, a_scale = _mxfp8_e4m3_quantize_torch(
+        torch.randn(num_tokens, n_groups * group_dim, device=device)
+    )
+    w, w_scale = _mxfp8_e4m3_quantize_torch(
+        torch.randn(n_groups * o_lora_rank, group_dim, device=device)
+    )
+    if block_scales:
+        w_scale = _as_block32_scale(w_scale[::32].repeat_interleave(32, dim=0))
+        assert w_scale is not None
+    wo_a = SimpleNamespace(weight=w, weight_scale=w_scale)
+
+    q, q_scale = rocm_mxfp8_wo_a_bmm_mxfp8(a, a_scale, wo_a, n_groups, o_lora_rank)
+
+    ref_q, ref_scale = mxfp8_e4m3_quantize(
+        rocm_mxfp8_wo_a_bmm(a, a_scale, wo_a, n_groups, o_lora_rank)
+    )
+    assert torch.equal(q.view(torch.uint8), ref_q.view(torch.uint8))
+    assert torch.equal(q_scale, ref_scale)
+
+
+@requires_gfx950
+@pytest.mark.parametrize(
+    "num_tokens, n_groups",
+    # On both sides of the narrowest wo_a tile that quantizes in its epilogue
+    # (T * G > 128), and at the per-rank decode sizes of DP8 at conc 256.
+    [(1, 8), (8, 8), (16, 8), (17, 8), (224, 8), (300, 8), (8, 4), (32, 4), (33, 4)],
+)
+@torch.inference_mode()
+def test_rocm_mxfp8_wo_a_bmm_o_proj_matches_wo_b_linear(
+    num_tokens: int, n_groups: int, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """The o_proj op hands wo_b its MXFP8 input from wo_a's epilogue only where
+    the wo_a tile can quantize; below that wo_b's own quantizing GEMM runs, so a
+    tiny batch keeps the block32 GEMM's fused quant. Either way the output is
+    bit-identical to wo_b's linear on the bf16 wo_a output."""
+    import vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm as block32
+    from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+        _as_block32_scale,
+    )
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        _mxfp8_e4m3_quantize_torch,
+    )
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        rocm_mxfp8_o_proj,
+        rocm_mxfp8_wo_a_bmm,
+    )
+
+    device = torch.device("cuda")
+    torch.manual_seed(7)
+    group_dim, o_lora_rank, n = 2048, 128, 512
+
+    def block_scaled(rows: int, cols: int) -> SimpleNamespace:
+        w, w_scale = _mxfp8_e4m3_quantize_torch(torch.randn(rows, cols, device=device))
+        w_scale = _as_block32_scale(w_scale[::32].repeat_interleave(32, dim=0))
+        assert w_scale is not None
+        return SimpleNamespace(weight=w, weight_scale=w_scale)
+
+    a, a_scale = _mxfp8_e4m3_quantize_torch(
+        torch.randn(num_tokens, n_groups * group_dim, device=device)
+    )
+    wo_a = block_scaled(n_groups * o_lora_rank, group_dim)
+    wo_b = block_scaled(n, n_groups * o_lora_rank)
+    expected = block32.rocm_mxfp8_block32_gemm_quant(
+        rocm_mxfp8_wo_a_bmm(a, a_scale, wo_a, n_groups, o_lora_rank),
+        wo_b.weight,
+        wo_b.weight_scale,
+        torch.bfloat16,
+    )
+
+    calls: list[str] = []
+
+    def counted(name: str):
+        fn = getattr(block32, name)
+
+        def wrapper(*args):
+            calls.append(name)
+            return fn(*args)
+
+        return wrapper
+
+    for name in ("rocm_mxfp8_block32_gemm", "rocm_mxfp8_block32_gemm_quant"):
+        monkeypatch.setattr(block32, name, counted(name))
+    out = rocm_mxfp8_o_proj(
+        a, a_scale, wo_a, n_groups, o_lora_rank, wo_b, torch.bfloat16
+    )
+
+    fused = num_tokens * n_groups > 128
+    assert calls == [
+        "rocm_mxfp8_block32_gemm" if fused else "rocm_mxfp8_block32_gemm_quant"
+    ]
+    assert torch.equal(out, expected)
+
+
+@pytest.mark.parametrize(
+    "case, expected",
+    [
+        ("block32_gemm", True),
+        ("reduce_at_tp1", True),
+        ("other_mxfp8_kernel", False),
+        ("per_row_scale", False),
+        ("bf16_weight", False),
+        ("bias", False),
+        ("tp_reduce", False),
+        ("bpreshuffled_scale", False),
+    ],
+)
+def test_wo_b_takes_wo_a_mxfp8_gate(case: str, expected: bool) -> None:
+    """wo_a hands wo_b its MXFP8 output only when wo_b's linear is the block32
+    MXFP8 GEMM on a 32x32 block-scaled e4m3 weight, without a bias or a
+    tensor-parallel reduce. Another kernel that keeps the same weight dtypes
+    and shapes does not get the fused input."""
+    from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+        RocmDotScaledMxfp8LinearKernel,
+    )
+    from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention
+
+    class OtherMxfp8Kernel(RocmDotScaledMxfp8LinearKernel):
+        pass
+
+    n, k = 64, 128
+    wo_b = SimpleNamespace(
+        quant_method=SimpleNamespace(
+            kernel=object.__new__(RocmDotScaledMxfp8LinearKernel)
+        ),
+        weight=torch.zeros(n, k, dtype=torch.float8_e4m3fn),
+        weight_scale=torch.zeros(n // 32, k // 32, dtype=torch.uint8),
+        bias=None,
+        reduce_results=False,
+        tp_size=1,
+    )
+    attn = SimpleNamespace(_wo_b_scale=None, wo_b=wo_b)
+    if case == "reduce_at_tp1":
+        wo_b.reduce_results = True
+    elif case == "other_mxfp8_kernel":
+        wo_b.quant_method.kernel = object.__new__(OtherMxfp8Kernel)
+    elif case == "per_row_scale":
+        wo_b.weight_scale = torch.zeros(n, k // 32, dtype=torch.uint8)
+    elif case == "bf16_weight":
+        wo_b.weight = torch.zeros(n, k, dtype=torch.bfloat16)
+    elif case == "bias":
+        wo_b.bias = torch.zeros(n, dtype=torch.bfloat16)
+    elif case == "tp_reduce":
+        wo_b.reduce_results, wo_b.tp_size = True, 2
+    elif case == "bpreshuffled_scale":
+        attn._wo_b_scale = torch.ones(1)
+
+    assert DeepseekV41ROCMAiterMLAAttention._wo_b_takes_wo_a_mxfp8(attn) == expected
+
+
 @requires_gfx950
 @pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65])
 @torch.inference_mode()

tests/kernels/quantization/test_rocm_mxfp8_linear.py

@@ -280,17 +280,22 @@ def test_rocm_mxfp8_linear_handles_extreme_magnitudes():
     assert (out[0] == 0).all(), "an all-zero activation row must give a zero row"
 
 
-# (name, N, K) of the DeepSeek-V4.1-Flash MXFP8 linears at TP=2 and TP=4. The
+# (name, N, K) of the DeepSeek-V4.1-Flash MXFP8 linears at TP=1, 2 and 4. The
 # checkpoint stores their scales in 32x32 blocks; shared_down at TP=4 has
 # K % 128 != 0, which only the block-scaled GEMM runs natively.
 V41_BLOCK32_SHAPES = [
     ("fused_wqa_wkv", 1280 + 512, 5120),
+    ("wq_b_tp1", 64 * 512, 1280),
+    ("wo_b_tp1", 5120, 8192),
+    ("shared_gate_up_tp1", 2 * 2304, 5120),
+    ("shared_down_tp1", 5120, 2304),
     ("wq_b_tp2", 64 * 512 // 2, 1280),
     ("indexer_wq_b", 32 * 128, 1280),
     ("wo_b_tp2", 5120, 8192 // 2),
     ("shared_gate_up_tp2", 2 * 2304 // 2, 5120),
     ("shared_down_tp2", 5120, 2304 // 2),
     ("shared_down_tp4", 5120, 2304 // 4),
+    ("dspark_main_proj", 5120, 3 * 5120),
 ]
 
 
@@ -311,7 +316,7 @@ def _make_block32_layer(n: int, k: int, device: str) -> torch.nn.Module:
 
 
 @pytest.mark.parametrize("name,n,k", V41_BLOCK32_SHAPES, ids=lambda v: str(v))
-@pytest.mark.parametrize("num_tokens", [1, 7, 32, 33, 192, 1024])
+@pytest.mark.parametrize("num_tokens", [1, 7, 32, 33, 96, 192, 300, 1024, 2100])
 @torch.inference_mode()
 def test_rocm_mxfp8_linear_block32_matches_dequant_reference(name, n, k, num_tokens):
     """32x32 block scales stay compact and the block-scaled GEMM is exact.

vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

@@ -456,7 +456,7 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf
 
 
 # (N, K) -> [(max M, config), ...] for the DeepSeek-V4.1-Flash MXFP8
-# linears at TP2/TP4, tuned on MI355X under HIP graphs. A tier serves
+# linears at TP1/TP2/TP4, tuned on MI355X under HIP graphs. A tier serves
 # every M up to its bound; larger M reuse the last tier.
 _TUNED: dict[tuple[int, int], list[tuple[int, _Config]]] = {
     # shared expert gate_up, TP4
@@ -550,6 +550,20 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf
         (4096, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
         (8192, _tiled(128, 128, 128, 1, True, 4, 2, 16, 0)),
     ],
+    # shared expert gate_up, TP1 (DP attention), tuned at the CUDA graph sizes 48-4096
+    (4608, 5120): [
+        (8, _packed(8, 16, 256, 2, 2, 3, 16, 0)),
+        (32, _packed(32, 32, 256, 1, 2, 3, 16, 0)),
+        (64, _tiled(64, 64, 512, 1, False, 4, 2, 32, 0, 2)),
+        (96, _tiled(32, 64, 512, 1, True, 8, 2, 16, 2)),
+        (128, _tiled(64, 64, 256, 1, False, 4, 2, 32, 2, 2)),
+        (192, _tiled(64, 64, 512, 1, True, 4, 2, 32, 2)),
+        (448, _tiled(64, 64, 256, 1, True, 4, 2, 32, 2)),
+        (896, _tiled(128, 128, 256, 1, True, 4, 2, 16, 0)),
+        (1792, _tiled(128, 128, 128, 1, True, 4, 2, 32, 2)),
+        (2560, _tiled(128, 128, 256, 1, True, 8, 2, 32, 0)),
+        (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 2)),
+    ],
     # shared expert down, TP4
     (5120, 576): [
         (8, _packed(8, 32, 256, 2, 2, 2, 16, 1)),
@@ -607,6 +621,20 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf
         (4096, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
         (8192, _tiled(128, 128, 128, 1, True, 2, 2, 32, 1)),
     ],
+    # shared expert down, TP1 (DP attention), tuned at the CUDA graph sizes 48-4096
+    (5120, 2304): [
+        (8, _packed(8, 16, 256, 2, 2, 3, 16, 0)),
+        (32, _packed(32, 32, 256, 1, 2, 3, 16, 0)),
+        (64, _tiled(64, 64, 256, 1, True, 8, 3, 16, 0, 4)),
+        (96, _tiled(32, 64, 512, 1, True, 8, 3, 16, 2)),
+        (192, _tiled(64, 64, 512, 1, True, 4, 2, 32, 2)),
+        (288, _tiled(64, 64, 256, 1, True, 4, 2, 32, 2)),
+        (384, _tiled(128, 64, 256, 1, True, 8, 2, 32, 0)),
+        (768, _tiled(128, 128, 256, 1, True, 8, 2, 32, 2)),
+        (2048, _tiled(128, 128, 128, 1, True, 4, 2, 32, 2)),
+        (3072, _tiled(128, 256, 128, 1, True, 8, 3, 32, 2)),
+        (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
+    ],
     # wo_b, TP2
     (5120, 4096): [
         (4, _packed(4, 32, 512, 4, 2, 2, 16, 0)),
@@ -629,6 +657,34 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf
         (4096, _tiled(128, 128, 128, 1, True, 8, 2, 16, 1)),
         (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
     ],
+    # wo_b, TP1 (DP attention), tuned at the CUDA graph sizes 48-4096
+    (5120, 8192): [
+        (8, _packed(8, 16, 256, 2, 2, 3, 16, 0)),
+        (32, _packed(32, 32, 256, 1, 2, 3, 16, 0)),
+        (48, _tiled(64, 64, 256, 1, True, 4, 2, 16, 2, 4)),
+        (64, _tiled(64, 64, 512, 1, True, 4, 2, 32, 2, 2)),
+        (96, _tiled(32, 64, 512, 1, True, 8, 2, 16, 0)),
+        (128, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0, 2)),
+        (192, _tiled(64, 64, 512, 1, True, 4, 2, 16, 0)),
+        (384, _tiled(128, 128, 256, 2, True, 4, 2, 16, 0)),
+        (768, _tiled(128, 128, 256, 1, True, 4, 2, 16, 0)),
+        (1024, _tiled(128, 128, 256, 2, True, 4, 2, 16, 0)),
+        (1536, _tiled(128, 128, 256, 1, True, 8, 2, 32, 0)),
+        (2048, _tiled(128, 128, 256, 1, True, 8, 2, 32, 2)),
+        (3072, _tiled(128, 128, 128, 1, True, 4, 2, 16, 0)),
+        (8192, _tiled(128, 128, 256, 1, True, 8, 2, 32, 2)),
+    ],
+    # DSpark main_proj (15360 -> 5120, replicated): 7 rows a request per step
+    (5120, 15360): [
+        (8, _packed(8, 16, 256, 2, 2, 3, 16, 0)),
+        (32, _packed(32, 32, 256, 1, 2, 3, 16, 0)),
+        (56, _tiled(64, 64, 512, 1, True, 4, 2, 16, 2, 2)),
+        (126, _tiled(128, 128, 256, 1, True, 4, 2, 16, 2, 4)),
+        (133, _tiled(64, 64, 256, 1, True, 4, 2, 16, 2, 2)),
+        (168, _tiled(64, 64, 512, 1, True, 4, 2, 16, 2)),
+        (224, _tiled(128, 128, 256, 1, True, 4, 2, 16, 0, 2)),
+        (8192, _tiled(128, 128, 256, 1, True, 4, 2, 16, 2, 2)),
+    ],
     # wq_b, TP4
     (8192, 1280): [
         (1, _packed(8, 32, 256, 2, 4, 3, 16, 1)),
@@ -663,6 +719,20 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf
         (2048, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
         (8192, _tiled(128, 256, 128, 1, True, 2, 2, 32, 0)),
     ],
+    # wq_b, TP1 (DP attention), tuned at the CUDA graph sizes 48-4096
+    (32768, 1280): [
+        (8, _packed(8, 16, 256, 2, 2, 3, 16, 0)),
+        (32, _packed(32, 32, 256, 1, 2, 3, 16, 0)),
+        (56, _tiled(64, 64, 256, 1, True, 4, 2, 32, 2)),
+        (64, _tiled(64, 64, 256, 1, True, 4, 2, 16, 2)),
+        (128, _tiled(128, 128, 256, 1, False, 8, 2, 32, 0)),
+        (160, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)),
+        (224, _tiled(128, 128, 128, 1, True, 4, 2, 32, 2)),
+        (288, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
+        (448, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)),
+        (1280, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
+        (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 2)),
+    ],
 }