[ROCm][Perf][MiniMax-M3.1] Pick the NVFP4 sparse decode splits per XCD, not per CU

.buildkite/test_areas/kernels.yaml

@@ -698,9 +698,11 @@ steps:
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py
       - tests/kernels/attention/test_minimax_m3_nvfp4_index.py
+      - tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py
       commands:
       - pytest -v -s tests/kernels/attention/test_minimax_m3.py
           tests/kernels/attention/test_minimax_m3_nvfp4_index.py
+          tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py
 
 - label: ":nvidia: (H100) MiniMax-M3 Sparse Prefill Kernels"
   key: kernels-minimax-m3-sparse-prefill-h100

tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py

@@ -241,3 +241,91 @@ 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 test_xcd_splits_table():
+    """Splits of the XCD-aware grid: 4 up to 8 requests (one per XCD) and for
+    17-24, 2 for 9-16, the cost model beyond and off the tuned shape."""
+    from vllm.models.minimax_m3.amd.ops.nvfp4_sparse_decode import host
+
+    got = [host._xcd_splits(n, 8, 16, 8, 1) for n in (1, 5, 8, 9, 16, 17, 24, 25)]
+    assert got == [4, 4, 4, 2, 2, 4, 4, None]
+    assert host._xcd_splits(5, 8, 16, 4, 1) is None  # DSpark 3: not tuned
+    assert host._xcd_splits(5, 0, 16, 8, 1) is None  # no XCD-aware grid
+    assert host._xcd_splits(5, 8, 32, 8, 1) is None  # other top-k width
+    assert host._xcd_splits(5, 8, 16, 8, 2) is None  # TP2: over-splits
+    assert host._xcd_splits(5, 8, 16, 8, 4) is None  # TP1: over-splits
+
+
+@pytest.mark.parametrize("num_reqs", [1, 3, 5, 8, 9, 12, 16, 17, 24, 40])
+@pytest.mark.parametrize("ctx", [30000, 500000])
+def test_default_splits_match_unsplit(num_reqs, ctx):
+    """The default split count agrees with the unsplit kernel (the summation
+    order is the only difference) and repeats bitwise, at DSpark 7 batch sizes."""
+    torch.manual_seed(num_reqs)
+    lens = torch.randint(ctx * 6 // 10, ctx * 14 // 10 + 1, (num_reqs,)).tolist()
+    table, num_blocks = _block_table(lens)
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    topk = _selections(lens, 8)
+    cache = torch.randint(
+        0, 256, (num_blocks, 2, BLOCK, 72), dtype=torch.uint8, device="cuda"
+    )
+    for side in range(2):
+        sf = cache.view(num_blocks, -1)[:, side * 9216 + 8192 : (side + 1) * 9216]
+        sf.masked_fill_((sf & 0x7F) == 0x7F, 0x38)  # no NaN scales
+    q = (torch.randn(num_reqs * 8, GROUP, D, device="cuda") * 0.05).bfloat16()
+    out = _decode(q, cache, topk, table, seq_lens, 8)
+    unsplit = _decode(q, cache, topk, table, seq_lens, 8, splits=1)
+    assert torch.isfinite(out).all()
+    assert _rel(out, unsplit) < 3e-3
+    assert torch.equal(out, _decode(q, cache, topk, table, seq_lens, 8))
+
+
+@pytest.mark.parametrize("num_kv_heads", [1, 2, 4])
+@pytest.mark.parametrize("num_reqs", [1, 5, 8, 9, 16, 17, 24, 40])
+def test_default_splits_are_chosen_per_xcd(monkeypatch, num_kv_heads, num_reqs):
+    """The default split count is the XCD table at one KV head per rank and the
+    cost model otherwise, counted over every request including padding
+    (seq_len 0), and agrees with the unsplit kernel."""
+    from vllm.models.minimax_m3.amd.ops.nvfp4_sparse_decode import host
+
+    torch.manual_seed(num_reqs)
+    lens = [0 if r % 3 == 2 else 20000 + 700 * r for r in range(num_reqs)]
+    live = [max(s, 1) for s in lens]
+    table, num_blocks = _block_table(live)
+    seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda")
+    topk = _selections(live, 8, num_kv_heads)
+    cache = _nvfp4_cache(num_blocks, _scales(), num_kv_heads)
+    q = (
+        torch.randn(num_reqs * 8, GROUP * num_kv_heads, D, device="cuda") * 0.05
+    ).bfloat16()
+
+    chosen = []
+    real = host._get_kernel
+
+    def record(*args):
+        chosen.append(args)
+        return real(*args)
+
+    monkeypatch.setattr(host, "_get_kernel", record)
+    out = _decode(q, cache, topk, table, seq_lens, 8)
+
+    xcds = host._xcd_count(q.device.index)
+    assert xcds == 8
+    rows = num_reqs * 8
+    expected = host._xcd_splits(num_reqs, xcds, 16, 8, num_kv_heads)
+    if expected is None:
+        expected = host.decode_splits(rows * num_kv_heads, 16, host.cu_count(q.device))
+    grid_x = -(-rows // (xcds * 8)) * xcds * 8
+    if grid_x * num_kv_heads > host.MAX_SPLIT_ROWS:
+        expected = 1
+    assert chosen == [(expected, 16, xcds)]
+    if num_kv_heads == 1:
+        table_splits = {1: 4, 5: 4, 8: 4, 9: 2, 16: 2, 17: 4, 24: 4}
+        assert expected == table_splits.get(num_reqs, expected)
+
+    unsplit = _decode(q, cache, topk, table, seq_lens, 8, splits=1)
+    real_rows = (seq_lens.repeat_interleave(8) > 0).cpu()
+    assert torch.isfinite(out).all()
+    assert (out[~real_rows] == 0).all()
+    assert _rel(out[real_rows], unsplit[real_rows]) < 3e-3

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

@@ -99,6 +99,29 @@ def decode_splits(num_rows: int, num_blocks: int, cus: int) -> int:
     return best
 
 
+def _xcd_splits(
+    num_reqs: int, xcds: int, topk: int, decode_query_len: int, num_kv_heads: int
+) -> int | None:
+    """Splits for the XCD-aware grid with up to three requests per XCD, or None.
+
+    A request's rows and splits all run on one XCD (32 CUs), so the cost model
+    of ``decode_splits``, which spreads the grid over every CU, mispredicts: at
+    5 requests it picks 6 splits for 15.7 us where 4 splits take 11.9 us.
+    Measured on MI355X (top-k 16, 8 rows a request, 30K-400K contexts, 1-24
+    requests; the kernel time does not depend on the context): 4 splits is
+    best up to 8 requests (11.4-12.4 us, one request per XCD) and for 17-24
+    requests (21 us), 2 splits for 9-16 requests (15.3-16 us). Wider grids use
+    the cost model. A request's KV heads also share its XCD, so with 2 or 4 KV
+    heads per rank (TP2, TP1) the table over-splits and the cost model is kept.
+    """
+    if xcds <= 0 or topk != 16 or decode_query_len != 8 or num_kv_heads != 1:
+        return None  # tuned for DSpark 7, one KV head per rank
+    per_xcd = -(-num_reqs // xcds)
+    if per_xcd == 2:
+        return 2
+    return 4 if per_xcd <= 3 else None
+
+
 def minimax_m3_nvfp4_sparse_decode_fly(
     q: torch.Tensor,  # [rows, H, 128] bf16
     kv_cache: torch.Tensor,  # [pages, 2 * Hkv, 128, 72] uint8
@@ -144,11 +167,14 @@ def minimax_m3_nvfp4_sparse_decode_fly(
         grid_x = rows
     if splits is None:
         wg_rows = grid_x * num_kv_heads
-        splits = (
-            1
-            if wg_rows > MAX_SPLIT_ROWS
-            else decode_splits(rows * num_kv_heads, topk, cu_count(q.device))
-        )
+        if wg_rows > MAX_SPLIT_ROWS:
+            splits = 1
+        else:
+            splits = _xcd_splits(
+                rows // decode_query_len, xcds, topk, decode_query_len, num_kv_heads
+            )
+            if splits is None:
+                splits = decode_splits(rows * num_kv_heads, topk, cu_count(q.device))
     assert splits == 1 or grid_x * num_kv_heads <= MAX_SPLIT_ROWS
     part, cnt = _split_buffers(q.device)
     unit = unit_scale(q.device)