[ROCm][Perf][MiniMax-M3.1] Tune MXFP8 linear tiles for 16 < M <= 256

tests/kernels/quantization/test_rocm_mxfp8_linear.py

@@ -2,12 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Native CDNA4 MXFP8 linear (``RocmDotScaledMxfp8LinearKernel``).
 
-DeepSeek V4.1 routes every attention linear through this kernel. It used to
-carry two BF16-dequantize workarounds on ROCm (one for all ``.attn.`` linears,
-one for the row-parallel ``wo_b``, whose comment blamed NaNs at small token
-counts). These tests pin the two properties those workarounds stood in for --
-finite output and MXFP8-level accuracy -- across the V4.1 attention shapes at
-TP=4 and the small token counts decode actually runs.
+The per-row-scale tests use synthetic weights with MiniMax-M3.1 and DeepSeek
+V4.1 matrix shapes to check finite output and MXFP8-level accuracy. V4.1's
+checkpoint has repeated 32-row scales and normally dispatches to the separate
+block-scaled GEMM; its synthetic per-row cases do not represent serving.
 
 They also cover the tiny-M fused quant+GEMM fastpath of
 ``rocm_mxfp8_block32_gemm_quant`` (at M <= 8 the activation quant runs inside
@@ -60,6 +58,11 @@
     ("shared_gate_up", 1536, 6144),
     ("shared_down", 6144, 768),
 ]
+# Wider local shapes at TP2 (fused qkv, dense-layer gate_up).
+M31_TP2_SHAPES = [
+    ("tp2_fused_qkv", 4992, 6144),
+    ("tp2_dense_gate_up", 12288, 6144),
+]
 
 
 def _make_layer(n: int, k: int, device: str) -> torch.nn.Module:
@@ -72,23 +75,38 @@ def _make_layer(n: int, k: int, device: str) -> torch.nn.Module:
 
 
 @pytest.mark.parametrize(
-    "n,k,selected,fallback",
+    "n,k,selected,mid_m",
     [
-        (2560, 6144, (16, 16, 1024, 2, 3), (16, 16, 1024, 2, 2)),
-        (1536, 6144, (16, 16, 1024, 2, 3), (16, 16, 1024, 2, 2)),
-        (6144, 768, (16, 16, 256, 2, 3), (16, 32, 256, 4, 3)),
+        (2560, 6144, (16, 16, 1024, 2, 3), (32, 16, 1024, 2, 3)),
+        (1536, 6144, (16, 16, 1024, 2, 3), (16, 16, 1024, 2, 3)),
+        (6144, 768, (16, 16, 256, 2, 3), (32, 16, 256, 2, 3)),
     ],
 )
-def test_rocm_mxfp8_linear_selects_m31_tp4_tiles(n, k, selected, fallback):
-    """The gfx950 small-M tuning is shape-scoped and ends after M=16."""
+def test_rocm_mxfp8_linear_selects_m31_tp4_tiles(n, k, selected, mid_m):
+    """The gfx950 small-M tuning is shape-scoped through M=16; from M=17 the
+    mid-M table takes over."""
     from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
     from vllm.platforms.rocm import on_gfx950
 
     if not on_gfx950():
         pytest.skip("MiniMax-M3.1 small-M tuning requires gfx950")
     for num_tokens in (1, 8, 16):
         assert rocm_native._select_cfg(num_tokens, n, k) == selected
-    assert rocm_native._select_cfg(17, n, k) == fallback
+    assert rocm_native._select_cfg(17, n, k) == mid_m
+
+
+@pytest.mark.parametrize("k", [768, 1280, 1536, 2048, 3072, 4096, 6144])
+@pytest.mark.parametrize("n", [1536, 2560, 3072, 6144, 12288])
+def test_gfx950_mid_m_tiles_divide_k(n, k):
+    """The mid-M table's BLOCK_K always divides K (the K-loop is unmasked)."""
+    from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+        _gfx950_mid_m_cfg,
+    )
+
+    for m in range(17, 257):
+        block_m, block_n, block_k, _, _ = _gfx950_mid_m_cfg(m, n, k)
+        assert k % block_k == 0 and block_k >= 128, (m, n, k, block_k)
+        assert block_m in (16, 32, 64) and block_n in (16, 32, 64)
 
 
 @pytest.mark.parametrize(
@@ -108,9 +126,13 @@ def test_rocm_mxfp8_linear_preserves_non_gfx950_tiles(monkeypatch, n, k, fallbac
 
 
 @pytest.mark.parametrize(
-    "name,n,k", V41_ATTN_SHAPES + M31_TP4_SHAPES, ids=lambda v: str(v)
+    "name,n,k",
+    V41_ATTN_SHAPES + M31_TP4_SHAPES + M31_TP2_SHAPES,
+    ids=lambda v: str(v),
+)
+@pytest.mark.parametrize(
+    "num_tokens", [1, 2, 7, 8, 16, 17, 31, 32, 48, 64, 96, 128, 192, 256, 1024]
 )
-@pytest.mark.parametrize("num_tokens", [1, 2, 7, 8, 16, 31, 32, 128, 1024])
 @torch.inference_mode()
 def test_rocm_mxfp8_linear_matches_dequant_reference(name, n, k, num_tokens):
     """Output is finite and matches an exact MXFP8 GEMM of the same inputs.

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

@@ -266,6 +266,98 @@ def _mxfp8_dot_scaled_gemm(
         _BK256_STAGES = 2
 
 
+# gfx950 tiles for 16 < M <= 256 (decode batches carrying draft tokens, small
+# prefill chunks) with local N > 1280: (BLOCK_M, BLOCK_N, preferred BLOCK_K,
+# num_warps, num_stages) per M bucket and local-N bucket (<= 1536, <= 2560,
+# <= 3072, <= 6144, wider), for K >= 1024. Tuned for MiniMax-M3.1-Flash
+# TP2/TP4 linears on MI355X.
+_GFX950_MID_M_BUCKETS = (32, 48, 64, 96, 128, 192, 256)
+_GFX950_MID_N_BUCKETS = (1536, 2560, 3072, 6144)
+_GFX950_MID_M_TILES = {
+    32: (
+        (16, 16, 1024, 2, 3),
+        (32, 16, 1024, 2, 3),
+        (32, 16, 1024, 2, 3),
+        (32, 32, 512, 2, 3),
+        (32, 16, 512, 2, 2),
+    ),
+    48: (
+        (32, 16, 1024, 2, 3),
+        # M3.1 TP4 qkv (N=2560, K=6144) retains the narrow tile at M=48.
+        (16, 16, 1024, 2, 2),
+        (32, 32, 512, 2, 3),
+        (64, 32, 512, 2, 3),
+        (64, 64, 512, 4, 2),
+    ),
+    64: (
+        (32, 16, 1024, 2, 3),
+        (32, 32, 512, 2, 3),
+        (32, 32, 512, 2, 3),
+        (64, 32, 512, 2, 3),
+        (64, 64, 512, 4, 2),
+    ),
+    96: (
+        (32, 32, 512, 2, 3),
+        (64, 32, 512, 2, 3),
+        (64, 32, 512, 2, 3),
+        (64, 64, 512, 4, 2),
+        (64, 64, 256, 4, 2),
+    ),
+    128: (
+        (32, 32, 512, 2, 3),
+        (64, 32, 512, 2, 3),
+        (64, 32, 512, 2, 3),
+        (64, 64, 512, 4, 2),
+        (64, 64, 256, 4, 2),
+    ),
+    192: (
+        (64, 32, 512, 2, 3),
+        (64, 64, 512, 4, 2),
+        (64, 64, 512, 4, 2),
+        (64, 64, 256, 4, 2),
+        (64, 64, 256, 4, 2),
+    ),
+    256: (
+        (64, 32, 512, 2, 3),
+        (64, 64, 512, 4, 2),
+        (64, 64, 512, 4, 2),
+        (64, 64, 256, 4, 2),
+        (64, 64, 512, 4, 2),
+    ),
+}
+# The same for short K (< 1024, e.g. a shared-expert down projection).
+_GFX950_MID_M_SHORT_K_TILES = {
+    32: (32, 16, 256, 2, 3),
+    48: (64, 32, 256, 4, 3),
+    64: (64, 32, 256, 4, 3),
+    96: (64, 64, 256, 4, 3),
+    128: (64, 64, 256, 4, 3),
+    192: (64, 32, 256, 2, 2),
+    256: (64, 32, 256, 2, 2),
+}
+
+
+def _gfx950_mid_m_cfg(M, N, K):
+    m_bucket = next(b for b in _GFX950_MID_M_BUCKETS if b >= M)
+    if K < 1024:
+        block_m, block_n, block_k, num_warps, num_stages = _GFX950_MID_M_SHORT_K_TILES[
+            m_bucket
+        ]
+    else:
+        n_idx = next(
+            (i for i, b in enumerate(_GFX950_MID_N_BUCKETS) if b >= N),
+            len(_GFX950_MID_N_BUCKETS),
+        )
+        block_m, block_n, block_k, num_warps, num_stages = _GFX950_MID_M_TILES[
+            m_bucket
+        ][n_idx]
+    # The K-loop is unmasked: shrink BLOCK_K until it divides K (the entry
+    # requires K % 128 == 0).
+    while block_k > 128 and K % block_k:
+        block_k //= 2
+    return block_m, block_n, block_k, num_warps, num_stages
+
+
 def _select_cfg(M, N, K):
     """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) — graph-tuned on gfx950.
 
@@ -282,6 +374,14 @@ def _select_cfg(M, N, K):
             return 16, 16, 1024, 2, 3
         if K == 768 and N == 6144:
             return 16, 16, 256, 2, 3
+    # M3.1 TP2 qkv regresses with the M192 tile; use the fallback below.
+    if (
+        _ON_GFX950
+        and 16 < M <= 256
+        and N > 1280
+        and not (128 < M <= 192 and N == 4992 and K == 6144)
+    ):
+        return _gfx950_mid_m_cfg(M, N, K)
     if M <= 64:
         # decode (M in {1,32,64}): tiny-M GEMV is weight-BW + GPU-OCCUPANCY bound. The
         # lever is NARROW BLOCK_N=16 (maximize N-tiles so more CUs stream the weight in