[Perf][MiniMax-M3] Triton indexer: decode grid retune + SM12.0 split-K (#56151)
Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com> Signed-off-by: Zijing Liu <liuzijing2014@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: Zijing Liu <liuzijing2014@gmail.com> Co-authored-by: Kyle Liang <199154931+kyleliang-nv@users.noreply.github.com>
Venky committed
Sep 30, 2026 at 12:33 UTC
84f738f1878ff2ab7664f1a4464b7f9f3cdb063d
2 files changed
+38
-7
tests/kernels/attention/test_minimax_m3.py
+10
-2
@@ -614,14 +614,22 @@ def _reference_decode_index_score(
614
return out
615
616
617
-def test_prefill_index_topk_correctness():
617
+# 3600: 29 blocks over 12 splits of 3, so the last busy split has 2 blocks.
618
+@pytest.mark.parametrize("prefix_len", [1024, 3600])
619
+@pytest.mark.parametrize("force_split_k", [False, True])
620
+def test_prefill_index_topk_correctness(monkeypatch, prefix_len, force_split_k):
621
+ if force_split_k:
622
+ # Take the SM12.0-only split-K launch on any CUDA GPU.
623
+ monkeypatch.setattr(
624
+ current_platform, "is_device_capability", lambda cap, *_: cap == (12, 0)
625
+ )
626
topk = 6
627
init_blocks = 0
628
local_blocks = 1
629
num_idx_heads = 2
630
head_dim = 16
631
q_lens = torch.tensor((4, 3), device="cuda", dtype=torch.int32)
624
- prefix_lens = torch.tensor((0, 1024), device="cuda", dtype=torch.int32)
632
+ prefix_lens = torch.tensor((0, prefix_len), device="cuda", dtype=torch.int32)
633
seq_lens = prefix_lens + q_lens
634
batch = q_lens.numel()
635
max_seq_len = seq_lens.max().item()
vllm/models/minimax_m3/common/ops/index_topk.py
+28
-5
@@ -101,6 +101,7 @@ def _index_block_score_kernel(
101
stride_bt_b,
102
BLOCK_SIZE_Q: tl.constexpr,
103
BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
104
+ USE_SPLIT_K: tl.constexpr,
105
):
106
pid_q = tl.program_id(0)
107
pid_bh = tl.program_id(1)
@@ -136,8 +137,19 @@ def _index_block_score_kernel(
137
bt_row = block_table_ptr + pid_b * stride_bt_b
138
# Causal window: only blocks up to the last query token's position.
139
hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q)
139
- for i in tl.range(0, hi, BLOCK_SIZE_K):
140
- blk = i // BLOCK_SIZE_K
140
+ num_blocks = tl.cdiv(hi, BLOCK_SIZE_K)
141
+ if USE_SPLIT_K:
142
+ pid_k = tl.program_id(2)
143
+ blocks_per_split = tl.cdiv(num_blocks, tl.num_programs(2))
144
+ block_start = pid_k * blocks_per_split
145
+ block_end = tl.minimum(block_start + blocks_per_split, num_blocks)
146
+ if block_start >= block_end:
147
+ return
148
+ else:
149
+ block_start = 0
150
+ block_end = num_blocks
151
+ for blk in tl.range(block_start, block_end):
152
+ i = blk * BLOCK_SIZE_K
153
page = tl.load(bt_row + blk).to(tl.int64)
154
pos = i + off_k
155
# index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed)
@@ -686,7 +698,17 @@ def minimax_m3_index_score(
698
device=idx_q.device,
699
)
700
BLOCK_SIZE_Q = 64
689
- grid_score = (triton.cdiv(max_query_len, BLOCK_SIZE_Q), batch * num_idx_heads)
701
+ n_q_tiles = triton.cdiv(max_query_len, BLOCK_SIZE_Q)
702
+ SCORE_TARGET_GRID = 48
703
+ split_k = max(
704
+ 1,
705
+ min(max_block, SCORE_TARGET_GRID // max(1, n_q_tiles * batch * num_idx_heads)),
706
+ )
707
+ if not (
708
+ current_platform.is_cuda() and current_platform.is_device_capability((12, 0))
709
+ ):
710
+ split_k = 1
711
+ grid_score = (n_q_tiles, batch * num_idx_heads, split_k)
712
_index_block_score_kernel[grid_score](
713
idx_q,
714
index_kv_cache,
@@ -709,6 +731,7 @@ def minimax_m3_index_score(
731
block_table.stride(0),
732
BLOCK_SIZE_Q=BLOCK_SIZE_Q,
733
BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
734
+ USE_SPLIT_K=split_k > 1,
735
)
736
return score
737
@@ -822,7 +845,7 @@ def minimax_m3_index_decode_score(
845
)
846
# split-K over seq blocks; chunk count depends only on shape constants so
847
# the grid is fixed within a cuda graph.
825
- TARGET_GRID = 512
848
+ TARGET_GRID = 4096
849
MAX_NUM_KV_CHUNKS = 256
850
# Use the configured max decode length to avoid Triton recompiles when
851
# switching between qlen=1 and spec-decode verification batches.
@@ -921,7 +944,7 @@ def minimax_m3_index_decode(
944
)
945
# Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts
946
# pow2(num_topk_chunks * pow2(topk)) candidates.
924
- TOPK_TARGET_GRID = 64
947
+ TOPK_TARGET_GRID = 512
948
MAX_NUM_TOPK_CHUNKS = 16
949
topk_target = max(
950
1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads))