[ROCm][Perf] Tune MiniMax-M3 decode top-k for short contexts (#55235)
Signed-off-by: fai <fangzhouai@gmail.com>
Fangzhou Ai committed
Sep 13, 2026 at 22:36 UTC
dfa1984e58dc898610a320de26f94b6039fccec3
2 files changed
+90
-48
tests/kernels/attention/test_minimax_m3.py
+63
-41
@@ -1025,17 +1025,20 @@ def test_decode_index_topk_correctness(
1025
"expected",
1026
),
1027
[
1028
- (0, 24, 1, 16, True, (2, False, False)),
1029
- (1, 1, 1, 16, True, (16, True, True)),
1030
- (512, 4, 1, 16, True, (16, True, True)),
1031
- (513, 24, 4, 16, True, (16, True, True)),
1032
- (2760, 5, 2, 16, True, (16, True, True)),
1033
- (8192, 64, 4, 16, True, (16, True, True)),
1034
- (8193, 4, 1, 16, True, (16, False, False)),
1035
- (8193, 16, 1, 16, True, (4, False, False)),
1036
- (65534, 64, 1, 16, True, (1, False, False)),
1037
- (513, 16, 1, 13, True, (4, False, False)),
1038
- (513, 16, 4, 16, False, (1, False, False)),
1028
+ (0, 24, 1, 16, True, (2, 512, 8, 2, False, False)),
1029
+ (1, 1, 1, 16, True, (1, 128, 2, 1, True, True)),
1030
+ (74, 4, 1, 16, True, (1, 128, 2, 1, True, True)),
1031
+ (128, 4, 1, 16, True, (1, 128, 2, 1, True, True)),
1032
+ (129, 4, 1, 16, True, (16, 512, 4, 2, True, True)),
1033
+ (512, 4, 1, 16, True, (16, 512, 4, 2, True, True)),
1034
+ (513, 24, 4, 16, True, (16, 512, 4, 2, True, True)),
1035
+ (2760, 5, 2, 16, True, (16, 512, 4, 2, True, True)),
1036
+ (8192, 64, 4, 16, True, (16, 512, 4, 2, True, True)),
1037
+ (8193, 4, 1, 16, True, (16, 512, 8, 2, False, False)),
1038
+ (8193, 16, 1, 16, True, (4, 512, 8, 2, False, False)),
1039
+ (65534, 64, 1, 16, True, (1, 512, 8, 2, False, False)),
1040
+ (513, 16, 1, 13, True, (4, 512, 8, 2, False, False)),
1041
+ (513, 16, 4, 16, False, (1, 512, 8, 2, False, False)),
1042
],
1043
)
1044
def test_amd_decode_topk_launch_policy(
@@ -1044,9 +1047,9 @@ def test_amd_decode_topk_launch_policy(
1047
num_idx_heads: int,
1048
topk: int,
1049
is_gfx950: bool,
1047
- expected: tuple[int, bool, bool],
1050
+ expected: tuple[int, int, int, int, bool, bool],
1051
):
1049
- """All supported head and query layouts use adaptive selection."""
1052
+ """The short-context specialization has strict dispatch boundaries."""
1053
from vllm.models.minimax_m3.amd.ops.index_topk import (
1054
_decode_topk_launch_policy,
1055
)
@@ -1073,19 +1076,23 @@ def test_amd_decode_topk_launch_policy(
1076
"single_tile_guaranteed",
1077
"adaptive_final_merge",
1078
"max_blocks",
1079
+ "block_size_k",
1080
+ "num_warps",
1081
+ "num_stages",
1082
"topk",
1083
"num_idx_heads",
1084
),
1085
[
1080
- (1, False, False, 65534, 16, 1),
1081
- (2, False, False, 8193, 16, 1),
1082
- (4, False, False, 8193, 16, 1),
1083
- (8, False, False, 8193, 16, 1),
1084
- (16, False, False, 8193, 16, 1),
1085
- (16, False, False, 65534, 16, 1),
1086
- (2, True, False, 513, 13, 1),
1087
- (4, False, True, 8192, 16, 2),
1088
- (16, True, True, 8192, 16, 4),
1086
+ (1, False, False, 65534, 512, 8, 2, 16, 1),
1087
+ (2, False, False, 8193, 512, 8, 2, 16, 1),
1088
+ (4, False, False, 8193, 512, 8, 2, 16, 1),
1089
+ (8, False, False, 8193, 512, 8, 2, 16, 1),
1090
+ (16, False, False, 8193, 512, 8, 2, 16, 1),
1091
+ (16, False, False, 65534, 512, 8, 2, 16, 1),
1092
+ (2, True, False, 513, 512, 8, 2, 13, 1),
1093
+ (4, False, True, 8192, 512, 4, 2, 16, 2),
1094
+ (16, True, True, 8192, 512, 4, 2, 16, 4),
1095
+ (1, True, True, 74, 128, 2, 1, 16, 1),
1096
],
1097
)
1098
def test_amd_decode_fused_topk_total_order_and_replay(
@@ -1093,10 +1100,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
1100
single_tile_guaranteed: bool,
1101
adaptive_final_merge: bool,
1102
max_blocks: int,
1103
+ block_size_k: int,
1104
+ num_warps: int,
1105
+ num_stages: int,
1106
topk: int,
1107
num_idx_heads: int,
1108
):
1099
- """Atomic multi-chunk selection is ordered and resets graph state."""
1109
+ """Selector specializations preserve ordering and graph state."""
1110
from vllm.models.minimax_m3.amd.ops.index_topk import (
1111
SPARSE_BLOCK_SIZE,
1112
_decode_topk_fused_kernel,
@@ -1104,18 +1114,25 @@ def test_amd_decode_fused_topk_total_order_and_replay(
1114
1115
set_random_seed(0)
1116
block_size_t = 1 << (topk - 1).bit_length()
1117
+ block_counts = tuple(
1118
+ count
1119
+ for count in (
1120
+ 15,
1121
+ 16,
1122
+ 17,
1123
+ block_size_k // 2 - 1,
1124
+ block_size_k // 2,
1125
+ block_size_k // 2 + 1,
1126
+ block_size_k - 1,
1127
+ block_size_k,
1128
+ block_size_k + 1,
1129
+ )
1130
+ if count <= max_blocks
1131
+ )
1132
long_seq_lens_list = [
1133
0,
1134
1,
1110
- 15 * SPARSE_BLOCK_SIZE,
1111
- 16 * SPARSE_BLOCK_SIZE,
1112
- 17 * SPARSE_BLOCK_SIZE,
1113
- 255 * SPARSE_BLOCK_SIZE,
1114
- 256 * SPARSE_BLOCK_SIZE,
1115
- 257 * SPARSE_BLOCK_SIZE,
1116
- 511 * SPARSE_BLOCK_SIZE,
1117
- 512 * SPARSE_BLOCK_SIZE,
1118
- 513 * SPARSE_BLOCK_SIZE,
1135
+ *(blocks * SPARSE_BLOCK_SIZE for blocks in block_counts),
1136
max_blocks * SPARSE_BLOCK_SIZE,
1137
]
1138
short_seq_lens_list = [0, 1, 127, 128, 129, 0, 1, 128, 129, 0, 1, 129]
@@ -1126,12 +1143,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
1143
boundary_limit = 4
1144
if boundary_limit:
1145
boundary_blocks = [
1129
- (active_chunks - 1) * 512 + 1
1146
+ (active_chunks - 1) * block_size_k + 1
1147
for active_chunks in range(3, boundary_limit + 1)
1148
]
1132
- assert [(num_blocks + 511) // 512 for num_blocks in boundary_blocks] == list(
1133
- range(3, boundary_limit + 1)
1134
- )
1149
+ assert [
1150
+ (num_blocks + block_size_k - 1) // block_size_k
1151
+ for num_blocks in boundary_blocks
1152
+ ] == list(range(3, boundary_limit + 1))
1153
long_seq_lens_list.extend(
1154
num_blocks * SPARSE_BLOCK_SIZE for num_blocks in boundary_blocks
1155
)
@@ -1139,6 +1157,8 @@ def test_amd_decode_fused_topk_total_order_and_replay(
1157
num_blocks * SPARSE_BLOCK_SIZE for num_blocks in reversed(boundary_blocks)
1158
)
1159
batch = len(long_seq_lens_list)
1160
+ repeats = (batch + len(short_seq_lens_list) - 1) // len(short_seq_lens_list)
1161
+ short_seq_lens_list = (short_seq_lens_list * repeats)[:batch]
1162
score = torch.randn((num_idx_heads, batch, max_blocks), device="cuda")
1163
score[:, :, 0] = float("nan")
1164
score[:, :, 1] = float("inf")
@@ -1147,8 +1167,9 @@ def test_amd_decode_fused_topk_total_order_and_replay(
1167
score[:, :, 4] = -0.0
1168
score[:, :, 5] = 7.0
1169
score[:, :, 6] = -1e30
1150
- score[:, :, 256] = 7.0
1151
- score[:, :, 512] = 7.0
1170
+ for index in (block_size_k // 2, block_size_k):
1171
+ if index < max_blocks:
1172
+ score[:, :, index] = 7.0
1173
1174
seq_lens = torch.empty(batch, device="cuda", dtype=torch.int32)
1175
partial = torch.full(
@@ -1198,13 +1219,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
1219
pages_per_sparse_block=8,
1220
block_page_stride=8,
1221
NUM_TOPK_CHUNKS=num_topk_chunks,
1201
- BLOCK_SIZE_K=512,
1222
+ BLOCK_SIZE_K=block_size_k,
1223
BLOCK_SIZE_T=block_size_t,
1224
EMIT_SPARSE_TABLE=False,
1225
SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
1226
ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
1206
- num_warps=4 if adaptive_final_merge else 8,
1207
- num_stages=2,
1227
+ num_warps=num_warps,
1228
+ num_stages=num_stages,
1229
)
1230
1231
score_cpu = score.cpu()
@@ -1251,6 +1272,7 @@ def test_amd_decode_fused_topk_total_order_and_replay(
1272
(16, False, False, 8193),
1273
(4, False, True, 8192),
1274
(16, True, True, 8192),
1275
+ (1, True, True, 74),
1276
):
1277
graph = torch.cuda.CUDAGraph()
1278
torch.accelerator.synchronize()
vllm/models/minimax_m3/amd/ops/index_topk.py
+27
-7
@@ -37,6 +37,7 @@ MIN_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS = 20
37
MAX_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS = 64
38
MAX_DECODE_SCORE_KV_CHUNKS = 256
39
DECODE_TOPK_BLOCKS_PER_CHUNK = 512
40
+DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK = 128
41
MAX_DECODE_TOPK_FAST_CHUNKS = 16
42
DECODE_TOPK_TARGET_GRID = 64
43
@@ -103,14 +104,23 @@ def _decode_topk_launch_policy(
104
topk: int,
105
*,
106
is_gfx950: bool,
106
-) -> tuple[int, bool, bool]:
107
- """Choose the graph grid and bounded-context selector specialization."""
107
+) -> tuple[int, int, int, int, bool, bool]:
108
+ """Choose the selector grid and compile-time launch configuration."""
109
if (
110
is_gfx950
111
and topk == 16
112
and 0 < max_block <= MAX_DECODE_TOPK_FAST_CHUNKS * DECODE_TOPK_BLOCKS_PER_CHUNK
113
):
113
- return MAX_DECODE_TOPK_FAST_CHUNKS, True, True
114
+ if max_block <= DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK:
115
+ return 1, DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK, 2, 1, True, True
116
+ return (
117
+ MAX_DECODE_TOPK_FAST_CHUNKS,
118
+ DECODE_TOPK_BLOCKS_PER_CHUNK,
119
+ 4,
120
+ 2,
121
+ True,
122
+ True,
123
+ )
124
125
target = max(
126
1,
@@ -119,7 +129,14 @@ def _decode_topk_launch_policy(
129
DECODE_TOPK_TARGET_GRID // max(1, total_q * num_idx_heads),
130
),
131
)
122
- return 1 << (target.bit_length() - 1), False, False
132
+ return (
133
+ 1 << (target.bit_length() - 1),
134
+ DECODE_TOPK_BLOCKS_PER_CHUNK,
135
+ 8,
136
+ 2,
137
+ False,
138
+ False,
139
+ )
140
141
142
# ---------------------------------------------------------------------------
@@ -1525,6 +1542,9 @@ def minimax_m3_index_decode(
1542
# only the number of chunks needed for its live context.
1543
(
1544
num_topk_chunks,
1545
+ topk_block_size,
1546
+ selector_num_warps,
1547
+ selector_num_stages,
1548
single_tile_guaranteed,
1549
adaptive_final_merge,
1550
) = _decode_topk_launch_policy(
@@ -1608,12 +1628,12 @@ def minimax_m3_index_decode(
1628
pages_per_sparse_block=PAGES_PER_SPARSE_BLOCK,
1629
block_page_stride=selector_block_page_stride,
1630
NUM_TOPK_CHUNKS=num_topk_chunks,
1611
- BLOCK_SIZE_K=512,
1631
+ BLOCK_SIZE_K=topk_block_size,
1632
BLOCK_SIZE_T=block_size_t,
1633
EMIT_SPARSE_TABLE=emit_sparse_table,
1634
SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
1635
ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
1616
- num_warps=4 if adaptive_final_merge else 8,
1617
- num_stages=2,
1636
+ num_warps=selector_num_warps,
1637
+ num_stages=selector_num_stages,
1638
)
1639
return topk_idx