[ROCm][Perf] Split MiniMax-M3 prefill index-score K loop

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

Author: akii96

tests/kernels/attention/test_minimax_m3.py

@@ -124,6 +124,46 @@ def test_sparse_kernels_recognize_fp8_dtypes(dtype: torch.dtype):
 
 
 # Index top-k kernels.
+def _assert_prefill_index_scores(
+    actual: torch.Tensor,
+    idx_q: torch.Tensor,
+    index_kv_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    q_lens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    block_size_q: int,
+) -> None:
+    q_start = 0
+    for req_id, (q_len, seq_len, prefix_len) in enumerate(
+        zip(q_lens.tolist(), seq_lens.tolist(), prefix_lens.tolist())
+    ):
+        q = idx_q[q_start : q_start + q_len]
+        num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+        pages = block_table[req_id, :num_blocks]
+        k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
+        expected = torch.einsum("qhd,kd->hqk", q.float(), k.float())
+
+        q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
+        k_pos = torch.arange(k.shape[0], device=idx_q.device)
+        expected.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf"))
+        expected = (
+            expected.reshape(idx_q.shape[1], q_len, num_blocks, BLOCK_SIZE)
+            .max(dim=3)
+            .values
+        )
+
+        for local_q in range(q_len):
+            q_block_end = min(q_len, (local_q // block_size_q + 1) * block_size_q)
+            hi = min(seq_len, prefix_len + q_block_end)
+            written_blocks = (hi + BLOCK_SIZE - 1) // BLOCK_SIZE
+            torch.testing.assert_close(
+                actual[:, q_start + local_q, :written_blocks],
+                expected[:, local_q, :written_blocks],
+            )
+        q_start += q_len
+
+
 def _reference_index_topk(
     idx_q: torch.Tensor,
     index_kv_cache: torch.Tensor,
@@ -222,14 +262,34 @@ def _reference_decode_index_score(
     return out
 
 
-def test_prefill_index_topk_correctness():
+@pytest.mark.parametrize("long_context", [False, True])
+def test_prefill_index_topk_correctness(long_context: bool):
+    if current_platform.is_rocm():
+        from vllm.models.minimax_m3.amd.ops.index_topk import (
+            minimax_m3_index_score as amd_index_score,
+        )
+
+        index_score = amd_index_score
+    else:
+        index_score = minimax_m3_index_score
+
+    if long_context:
+        if not current_platform.is_rocm():
+            pytest.skip("The split-K index-score path is ROCm-specific.")
+        from vllm.platforms.rocm import on_gfx942
+
+        if not on_gfx942():
+            pytest.skip("The split-K index-score path is enabled on gfx942.")
+
     topk = 6
     init_blocks = 0
     local_blocks = 1
     num_idx_heads = 2
     head_dim = 16
-    q_lens = torch.tensor((4, 3), device="cuda", dtype=torch.int32)
-    prefix_lens = torch.tensor((0, 1024), device="cuda", dtype=torch.int32)
+    q_lens_values = (128, 129) if long_context else (4, 3)
+    prefix_lens_values = (8192, 16384) if long_context else (0, 1024)
+    q_lens = torch.tensor(q_lens_values, device="cuda", dtype=torch.int32)
+    prefix_lens = torch.tensor(prefix_lens_values, device="cuda", dtype=torch.int32)
     seq_lens = prefix_lens + q_lens
     batch = q_lens.numel()
     max_seq_len = seq_lens.max().item()
@@ -242,13 +302,15 @@ def test_prefill_index_topk_correctness():
         batch, max_blocks
     )
     idx_q = torch.ones(q_lens.sum().item(), num_idx_heads, head_dim, device="cuda")
-    index_kv_cache = torch.empty(num_pages, BLOCK_SIZE, head_dim, device="cuda")
-    for req_id in range(batch):
-        for block_id in range(max_blocks):
-            page = block_table[req_id, block_id]
-            index_kv_cache[page].fill_(block_id + 1)
+    block_values = torch.empty(num_pages, device="cuda")
+    block_values[block_table] = torch.arange(
+        1, max_blocks + 1, device="cuda", dtype=torch.float32
+    ).expand(batch, -1)
+    index_kv_cache = (
+        block_values[:, None, None].expand(-1, BLOCK_SIZE, head_dim).contiguous()
+    )
 
-    score = minimax_m3_index_score(
+    score = index_score(
         idx_q,
         index_kv_cache,
         block_table,
@@ -259,6 +321,16 @@ def test_prefill_index_topk_correctness():
         max_seq_len=max_seq_len,
         num_kv_heads=num_idx_heads,
     )
+    _assert_prefill_index_scores(
+        score,
+        idx_q,
+        index_kv_cache,
+        block_table,
+        q_lens,
+        seq_lens,
+        prefix_lens,
+        block_size_q=128 if long_context else 64,
+    )
     actual = minimax_m3_index_topk(
         score,
         cu_seqlens,

vllm/models/minimax_m3/amd/ops/index_topk.py

@@ -17,12 +17,20 @@
 import torch
 
 from vllm.platforms import current_platform
+from vllm.platforms.rocm import on_gfx942
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import round_up
+from vllm.utils.platform_utils import num_compute_units
 
 # One sparse block == one KV page.
 SPARSE_BLOCK_SIZE = 128
 
+# gfx942 split-K prefill tuning. Target bounded K-loop lengths while retaining
+# enough work per split to amortize launch overhead.
+_INDEX_SCORE_TARGET_K_TILES_PER_SPLIT = 64
+_INDEX_SCORE_MAX_SPLITS = 16
+_INDEX_SCORE_MIN_SEQ_LEN = _INDEX_SCORE_TARGET_K_TILES_PER_SPLIT * SPARSE_BLOCK_SIZE
+
 
 # ---------------------------------------------------------------------------
 # Bitonic top-k helpers (layout-agnostic).
@@ -101,8 +109,10 @@ def _index_block_score_kernel(
     stride_bt_b,
     BLOCK_SIZE_Q: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,  # == SPARSE_BLOCK_SIZE (128)
+    NUM_K_SPLITS: tl.constexpr,
 ):
-    pid_q = tl.program_id(0)
+    pid_q = tl.program_id(0) // NUM_K_SPLITS
+    pid_split = tl.program_id(0) % NUM_K_SPLITS
     pid_bh = tl.program_id(1)
     pid_b = pid_bh // num_idx_heads
     pid_h = pid_bh % num_idx_heads
@@ -114,6 +124,16 @@ def _index_block_score_kernel(
     if BLOCK_SIZE_Q * pid_q >= q_len:
         return
 
+    # Split the state-free K loop across programs. Every score element still
+    # has exactly one writer, so no reduction or atomic operation is needed.
+    hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q)
+    num_k_tiles = tl.cdiv(hi, BLOCK_SIZE_K)
+    tiles_per_split = tl.cdiv(num_k_tiles, NUM_K_SPLITS)
+    tile_lo = pid_split * tiles_per_split
+    tile_hi = min(num_k_tiles, tile_lo + tiles_per_split)
+    if tile_lo >= tile_hi:
+        return
+
     q_ptrs = tl.make_block_ptr(
         base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h,
         shape=(q_len, head_dim),
@@ -130,10 +150,8 @@ def _index_block_score_kernel(
     off_d = tl.arange(0, head_dim)
     # Block table row for this request.
     bt_row = block_table_ptr + pid_b * stride_bt_b
-    # Causal window: only blocks up to the last query token's position.
-    hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q)
-    for i in tl.range(0, hi, BLOCK_SIZE_K):
-        blk = i // BLOCK_SIZE_K
+    for blk in tl.range(tile_lo, tile_hi):
+        i = blk * BLOCK_SIZE_K
         page = tl.load(bt_row + blk).to(tl.int64)
         pos = i + off_k
         # index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed)
@@ -651,6 +669,39 @@ def _topk_index_merge_kernel(
 # ---------------------------------------------------------------------------
 # Python wrappers
 # ---------------------------------------------------------------------------
+def _index_score_launch_config(
+    max_query_len: int,
+    max_seq_len: int,
+    batch: int,
+    num_idx_heads: int,
+) -> tuple[int, int, dict]:
+    """Select a split-K launch only for shapes with useful parallel work."""
+    if not on_gfx942() or max_query_len < 128 or max_seq_len < _INDEX_SCORE_MIN_SEQ_LEN:
+        return 64, 1, {}
+
+    block_size_q = 128
+    num_k_tiles = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
+    q_programs = triton.cdiv(max_query_len, block_size_q) * batch * num_idx_heads
+
+    # Aim for two programs per CU when the split cap and K work permit, and
+    # avoid leaving very long K loops. Power-of-two splits limit compilations.
+    occupancy_splits = triton.cdiv(2 * num_compute_units(), q_programs)
+    work_splits = triton.cdiv(num_k_tiles, _INDEX_SCORE_TARGET_K_TILES_PER_SPLIT)
+    num_splits = triton.next_power_of_2(max(occupancy_splits, work_splits))
+    num_splits = min(_INDEX_SCORE_MAX_SPLITS, num_k_tiles, num_splits)
+    if num_splits == 1:
+        return 64, 1, {}
+
+    return (
+        block_size_q,
+        num_splits,
+        {
+            "num_warps": 4,
+            "num_stages": 1,
+        },
+    )
+
+
 @torch.no_grad()
 def minimax_m3_index_score(
     idx_q: torch.Tensor,  # [total_q, num_idx_heads, head_dim]
@@ -682,8 +733,16 @@ def minimax_m3_index_score(
         dtype=torch.float32,
         device=idx_q.device,
     )
-    BLOCK_SIZE_Q = 64
-    grid_score = (triton.cdiv(max_query_len, BLOCK_SIZE_Q), batch * num_idx_heads)
+    block_size_q, num_k_splits, launch_kwargs = _index_score_launch_config(
+        max_query_len,
+        max_seq_len,
+        batch,
+        num_idx_heads,
+    )
+    grid_score = (
+        triton.cdiv(max_query_len, block_size_q) * num_k_splits,
+        batch * num_idx_heads,
+    )
     _index_block_score_kernel[grid_score](
         idx_q,
         index_kv_cache,
@@ -704,8 +763,10 @@ def minimax_m3_index_score(
         score.stride(1),
         score.stride(2),
         block_table.stride(0),
-        BLOCK_SIZE_Q=BLOCK_SIZE_Q,
+        BLOCK_SIZE_Q=block_size_q,
         BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
+        NUM_K_SPLITS=num_k_splits,
+        **launch_kwargs,
     )
     return score