@tduncan / Nvidia-Optimization-Stack-VLLM / commits / cee0f92c02

[ROCm][Perf] Optimize MiniMax-M3 decode indexer and top-k (#54682)

Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Fangzhou Ai committed Sep 3, 2026 at 16:20 UTC cee0f92c02112ace7120da45896d27e0fe95ea1e
6 files changed +2010 -327
tests/kernels/attention/test_minimax_m3.py
+840
@@ -21,6 +21,7 @@ from vllm.models.minimax_m3.common.ops.sparse_attn import (
21 )
22 from vllm.models.minimax_m3.common.sparse_attention import MiniMaxM3SparseTritonImpl
23 from vllm.platforms import current_platform
24 +from vllm.utils.torch_utils import set_random_seed
25 from vllm.v1.attention.backends.utils import record_kv_cache_layout
26 from vllm.v1.kv_cache_interface import (
27 FullAttentionSpec,
@@ -92,6 +93,382 @@ SM_SCALE = HEAD_DIM**-0.5
93 TOPK = 16
94
95
96 +@pytest.mark.parametrize(
97 + (
98 + "num_reqs",
99 + "head_dim",
100 + "query_dtype",
101 + "cache_dtype",
102 + "is_gfx950",
103 + "expected",
104 + ),
105 + [
106 + (19, 128, torch.bfloat16, torch.bfloat16, True, (16, False)),
107 + (20, 128, torch.bfloat16, torch.bfloat16, True, (32, True)),
108 + (32, 128, torch.bfloat16, torch.bfloat16, True, (32, True)),
109 + (33, 128, torch.bfloat16, torch.bfloat16, True, (16, True)),
110 + (64, 128, torch.bfloat16, torch.bfloat16, True, (16, True)),
111 + (65, 128, torch.bfloat16, torch.bfloat16, True, (4, False)),
112 + (20, 128, torch.bfloat16, torch.bfloat16, False, (16, False)),
113 + (20, 64, torch.bfloat16, torch.bfloat16, True, (16, False)),
114 + (20, 128, torch.float16, torch.bfloat16, True, (16, False)),
115 + (20, 128, torch.bfloat16, torch.float16, True, (16, False)),
116 + ],
117 +)
118 +def test_amd_decode_score_split_launch_policy(
119 + num_reqs: int,
120 + head_dim: int,
121 + query_dtype: torch.dtype,
122 + cache_dtype: torch.dtype,
123 + is_gfx950: bool,
124 + expected: tuple[int, bool],
125 +):
126 + from vllm.models.minimax_m3.amd.ops.index_topk import (
127 + _decode_score_split_launch_policy,
128 + )
129 +
130 + assert (
131 + _decode_score_split_launch_policy(
132 + num_reqs,
133 + head_dim,
134 + query_dtype,
135 + cache_dtype,
136 + is_gfx950=is_gfx950,
137 + )
138 + == expected
139 + )
140 +
141 +
142 +@pytest.mark.parametrize("num_reqs", range(1, 12))
143 +def test_amd_balanced_decode_score_mapping_exact_coverage(num_reqs: int):
144 + from vllm.models.minimax_m3.amd.ops.index_topk import (
145 + _decode_score_program_budget,
146 + )
147 +
148 + budget = _decode_score_program_budget(
149 + num_reqs,
150 + 128,
151 + torch.bfloat16,
152 + torch.bfloat16,
153 + is_gfx950=True,
154 + )
155 + assert budget is not None
156 + patterns = (
157 + (0,) * num_reqs,
158 + tuple((request * 7919 + 17) % 65535 for request in range(num_reqs)),
159 + (65534,) * num_reqs,
160 + )
161 + for request_blocks in patterns:
162 + width = max(1, (sum(request_blocks) + budget - 1) // budget)
163 + grid = budget + len(request_blocks) - 1
164 + intervals = []
165 + for request, blocks in enumerate(request_blocks):
166 + programs = (blocks + width - 1) // width
167 + for local_program in range(programs):
168 + begin = local_program * width
169 + intervals.append((request, begin, min(begin + width, blocks)))
170 +
171 + assert len(intervals) <= grid
172 + by_request: list[list[tuple[int, int]]] = [[] for _ in request_blocks]
173 + for request, begin, end in intervals:
174 + assert 0 <= request < len(request_blocks)
175 + assert 0 <= begin < end <= request_blocks[request]
176 + by_request[request].append((begin, end))
177 + for blocks, request_intervals in zip(
178 + request_blocks,
179 + by_request,
180 + strict=True,
181 + ):
182 + cursor = 0
183 + for begin, end in request_intervals:
184 + assert begin == cursor
185 + cursor = end
186 + assert cursor == blocks
187 +
188 +
189 +@pytest.mark.skipif(
190 + not current_platform.is_rocm(),
191 + reason="MiniMax-M3 balanced decode score is ROCm-only",
192 +)
193 +@pytest.mark.parametrize(
194 + (
195 + "num_reqs",
196 + "max_block",
197 + "num_idx_heads",
198 + "decode_query_len",
199 + "max_decode_query_len",
200 + "init_blocks",
201 + "local_blocks",
202 + ),
203 + [
204 + (1, 37, 1, 1, 1, 1, 0),
205 + (8, 37, 1, 1, 4, 2, 1),
206 + (9, 37, 1, 4, 4, 0, 2),
207 + (11, 37, 1, 3, 5, 1, 3),
208 + (1, 37, 2, 1, 1, 2, 0),
209 + (8, 37, 2, 1, 4, 0, 1),
210 + (9, 37, 2, 4, 4, 1, 2),
211 + (11, 37, 2, 3, 5, 2, 3),
212 + (1, 37, 4, 1, 1, 0, 0),
213 + (8, 37, 4, 1, 4, 1, 1),
214 + (9, 37, 4, 4, 4, 2, 2),
215 + (11, 37, 4, 3, 5, 0, 3),
216 + (20, 65, 1, 4, 4, 1, 3),
217 + (36, 65, 1, 4, 4, 2, 2),
218 + (64, 65, 1, 4, 4, 0, 1),
219 + ],
220 +)
221 +def test_amd_optimized_decode_score_bitwise_and_graph_replay(
222 + num_reqs: int,
223 + max_block: int,
224 + num_idx_heads: int,
225 + decode_query_len: int,
226 + max_decode_query_len: int,
227 + init_blocks: int,
228 + local_blocks: int,
229 +):
230 + from vllm.models.minimax_m3.amd.ops.index_topk import (
231 + _decode_index_score_balanced_kernel,
232 + _decode_index_score_kernel,
233 + _decode_score_program_budget,
234 + _decode_score_split_launch_policy,
235 + )
236 +
237 + set_random_seed(0)
238 + head_dim = 128
239 + block_size_q = 1 << (max_decode_query_len - 1).bit_length()
240 + score_stride = (max_block + 15) // 16 * 16
241 + num_pages = num_reqs * max_block
242 + block_table = torch.arange(
243 + num_pages,
244 + device="cuda",
245 + dtype=torch.int32,
246 + ).reshape(num_reqs, max_block)
247 + idx_q = torch.randn(
248 + (num_reqs * decode_query_len, num_idx_heads, head_dim),
249 + device="cuda",
250 + dtype=torch.bfloat16,
251 + )
252 + index_kv_cache = torch.randn(
253 + (num_pages, BLOCK_SIZE, head_dim),
254 + device="cuda",
255 + dtype=torch.bfloat16,
256 + )
257 + seq_lens = torch.empty(num_reqs, device="cuda", dtype=torch.int32)
258 + deployed_score = torch.empty(
259 + (num_idx_heads, num_reqs * decode_query_len, score_stride),
260 + device="cuda",
261 + dtype=torch.float32,
262 + )
263 + optimized_score = torch.empty_like(deployed_score)
264 + budget = _decode_score_program_budget(
265 + num_reqs,
266 + head_dim,
267 + idx_q.dtype,
268 + index_kv_cache.dtype,
269 + is_gfx950=True,
270 + )
271 + optimized_chunks = None
272 + if budget is None:
273 + optimized_chunks, use_high_batch_config = _decode_score_split_launch_policy(
274 + num_reqs,
275 + head_dim,
276 + idx_q.dtype,
277 + index_kv_cache.dtype,
278 + is_gfx950=True,
279 + )
280 + assert use_high_batch_config
281 + target = max(1, min(256, 512 // num_reqs))
282 + deployed_chunks = 1 << (target.bit_length() - 1)
283 +
284 + def launch_deployed() -> None:
285 + _decode_index_score_kernel[(num_reqs, deployed_chunks)](
286 + idx_q,
287 + index_kv_cache,
288 + deployed_score,
289 + block_table,
290 + seq_lens,
291 + num_idx_heads,
292 + head_dim,
293 + init_blocks,
294 + local_blocks,
295 + decode_query_len,
296 + idx_q.stride(0),
297 + idx_q.stride(1),
298 + idx_q.stride(2),
299 + index_kv_cache.stride(0),
300 + index_kv_cache.stride(1),
301 + index_kv_cache.stride(2),
302 + deployed_score.stride(0),
303 + deployed_score.stride(1),
304 + deployed_score.stride(2),
305 + block_table.stride(0),
306 + BLOCK_SIZE_K=BLOCK_SIZE,
307 + BLOCK_SIZE_Q=block_size_q,
308 + num_kv_chunks=deployed_chunks,
309 + USE_PDL=False,
310 + )
311 +
312 + def launch_optimized() -> None:
313 + if budget is not None:
314 + _decode_index_score_balanced_kernel[(budget + num_reqs - 1,)](
315 + idx_q,
316 + index_kv_cache,
317 + optimized_score,
318 + block_table,
319 + seq_lens,
320 + num_idx_heads,
321 + head_dim,
322 + init_blocks,
323 + local_blocks,
324 + num_reqs,
325 + budget,
326 + decode_query_len,
327 + idx_q.stride(0),
328 + idx_q.stride(1),
329 + idx_q.stride(2),
330 + index_kv_cache.stride(0),
331 + index_kv_cache.stride(1),
332 + index_kv_cache.stride(2),
333 + optimized_score.stride(0),
334 + optimized_score.stride(1),
335 + optimized_score.stride(2),
336 + block_table.stride(0),
337 + BLOCK_SIZE_K=BLOCK_SIZE,
338 + BLOCK_SIZE_Q=block_size_q,
339 + num_warps=2,
340 + num_stages=1,
341 + )
342 + else:
343 + assert optimized_chunks is not None
344 + _decode_index_score_kernel[(num_reqs, optimized_chunks)](
345 + idx_q,
346 + index_kv_cache,
347 + optimized_score,
348 + block_table,
349 + seq_lens,
350 + num_idx_heads,
351 + head_dim,
352 + init_blocks,
353 + local_blocks,
354 + decode_query_len,
355 + idx_q.stride(0),
356 + idx_q.stride(1),
357 + idx_q.stride(2),
358 + index_kv_cache.stride(0),
359 + index_kv_cache.stride(1),
360 + index_kv_cache.stride(2),
361 + optimized_score.stride(0),
362 + optimized_score.stride(1),
363 + optimized_score.stride(2),
364 + block_table.stride(0),
365 + BLOCK_SIZE_K=BLOCK_SIZE,
366 + BLOCK_SIZE_Q=block_size_q,
367 + num_kv_chunks=optimized_chunks,
368 + USE_PDL=False,
369 + num_warps=2,
370 + num_stages=1,
371 + )
372 +
373 + def score_masks(
374 + expected_seq_lens: torch.Tensor,
375 + ) -> tuple[torch.Tensor, torch.Tensor]:
376 + query_offsets = torch.arange(decode_query_len, device="cuda")
377 + kv_lens = torch.clamp(
378 + expected_seq_lens[:, None] - decode_query_len + query_offsets + 1,
379 + min=0,
380 + )
381 + query_blocks = (kv_lens + BLOCK_SIZE - 1) // BLOCK_SIZE
382 + block_ids = torch.arange(score_stride, device="cuda")
383 + topk_mask = block_ids[None, :] < query_blocks.reshape(-1, 1)
384 + request_blocks = query_blocks.max(dim=1).values
385 + write_mask = (
386 + block_ids[None, :]
387 + < request_blocks.repeat_interleave(decode_query_len)[:, None]
388 + )
389 + return write_mask, topk_mask
390 +
391 + def packed_keys(score: torch.Tensor, live: torch.Tensor) -> torch.Tensor:
392 + normalized = torch.where(torch.isnan(score), -1e30, score)
393 + normalized = torch.where(normalized == 0.0, 0.0, normalized)
394 + bits = normalized.contiguous().view(torch.int32).to(torch.int64)
395 + bits &= 0xFFFFFFFF
396 + ordered = bits ^ torch.where(
397 + bits >> 31 != 0,
398 + 0xFFFFFFFF,
399 + 0x80000000,
400 + )
401 + index = torch.arange(score_stride, device="cuda") + 1
402 + keys = (1 << 48) | (ordered << 16) | (0xFFFF - index)
403 + return torch.where(live, keys, 0)
404 +
405 + def assert_exact(expected_seq_lens: torch.Tensor) -> None:
406 + write_mask, topk_mask = score_masks(expected_seq_lens)
407 + write_mask = write_mask.unsqueeze(0).expand_as(deployed_score)
408 + topk_mask = topk_mask.unsqueeze(0)
409 + deployed = deployed_score
410 + optimized = optimized_score
411 + assert torch.equal(
412 + deployed[write_mask].contiguous().view(torch.int32),
413 + optimized[write_mask].contiguous().view(torch.int32),
414 + )
415 + assert torch.isnan(deployed[~write_mask]).all()
416 + assert torch.isnan(optimized[~write_mask]).all()
417 + deployed_keys = packed_keys(deployed, topk_mask)
418 + optimized_keys = packed_keys(optimized, topk_mask)
419 + assert torch.equal(deployed_keys, optimized_keys)
420 + deployed_topk = torch.topk(deployed_keys, 16, dim=-1, sorted=True)
421 + optimized_topk = torch.topk(optimized_keys, 16, dim=-1, sorted=True)
422 + assert torch.equal(deployed_topk.values, optimized_topk.values)
423 + assert torch.equal(deployed_topk.indices, optimized_topk.indices)
424 +
425 + uniform = torch.full(
426 + (num_reqs,),
427 + max_block * BLOCK_SIZE,
428 + device="cuda",
429 + dtype=torch.int32,
430 + )
431 + one_active = torch.zeros(num_reqs, device="cuda", dtype=torch.int32)
432 + one_active[0] = max(1, max_block // 2) * BLOCK_SIZE
433 + patterns = [uniform, one_active]
434 + non_aligned = torch.tensor(
435 + [(max_block - 1) * BLOCK_SIZE + request % 3 + 1 for request in range(num_reqs)],
436 + device="cuda",
437 + dtype=torch.int32,
438 + )
439 + patterns.append(non_aligned)
440 +
441 + seq_lens.copy_(uniform)
442 + deployed_score.fill_(float("nan"))
443 + optimized_score.fill_(float("nan"))
444 + launch_deployed()
445 + launch_optimized()
446 + assert_exact(uniform)
447 +
448 + graph = torch.cuda.CUDAGraph()
449 + with torch.cuda.graph(graph):
450 + launch_optimized()
451 + addresses = (
452 + seq_lens.data_ptr(),
453 + optimized_score.data_ptr(),
454 + block_table.data_ptr(),
455 + index_kv_cache.data_ptr(),
456 + )
457 + for pattern in patterns:
458 + seq_lens.copy_(pattern)
459 + deployed_score.fill_(float("nan"))
460 + optimized_score.fill_(float("nan"))
461 + launch_deployed()
462 + graph.replay()
463 + assert_exact(pattern)
464 + assert addresses == (
465 + seq_lens.data_ptr(),
466 + optimized_score.data_ptr(),
467 + block_table.data_ptr(),
468 + index_kv_cache.data_ptr(),
469 + )
470 +
471 +
472 @pytest.mark.parametrize(
473 ("kv_cache_dtype", "expected_dtype"),
474 [
@@ -638,6 +1015,338 @@ def test_decode_index_topk_correctness(
1015 _assert_topk_indices_equal_unordered(actual, expected)
1016
1017
1018 +@pytest.mark.parametrize(
1019 + (
1020 + "max_block",
1021 + "total_q",
1022 + "num_idx_heads",
1023 + "topk",
1024 + "is_gfx950",
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)),
1039 + ],
1040 +)
1041 +def test_amd_decode_topk_launch_policy(
1042 + max_block: int,
1043 + total_q: int,
1044 + num_idx_heads: int,
1045 + topk: int,
1046 + is_gfx950: bool,
1047 + expected: tuple[int, bool, bool],
1048 +):
1049 + """All supported head and query layouts use adaptive selection."""
1050 + from vllm.models.minimax_m3.amd.ops.index_topk import (
1051 + _decode_topk_launch_policy,
1052 + )
1053 +
1054 + assert (
1055 + _decode_topk_launch_policy(
1056 + max_block,
1057 + total_q,
1058 + num_idx_heads,
1059 + topk,
1060 + is_gfx950=is_gfx950,
1061 + )
1062 + == expected
1063 + )
1064 +
1065 +
1066 +@pytest.mark.skipif(
1067 + not current_platform.is_rocm(),
1068 + reason="MiniMax-M3 fused decode top-k is ROCm-only",
1069 +)
1070 +@pytest.mark.parametrize(
1071 + (
1072 + "num_topk_chunks",
1073 + "single_tile_guaranteed",
1074 + "adaptive_final_merge",
1075 + "max_blocks",
1076 + "topk",
1077 + "num_idx_heads",
1078 + ),
1079 + [
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),
1089 + ],
1090 +)
1091 +def test_amd_decode_fused_topk_total_order_and_replay(
1092 + num_topk_chunks: int,
1093 + single_tile_guaranteed: bool,
1094 + adaptive_final_merge: bool,
1095 + max_blocks: int,
1096 + topk: int,
1097 + num_idx_heads: int,
1098 +):
1099 + """Atomic multi-chunk selection is ordered and resets graph state."""
1100 + from vllm.models.minimax_m3.amd.ops.index_topk import (
1101 + SPARSE_BLOCK_SIZE,
1102 + _decode_topk_fused_kernel,
1103 + )
1104 +
1105 + set_random_seed(0)
1106 + block_size_t = 1 << (topk - 1).bit_length()
1107 + long_seq_lens_list = [
1108 + 0,
1109 + 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,
1119 + max_blocks * SPARSE_BLOCK_SIZE,
1120 + ]
1121 + short_seq_lens_list = [0, 1, 127, 128, 129, 0, 1, 128, 129, 0, 1, 129]
1122 + boundary_limit = 0
1123 + if num_topk_chunks == 16 and max_blocks >= 8192:
1124 + boundary_limit = 15
1125 + elif num_topk_chunks == 4 and adaptive_final_merge:
1126 + boundary_limit = 4
1127 + if boundary_limit:
1128 + boundary_blocks = [
1129 + (active_chunks - 1) * 512 + 1
1130 + for active_chunks in range(3, boundary_limit + 1)
1131 + ]
1132 + assert [(num_blocks + 511) // 512 for num_blocks in boundary_blocks] == list(
1133 + range(3, boundary_limit + 1)
1134 + )
1135 + long_seq_lens_list.extend(
1136 + num_blocks * SPARSE_BLOCK_SIZE for num_blocks in boundary_blocks
1137 + )
1138 + short_seq_lens_list.extend(
1139 + num_blocks * SPARSE_BLOCK_SIZE for num_blocks in reversed(boundary_blocks)
1140 + )
1141 + batch = len(long_seq_lens_list)
1142 + score = torch.randn((num_idx_heads, batch, max_blocks), device="cuda")
1143 + score[:, :, 0] = float("nan")
1144 + score[:, :, 1] = float("inf")
1145 + score[:, :, 2] = float("-inf")
1146 + score[:, :, 3] = 0.0
1147 + score[:, :, 4] = -0.0
1148 + score[:, :, 5] = 7.0
1149 + score[:, :, 6] = -1e30
1150 + score[:, :, 256] = 7.0
1151 + score[:, :, 512] = 7.0
1152 +
1153 + seq_lens = torch.empty(batch, device="cuda", dtype=torch.int32)
1154 + partial = torch.full(
1155 + (num_topk_chunks, num_idx_heads, batch, block_size_t),
1156 + -1,
1157 + device="cuda",
1158 + dtype=torch.int64,
1159 + )
1160 + counter = torch.zeros((num_idx_heads, batch), device="cuda", dtype=torch.int32)
1161 + actual = torch.full(
1162 + (num_idx_heads, batch, topk), -2, device="cuda", dtype=torch.int32
1163 + )
1164 + dummy_block_table = torch.zeros(
1165 + (batch, max_blocks), device="cuda", dtype=torch.int32
1166 + )
1167 + addresses = (partial.data_ptr(), counter.data_ptr(), actual.data_ptr())
1168 + long_seq_lens = torch.tensor(long_seq_lens_list, device="cuda", dtype=torch.int32)
1169 + short_seq_lens = torch.tensor(short_seq_lens_list, device="cuda", dtype=torch.int32)
1170 +
1171 + def launch() -> None:
1172 + _decode_topk_fused_kernel[(batch, num_idx_heads, num_topk_chunks)](
1173 + score,
1174 + partial,
1175 + counter,
1176 + actual,
1177 + seq_lens,
1178 + dummy_block_table,
1179 + actual,
1180 + seq_lens,
1181 + 1,
1182 + score.stride(0),
1183 + score.stride(1),
1184 + score.stride(2),
1185 + partial.stride(0),
1186 + partial.stride(1),
1187 + partial.stride(2),
1188 + partial.stride(3),
1189 + counter.stride(0),
1190 + counter.stride(1),
1191 + actual.stride(0),
1192 + actual.stride(1),
1193 + actual.stride(2),
1194 + dummy_block_table.stride(0),
1195 + actual.stride(1),
1196 + topk=topk,
1197 + block_size=SPARSE_BLOCK_SIZE,
1198 + pages_per_sparse_block=8,
1199 + block_page_stride=8,
1200 + NUM_TOPK_CHUNKS=num_topk_chunks,
1201 + BLOCK_SIZE_K=512,
1202 + BLOCK_SIZE_T=block_size_t,
1203 + EMIT_SPARSE_TABLE=False,
1204 + SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
1205 + ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
1206 + num_warps=4 if adaptive_final_merge else 8,
1207 + num_stages=2,
1208 + )
1209 +
1210 + score_cpu = score.cpu()
1211 +
1212 + def reference(seq_lens_list: list[int]) -> torch.Tensor:
1213 + expected = torch.full((num_idx_heads, batch, topk), -1, dtype=torch.int32)
1214 + for head in range(num_idx_heads):
1215 + for row, seq_len in enumerate(seq_lens_list):
1216 + num_blocks = (seq_len + SPARSE_BLOCK_SIZE - 1) // SPARSE_BLOCK_SIZE
1217 + values = score_cpu[head, row, :num_blocks].tolist()
1218 + values = [-1e30 if value != value else value for value in values]
1219 + selected = sorted(
1220 + range(num_blocks),
1221 + key=lambda index: (-values[index], index),
1222 + )[:topk]
1223 + expected[head, row, : len(selected)] = torch.tensor(selected)
1224 + return expected
1225 +
1226 + seq_lens.copy_(long_seq_lens)
1227 + launch()
1228 + torch.accelerator.synchronize()
1229 + assert torch.equal(actual.cpu(), reference(long_seq_lens_list))
1230 + assert torch.count_nonzero(counter).item() == 0
1231 +
1232 + seq_lens.copy_(short_seq_lens)
1233 + launch()
1234 + torch.accelerator.synchronize()
1235 + assert torch.equal(actual.cpu(), reference(short_seq_lens_list))
1236 + assert torch.count_nonzero(counter).item() == 0
1237 +
1238 + replay_count = 8 if max_blocks == 65534 else 128
1239 + for replay in range(replay_count):
1240 + seq_lens.copy_(long_seq_lens if replay % 2 == 0 else short_seq_lens)
1241 + launch()
1242 + seq_lens.copy_(long_seq_lens)
1243 + launch()
1244 + torch.accelerator.synchronize()
1245 + assert torch.equal(actual.cpu(), reference(long_seq_lens_list))
1246 + assert torch.count_nonzero(counter).item() == 0
1247 + assert addresses == (partial.data_ptr(), counter.data_ptr(), actual.data_ptr())
1248 +
1249 + if (num_topk_chunks, single_tile_guaranteed, adaptive_final_merge, max_blocks) in (
1250 + (4, False, False, 8193),
1251 + (16, False, False, 8193),
1252 + (4, False, True, 8192),
1253 + (16, True, True, 8192),
1254 + ):
1255 + graph = torch.cuda.CUDAGraph()
1256 + torch.accelerator.synchronize()
1257 + with torch.cuda.graph(graph):
1258 + launch()
1259 + seq_lens.copy_(short_seq_lens)
1260 + graph.replay()
1261 + torch.accelerator.synchronize()
1262 + assert torch.equal(actual.cpu(), reference(short_seq_lens_list))
1263 + assert torch.count_nonzero(counter).item() == 0
1264 + seq_lens.copy_(long_seq_lens)
1265 + graph.replay()
1266 + torch.accelerator.synchronize()
1267 + assert torch.equal(actual.cpu(), reference(long_seq_lens_list))
1268 + assert torch.count_nonzero(counter).item() == 0
1269 + assert addresses == (partial.data_ptr(), counter.data_ptr(), actual.data_ptr())
1270 +
1271 +
1272 +@pytest.mark.skipif(
1273 + not current_platform.is_rocm(),
1274 + reason="MiniMax-M3 fused decode top-k is ROCm-only",
1275 +)
1276 +@pytest.mark.parametrize(
1277 + ("num_idx_heads", "decode_query_len", "max_decode_query_len"),
1278 + [(1, 1, 1), (2, 1, 4), (4, 4, 4), (4, 3, 8)],
1279 +)
1280 +def test_amd_decode_index_topk_end_to_end(
1281 + num_idx_heads: int,
1282 + decode_query_len: int,
1283 + max_decode_query_len: int,
1284 +):
1285 + """ROCm BF16 score-to-selector parity crosses the atomic boundary."""
1286 + from vllm.models.minimax_m3.amd.ops.index_topk import (
1287 + minimax_m3_index_decode as amd_minimax_m3_index_decode,
1288 + )
1289 +
1290 + torch.manual_seed(0)
1291 + topk = 16
1292 + head_dim = 16
1293 + seq_lens = torch.tensor([65537], device="cuda", dtype=torch.int32)
1294 + q_lens = torch.tensor([decode_query_len], device="cuda", dtype=torch.int32)
1295 + prefix_lens = seq_lens - q_lens
1296 + max_seq_len = int(seq_lens.max())
1297 + max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
1298 + block_table = torch.randperm(
1299 + max_blocks, device="cuda", dtype=torch.int32
1300 + ).unsqueeze(0)
1301 + idx_q = torch.randn(
1302 + decode_query_len,
1303 + num_idx_heads,
1304 + head_dim,
1305 + device="cuda",
1306 + dtype=torch.bfloat16,
1307 + )
1308 + index_kv_cache = torch.randn(
1309 + max_blocks,
1310 + BLOCK_SIZE,
1311 + head_dim,
1312 + device="cuda",
1313 + dtype=torch.bfloat16,
1314 + )
1315 + completion_counter = torch.zeros(
1316 + (num_idx_heads, decode_query_len), device="cuda", dtype=torch.int32
1317 + )
1318 +
1319 + actual = amd_minimax_m3_index_decode(
1320 + idx_q,
1321 + index_kv_cache,
1322 + block_table,
1323 + seq_lens,
1324 + max_seq_len=max_seq_len,
1325 + topk=topk,
1326 + init_blocks=0,
1327 + local_blocks=1,
1328 + num_kv_heads=num_idx_heads,
1329 + decode_query_len=decode_query_len,
1330 + max_decode_query_len=max_decode_query_len,
1331 + completion_counter=completion_counter,
1332 + )
1333 + expected = _reference_index_topk(
1334 + idx_q,
1335 + index_kv_cache,
1336 + block_table,
1337 + q_lens,
1338 + seq_lens,
1339 + prefix_lens,
1340 + topk,
1341 + 0,
1342 + 1,
1343 + )
1344 + torch.accelerator.synchronize()
1345 +
1346 + _assert_topk_indices_equal_unordered(actual, expected)
1347 + assert torch.count_nonzero(completion_counter).item() == 0
1348 +
1349 +
1350 @pytest.mark.skipif(
1351 not current_platform.is_device_capability_family(100),
1352 reason="fp8 e4m3 indexer cache is the SM100 (MSA) path.",
@@ -1250,6 +1959,137 @@ def test_aiter_decode_sparse_block_table_supports_spec_decode():
1959 assert torch.equal(sparse_bt, expected_bt)
1960
1961
1962 +@pytest.mark.skipif(
1963 + not current_platform.is_rocm(),
1964 + reason="MiniMax-M3 fused decode top-k is ROCm-only",
1965 +)
1966 +@pytest.mark.parametrize("block_page_stride", [8, 16])
1967 +@pytest.mark.parametrize(
1968 + "seq_lens_list",
1969 + [
1970 + [4097],
1971 + [65537],
1972 + [1, 129, 4097, 0],
1973 + [127, 128, 129, 256, 4097, 0],
1974 + [1, 127, 128, 129, 255, 256, 257, 4097, 0],
1975 + ],
1976 +)
1977 +def test_amd_decode_fused_topk_emits_sparse_table(
1978 + block_page_stride: int, seq_lens_list: list[int]
1979 +):
1980 + """The fused selector must match top-k followed by table construction."""
1981 + from vllm.models.minimax_m3.amd.ops.index_topk import (
1982 + minimax_m3_index_decode as amd_minimax_m3_index_decode,
1983 + )
1984 + from vllm.models.minimax_m3.amd.ops.sparse_pa import (
1985 + PAGES_PER_SPARSE_BLOCK,
1986 + minimax_m3_build_sparse_block_table_decode,
1987 + )
1988 +
1989 + torch.manual_seed(0)
1990 + decode_query_len = 4
1991 + seq_lens = torch.tensor(seq_lens_list, device="cuda", dtype=torch.int32)
1992 + batch = seq_lens.numel()
1993 + total_q = batch * decode_query_len
1994 + max_seq_len = int(seq_lens.max())
1995 + max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
1996 +
1997 + index_block_table = torch.zeros(
1998 + (batch, max_blocks), device="cuda", dtype=torch.int32
1999 + )
2000 + for req_id, seq_len in enumerate(seq_lens.tolist()):
2001 + num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
2002 + pages = torch.arange(max_blocks, device="cuda", dtype=torch.int32)
2003 + index_block_table[req_id, :num_blocks] = torch.roll(pages, req_id)[:num_blocks]
2004 +
2005 + # The index cache and main-attention cache have independent physical page
2006 + # mappings in production. Disjoint values make using the wrong table fail.
2007 + attention_block_table = torch.full_like(index_block_table, -31337)
2008 + for req_id, seq_len in enumerate(seq_lens.tolist()):
2009 + num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
2010 + first_page = 10_000 + req_id * max_blocks
2011 + attention_block_table[req_id, :num_blocks] = torch.arange(
2012 + first_page + num_blocks - 1,
2013 + first_page - 1,
2014 + -1,
2015 + device="cuda",
2016 + dtype=torch.int32,
2017 + )
2018 + assert not torch.equal(index_block_table, attention_block_table)
2019 +
2020 + idx_q = torch.randn(total_q, 1, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
2021 + index_kv_cache = torch.randn(
2022 + max_blocks,
2023 + BLOCK_SIZE,
2024 + HEAD_DIM,
2025 + device="cuda",
2026 + dtype=torch.bfloat16,
2027 + )
2028 + decode_kwargs = {
2029 + "max_seq_len": max_seq_len,
2030 + "topk": TOPK,
2031 + "init_blocks": 0,
2032 + "local_blocks": 1,
2033 + "num_kv_heads": 1,
2034 + "decode_query_len": decode_query_len,
2035 + "max_decode_query_len": decode_query_len,
2036 + }
2037 + completion_counter = torch.zeros((1, total_q), device="cuda", dtype=torch.int32)
2038 +
2039 + expected_topk = amd_minimax_m3_index_decode(
2040 + idx_q,
2041 + index_kv_cache,
2042 + index_block_table,
2043 + seq_lens,
2044 + completion_counter=completion_counter,
2045 + **decode_kwargs,
2046 + )
2047 + expected_sparse_bt, expected_sparse_ctx = (
2048 + minimax_m3_build_sparse_block_table_decode(
2049 + expected_topk,
2050 + attention_block_table,
2051 + seq_lens,
2052 + decode_query_len,
2053 + block_page_stride,
2054 + )
2055 + )
2056 +
2057 + topk_out = torch.full_like(expected_topk, -27182)
2058 + sparse_bt_out = torch.full(
2059 + (total_q, TOPK * PAGES_PER_SPARSE_BLOCK),
2060 + -27182,
2061 + device="cuda",
2062 + dtype=torch.int32,
2063 + )
2064 + sparse_ctx_out = torch.full((total_q,), -27182, device="cuda", dtype=torch.int32)
2065 + actual_topk = amd_minimax_m3_index_decode(
2066 + idx_q,
2067 + index_kv_cache,
2068 + index_block_table,
2069 + seq_lens,
2070 + out=topk_out,
2071 + attention_block_table=attention_block_table,
2072 + sparse_block_table_out=sparse_bt_out,
2073 + sparse_context_lens_out=sparse_ctx_out,
2074 + block_page_stride=block_page_stride,
2075 + completion_counter=completion_counter,
2076 + **decode_kwargs,
2077 + )
2078 + torch.accelerator.synchronize()
2079 +
2080 + assert actual_topk.data_ptr() == topk_out.data_ptr()
2081 + assert torch.equal(actual_topk, expected_topk)
2082 + assert torch.equal(sparse_bt_out, expected_sparse_bt)
2083 + assert torch.equal(sparse_ctx_out, expected_sparse_ctx)
2084 + assert torch.count_nonzero(completion_counter).item() == 0
2085 +
2086 + if seq_lens_list[-1] == 0:
2087 + padded_rows = slice(total_q - decode_query_len, total_q)
2088 + assert torch.all(actual_topk[:, padded_rows] == -1)
2089 + assert torch.count_nonzero(sparse_bt_out[padded_rows]) == 0
2090 + assert torch.count_nonzero(sparse_ctx_out[padded_rows]) == 0
2091 +
2092 +
2093 @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
2094 @pytest.mark.parametrize(
2095 "seq_lens_list",
vllm/models/minimax_m3/amd/model.py
+46 -1
@@ -92,6 +92,13 @@ from vllm.models.minimax_m3.amd.ops import (
92 gemma_rmsnorm,
93 swiglu_oai_split,
94 )
95 +from vllm.models.minimax_m3.amd.ops.sparse_pa import (
96 + minimax_m3_alloc_sparse_block_table,
97 + minimax_m3_sparse_block_page_stride,
98 +)
99 +from vllm.models.minimax_m3.amd.sparse_attention_msa import (
100 + MiniMaxM3SparseAiterPAImpl,
101 +)
102 from vllm.models.minimax_m3.common.indexer import MiniMaxM3Indexer
103 from vllm.models.minimax_m3.common.mm_preprocess import (
104 MiniMaxM3VLDummyInputsBuilder,
@@ -701,6 +708,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
708 self.kv_cache_k = torch.tensor([])
709 self.kv_cache_v = torch.tensor([])
710 self._aiter_sparse_pa_cache_data_ptr = 0
711 + self._aiter_sparse_pa_block_page_stride = 0
712 # Self-contained nn.Module: owns its side cache, selects its impl in init
713 # (Triton on ROCm, where the SM100 gate is always False).
714 self.indexer = MiniMaxM3Indexer(
@@ -796,6 +804,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
804 self.head_dim,
805 x,
806 )[v_page_offset:]
807 + self._aiter_sparse_pa_block_page_stride = minimax_m3_sparse_block_page_stride(
808 + self.kv_cache_k, self.kv_cache_v
809 + )
810 self._aiter_sparse_pa_cache_data_ptr = self.kv_cache.data_ptr()
811
812 def get_aiter_sparse_pa_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]:
@@ -1020,9 +1031,43 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
1031 # top-k into the shared ``topk_indices_buffer``; the attend reads it back.
1032 # When skip_index_topk is set (ATOM index_topk_freq), reuse the selection
1033 # the preceding compute layer wrote into the shared buffer this forward.
1034 + decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None
1035 if not self.skip_index_topk:
1036 assert index_query is not None
1025 - self.indexer(index_query)
1037 + attn_metadata = get_forward_context().attn_metadata
1038 + if self.use_aiter_sparse_pa and isinstance(attn_metadata, dict):
1039 + main_md = attn_metadata[self.layer_name]
1040 + assert isinstance(main_md, MiniMaxM3SparseMetadata)
1041 + if main_md.num_decodes > 0:
1042 + d = main_md.decode
1043 + assert d is not None
1044 + topk = self.topk_indices_buffer
1045 + assert topk is not None
1046 + decode_sparse_table = minimax_m3_alloc_sparse_block_table(
1047 + topk[:, : main_md.num_decode_tokens, :]
1048 + )
1049 + block_page_stride = self._aiter_sparse_pa_block_page_stride
1050 + assert block_page_stride > 0
1051 + self.indexer(
1052 + index_query,
1053 + attention_block_table=d.block_table,
1054 + sparse_block_table_out=decode_sparse_table[0],
1055 + sparse_context_lens_out=decode_sparse_table[1],
1056 + block_page_stride=block_page_stride,
1057 + )
1058 + else:
1059 + self.indexer(index_query)
1060 + else:
1061 + self.indexer(index_query)
1062 + if self.use_aiter_sparse_pa:
1063 + assert isinstance(self.impl, MiniMaxM3SparseAiterPAImpl)
1064 + return self.impl.forward(
1065 + self,
1066 + query,
1067 + self.kv_cache,
1068 + output,
1069 + decode_sparse_table=decode_sparse_table,
1070 + )
1071 return self.impl.forward(self, query, self.kv_cache, output)
1072
1073
vllm/models/minimax_m3/amd/ops/index_topk.py
+986 -306
@@ -7,21 +7,119 @@ then the top-k blocks (plus forced init/local blocks) are selected per query
7 token. Adapted to vLLM's paged KV cache: the KV page size is forced to equal the
8 sparse block size (128), so one sparse block maps to exactly one page.
9
10 -Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (single head).
10 +Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (one shared
11 +key vector per token).
12
13 Only the paths MiniMax M3 uses are implemented: score_type="max", index value
13 -disabled (score-only indexer), single shared index head. The selected block ids
14 -feed the block-sparse attention kernels in ``sparse_attn``.
14 +disabled (score-only indexer), and shared index keys. Each local index-query
15 +head selects its own block ids for the block-sparse attention kernels in
16 +``sparse_attn``.
17 """
18
19 import torch
20
21 +from vllm.models.minimax_m3.amd.ops.sparse_pa import (
22 + PAGES_PER_SPARSE_BLOCK,
23 + _write_sparse_block_table_row_from_values,
24 +)
25 from vllm.platforms import current_platform
26 from vllm.triton_utils import tl, triton
27 from vllm.utils.math_utils import round_up
28
29 # One sparse block == one KV page.
30 SPARSE_BLOCK_SIZE = 128
31 +DECODE_SCORE_BALANCED_PROGRAM_BUDGET = 1024
32 +DECODE_SCORE_HIGH_BATCH_PROGRAM_BUDGET = 768
33 +MAX_DECODE_SCORE_BALANCED_REQUESTS = 11
34 +DECODE_SCORE_DEFAULT_TARGET_GRID = 512
35 +DECODE_SCORE_GFX950_HIGH_BATCH_TARGET_GRID = 1024
36 +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 +MAX_DECODE_TOPK_FAST_CHUNKS = 16
41 +DECODE_TOPK_TARGET_GRID = 64
42 +
43 +
44 +def _decode_score_program_budget(
45 + num_reqs: int,
46 + head_dim: int,
47 + query_dtype: torch.dtype,
48 + cache_dtype: torch.dtype,
49 + *,
50 + is_gfx950: bool,
51 +) -> int | None:
52 + if not (
53 + is_gfx950
54 + and head_dim == 128
55 + and query_dtype == torch.bfloat16
56 + and cache_dtype == torch.bfloat16
57 + ):
58 + return None
59 + if 1 <= num_reqs <= 8:
60 + return DECODE_SCORE_BALANCED_PROGRAM_BUDGET
61 + if 9 <= num_reqs <= MAX_DECODE_SCORE_BALANCED_REQUESTS:
62 + return DECODE_SCORE_HIGH_BATCH_PROGRAM_BUDGET
63 + return None
64 +
65 +
66 +def _decode_score_split_launch_policy(
67 + num_reqs: int,
68 + head_dim: int,
69 + query_dtype: torch.dtype,
70 + cache_dtype: torch.dtype,
71 + *,
72 + is_gfx950: bool,
73 +) -> tuple[int, bool]:
74 + """Choose the generic split-K launch and high-batch specialization."""
75 + use_high_batch_config = (
76 + is_gfx950
77 + and MIN_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS
78 + <= num_reqs
79 + <= MAX_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS
80 + and head_dim == 128
81 + and query_dtype == torch.bfloat16
82 + and cache_dtype == torch.bfloat16
83 + )
84 + target_grid = (
85 + DECODE_SCORE_GFX950_HIGH_BATCH_TARGET_GRID
86 + if use_high_batch_config
87 + else DECODE_SCORE_DEFAULT_TARGET_GRID
88 + )
89 + target = max(
90 + 1,
91 + min(
92 + MAX_DECODE_SCORE_KV_CHUNKS,
93 + target_grid // max(1, num_reqs),
94 + ),
95 + )
96 + return 1 << (target.bit_length() - 1), use_high_batch_config
97 +
98 +
99 +def _decode_topk_launch_policy(
100 + max_block: int,
101 + total_q: int,
102 + num_idx_heads: int,
103 + topk: int,
104 + *,
105 + is_gfx950: bool,
106 +) -> tuple[int, bool, bool]:
107 + """Choose the graph grid and bounded-context selector specialization."""
108 + if (
109 + is_gfx950
110 + and topk == 16
111 + and 0 < max_block <= MAX_DECODE_TOPK_FAST_CHUNKS * DECODE_TOPK_BLOCKS_PER_CHUNK
112 + ):
113 + return MAX_DECODE_TOPK_FAST_CHUNKS, True, True
114 +
115 + target = max(
116 + 1,
117 + min(
118 + MAX_DECODE_TOPK_FAST_CHUNKS,
119 + DECODE_TOPK_TARGET_GRID // max(1, total_q * num_idx_heads),
120 + ),
121 + )
122 + return 1 << (target.bit_length() - 1), False, False
123
124
125 # ---------------------------------------------------------------------------
@@ -397,255 +495,670 @@ def _decode_index_score_kernel(
495 )
496
497
400 -# ---------------------------------------------------------------------------
401 -# Decode top-k (split-K): per-chunk partial top-k + merge. Forced init/local
402 -# blocks are already encoded in the scores.
403 -# ---------------------------------------------------------------------------
404 -@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])})
405 -@triton.autotune(
406 - configs=[
407 - triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2),
408 - triton.Config({"BLOCK_SIZE_K": 256}, num_warps=4, num_stages=2),
409 - triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2),
410 - triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=3),
411 - triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2),
412 - ],
413 - key=["topk"],
414 -)
415 -@triton.jit(do_not_specialize=["chunk_blocks", "decode_query_len"])
416 -def _topk_index_partial_kernel(
417 - s_ptr, # score: [num_idx_heads, total_q, max_block]
418 - ts_partial_ptr, # partial scores out: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T]
419 - ti_partial_ptr, # partial idx out (1-indexed global, 0=invalid): same shape
420 - seq_lens, # [num_reqs]
421 - block_size: tl.constexpr, # sparse block size (128)
422 - topk: tl.constexpr,
423 - chunk_blocks, # how many score-blocks each chunk owns
498 +@triton.jit
499 +def _ceil_div_nonnegative(value, divisor):
500 + return value // divisor + tl.where(value % divisor != 0, 1, 0)
501 +
502 +
503 +@triton.jit
504 +def _decode_index_score_mapped_range(
505 + q_ptr,
506 + ik_cache_ptr,
507 + score_ptr,
508 + block_table_ptr,
509 + pid_r,
510 + seq_len,
511 + chunk_start_block,
512 + blocks_per_program,
513 + num_idx_heads: tl.constexpr,
514 + head_dim: tl.constexpr,
515 + init_blocks,
516 + local_blocks,
517 decode_query_len,
518 + stride_q_n,
519 + stride_q_h,
520 + stride_q_d,
521 + stride_ik_blk,
522 + stride_ik_pos,
523 + stride_ik_d,
524 stride_s_h,
426 - stride_s_b,
525 + stride_s_n,
526 stride_s_k,
428 - stride_ts_c,
429 - stride_ts_h,
430 - stride_ts_b,
431 - stride_ts_t,
432 - stride_ti_c,
433 - stride_ti_h,
434 - stride_ti_b,
435 - stride_ti_t,
527 + stride_bt_b,
528 BLOCK_SIZE_K: tl.constexpr,
437 - BLOCK_SIZE_T: tl.constexpr,
438 - USE_PDL: tl.constexpr,
529 + BLOCK_SIZE_Q: tl.constexpr,
530 ):
440 - tl.static_assert(topk < BLOCK_SIZE_K)
441 - pid_b = tl.program_id(0) # flattened query-token id
442 - pid_h = tl.program_id(1)
443 - pid_chunk = tl.program_id(2)
444 - req_id = pid_b // decode_query_len
445 - q_offset = pid_b - req_id * decode_query_len
446 -
447 - if USE_PDL:
448 - tl.extra.cuda.gdc_wait()
531 + BLOCK_SIZE_HQ: tl.constexpr = num_idx_heads * BLOCK_SIZE_Q
532 + hq_offsets = tl.arange(0, BLOCK_SIZE_HQ)
533 + h_offsets = hq_offsets // BLOCK_SIZE_Q
534 + q_offsets = hq_offsets % BLOCK_SIZE_Q
535 + q_mask = q_offsets < decode_query_len
536 + q_ids = pid_r * decode_query_len + q_offsets
537
450 - seq_len = tl.load(seq_lens + req_id)
451 - query_pos = seq_len - decode_query_len + q_offset
452 - # Full-CG padding uses zero-length request rows. Clamp to an empty
453 - # attention range instead of letting padded rows produce negative lengths.
538 + query_pos = seq_len - decode_query_len + q_offsets
539 kv_len = tl.maximum(query_pos + 1, 0)
455 - num_blocks = (kv_len + block_size - 1) // block_size
456 -
457 - # Slice this chunk owns within [0, num_blocks).
458 - chunk_start = pid_chunk * chunk_blocks
459 - chunk_end = tl.minimum(chunk_start + chunk_blocks, num_blocks)
460 - chunk_actual = tl.maximum(chunk_end - chunk_start, 0)
461 -
462 - off_k = tl.arange(0, BLOCK_SIZE_K)
463 - off_t = tl.arange(0, BLOCK_SIZE_T)
464 -
465 - s_ptrs = (
466 - s_ptr
467 - + pid_b * stride_s_b
468 - + pid_h * stride_s_h
469 - + (chunk_start + off_k) * stride_s_k
540 + num_blocks_q = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
541 + kv_len_max = tl.max(tl.where(q_mask, kv_len, 0), axis=0)
542 + num_blocks = (kv_len_max + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
543 + chunk_end_block = tl.minimum(
544 + chunk_start_block + blocks_per_program,
545 + num_blocks,
546 )
547
472 - topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32)
473 - topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32)
474 - left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2
475 -
476 - # Streaming top-K within this chunk. tl.range(0, 0) is a no-op so empty
477 - # chunks (chunk_actual == 0) skip the body and store sentinel -1e30 / 0.
478 - for i in tl.range(0, chunk_actual, BLOCK_SIZE_K):
479 - mask = off_k < chunk_actual - i
480 - score = tl.load(s_ptrs, mask=mask, other=-1e30).to(tl.float32)
481 - score = tl.where(score != score, -1e30, score)
482 - s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K
483 - topk_score, last_topk_score = score, topk_score
484 - topk_idx, last_topk_idx = (
485 - tl.where(mask, chunk_start + i + off_k + 1, 0), # 1-indexed global
486 - topk_idx,
548 + if chunk_start_block < chunk_end_block:
549 + off_k = tl.arange(0, BLOCK_SIZE_K)
550 + off_d = tl.arange(0, head_dim)
551 + bt_row = block_table_ptr + pid_r * stride_bt_b
552 + local_start = tl.maximum(0, num_blocks_q - local_blocks)
553 + q = tl.load(
554 + q_ptr
555 + + q_ids[None, :] * stride_q_n
556 + + h_offsets[None, :] * stride_q_h
557 + + off_d[:, None] * stride_q_d,
558 + mask=q_mask[None, :],
559 + other=0.0,
560 )
488 - n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
489 - for j in tl.static_range(1, n_dims):
490 - topk_score, topk_idx = _bitonic_merge(
491 - topk_score, topk_idx.to(tl.int32), j, 2, n_dims
492 - )
493 - if i != 0:
494 - topk_score, topk_idx = _bitonic_merge(
495 - topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims
561 + for blk in tl.range(chunk_start_block, chunk_end_block):
562 + page = tl.load(bt_row + blk).to(tl.int64)
563 + pos = blk * BLOCK_SIZE_K + off_k
564 + pos_mask = pos[:, None] < kv_len[None, :]
565 + k = tl.load(
566 + ik_cache_ptr
567 + + page * stride_ik_blk
568 + + off_k[:, None] * stride_ik_pos
569 + + off_d * stride_ik_d,
570 )
497 - topk_score_new = last_topk_score * left_half_mask + topk_score * (
498 - 1 - left_half_mask
499 - )
500 - topk_idx_new = last_topk_idx * left_half_mask + topk_idx * (
501 - 1 - left_half_mask
571 + if BLOCK_SIZE_HQ == 1:
572 + q_vec = tl.sum(q, axis=1).to(tl.float32)
573 + kq = tl.sum(
574 + k.to(tl.float32) * q_vec[None, :],
575 + axis=1,
576 + )[:, None]
577 + else:
578 + kq = tl.dot(k, q, out_dtype=tl.float32)
579 + kq = tl.where(
580 + pos_mask & q_mask[None, :],
581 + kq,
582 + float("-inf"),
583 )
503 - topk_score, topk_idx = _bitonic_merge(
504 - topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims
584 + score = tl.max(kq, axis=0)
585 + is_visible_block = blk < num_blocks_q
586 + is_init = (blk < init_blocks) & is_visible_block
587 + is_local = (blk >= local_start) & is_visible_block
588 + score = tl.where(
589 + is_local,
590 + 1e29,
591 + tl.where(is_init, 1e30, score),
592 )
506 - else:
507 - topk_score, topk_idx = _bitonic_merge(
508 - topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims
593 + tl.store(
594 + score_ptr
595 + + h_offsets * stride_s_h
596 + + q_ids * stride_s_n
597 + + blk * stride_s_k,
598 + score,
599 + mask=q_mask,
600 )
601
511 - if USE_PDL:
512 - tl.extra.cuda.gdc_launch_dependents()
602
514 - # Extract first BLOCK_SIZE_T entries (top-K of this chunk after the sort).
515 - topk_mask_extract = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
516 - final_score = tl.sum(
517 - topk_mask_extract[:, None]
518 - * tl.reshape(topk_score, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
519 - axis=0,
520 - )
521 - final_idx = tl.sum(
522 - topk_mask_extract[:, None]
523 - * tl.reshape(topk_idx, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
524 - axis=0,
603 +@triton.jit(do_not_specialize=["score_program_budget", "decode_query_len"])
604 +def _decode_index_score_balanced_kernel(
605 + q_ptr,
606 + ik_cache_ptr,
607 + score_ptr,
608 + block_table_ptr,
609 + seq_lens,
610 + num_idx_heads: tl.constexpr,
611 + head_dim: tl.constexpr,
612 + init_blocks,
613 + local_blocks,
614 + NUM_REQUESTS: tl.constexpr,
615 + score_program_budget,
616 + decode_query_len,
617 + stride_q_n,
618 + stride_q_h,
619 + stride_q_d,
620 + stride_ik_blk,
621 + stride_ik_pos,
622 + stride_ik_d,
623 + stride_s_h,
624 + stride_s_n,
625 + stride_s_k,
626 + stride_bt_b,
627 + BLOCK_SIZE_K: tl.constexpr,
628 + BLOCK_SIZE_Q: tl.constexpr,
629 +):
630 + pid = tl.program_id(0)
631 + total_blocks = tl.zeros((), dtype=tl.int32)
632 + for req in tl.static_range(0, NUM_REQUESTS):
633 + req_seq_len = tl.maximum(tl.load(seq_lens + req), 0)
634 + total_blocks += _ceil_div_nonnegative(
635 + req_seq_len,
636 + BLOCK_SIZE_K,
637 + )
638 + blocks_per_program = tl.maximum(
639 + 1,
640 + _ceil_div_nonnegative(
641 + total_blocks,
642 + score_program_budget,
643 + ),
644 )
645
527 - # Always write all BLOCK_SIZE_T slots — invalid slots carry -1e30 / 0
528 - # sentinels and lose to real scores in the merge stage.
529 - ts_ptrs = (
530 - ts_partial_ptr
531 - + pid_chunk * stride_ts_c
532 - + pid_b * stride_ts_b
533 - + pid_h * stride_ts_h
534 - + off_t * stride_ts_t
646 + program_offset = tl.zeros((), dtype=tl.int32)
647 + pid_r = tl.zeros((), dtype=tl.int32)
648 + chunk_start_block = tl.zeros((), dtype=tl.int32)
649 + selected_seq_len = tl.zeros((), dtype=tl.int32)
650 + owns_work = tl.zeros((), dtype=tl.int32)
651 + for req in tl.static_range(0, NUM_REQUESTS):
652 + req_seq_len = tl.maximum(tl.load(seq_lens + req), 0)
653 + req_num_blocks = _ceil_div_nonnegative(
654 + req_seq_len,
655 + BLOCK_SIZE_K,
656 + )
657 + req_programs = _ceil_div_nonnegative(
658 + req_num_blocks,
659 + blocks_per_program,
660 + )
661 + owns_request = (pid >= program_offset) & (pid < program_offset + req_programs)
662 + pid_r = tl.where(owns_request, req, pid_r)
663 + local_program = tl.where(
664 + owns_request,
665 + pid - program_offset,
666 + 0,
667 + )
668 + chunk_start_block = tl.where(
669 + owns_request,
670 + local_program * blocks_per_program,
671 + chunk_start_block,
672 + )
673 + selected_seq_len = tl.where(
674 + owns_request,
675 + req_seq_len,
676 + selected_seq_len,
677 + )
678 + owns_work += owns_request.to(tl.int32)
679 + program_offset += req_programs
680 + if owns_work == 0:
681 + return
682 +
683 + _decode_index_score_mapped_range(
684 + q_ptr,
685 + ik_cache_ptr,
686 + score_ptr,
687 + block_table_ptr,
688 + pid_r,
689 + selected_seq_len,
690 + chunk_start_block,
691 + blocks_per_program,
692 + num_idx_heads,
693 + head_dim,
694 + init_blocks,
695 + local_blocks,
696 + decode_query_len,
697 + stride_q_n,
698 + stride_q_h,
699 + stride_q_d,
700 + stride_ik_blk,
701 + stride_ik_pos,
702 + stride_ik_d,
703 + stride_s_h,
704 + stride_s_n,
705 + stride_s_k,
706 + stride_bt_b,
707 + BLOCK_SIZE_K,
708 + BLOCK_SIZE_Q,
709 )
536 - ti_ptrs = (
537 - ti_partial_ptr
538 - + pid_chunk * stride_ti_c
539 - + pid_b * stride_ti_b
540 - + pid_h * stride_ti_h
541 - + off_t * stride_ti_t
710 +
711 +
712 +# ---------------------------------------------------------------------------
713 +# Decode top-k. Each CTA keeps a packed score/id top-k. Short rows emit
714 +# directly; long rows synchronize their live CTAs and the last arrival merges
715 +# the partial lists and optionally builds the sparse page table.
716 +# ---------------------------------------------------------------------------
717 +@triton.jit
718 +def _decode_topk_key(score, index, valid):
719 + score = tl.where(score != score, -1e30, score)
720 + score = tl.where(score == 0.0, 0.0, score)
721 + bits = score.to(tl.uint32, bitcast=True)
722 + ordered = bits ^ tl.where(bits >> 31 != 0, 0xFFFFFFFF, 0x80000000)
723 + tie = (0xFFFF - index).to(tl.int64)
724 + key = (1 << 48) | (ordered.to(tl.int64) << 16) | tie
725 + return tl.where(valid, key, 0)
726 +
727 +
728 +@triton.jit
729 +def _store_decode_topk(
730 + winners,
731 + topk_ptr,
732 + attention_block_table_ptr,
733 + sparse_bt_ptr,
734 + sparse_ctx_ptr,
735 + pid_b,
736 + pid_h,
737 + req_id,
738 + query_pos,
739 + num_blocks,
740 + topk: tl.constexpr,
741 + block_size: tl.constexpr,
742 + pages_per_sparse_block: tl.constexpr,
743 + block_page_stride: tl.constexpr,
744 + stride_topk_h,
745 + stride_topk_b,
746 + stride_topk_t,
747 + stride_attention_bt_b,
748 + stride_sparse_bt_b,
749 + BLOCK_SIZE_T: tl.constexpr,
750 + EMIT_SPARSE_TABLE: tl.constexpr,
751 +):
752 + off_t = tl.arange(0, BLOCK_SIZE_T)
753 + topk_idx = (0xFFFF - (winners & 0xFFFF)).to(tl.int32) - 1
754 + topk_idx = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx, -1)
755 + tl.store(
756 + topk_ptr
757 + + pid_h * stride_topk_h
758 + + pid_b * stride_topk_b
759 + + off_t * stride_topk_t,
760 + topk_idx,
761 + mask=off_t < topk,
762 )
543 - tl.store(ts_ptrs, final_score)
544 - tl.store(ti_ptrs, final_idx)
763 + if EMIT_SPARSE_TABLE:
764 + _write_sparse_block_table_row_from_values(
765 + topk_idx,
766 + attention_block_table_ptr + req_id * stride_attention_bt_b,
767 + sparse_bt_ptr + pid_b * stride_sparse_bt_b,
768 + sparse_ctx_ptr + pid_b,
769 + query_pos,
770 + topk,
771 + block_size,
772 + pages_per_sparse_block,
773 + block_page_stride,
774 + BLOCK_SIZE_T,
775 + )
776
777
547 -@triton.heuristics(
548 - {
549 - "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]),
550 - "BLOCK_SIZE_K": lambda args: triton.next_power_of_2(
551 - args["num_topk_chunks"] * triton.next_power_of_2(args["topk"])
552 - ),
553 - }
554 -)
555 -@triton.jit(do_not_specialize=["num_topk_chunks", "decode_query_len"])
556 -def _topk_index_merge_kernel(
557 - ts_partial_ptr, # partial scores: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T]
558 - ti_partial_ptr, # partial idx (1-indexed global, 0=invalid): same shape
559 - ti_final_ptr, # final idx (0-indexed, -1=invalid): [num_idx_heads, total_q, topk]
560 - seq_lens, # [num_reqs]
561 - block_size: tl.constexpr, # sparse block size (128)
778 +@triton.jit
779 +def _merge_store_decode_topk(
780 + partial_ptr,
781 + counter,
782 + topk_ptr,
783 + attention_block_table_ptr,
784 + sparse_bt_ptr,
785 + sparse_ctx_ptr,
786 + active_chunks,
787 + pid_b,
788 + pid_h,
789 + req_id,
790 + query_pos,
791 + num_blocks,
792 + stride_partial_c,
793 + stride_partial_h,
794 + stride_partial_b,
795 + stride_partial_t,
796 + stride_topk_h,
797 + stride_topk_b,
798 + stride_topk_t,
799 + stride_attention_bt_b,
800 + stride_sparse_bt_b,
801 topk: tl.constexpr,
802 + block_size: tl.constexpr,
803 + pages_per_sparse_block: tl.constexpr,
804 + block_page_stride: tl.constexpr,
805 + BLOCK_SIZE_T: tl.constexpr,
806 + MERGE_CHUNKS: tl.constexpr,
807 + EMIT_SPARSE_TABLE: tl.constexpr,
808 +):
809 + candidate_width: tl.constexpr = MERGE_CHUNKS * BLOCK_SIZE_T
810 + off_candidate = tl.arange(0, candidate_width)
811 + candidate_chunk = off_candidate // BLOCK_SIZE_T
812 + candidate_slot = off_candidate % BLOCK_SIZE_T
813 + candidates = tl.load(
814 + partial_ptr
815 + + candidate_chunk * stride_partial_c
816 + + pid_h * stride_partial_h
817 + + pid_b * stride_partial_b
818 + + candidate_slot * stride_partial_t,
819 + mask=candidate_chunk < active_chunks,
820 + other=0,
821 + volatile=True,
822 + )
823 + winners = tl.topk(candidates, BLOCK_SIZE_T)
824 + _store_decode_topk(
825 + winners,
826 + topk_ptr,
827 + attention_block_table_ptr,
828 + sparse_bt_ptr,
829 + sparse_ctx_ptr,
830 + pid_b,
831 + pid_h,
832 + req_id,
833 + query_pos,
834 + num_blocks,
835 + topk,
836 + block_size,
837 + pages_per_sparse_block,
838 + block_page_stride,
839 + stride_topk_h,
840 + stride_topk_b,
841 + stride_topk_t,
842 + stride_attention_bt_b,
843 + stride_sparse_bt_b,
844 + BLOCK_SIZE_T,
845 + EMIT_SPARSE_TABLE,
846 + )
847 + tl.atomic_xchg(counter, 0, sem="acq_rel", scope="gpu")
848 +
849 +
850 +@triton.jit(do_not_specialize=["decode_query_len"])
851 +def _decode_topk_fused_kernel(
852 + score_ptr, # [num_idx_heads, total_q, max_block]
853 + partial_ptr, # [NUM_TOPK_CHUNKS, num_idx_heads, total_q, BLOCK_SIZE_T]
854 + counter_ptr, # [num_idx_heads, total_q]
855 + topk_ptr, # [num_idx_heads, total_q, topk]
856 + seq_lens_ptr, # [num_reqs]
857 + attention_block_table_ptr, # [num_reqs, max_blocks]
858 + sparse_bt_ptr, # [total_q, topk * PAGES_PER_SPARSE_BLOCK]
859 + sparse_ctx_ptr, # [total_q]
860 decode_query_len,
564 - stride_ts_c,
565 - stride_ts_h,
566 - stride_ts_b,
567 - stride_ts_t,
568 - stride_ti_c,
569 - stride_ti_h,
570 - stride_ti_b,
571 - stride_ti_t,
572 - stride_tif_h,
573 - stride_tif_b,
574 - stride_tif_t,
575 - num_topk_chunks,
861 + stride_s_h,
862 + stride_s_b,
863 + stride_s_k,
864 + stride_partial_c,
865 + stride_partial_h,
866 + stride_partial_b,
867 + stride_partial_t,
868 + stride_counter_h,
869 + stride_counter_b,
870 + stride_topk_h,
871 + stride_topk_b,
872 + stride_topk_t,
873 + stride_attention_bt_b,
874 + stride_sparse_bt_b,
875 + topk: tl.constexpr,
876 + block_size: tl.constexpr,
877 + pages_per_sparse_block: tl.constexpr,
878 + block_page_stride: tl.constexpr,
879 + NUM_TOPK_CHUNKS: tl.constexpr,
880 BLOCK_SIZE_K: tl.constexpr,
881 BLOCK_SIZE_T: tl.constexpr,
578 - USE_PDL: tl.constexpr,
882 + EMIT_SPARSE_TABLE: tl.constexpr,
883 + SINGLE_TILE_GUARANTEED: tl.constexpr,
884 + ADAPTIVE_FINAL_MERGE: tl.constexpr,
885 ):
580 - pid_b = tl.program_id(0) # flattened query-token id
886 + tl.static_assert(topk <= BLOCK_SIZE_T)
887 + tl.static_assert(BLOCK_SIZE_T <= BLOCK_SIZE_K)
888 + pid_b = tl.program_id(0)
889 pid_h = tl.program_id(1)
890 + pid_chunk = tl.program_id(2)
891 req_id = pid_b // decode_query_len
892 q_offset = pid_b - req_id * decode_query_len
893
585 - if USE_PDL:
586 - tl.extra.cuda.gdc_wait()
587 - tl.extra.cuda.gdc_launch_dependents()
588 -
589 - seq_len = tl.load(seq_lens + req_id)
894 + seq_len = tl.load(seq_lens_ptr + req_id)
895 query_pos = seq_len - decode_query_len + q_offset
591 - # Full-CG padding uses zero-length request rows. Clamp to an empty
592 - # attention range instead of letting padded rows produce negative lengths.
896 kv_len = tl.maximum(query_pos + 1, 0)
897 num_blocks = (kv_len + block_size - 1) // block_size
898
596 - # Load NUM_TOPK_CHUNKS * BLOCK_SIZE_T candidates, padded to BLOCK_SIZE_K.
597 - # Candidate at flat position p comes from chunk = p // BLOCK_SIZE_T,
598 - # in_chunk = p % BLOCK_SIZE_T.
599 - off = tl.arange(0, BLOCK_SIZE_K)
600 - chunk_idx = off // BLOCK_SIZE_T
601 - in_chunk_idx = off % BLOCK_SIZE_T
602 - valid = chunk_idx < num_topk_chunks
603 -
604 - score_offset = (
605 - chunk_idx * stride_ts_c
606 - + pid_h * stride_ts_h
607 - + pid_b * stride_ts_b
608 - + in_chunk_idx * stride_ts_t
609 - )
610 - idx_offset = (
611 - chunk_idx * stride_ti_c
612 - + pid_h * stride_ti_h
613 - + pid_b * stride_ti_b
614 - + in_chunk_idx * stride_ti_t
615 - )
899 + desired_chunks = (num_blocks + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
900 + active_chunks = tl.maximum(1, tl.minimum(NUM_TOPK_CHUNKS, desired_chunks))
901 + if pid_chunk >= active_chunks:
902 + return
903
617 - score = tl.load(ts_partial_ptr + score_offset, mask=valid, other=-1e30).to(
618 - tl.float32
619 - )
620 - score = tl.where(score != score, -1e30, score)
621 - idx = tl.load(ti_partial_ptr + idx_offset, mask=valid, other=0).to(tl.int32)
622 -
623 - # Full bitonic descending sort of BLOCK_SIZE_K items.
624 - n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
625 - for j in tl.static_range(1, n_dims):
626 - score, idx = _bitonic_merge(score, idx.to(tl.int32), j, 2, n_dims)
627 - score, idx = _bitonic_merge(score, idx.to(tl.int32), n_dims, True, n_dims)
628 -
629 - # Extract first BLOCK_SIZE_T positions — these are the global top-K.
630 - extract_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
631 - topk_idx_final = tl.sum(
632 - extract_mask[:, None]
633 - * tl.reshape(idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
634 - axis=0,
635 - )
904 + blocks_per_chunk = num_blocks // active_chunks
905 + extra_chunks = num_blocks - blocks_per_chunk * active_chunks
906 + chunk_start = pid_chunk * blocks_per_chunk + tl.minimum(pid_chunk, extra_chunks)
907 + chunk_blocks = blocks_per_chunk + (pid_chunk < extra_chunks)
908
909 + off_k = tl.arange(0, BLOCK_SIZE_K)
910 off_t = tl.arange(0, BLOCK_SIZE_T)
638 - tif_ptrs = (
639 - ti_final_ptr
640 - + pid_h * stride_tif_h
641 - + pid_b * stride_tif_b
642 - + off_t * stride_tif_t
643 - )
644 - store_mask = off_t < topk
645 - topk_idx_final = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx_final, -1)
911 + if SINGLE_TILE_GUARANTEED:
912 + valid = off_k < chunk_blocks
913 + score = tl.load(
914 + score_ptr
915 + + pid_h * stride_s_h
916 + + pid_b * stride_s_b
917 + + (chunk_start + off_k) * stride_s_k,
918 + mask=valid,
919 + other=-1e30,
920 + ).to(tl.float32)
921 + index = (chunk_start + off_k + 1).to(tl.int32)
922 + local_topk = tl.topk(
923 + _decode_topk_key(score, index, valid),
924 + BLOCK_SIZE_T,
925 + )
926 + else:
927 + score_ptrs = (
928 + score_ptr
929 + + pid_h * stride_s_h
930 + + pid_b * stride_s_b
931 + + (chunk_start + off_k) * stride_s_k
932 + )
933 + valid = off_k < chunk_blocks
934 + score = tl.load(score_ptrs, mask=valid, other=-1e30).to(tl.float32)
935 + index = (chunk_start + off_k + 1).to(tl.int32)
936 + local_topk = tl.topk(
937 + _decode_topk_key(score, index, valid),
938 + BLOCK_SIZE_T,
939 + )
940 + score_ptrs += BLOCK_SIZE_K * stride_s_k
941 + for offset in tl.range(BLOCK_SIZE_K, chunk_blocks, BLOCK_SIZE_K):
942 + valid = off_k < chunk_blocks - offset
943 + score = tl.load(score_ptrs, mask=valid, other=-1e30).to(tl.float32)
944 + index = (chunk_start + offset + off_k + 1).to(tl.int32)
945 + tile_topk = tl.topk(
946 + _decode_topk_key(score, index, valid),
947 + BLOCK_SIZE_T,
948 + )
949 + local_topk = tl.topk(
950 + tl.cat(local_topk, tile_topk, can_reorder=True),
951 + BLOCK_SIZE_T,
952 + )
953 + score_ptrs += BLOCK_SIZE_K * stride_s_k
954 +
955 + if active_chunks == 1:
956 + _store_decode_topk(
957 + local_topk,
958 + topk_ptr,
959 + attention_block_table_ptr,
960 + sparse_bt_ptr,
961 + sparse_ctx_ptr,
962 + pid_b,
963 + pid_h,
964 + req_id,
965 + query_pos,
966 + num_blocks,
967 + topk,
968 + block_size,
969 + pages_per_sparse_block,
970 + block_page_stride,
971 + stride_topk_h,
972 + stride_topk_b,
973 + stride_topk_t,
974 + stride_attention_bt_b,
975 + stride_sparse_bt_b,
976 + BLOCK_SIZE_T,
977 + EMIT_SPARSE_TABLE,
978 + )
979 + return
980 +
981 tl.store(
647 - tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty), mask=store_mask
982 + partial_ptr
983 + + pid_chunk * stride_partial_c
984 + + pid_h * stride_partial_h
985 + + pid_b * stride_partial_b
986 + + off_t * stride_partial_t,
987 + local_topk,
988 + )
989 + counter = counter_ptr + pid_h * stride_counter_h + pid_b * stride_counter_b
990 + arrival = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu")
991 + if arrival != active_chunks - 1:
992 + return
993 +
994 + if ADAPTIVE_FINAL_MERGE:
995 + if active_chunks <= 2:
996 + _merge_store_decode_topk(
997 + partial_ptr,
998 + counter,
999 + topk_ptr,
1000 + attention_block_table_ptr,
1001 + sparse_bt_ptr,
1002 + sparse_ctx_ptr,
1003 + active_chunks,
1004 + pid_b,
1005 + pid_h,
1006 + req_id,
1007 + query_pos,
1008 + num_blocks,
1009 + stride_partial_c,
1010 + stride_partial_h,
1011 + stride_partial_b,
1012 + stride_partial_t,
1013 + stride_topk_h,
1014 + stride_topk_b,
1015 + stride_topk_t,
1016 + stride_attention_bt_b,
1017 + stride_sparse_bt_b,
1018 + topk,
1019 + block_size,
1020 + pages_per_sparse_block,
1021 + block_page_stride,
1022 + BLOCK_SIZE_T,
1023 + 2,
1024 + EMIT_SPARSE_TABLE,
1025 + )
1026 + return
1027 + if active_chunks <= 4:
1028 + _merge_store_decode_topk(
1029 + partial_ptr,
1030 + counter,
1031 + topk_ptr,
1032 + attention_block_table_ptr,
1033 + sparse_bt_ptr,
1034 + sparse_ctx_ptr,
1035 + active_chunks,
1036 + pid_b,
1037 + pid_h,
1038 + req_id,
1039 + query_pos,
1040 + num_blocks,
1041 + stride_partial_c,
1042 + stride_partial_h,
1043 + stride_partial_b,
1044 + stride_partial_t,
1045 + stride_topk_h,
1046 + stride_topk_b,
1047 + stride_topk_t,
1048 + stride_attention_bt_b,
1049 + stride_sparse_bt_b,
1050 + topk,
1051 + block_size,
1052 + pages_per_sparse_block,
1053 + block_page_stride,
1054 + BLOCK_SIZE_T,
1055 + 4,
1056 + EMIT_SPARSE_TABLE,
1057 + )
1058 + return
1059 + if active_chunks <= 8:
1060 + _merge_store_decode_topk(
1061 + partial_ptr,
1062 + counter,
1063 + topk_ptr,
1064 + attention_block_table_ptr,
1065 + sparse_bt_ptr,
1066 + sparse_ctx_ptr,
1067 + active_chunks,
1068 + pid_b,
1069 + pid_h,
1070 + req_id,
1071 + query_pos,
1072 + num_blocks,
1073 + stride_partial_c,
1074 + stride_partial_h,
1075 + stride_partial_b,
1076 + stride_partial_t,
1077 + stride_topk_h,
1078 + stride_topk_b,
1079 + stride_topk_t,
1080 + stride_attention_bt_b,
1081 + stride_sparse_bt_b,
1082 + topk,
1083 + block_size,
1084 + pages_per_sparse_block,
1085 + block_page_stride,
1086 + BLOCK_SIZE_T,
1087 + 8,
1088 + EMIT_SPARSE_TABLE,
1089 + )
1090 + return
1091 + _merge_store_decode_topk(
1092 + partial_ptr,
1093 + counter,
1094 + topk_ptr,
1095 + attention_block_table_ptr,
1096 + sparse_bt_ptr,
1097 + sparse_ctx_ptr,
1098 + active_chunks,
1099 + pid_b,
1100 + pid_h,
1101 + req_id,
1102 + query_pos,
1103 + num_blocks,
1104 + stride_partial_c,
1105 + stride_partial_h,
1106 + stride_partial_b,
1107 + stride_partial_t,
1108 + stride_topk_h,
1109 + stride_topk_b,
1110 + stride_topk_t,
1111 + stride_attention_bt_b,
1112 + stride_sparse_bt_b,
1113 + topk,
1114 + block_size,
1115 + pages_per_sparse_block,
1116 + block_page_stride,
1117 + BLOCK_SIZE_T,
1118 + 16,
1119 + EMIT_SPARSE_TABLE,
1120 + )
1121 + return
1122 +
1123 + candidate_width: tl.constexpr = NUM_TOPK_CHUNKS * BLOCK_SIZE_T
1124 + off_candidate = tl.arange(0, candidate_width)
1125 + candidate_chunk = off_candidate // BLOCK_SIZE_T
1126 + candidate_slot = off_candidate % BLOCK_SIZE_T
1127 + candidates = tl.load(
1128 + partial_ptr
1129 + + candidate_chunk * stride_partial_c
1130 + + pid_h * stride_partial_h
1131 + + pid_b * stride_partial_b
1132 + + candidate_slot * stride_partial_t,
1133 + mask=candidate_chunk < active_chunks,
1134 + other=0,
1135 + volatile=True,
1136 )
1137 + winners = tl.topk(candidates, BLOCK_SIZE_T)
1138 + _store_decode_topk(
1139 + winners,
1140 + topk_ptr,
1141 + attention_block_table_ptr,
1142 + sparse_bt_ptr,
1143 + sparse_ctx_ptr,
1144 + pid_b,
1145 + pid_h,
1146 + req_id,
1147 + query_pos,
1148 + num_blocks,
1149 + topk,
1150 + block_size,
1151 + pages_per_sparse_block,
1152 + block_page_stride,
1153 + stride_topk_h,
1154 + stride_topk_b,
1155 + stride_topk_t,
1156 + stride_attention_bt_b,
1157 + stride_sparse_bt_b,
1158 + BLOCK_SIZE_T,
1159 + EMIT_SPARSE_TABLE,
1160 + )
1161 + tl.atomic_xchg(counter, 0, sem="acq_rel", scope="gpu")
1162
1163
1164 # ---------------------------------------------------------------------------
@@ -777,21 +1290,118 @@ def minimax_m3_index_decode(
1290 decode_query_len: int,
1291 max_decode_query_len: int,
1292 out: torch.Tensor | None = None,
1293 + *,
1294 + attention_block_table: torch.Tensor | None = None,
1295 + sparse_block_table_out: torch.Tensor | None = None,
1296 + sparse_context_lens_out: torch.Tensor | None = None,
1297 + block_page_stride: int | None = None,
1298 + completion_counter: torch.Tensor | None = None,
1299 ) -> torch.Tensor:
781 - """Decode index block-score + top-k, both split-K (cudagraph-safe).
1300 + """Decode index block-score followed by fused adaptive top-k selection.
1301
1302 Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).
1303 When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into
1304 ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating.
1305 + The optional sparse-table arguments fuse current-layer table construction
1306 + into the selector. They must be provided together. ``completion_counter``
1307 + provides stable per-query synchronization storage for CUDA graphs. It must
1308 + be zero before its first launch and must not be shared by overlapping
1309 + selector invocations; every completed launch resets its active entries.
1310 """
1311 total_q, num_idx_heads, head_dim = idx_q.shape
1312 assert num_idx_heads == num_kv_heads, (
1313 "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
1314 )
791 - assert decode_query_len <= max_decode_query_len
1315 + assert 0 < decode_query_len <= max_decode_query_len
1316 assert total_q == seq_lens.shape[0] * decode_query_len
1317 batch = total_q
1318 + emit_sparse_table = attention_block_table is not None
1319 + if emit_sparse_table and (
1320 + sparse_block_table_out is None
1321 + or sparse_context_lens_out is None
1322 + or block_page_stride is None
1323 + ):
1324 + raise ValueError(
1325 + "MiniMax-M3 fused decode sparse-table arguments must be provided together"
1326 + )
1327 + if not emit_sparse_table and (
1328 + sparse_block_table_out is not None
1329 + or sparse_context_lens_out is not None
1330 + or block_page_stride is not None
1331 + ):
1332 + raise ValueError(
1333 + "MiniMax-M3 fused decode sparse-table arguments must be provided together"
1334 + )
1335 + if emit_sparse_table:
1336 + assert attention_block_table is not None
1337 + assert sparse_block_table_out is not None
1338 + assert sparse_context_lens_out is not None
1339 + assert block_page_stride is not None
1340 + if num_idx_heads != 1:
1341 + raise ValueError(
1342 + "MiniMax-M3 fused decode sparse-table construction requires "
1343 + f"one index head, got {num_idx_heads}"
1344 + )
1345 + if block_page_stride not in (
1346 + PAGES_PER_SPARSE_BLOCK,
1347 + 2 * PAGES_PER_SPARSE_BLOCK,
1348 + ):
1349 + raise ValueError(
1350 + "MiniMax-M3 fused decode sparse-table page stride must be "
1351 + f"{PAGES_PER_SPARSE_BLOCK} or {2 * PAGES_PER_SPARSE_BLOCK}, "
1352 + f"got {block_page_stride}"
1353 + )
1354 + expected_sparse_shape = (total_q, topk * PAGES_PER_SPARSE_BLOCK)
1355 + if sparse_block_table_out.shape != expected_sparse_shape:
1356 + raise ValueError(
1357 + "MiniMax-M3 fused decode sparse block table has shape "
1358 + f"{tuple(sparse_block_table_out.shape)}, expected "
1359 + f"{expected_sparse_shape}"
1360 + )
1361 + if sparse_context_lens_out.shape != (total_q,):
1362 + raise ValueError(
1363 + "MiniMax-M3 fused decode sparse context lengths have shape "
1364 + f"{tuple(sparse_context_lens_out.shape)}, expected {(total_q,)}"
1365 + )
1366 + if (
1367 + attention_block_table.dim() != 2
1368 + or attention_block_table.shape[0] != seq_lens.shape[0]
1369 + ):
1370 + raise ValueError(
1371 + "MiniMax-M3 fused decode attention block table must be "
1372 + "[num_requests, max_blocks]"
1373 + )
1374 + int32_tensors = (
1375 + attention_block_table,
1376 + sparse_block_table_out,
1377 + sparse_context_lens_out,
1378 + )
1379 + if any(tensor.dtype != torch.int32 for tensor in int32_tensors):
1380 + raise ValueError("MiniMax-M3 fused decode sparse tables require int32")
1381 + if any(tensor.device != idx_q.device for tensor in int32_tensors):
1382 + raise ValueError(
1383 + "MiniMax-M3 fused decode sparse tables must share the query device"
1384 + )
1385 + if (
1386 + attention_block_table.stride(1) != 1
1387 + or sparse_block_table_out.stride(1) != 1
1388 + or sparse_context_lens_out.stride(0) != 1
1389 + ):
1390 + raise ValueError(
1391 + "MiniMax-M3 fused decode sparse tables require contiguous rows"
1392 + )
1393 max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
1394 + if max_block >= 0xFFFF:
1395 + raise ValueError(
1396 + "MiniMax-M3 decode top-k supports fewer than 65535 sparse blocks"
1397 + )
1398 + if emit_sparse_table:
1399 + assert attention_block_table is not None
1400 + if attention_block_table.shape[1] < max_block:
1401 + raise ValueError(
1402 + "MiniMax-M3 fused decode attention block table is shorter "
1403 + f"than the required {max_block} sparse blocks"
1404 + )
1405 use_pdl = current_platform.is_arch_support_pdl()
1406 # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
1407 # SM9+); this ROCm Triton rejects it even when False ("Keyword argument
@@ -800,9 +1410,13 @@ def minimax_m3_index_decode(
1410 pdl_kwargs: dict[str, bool | int] = {}
1411 if use_pdl:
1412 pdl_kwargs.update({"launch_pdl": True})
803 - # TP=1 spec decode scores a wide 4-head x 4-position query tile per K block;
804 - # reduce stages to ease memory/register pressure. Keep no-spec and TP=4
805 - # single-head codegen unchanged.
1413 + is_gfx950 = False
1414 + if current_platform.is_rocm():
1415 + from vllm.platforms.rocm import on_gfx950
1416 +
1417 + is_gfx950 = on_gfx950()
1418 + # Multi-head spec decode scores a wider head-position tile per K block;
1419 + # reduce stages to ease memory/register pressure on the fallback path.
1420 score_kwargs = pdl_kwargs.copy()
1421 if num_idx_heads > 1 and max_decode_query_len > 1:
1422 score_kwargs.update({"num_warps": 4, "num_stages": 2})
@@ -814,47 +1428,90 @@ def minimax_m3_index_decode(
1428 dtype=torch.float32,
1429 device=idx_q.device,
1430 )
817 - # split-K over seq blocks; chunk count depends only on shape constants so
818 - # the grid is fixed within a cuda graph.
819 - TARGET_GRID = 512
820 - MAX_NUM_KV_CHUNKS = 256
1431 # Use the configured max decode length to avoid Triton recompiles when
1432 # switching between qlen=1 and spec-decode verification batches.
1433 BLOCK_SIZE_Q = triton.next_power_of_2(max_decode_query_len)
824 - score_ctas_per_chunk = seq_lens.shape[0]
825 - target = max(
826 - 1,
827 - min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, score_ctas_per_chunk)),
828 - )
829 - num_kv_chunks = 1 << (target.bit_length() - 1)
830 - grid_score = (seq_lens.shape[0], num_kv_chunks)
831 - _decode_index_score_kernel[grid_score](
832 - idx_q,
833 - index_kv_cache,
834 - score,
835 - block_table,
836 - seq_lens,
837 - num_idx_heads,
1434 + num_reqs = seq_lens.shape[0]
1435 + score_program_budget = _decode_score_program_budget(
1436 + num_reqs,
1437 head_dim,
839 - init_blocks,
840 - local_blocks,
841 - decode_query_len,
842 - idx_q.stride(0),
843 - idx_q.stride(1),
844 - idx_q.stride(2),
845 - index_kv_cache.stride(0),
846 - index_kv_cache.stride(1),
847 - index_kv_cache.stride(2),
848 - score.stride(0),
849 - score.stride(1),
850 - score.stride(2),
851 - block_table.stride(0),
852 - BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
853 - BLOCK_SIZE_Q=BLOCK_SIZE_Q,
854 - num_kv_chunks=num_kv_chunks,
855 - USE_PDL=use_pdl,
856 - **score_kwargs,
1438 + idx_q.dtype,
1439 + index_kv_cache.dtype,
1440 + is_gfx950=is_gfx950,
1441 )
1442 + grid_score: tuple[int, ...]
1443 + if score_program_budget is not None:
1444 + grid_score = (score_program_budget + num_reqs - 1,)
1445 + _decode_index_score_balanced_kernel[grid_score](
1446 + idx_q,
1447 + index_kv_cache,
1448 + score,
1449 + block_table,
1450 + seq_lens,
1451 + num_idx_heads,
1452 + head_dim,
1453 + init_blocks,
1454 + local_blocks,
1455 + num_reqs,
1456 + score_program_budget,
1457 + decode_query_len,
1458 + idx_q.stride(0),
1459 + idx_q.stride(1),
1460 + idx_q.stride(2),
1461 + index_kv_cache.stride(0),
1462 + index_kv_cache.stride(1),
1463 + index_kv_cache.stride(2),
1464 + score.stride(0),
1465 + score.stride(1),
1466 + score.stride(2),
1467 + block_table.stride(0),
1468 + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
1469 + BLOCK_SIZE_Q=BLOCK_SIZE_Q,
1470 + num_warps=2,
1471 + num_stages=1,
1472 + )
1473 + else:
1474 + # Increase independent work for the measured high-batch gfx950 decode
1475 + # range while preserving the deployed split for every other shape.
1476 + num_kv_chunks, use_high_batch_config = _decode_score_split_launch_policy(
1477 + num_reqs,
1478 + head_dim,
1479 + idx_q.dtype,
1480 + index_kv_cache.dtype,
1481 + is_gfx950=is_gfx950,
1482 + )
1483 + if use_high_batch_config and not (
1484 + num_idx_heads > 1 and max_decode_query_len > 1
1485 + ):
1486 + score_kwargs.update({"num_warps": 2, "num_stages": 1})
1487 + grid_score = (num_reqs, num_kv_chunks)
1488 + _decode_index_score_kernel[grid_score](
1489 + idx_q,
1490 + index_kv_cache,
1491 + score,
1492 + block_table,
1493 + seq_lens,
1494 + num_idx_heads,
1495 + head_dim,
1496 + init_blocks,
1497 + local_blocks,
1498 + decode_query_len,
1499 + idx_q.stride(0),
1500 + idx_q.stride(1),
1501 + idx_q.stride(2),
1502 + index_kv_cache.stride(0),
1503 + index_kv_cache.stride(1),
1504 + index_kv_cache.stride(2),
1505 + score.stride(0),
1506 + score.stride(1),
1507 + score.stride(2),
1508 + block_table.stride(0),
1509 + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
1510 + BLOCK_SIZE_Q=BLOCK_SIZE_Q,
1511 + num_kv_chunks=num_kv_chunks,
1512 + USE_PDL=use_pdl,
1513 + **score_kwargs,
1514 + )
1515
1516 if out is not None:
1517 topk_idx = out[:, :total_q, :]
@@ -864,76 +1521,99 @@ def minimax_m3_index_decode(
1521 dtype=torch.int32,
1522 device=idx_q.device,
1523 )
867 - # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts
868 - # pow2(num_topk_chunks * pow2(topk)) candidates.
869 - TOPK_TARGET_GRID = 64
870 - MAX_NUM_TOPK_CHUNKS = 16
871 - topk_target = max(
872 - 1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads))
873 - )
874 - num_topk_chunks = 1 << (topk_target.bit_length() - 1)
875 - block_size_t = triton.next_power_of_2(topk)
876 - chunk_blocks = (max_block + num_topk_chunks - 1) // num_topk_chunks
877 - topk_score_partial = torch.empty(
1524 + # The launch grid remains shape-constant for CUDA graphs. Each query uses
1525 + # only the number of chunks needed for its live context.
1526 + (
1527 num_topk_chunks,
879 - num_idx_heads,
1528 + single_tile_guaranteed,
1529 + adaptive_final_merge,
1530 + ) = _decode_topk_launch_policy(
1531 + max_block,
1532 batch,
881 - block_size_t,
882 - dtype=torch.float32,
883 - device=idx_q.device,
1533 + num_idx_heads,
1534 + topk,
1535 + is_gfx950=is_gfx950,
1536 )
885 - topk_idx_partial = torch.empty(
1537 + block_size_t = triton.next_power_of_2(topk)
1538 + topk_partial = torch.empty(
1539 num_topk_chunks,
1540 num_idx_heads,
1541 batch,
1542 block_size_t,
890 - dtype=torch.int32,
1543 + dtype=torch.int64,
1544 device=idx_q.device,
1545 )
893 - _topk_index_partial_kernel[(batch, num_idx_heads, num_topk_chunks)](
1546 + if completion_counter is None:
1547 + active_counter = torch.zeros(
1548 + (num_idx_heads, batch),
1549 + dtype=torch.int32,
1550 + device=idx_q.device,
1551 + )
1552 + else:
1553 + if (
1554 + completion_counter.dim() != 2
1555 + or completion_counter.shape[0] != num_idx_heads
1556 + or completion_counter.shape[1] < batch
1557 + or completion_counter.dtype != torch.int32
1558 + or completion_counter.device != idx_q.device
1559 + or not completion_counter.is_contiguous()
1560 + ):
1561 + raise ValueError(
1562 + "MiniMax-M3 completion counter must be contiguous int32 "
1563 + "[num_idx_heads, >=total_q] on the query device"
1564 + )
1565 + active_counter = completion_counter[:, :batch]
1566 +
1567 + selector_attention_block_table = block_table
1568 + selector_sparse_block_table = topk_idx
1569 + selector_sparse_context_lens = seq_lens
1570 + selector_block_page_stride = PAGES_PER_SPARSE_BLOCK
1571 + selector_sparse_block_stride = topk_idx.stride(1)
1572 + if emit_sparse_table:
1573 + assert attention_block_table is not None
1574 + assert sparse_block_table_out is not None
1575 + assert sparse_context_lens_out is not None
1576 + assert block_page_stride is not None
1577 + selector_attention_block_table = attention_block_table
1578 + selector_sparse_block_table = sparse_block_table_out
1579 + selector_sparse_context_lens = sparse_context_lens_out
1580 + selector_block_page_stride = block_page_stride
1581 + selector_sparse_block_stride = sparse_block_table_out.stride(0)
1582 + _decode_topk_fused_kernel[(batch, num_idx_heads, num_topk_chunks)](
1583 score,
895 - topk_score_partial,
896 - topk_idx_partial,
1584 + topk_partial,
1585 + active_counter,
1586 + topk_idx,
1587 seq_lens,
898 - SPARSE_BLOCK_SIZE,
899 - topk,
900 - chunk_blocks,
1588 + selector_attention_block_table,
1589 + selector_sparse_block_table,
1590 + selector_sparse_context_lens,
1591 decode_query_len,
1592 score.stride(0),
1593 score.stride(1),
1594 score.stride(2),
905 - topk_score_partial.stride(0),
906 - topk_score_partial.stride(1),
907 - topk_score_partial.stride(2),
908 - topk_score_partial.stride(3),
909 - topk_idx_partial.stride(0),
910 - topk_idx_partial.stride(1),
911 - topk_idx_partial.stride(2),
912 - topk_idx_partial.stride(3),
913 - USE_PDL=use_pdl,
914 - **pdl_kwargs,
915 - )
916 - _topk_index_merge_kernel[(batch, num_idx_heads)](
917 - topk_score_partial,
918 - topk_idx_partial,
919 - topk_idx,
920 - seq_lens,
921 - SPARSE_BLOCK_SIZE,
922 - topk,
923 - decode_query_len,
924 - topk_score_partial.stride(0),
925 - topk_score_partial.stride(1),
926 - topk_score_partial.stride(2),
927 - topk_score_partial.stride(3),
928 - topk_idx_partial.stride(0),
929 - topk_idx_partial.stride(1),
930 - topk_idx_partial.stride(2),
931 - topk_idx_partial.stride(3),
1595 + topk_partial.stride(0),
1596 + topk_partial.stride(1),
1597 + topk_partial.stride(2),
1598 + topk_partial.stride(3),
1599 + active_counter.stride(0),
1600 + active_counter.stride(1),
1601 topk_idx.stride(0),
1602 topk_idx.stride(1),
1603 topk_idx.stride(2),
935 - num_topk_chunks=num_topk_chunks,
936 - USE_PDL=use_pdl,
937 - **pdl_kwargs,
1604 + selector_attention_block_table.stride(0),
1605 + selector_sparse_block_stride,
1606 + topk=topk,
1607 + block_size=SPARSE_BLOCK_SIZE,
1608 + pages_per_sparse_block=PAGES_PER_SPARSE_BLOCK,
1609 + block_page_stride=selector_block_page_stride,
1610 + NUM_TOPK_CHUNKS=num_topk_chunks,
1611 + BLOCK_SIZE_K=512,
1612 + BLOCK_SIZE_T=block_size_t,
1613 + EMIT_SPARSE_TABLE=emit_sparse_table,
1614 + SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
1615 + ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
1616 + num_warps=4 if adaptive_final_merge else 8,
1617 + num_stages=2,
1618 )
1619 return topk_idx
vllm/models/minimax_m3/amd/ops/sparse_pa.py
+60 -19
@@ -32,14 +32,13 @@ def _is_fp8_kv_cache_tensor(kv_cache: torch.Tensor) -> bool:
32
33
34 @triton.jit
35 -def _write_sparse_block_table_row(
36 - topk_row, # [topk] int32, selected logical block ids for this query
35 +def _write_sparse_block_table_row_from_values(
36 + blk, # [BLOCK_SIZE_T] int32 selected logical block ids for this query
37 bt_row, # [max_blocks] int32, this request's logical page table
38 sbt_row, # [topk * PAGES_PER_BLOCK] int32, physical 16-page table
39 ctx_ptr, # int32, this query's attended context length
40 abs_pos, # absolute position of this query token, may be negative
41 max_topk,
42 - stride_topk_k,
42 SPARSE_BLOCK_SIZE_C: tl.constexpr,
43 PAGES_PER_BLOCK: tl.constexpr,
44 BLOCK_PAGE_STRIDE: tl.constexpr,
@@ -56,8 +55,7 @@ def _write_sparse_block_table_row(
55 self_blk = abs_pos // SPARSE_BLOCK_SIZE_C
56
57 off_t = tl.arange(0, BLOCK_SIZE_T)
59 - blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
60 - valid = (causal_len > 0) & (blk >= 0) & (blk <= self_blk)
58 + valid = (off_t < max_topk) & (causal_len > 0) & (blk >= 0) & (blk <= self_blk)
59 is_tail = valid & (blk == self_blk)
60 is_full = valid & (blk < self_blk)
61
@@ -89,6 +87,36 @@ def _write_sparse_block_table_row(
87 tl.store(ctx_ptr, ctx)
88
89
90 +@triton.jit
91 +def _write_sparse_block_table_row(
92 + topk_row, # [topk] int32, selected logical block ids for this query
93 + bt_row, # [max_blocks] int32, this request's logical page table
94 + sbt_row, # [topk * PAGES_PER_BLOCK] int32, physical 16-page table
95 + ctx_ptr, # int32, this query's attended context length
96 + abs_pos, # absolute position of this query token, may be negative
97 + max_topk,
98 + stride_topk_k,
99 + SPARSE_BLOCK_SIZE_C: tl.constexpr,
100 + PAGES_PER_BLOCK: tl.constexpr,
101 + BLOCK_PAGE_STRIDE: tl.constexpr,
102 + BLOCK_SIZE_T: tl.constexpr,
103 +):
104 + off_t = tl.arange(0, BLOCK_SIZE_T)
105 + blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
106 + _write_sparse_block_table_row_from_values(
107 + blk,
108 + bt_row,
109 + sbt_row,
110 + ctx_ptr,
111 + abs_pos,
112 + max_topk,
113 + SPARSE_BLOCK_SIZE_C,
114 + PAGES_PER_BLOCK,
115 + BLOCK_PAGE_STRIDE,
116 + BLOCK_SIZE_T,
117 + )
118 +
119 +
120 @triton.jit
121 def _build_sparse_block_table_kernel(
122 topk_ptr, # [1, batch * DECODE_QUERY_LEN, topk] int32, logical block ids
@@ -169,7 +197,7 @@ def _build_sparse_block_table_prefill_kernel(
197 )
198
199
172 -def _alloc_sparse_block_table(
200 +def minimax_m3_alloc_sparse_block_table(
201 topk_idx: torch.Tensor,
202 ) -> tuple[torch.Tensor, torch.Tensor]:
203 assert topk_idx.shape[0] == 1, "AITER sparse PA requires num_kv_heads == 1"
@@ -190,7 +218,7 @@ def minimax_m3_build_sparse_block_table_prefill(
218 block_page_stride: int = PAGES_PER_SPARSE_BLOCK,
219 ) -> tuple[torch.Tensor, torch.Tensor]:
220 """Build one page-16 sparse block table row per prefill query token."""
193 - sparse_bt, sparse_ctx = _alloc_sparse_block_table(topk_idx)
221 + sparse_bt, sparse_ctx = minimax_m3_alloc_sparse_block_table(topk_idx)
222 topk = topk_idx.shape[-1]
223 _build_sparse_block_table_prefill_kernel[(topk_idx.shape[1],)](
224 topk_idx,
@@ -227,7 +255,7 @@ def minimax_m3_build_sparse_block_table_decode(
255 "MiniMax-M3 decode top-k rows must equal batch * decode_query_len: "
256 f"{total_q} != {expected_q}"
257 )
230 - sparse_bt, sparse_ctx = _alloc_sparse_block_table(topk_idx)
258 + sparse_bt, sparse_ctx = minimax_m3_alloc_sparse_block_table(topk_idx)
259 topk = topk_idx.shape[-1]
260 _build_sparse_block_table_kernel[(total_q,)](
261 topk_idx,
@@ -328,7 +356,9 @@ def _sides_are_packed(k_cache: torch.Tensor, v_cache: torch.Tensor) -> bool:
356 return v_cache.shape[0] != k_cache.shape[0]
357
358
331 -def _block_page_stride(k_cache: torch.Tensor, v_cache: torch.Tensor) -> int:
359 +def minimax_m3_sparse_block_page_stride(
360 + k_cache: torch.Tensor, v_cache: torch.Tensor
361 +) -> int:
362 """How many page ids one sparse block spans.
363
364 Each side owns ``PAGES_PER_SPARSE_BLOCK`` pages. When the sides are dense
@@ -384,20 +414,31 @@ def minimax_m3_sparse_attn_decode_aiter(
414 k_scale: torch.Tensor | None = None,
415 v_scale: torch.Tensor | None = None,
416 decode_query_len: int = 1,
417 + sparse_block_table: torch.Tensor | None = None,
418 + sparse_context_lens: torch.Tensor | None = None,
419 ) -> None:
388 - sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_decode(
389 - topk_idx,
390 - block_table,
391 - seq_lens,
392 - decode_query_len,
393 - _block_page_stride(k_cache, v_cache),
394 - )
420 + if (sparse_block_table is None) != (sparse_context_lens is None):
421 + raise ValueError(
422 + "MiniMax-M3 prepared sparse block table and context lengths "
423 + "must be provided together"
424 + )
425 + if sparse_block_table is None:
426 + sparse_block_table, sparse_context_lens = (
427 + minimax_m3_build_sparse_block_table_decode(
428 + topk_idx,
429 + block_table,
430 + seq_lens,
431 + decode_query_len,
432 + minimax_m3_sparse_block_page_stride(k_cache, v_cache),
433 + )
434 + )
435 + assert sparse_context_lens is not None
436 _run_gluon_decode(
437 q,
438 k_cache,
439 v_cache,
399 - sparse_bt,
400 - sparse_ctx,
440 + sparse_block_table,
441 + sparse_context_lens,
442 num_kv_heads,
443 sm_scale,
444 output,
@@ -426,7 +467,7 @@ def minimax_m3_sparse_attn_prefill_aiter(
467 block_table,
468 query_req_id,
469 query_abs_pos,
429 - _block_page_stride(k_cache, v_cache),
470 + minimax_m3_sparse_block_page_stride(k_cache, v_cache),
471 )
472 _run_gluon_decode(
473 q,
vllm/models/minimax_m3/amd/sparse_attention_msa.py
+7
@@ -25,6 +25,7 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
25 output: torch.Tensor,
26 *,
27 query_fp8: torch.Tensor | None = None,
28 + decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None,
29 ) -> torch.Tensor:
30 from vllm.models.minimax_m3.amd.ops.sparse_pa import (
31 minimax_m3_sparse_attn_decode_aiter,
@@ -70,6 +71,12 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
71 k_scale=k_scale,
72 v_scale=v_scale,
73 decode_query_len=d.decode_query_len,
74 + sparse_block_table=(
75 + decode_sparse_table[0] if decode_sparse_table is not None else None
76 + ),
77 + sparse_context_lens=(
78 + decode_sparse_table[1] if decode_sparse_table is not None else None
79 + ),
80 )
81
82 if main_md.num_prefills > 0:
vllm/models/minimax_m3/common/indexer.py
+71 -1
@@ -369,6 +369,18 @@ class MiniMaxM3IndexerImpl(nn.Module):
369 # Shared, stable-address top-k output buffer (set by the model for the
370 # cudagraph-safe MSA impl); None -> impl allocates fresh (eager).
371 self.topk_indices_buffer = topk_indices_buffer
372 + topk_completion_counter = None
373 + if current_platform.is_rocm() and topk_indices_buffer is not None:
374 + topk_completion_counter = torch.zeros(
375 + topk_indices_buffer.shape[:2],
376 + dtype=torch.int32,
377 + device=topk_indices_buffer.device,
378 + )
379 + self.register_buffer(
380 + "topk_completion_counter",
381 + topk_completion_counter,
382 + persistent=False,
383 + )
384 # Owns the side cache (registers itself in the static forward context).
385 self.index_cache = MiniMaxM3IndexerCache(
386 head_dim=index_head_dim,
@@ -392,6 +404,11 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
404 def forward(
405 self,
406 index_query: torch.Tensor,
407 + *,
408 + attention_block_table: torch.Tensor | None = None,
409 + sparse_block_table_out: torch.Tensor | None = None,
410 + sparse_context_lens_out: torch.Tensor | None = None,
411 + block_page_stride: int | None = None,
412 ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
413 attn_metadata = get_forward_context().attn_metadata
414 if not isinstance(attn_metadata, dict):
@@ -417,6 +434,19 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
434 if index_md.num_decodes > 0:
435 d = index_md.decode
436 assert d is not None
437 + fused_sparse_kwargs = {}
438 + if attention_block_table is not None:
439 + fused_sparse_kwargs = {
440 + "attention_block_table": attention_block_table,
441 + "sparse_block_table_out": sparse_block_table_out,
442 + "sparse_context_lens_out": sparse_context_lens_out,
443 + "block_page_stride": block_page_stride,
444 + }
445 + decode_backend_kwargs = {}
446 + if current_platform.is_rocm():
447 + decode_backend_kwargs["completion_counter"] = (
448 + self.topk_completion_counter
449 + )
450 decode_topk = minimax_m3_index_decode(
451 iq[:nd],
452 kv,
@@ -430,6 +460,8 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
460 d.decode_query_len,
461 d.max_decode_query_len,
462 out=buf_htk,
463 + **fused_sparse_kwargs,
464 + **decode_backend_kwargs,
465 )
466 if index_md.num_prefills > 0:
467 p = index_md.prefill
@@ -567,5 +599,43 @@ class MiniMaxM3Indexer(nn.Module):
599 def forward(
600 self,
601 index_query: torch.Tensor,
602 + *,
603 + attention_block_table: torch.Tensor | None = None,
604 + sparse_block_table_out: torch.Tensor | None = None,
605 + sparse_context_lens_out: torch.Tensor | None = None,
606 + block_page_stride: int | None = None,
607 ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
571 - return self.impl(index_query)
608 + if attention_block_table is None and all(
609 + arg is None
610 + for arg in (
611 + sparse_block_table_out,
612 + sparse_context_lens_out,
613 + block_page_stride,
614 + )
615 + ):
616 + return self.impl(index_query)
617 + if attention_block_table is None or any(
618 + arg is None
619 + for arg in (
620 + sparse_block_table_out,
621 + sparse_context_lens_out,
622 + block_page_stride,
623 + )
624 + ):
625 + raise ValueError(
626 + "MiniMax-M3 fused decode sparse-table arguments must be "
627 + "provided together"
628 + )
629 + if not current_platform.is_rocm() or not isinstance(
630 + self.impl, MiniMaxM3IndexerTritonImpl
631 + ):
632 + raise ValueError(
633 + "MiniMax-M3 fused decode sparse-table construction is ROCm-only"
634 + )
635 + return self.impl(
636 + index_query,
637 + attention_block_table=attention_block_table,
638 + sparse_block_table_out=sparse_block_table_out,
639 + sparse_context_lens_out=sparse_context_lens_out,
640 + block_page_stride=block_page_stride,
641 + )