@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 2,711 lines 93.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Correctness tests for MiniMax M3 sparse prefill attention kernels."""
4
5 import pytest
6 import torch
7
8 from tests.v1.attention.utils import dense_kv_cache_views
9 from vllm import _custom_ops as ops
10 from vllm.config import CacheConfig
11 from vllm.config.cache import _layout_from_name
12 from vllm.models.minimax_m3.common.ops.index_topk import (
13 minimax_m3_index_decode,
14 minimax_m3_index_score,
15 minimax_m3_index_topk,
16 )
17 from vllm.models.minimax_m3.common.ops.sparse_attn import (
18 _FP8_DTYPES,
19 minimax_m3_sparse_attn,
20 minimax_m3_sparse_attn_decode,
21 )
22 from vllm.models.minimax_m3.common.sparse_attention import (
23 MiniMaxM3SparseTritonImpl,
24 minimax_m3_rebase_slots_to_page16,
25 )
26 from vllm.platforms import current_platform
27 from vllm.utils.torch_utils import set_random_seed
28 from vllm.v1.attention.backends.utils import record_kv_cache_layout
29 from vllm.v1.kv_cache_interface import (
30 FullAttentionSpec,
31 KVCacheLayout,
32 MLAAttentionSpec,
33 compute_layer_kv_cache_shape_bytes,
34 get_kv_quant_mode,
35 )
36
37 if not (current_platform.is_cuda() or current_platform.is_rocm()):
38 pytest.skip(
39 "MiniMax M3 attention kernels require CUDA or ROCm.",
40 allow_module_level=True,
41 )
42
43
44 @pytest.fixture
45 def kv_layout(request) -> KVCacheLayout:
46 """Resolve the requested layout name (legacy NHD/HND aliases included)."""
47 return _layout_from_name(request.param)
48
49
50 def _layer_stride_order(layout: KVCacheLayout, ndim: int) -> tuple[int, ...]:
51 """Per-layer physical stride order for the given layout; the 3-dim
52 indexer side cache (H=1) is contiguous, so identity."""
53 if ndim == 3:
54 return (0, 1, 2)
55 stride_order = layout.layer_view_order
56 assert len(stride_order) == ndim
57 return stride_order
58
59
60 def _main_spec() -> FullAttentionSpec:
61 return FullAttentionSpec(
62 block_size=BLOCK_SIZE,
63 num_kv_heads=NUM_KV_HEADS,
64 head_size=HEAD_DIM,
65 head_size_v=HEAD_DIM,
66 dtype=DTYPE,
67 )
68
69
70 def _main_kv_logical_shape(num_pages: int) -> tuple[int, ...]:
71 """Standardized per-layer logical shape [B, H, N, C] for the main cache,
72 derived the same way the production allocator does."""
73 shape_bytes = compute_layer_kv_cache_shape_bytes(_main_spec(), num_pages)
74 return (*shape_bytes[:-1], shape_bytes[-1] // DTYPE.itemsize)
75
76
77 def _allocate_main_kv_via_contract(
78 num_pages: int, layout: KVCacheLayout, device: torch.device | str = "cuda"
79 ) -> torch.Tensor:
80 """Build the main KV cache exactly as the production allocator does for the
81 given layout: allocate the physical (permuted) tensor, then expose the
82 inverse-permuted logical [B, H, N, C] view the backend sees."""
83 logical_shape = _main_kv_logical_shape(num_pages)
84 stride_order = _layer_stride_order(layout, len(logical_shape))
85 physical_shape = tuple(logical_shape[i] for i in stride_order)
86 inv_order = [stride_order.index(i) for i in range(len(stride_order))]
87 raw = torch.randn(physical_shape, device=device, dtype=DTYPE)
88 return raw.permute(*inv_order)
89
90
91 NUM_Q_HEADS = 32
92 NUM_KV_HEADS = 2
93 HEAD_DIM = 128
94 BLOCK_SIZE = 128
95 DTYPE = torch.bfloat16
96 SM_SCALE = HEAD_DIM**-0.5
97 TOPK = 16
98
99
100 @pytest.mark.parametrize(
101 (
102 "num_reqs",
103 "head_dim",
104 "query_dtype",
105 "cache_dtype",
106 "is_gfx950",
107 "expected",
108 ),
109 [
110 (19, 128, torch.bfloat16, torch.bfloat16, True, (16, False)),
111 (20, 128, torch.bfloat16, torch.bfloat16, True, (32, True)),
112 (32, 128, torch.bfloat16, torch.bfloat16, True, (32, True)),
113 (33, 128, torch.bfloat16, torch.bfloat16, True, (16, True)),
114 (64, 128, torch.bfloat16, torch.bfloat16, True, (16, True)),
115 (65, 128, torch.bfloat16, torch.bfloat16, True, (4, False)),
116 (20, 128, torch.bfloat16, torch.bfloat16, False, (16, False)),
117 (20, 64, torch.bfloat16, torch.bfloat16, True, (16, False)),
118 (20, 128, torch.float16, torch.bfloat16, True, (16, False)),
119 (20, 128, torch.bfloat16, torch.float16, True, (16, False)),
120 ],
121 )
122 def test_amd_decode_score_split_launch_policy(
123 num_reqs: int,
124 head_dim: int,
125 query_dtype: torch.dtype,
126 cache_dtype: torch.dtype,
127 is_gfx950: bool,
128 expected: tuple[int, bool],
129 ):
130 from vllm.models.minimax_m3.amd.ops.index_topk import (
131 _decode_score_split_launch_policy,
132 )
133
134 assert (
135 _decode_score_split_launch_policy(
136 num_reqs,
137 head_dim,
138 query_dtype,
139 cache_dtype,
140 is_gfx950=is_gfx950,
141 )
142 == expected
143 )
144
145
146 @pytest.mark.parametrize("num_reqs", range(1, 12))
147 def test_amd_balanced_decode_score_mapping_exact_coverage(num_reqs: int):
148 from vllm.models.minimax_m3.amd.ops.index_topk import (
149 _decode_score_program_budget,
150 )
151
152 budget = _decode_score_program_budget(
153 num_reqs,
154 128,
155 torch.bfloat16,
156 torch.bfloat16,
157 is_gfx950=True,
158 )
159 assert budget is not None
160 patterns = (
161 (0,) * num_reqs,
162 tuple((request * 7919 + 17) % 65535 for request in range(num_reqs)),
163 (65534,) * num_reqs,
164 )
165 for request_blocks in patterns:
166 width = max(1, (sum(request_blocks) + budget - 1) // budget)
167 grid = budget + len(request_blocks) - 1
168 intervals = []
169 for request, blocks in enumerate(request_blocks):
170 programs = (blocks + width - 1) // width
171 for local_program in range(programs):
172 begin = local_program * width
173 intervals.append((request, begin, min(begin + width, blocks)))
174
175 assert len(intervals) <= grid
176 by_request: list[list[tuple[int, int]]] = [[] for _ in request_blocks]
177 for request, begin, end in intervals:
178 assert 0 <= request < len(request_blocks)
179 assert 0 <= begin < end <= request_blocks[request]
180 by_request[request].append((begin, end))
181 for blocks, request_intervals in zip(
182 request_blocks,
183 by_request,
184 strict=True,
185 ):
186 cursor = 0
187 for begin, end in request_intervals:
188 assert begin == cursor
189 cursor = end
190 assert cursor == blocks
191
192
193 @pytest.mark.skipif(
194 not current_platform.is_rocm(),
195 reason="MiniMax-M3 balanced decode score is ROCm-only",
196 )
197 @pytest.mark.parametrize(
198 (
199 "num_reqs",
200 "max_block",
201 "num_idx_heads",
202 "decode_query_len",
203 "max_decode_query_len",
204 "init_blocks",
205 "local_blocks",
206 ),
207 [
208 (1, 37, 1, 1, 1, 1, 0),
209 (8, 37, 1, 1, 4, 2, 1),
210 (9, 37, 1, 4, 4, 0, 2),
211 (11, 37, 1, 3, 5, 1, 3),
212 (1, 37, 2, 1, 1, 2, 0),
213 (8, 37, 2, 1, 4, 0, 1),
214 (9, 37, 2, 4, 4, 1, 2),
215 (11, 37, 2, 3, 5, 2, 3),
216 (1, 37, 4, 1, 1, 0, 0),
217 (8, 37, 4, 1, 4, 1, 1),
218 (9, 37, 4, 4, 4, 2, 2),
219 (11, 37, 4, 3, 5, 0, 3),
220 (20, 65, 1, 4, 4, 1, 3),
221 (36, 65, 1, 4, 4, 2, 2),
222 (64, 65, 1, 4, 4, 0, 1),
223 ],
224 )
225 def test_amd_optimized_decode_score_bitwise_and_graph_replay(
226 num_reqs: int,
227 max_block: int,
228 num_idx_heads: int,
229 decode_query_len: int,
230 max_decode_query_len: int,
231 init_blocks: int,
232 local_blocks: int,
233 ):
234 from vllm.models.minimax_m3.amd.ops.index_topk import (
235 _decode_index_score_balanced_kernel,
236 _decode_index_score_kernel,
237 _decode_score_program_budget,
238 _decode_score_split_launch_policy,
239 )
240
241 set_random_seed(0)
242 head_dim = 128
243 block_size_q = 1 << (max_decode_query_len - 1).bit_length()
244 score_stride = (max_block + 15) // 16 * 16
245 num_pages = num_reqs * max_block
246 block_table = torch.arange(
247 num_pages,
248 device="cuda",
249 dtype=torch.int32,
250 ).reshape(num_reqs, max_block)
251 idx_q = torch.randn(
252 (num_reqs * decode_query_len, num_idx_heads, head_dim),
253 device="cuda",
254 dtype=torch.bfloat16,
255 )
256 index_kv_cache = torch.randn(
257 (num_pages, BLOCK_SIZE, head_dim),
258 device="cuda",
259 dtype=torch.bfloat16,
260 )
261 seq_lens = torch.empty(num_reqs, device="cuda", dtype=torch.int32)
262 deployed_score = torch.empty(
263 (num_idx_heads, num_reqs * decode_query_len, score_stride),
264 device="cuda",
265 dtype=torch.float32,
266 )
267 optimized_score = torch.empty_like(deployed_score)
268 budget = _decode_score_program_budget(
269 num_reqs,
270 head_dim,
271 idx_q.dtype,
272 index_kv_cache.dtype,
273 is_gfx950=True,
274 )
275 optimized_chunks = None
276 if budget is None:
277 optimized_chunks, use_high_batch_config = _decode_score_split_launch_policy(
278 num_reqs,
279 head_dim,
280 idx_q.dtype,
281 index_kv_cache.dtype,
282 is_gfx950=True,
283 )
284 assert use_high_batch_config
285 target = max(1, min(256, 512 // num_reqs))
286 deployed_chunks = 1 << (target.bit_length() - 1)
287
288 def launch_deployed() -> None:
289 _decode_index_score_kernel[(num_reqs, deployed_chunks)](
290 idx_q,
291 index_kv_cache,
292 deployed_score,
293 block_table,
294 seq_lens,
295 num_idx_heads,
296 head_dim,
297 init_blocks,
298 local_blocks,
299 decode_query_len,
300 idx_q.stride(0),
301 idx_q.stride(1),
302 idx_q.stride(2),
303 index_kv_cache.stride(0),
304 index_kv_cache.stride(1),
305 index_kv_cache.stride(2),
306 deployed_score.stride(0),
307 deployed_score.stride(1),
308 deployed_score.stride(2),
309 block_table.stride(0),
310 BLOCK_SIZE_K=BLOCK_SIZE,
311 BLOCK_SIZE_Q=block_size_q,
312 num_kv_chunks=deployed_chunks,
313 USE_PDL=False,
314 )
315
316 def launch_optimized() -> None:
317 if budget is not None:
318 _decode_index_score_balanced_kernel[(budget + num_reqs - 1,)](
319 idx_q,
320 index_kv_cache,
321 optimized_score,
322 block_table,
323 seq_lens,
324 num_idx_heads,
325 head_dim,
326 init_blocks,
327 local_blocks,
328 num_reqs,
329 budget,
330 decode_query_len,
331 idx_q.stride(0),
332 idx_q.stride(1),
333 idx_q.stride(2),
334 index_kv_cache.stride(0),
335 index_kv_cache.stride(1),
336 index_kv_cache.stride(2),
337 optimized_score.stride(0),
338 optimized_score.stride(1),
339 optimized_score.stride(2),
340 block_table.stride(0),
341 BLOCK_SIZE_K=BLOCK_SIZE,
342 BLOCK_SIZE_Q=block_size_q,
343 num_warps=2,
344 num_stages=1,
345 )
346 else:
347 assert optimized_chunks is not None
348 _decode_index_score_kernel[(num_reqs, optimized_chunks)](
349 idx_q,
350 index_kv_cache,
351 optimized_score,
352 block_table,
353 seq_lens,
354 num_idx_heads,
355 head_dim,
356 init_blocks,
357 local_blocks,
358 decode_query_len,
359 idx_q.stride(0),
360 idx_q.stride(1),
361 idx_q.stride(2),
362 index_kv_cache.stride(0),
363 index_kv_cache.stride(1),
364 index_kv_cache.stride(2),
365 optimized_score.stride(0),
366 optimized_score.stride(1),
367 optimized_score.stride(2),
368 block_table.stride(0),
369 BLOCK_SIZE_K=BLOCK_SIZE,
370 BLOCK_SIZE_Q=block_size_q,
371 num_kv_chunks=optimized_chunks,
372 USE_PDL=False,
373 num_warps=2,
374 num_stages=1,
375 )
376
377 def score_masks(
378 expected_seq_lens: torch.Tensor,
379 ) -> tuple[torch.Tensor, torch.Tensor]:
380 query_offsets = torch.arange(decode_query_len, device="cuda")
381 kv_lens = torch.clamp(
382 expected_seq_lens[:, None] - decode_query_len + query_offsets + 1,
383 min=0,
384 )
385 query_blocks = (kv_lens + BLOCK_SIZE - 1) // BLOCK_SIZE
386 block_ids = torch.arange(score_stride, device="cuda")
387 topk_mask = block_ids[None, :] < query_blocks.reshape(-1, 1)
388 request_blocks = query_blocks.max(dim=1).values
389 write_mask = (
390 block_ids[None, :]
391 < request_blocks.repeat_interleave(decode_query_len)[:, None]
392 )
393 return write_mask, topk_mask
394
395 def packed_keys(score: torch.Tensor, live: torch.Tensor) -> torch.Tensor:
396 normalized = torch.where(torch.isnan(score), -1e30, score)
397 normalized = torch.where(normalized == 0.0, 0.0, normalized)
398 bits = normalized.contiguous().view(torch.int32).to(torch.int64)
399 bits &= 0xFFFFFFFF
400 ordered = bits ^ torch.where(
401 bits >> 31 != 0,
402 0xFFFFFFFF,
403 0x80000000,
404 )
405 index = torch.arange(score_stride, device="cuda") + 1
406 keys = (1 << 48) | (ordered << 16) | (0xFFFF - index)
407 return torch.where(live, keys, 0)
408
409 def assert_exact(expected_seq_lens: torch.Tensor) -> None:
410 write_mask, topk_mask = score_masks(expected_seq_lens)
411 write_mask = write_mask.unsqueeze(0).expand_as(deployed_score)
412 topk_mask = topk_mask.unsqueeze(0)
413 deployed = deployed_score
414 optimized = optimized_score
415 assert torch.equal(
416 deployed[write_mask].contiguous().view(torch.int32),
417 optimized[write_mask].contiguous().view(torch.int32),
418 )
419 assert torch.isnan(deployed[~write_mask]).all()
420 assert torch.isnan(optimized[~write_mask]).all()
421 deployed_keys = packed_keys(deployed, topk_mask)
422 optimized_keys = packed_keys(optimized, topk_mask)
423 assert torch.equal(deployed_keys, optimized_keys)
424 deployed_topk = torch.topk(deployed_keys, 16, dim=-1, sorted=True)
425 optimized_topk = torch.topk(optimized_keys, 16, dim=-1, sorted=True)
426 assert torch.equal(deployed_topk.values, optimized_topk.values)
427 assert torch.equal(deployed_topk.indices, optimized_topk.indices)
428
429 uniform = torch.full(
430 (num_reqs,),
431 max_block * BLOCK_SIZE,
432 device="cuda",
433 dtype=torch.int32,
434 )
435 one_active = torch.zeros(num_reqs, device="cuda", dtype=torch.int32)
436 one_active[0] = max(1, max_block // 2) * BLOCK_SIZE
437 patterns = [uniform, one_active]
438 non_aligned = torch.tensor(
439 [(max_block - 1) * BLOCK_SIZE + request % 3 + 1 for request in range(num_reqs)],
440 device="cuda",
441 dtype=torch.int32,
442 )
443 patterns.append(non_aligned)
444
445 seq_lens.copy_(uniform)
446 deployed_score.fill_(float("nan"))
447 optimized_score.fill_(float("nan"))
448 launch_deployed()
449 launch_optimized()
450 assert_exact(uniform)
451
452 graph = torch.cuda.CUDAGraph()
453 with torch.cuda.graph(graph):
454 launch_optimized()
455 addresses = (
456 seq_lens.data_ptr(),
457 optimized_score.data_ptr(),
458 block_table.data_ptr(),
459 index_kv_cache.data_ptr(),
460 )
461 for pattern in patterns:
462 seq_lens.copy_(pattern)
463 deployed_score.fill_(float("nan"))
464 optimized_score.fill_(float("nan"))
465 launch_deployed()
466 graph.replay()
467 assert_exact(pattern)
468 assert addresses == (
469 seq_lens.data_ptr(),
470 optimized_score.data_ptr(),
471 block_table.data_ptr(),
472 index_kv_cache.data_ptr(),
473 )
474
475
476 @pytest.mark.parametrize(
477 ("kv_cache_dtype", "expected_dtype"),
478 [
479 ("fp8", current_platform.fp8_dtype()),
480 ("fp8_e4m3", current_platform.fp8_dtype()),
481 (
482 "fp8_e5m2",
483 torch.float8_e5m2fnuz
484 if current_platform.is_fp8_fnuz()
485 else torch.float8_e5m2,
486 ),
487 ],
488 )
489 def test_sparse_impl_uses_platform_fp8_dtype(
490 kv_cache_dtype: str,
491 expected_dtype: torch.dtype,
492 ):
493 impl = MiniMaxM3SparseTritonImpl(
494 num_heads=NUM_Q_HEADS,
495 head_size=HEAD_DIM,
496 scale=SM_SCALE,
497 num_kv_heads=NUM_KV_HEADS,
498 kv_cache_dtype=kv_cache_dtype,
499 topk_blocks=TOPK,
500 sparse_block_size=BLOCK_SIZE,
501 )
502 assert impl.kv_cache_fp8_dtype == expected_dtype
503
504
505 @pytest.mark.parametrize(
506 "dtype",
507 [
508 torch.float8_e4m3fn,
509 torch.float8_e4m3fnuz,
510 torch.float8_e5m2,
511 torch.float8_e5m2fnuz,
512 ],
513 )
514 def test_sparse_kernels_recognize_fp8_dtypes(dtype: torch.dtype):
515 assert dtype in _FP8_DTYPES
516
517
518 # Index top-k kernels.
519 def _reference_index_topk(
520 idx_q: torch.Tensor,
521 index_kv_cache: torch.Tensor,
522 block_table: torch.Tensor,
523 q_lens: torch.Tensor,
524 seq_lens: torch.Tensor,
525 prefix_lens: torch.Tensor,
526 topk: int,
527 init_blocks: int,
528 local_blocks: int,
529 sm_scale: float = 1.0,
530 ) -> torch.Tensor:
531 total_q, num_idx_heads, _ = idx_q.shape
532 out = torch.full(
533 (num_idx_heads, total_q, topk), -1, device=idx_q.device, dtype=torch.int32
534 )
535
536 q_start = 0
537 for req_id, (q_len, seq_len, prefix_len) in enumerate(
538 zip(q_lens.tolist(), seq_lens.tolist(), prefix_lens.tolist())
539 ):
540 q_end = q_start + q_len
541 q = idx_q[q_start:q_end]
542 num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
543 pages = block_table[req_id, :num_blocks]
544 k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
545 score = sm_scale * torch.einsum("qhd,kd->hqk", q.float(), k.float())
546
547 q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
548 k_pos = torch.arange(k.shape[0], device=idx_q.device)
549 score.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf"))
550 score = score.reshape(num_idx_heads, q_len, num_blocks, BLOCK_SIZE)
551 score_tensor = score.max(dim=3).values
552
553 valid_blocks = (q_pos + BLOCK_SIZE) // BLOCK_SIZE
554 for local_q, num_valid_blocks in enumerate(valid_blocks.tolist()):
555 end = min(init_blocks, num_valid_blocks)
556 score_tensor[:, local_q, :end] = 1e30
557 start = max(0, num_valid_blocks - local_blocks)
558 score_tensor[:, local_q, start:num_valid_blocks] = 1e29
559
560 k = min(topk, num_valid_blocks)
561 topk_idx = score_tensor[:, local_q].topk(k, dim=1).indices
562 out[:, q_start + local_q, :k] = topk_idx
563 q_start = q_end
564
565 return out
566
567
568 def _assert_topk_indices_equal_unordered(
569 actual: torch.Tensor,
570 expected: torch.Tensor,
571 ) -> None:
572 """Compare selected sparse blocks without requiring a deterministic order."""
573 assert actual.shape == expected.shape
574 actual_flat = actual.cpu().reshape(-1, actual.shape[-1]).tolist()
575 expected_flat = expected.cpu().reshape(-1, expected.shape[-1]).tolist()
576 for actual_row, expected_row in zip(actual_flat, expected_flat):
577 assert set(actual_row) == set(expected_row)
578
579
580 def _reference_decode_index_score(
581 idx_q: torch.Tensor,
582 index_kv_cache: torch.Tensor,
583 block_table: torch.Tensor,
584 seq_lens: torch.Tensor,
585 decode_query_len: int,
586 score_block_stride: int,
587 ) -> torch.Tensor:
588 total_q, num_idx_heads, _ = idx_q.shape
589 out = torch.full(
590 (num_idx_heads, total_q, score_block_stride),
591 -float("inf"),
592 device=idx_q.device,
593 dtype=torch.float32,
594 )
595 for req_id, seq_len in enumerate(seq_lens.tolist()):
596 num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
597 token_start = req_id * decode_query_len
598 q = idx_q[token_start : token_start + decode_query_len].float()
599 pages = block_table[req_id, :num_blocks]
600 k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1).float()
601 score = torch.einsum("qhd,kd->hqk", q, k)
602 q_pos = (
603 seq_len
604 - decode_query_len
605 + torch.arange(decode_query_len, device=idx_q.device)
606 )
607 k_pos = torch.arange(k.shape[0], device=idx_q.device)
608 score.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf"))
609 out[:, token_start : token_start + decode_query_len, :num_blocks] = (
610 score.reshape(num_idx_heads, decode_query_len, num_blocks, BLOCK_SIZE)
611 .max(dim=3)
612 .values
613 )
614 return out
615
616
617 # 3600: 29 blocks over 12 splits of 3, so the last busy split has 2 blocks.
618 @pytest.mark.parametrize("prefix_len", [1024, 3600])
619 @pytest.mark.parametrize("force_split_k", [False, True])
620 def test_prefill_index_topk_correctness(monkeypatch, prefix_len, force_split_k):
621 if force_split_k:
622 # Take the SM12.0-only split-K launch on any CUDA GPU.
623 monkeypatch.setattr(
624 current_platform, "is_device_capability", lambda cap, *_: cap == (12, 0)
625 )
626 topk = 6
627 init_blocks = 0
628 local_blocks = 1
629 num_idx_heads = 2
630 head_dim = 16
631 q_lens = torch.tensor((4, 3), device="cuda", dtype=torch.int32)
632 prefix_lens = torch.tensor((0, prefix_len), device="cuda", dtype=torch.int32)
633 seq_lens = prefix_lens + q_lens
634 batch = q_lens.numel()
635 max_seq_len = seq_lens.max().item()
636 max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
637 num_pages = batch * max_blocks
638
639 cu_seqlens = torch.zeros(batch + 1, device="cuda", dtype=torch.int32)
640 cu_seqlens[1:] = q_lens.cumsum(0)
641 block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
642 batch, max_blocks
643 )
644 idx_q = torch.ones(q_lens.sum().item(), num_idx_heads, head_dim, device="cuda")
645 index_kv_cache = torch.empty(num_pages, BLOCK_SIZE, head_dim, device="cuda")
646 for req_id in range(batch):
647 for block_id in range(max_blocks):
648 page = block_table[req_id, block_id]
649 index_kv_cache[page].fill_(block_id + 1)
650
651 score = minimax_m3_index_score(
652 idx_q,
653 index_kv_cache,
654 block_table,
655 cu_seqlens,
656 seq_lens,
657 prefix_lens,
658 max_query_len=q_lens.max().item(),
659 max_seq_len=max_seq_len,
660 num_kv_heads=num_idx_heads,
661 )
662 actual = minimax_m3_index_topk(
663 score,
664 cu_seqlens,
665 prefix_lens,
666 max_query_len=q_lens.max().item(),
667 topk=topk,
668 init_blocks=init_blocks,
669 local_blocks=local_blocks,
670 )
671 expected = _reference_index_topk(
672 idx_q,
673 index_kv_cache,
674 block_table,
675 q_lens,
676 seq_lens,
677 prefix_lens,
678 topk,
679 init_blocks,
680 local_blocks,
681 )
682 _assert_topk_indices_equal_unordered(actual, expected)
683
684
685 # MSA indexer (SM100): fmha_sm100 OnlyScore for the per-block scores, then the
686 # Triton minimax_m3_index_topk for selection (no sparse_topk_select). Uses a
687 # deterministic construction (idx_q == 1, distinct e4m3-exact per-block values)
688 # so scores are strictly monotonic in the block id -> exact top-k agreement.
689 def _fmha_indexer_topk(
690 idx_q: torch.Tensor, # [total_q, H, 128] bf16/e4m3
691 index_cache: torch.Tensor, # [num_pages, 128, 128] bf16/e4m3
692 block_table: torch.Tensor,
693 q_lens: torch.Tensor,
694 seq_lens: torch.Tensor,
695 prefix_lens: torch.Tensor,
696 sm_scale: float,
697 topk: int,
698 ) -> torch.Tensor:
699 """Replicate MiniMaxM3IndexerMSAImpl's score path (single decode/prefill side)."""
700 from vllm.third_party.fmha_sm100.api import _fmha_sm100, _fmha_sm100_plan
701
702 num_idx_heads, head_dim = idx_q.shape[1], idx_q.shape[2]
703 nvp = [(s + 127) // 128 for s in seq_lens.tolist()]
704 kv_indices = torch.cat([block_table[r, : nvp[r]] for r in range(len(nvp))]).to(
705 torch.int32
706 )
707
708 qo = q_lens.cpu().to(torch.int32)
709 kv = seq_lens.cpu().to(torch.int32)
710 plan = _fmha_sm100_plan(
711 qo,
712 kv,
713 num_idx_heads,
714 num_kv_heads=1,
715 qo_offset=kv - qo,
716 page_size=128,
717 output_maxscore=True,
718 causal=True,
719 num_kv_splits=1,
720 )
721 k_pages = index_cache.view(index_cache.shape[0], 1, 128, head_dim)
722 _, max_score = _fmha_sm100(
723 idx_q,
724 k_pages,
725 k_pages,
726 plan,
727 kv_indices=kv_indices,
728 output_o=False,
729 output_maxscore=True,
730 sm_scale=sm_scale,
731 )
732
733 batch = q_lens.numel()
734 cu = torch.zeros(batch + 1, dtype=torch.int32, device=idx_q.device)
735 cu[1:] = q_lens.to(torch.int32).cumsum(0)
736 # max_score [H, k_tiles, total_q] -> transpose to [H, total_q, k_tiles].
737 return minimax_m3_index_topk(
738 max_score.transpose(1, 2),
739 cu,
740 prefix_lens.to(torch.int32),
741 int(q_lens.max()),
742 topk,
743 0, # init_blocks
744 0, # local_blocks
745 )
746
747
748 # e4m3-exact, strictly-increasing per-block values: with idx_q == 1 (also exact)
749 # the per-block scores are exact and distinct in BOTH bf16 and e4m3, so the fp8
750 # score path selects the same top-k as the reference (no quantization ties).
751 _E4M3_EXACT_VALUES = [
752 *range(1, 17), # 1..16 (step 1)
753 *range(18, 33, 2), # 18..32 (step 2)
754 *range(36, 65, 4), # 36..64 (step 4)
755 *range(72, 129, 8), # 72..128 (step 8)
756 ]
757
758
759 @pytest.mark.skipif(
760 not current_platform.is_device_capability_family(100),
761 reason="fmha_sm100 indexer requires SM100 (Blackwell).",
762 )
763 @pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])
764 @pytest.mark.parametrize(
765 ("q_lens", "prefix_lens"),
766 [
767 ((4, 3), (2048, 2560)), # prefill: every token sees >= 16 causal blocks
768 ((1, 1, 1), (2048, 3000, 4096)), # decode: one query token per request
769 ],
770 )
771 def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):
772 torch.manual_seed(0)
773 num_idx_heads, head_dim = 4, HEAD_DIM
774 device = "cuda"
775
776 q_lens_t = torch.tensor(q_lens, device=device, dtype=torch.int32)
777 prefix_lens_t = torch.tensor(prefix_lens, device=device, dtype=torch.int32)
778 seq_lens = prefix_lens_t + q_lens_t
779 batch = len(q_lens)
780 max_blocks = (int(seq_lens.max()) + BLOCK_SIZE - 1) // BLOCK_SIZE
781 assert max_blocks <= len(_E4M3_EXACT_VALUES)
782 num_pages = batch * max_blocks
783 block_table = torch.randperm(num_pages, device=device, dtype=torch.int32).reshape(
784 batch, max_blocks
785 )
786
787 idx_q = torch.ones(
788 int(q_lens_t.sum()), num_idx_heads, head_dim, device=device, dtype=index_dtype
789 )
790 index_cache = torch.empty(
791 num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype
792 )
793 for r in range(batch):
794 for b in range(max_blocks):
795 index_cache[block_table[r, b]] = float(_E4M3_EXACT_VALUES[b])
796
797 sm_scale = head_dim**-0.5
798 actual = _fmha_indexer_topk(
799 idx_q,
800 index_cache,
801 block_table,
802 q_lens_t,
803 seq_lens,
804 prefix_lens_t,
805 sm_scale,
806 TOPK,
807 )
808 expected = _reference_index_topk(
809 idx_q,
810 index_cache,
811 block_table,
812 q_lens_t,
813 seq_lens,
814 prefix_lens_t,
815 TOPK,
816 init_blocks=0,
817 local_blocks=0,
818 sm_scale=sm_scale,
819 )
820 _assert_topk_indices_equal_unordered(actual, expected)
821
822
823 # Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha/CuteDSL score
824 # + unified top-k) and MiniMaxM3IndexerTritonImpl through their real metadata
825 # builders on the SAME CommonAttentionMetadata + index cache, and assert the
826 # selected blocks agree. This exercises all the metadata the impl/kernels consume
827 # (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather,
828 # decode_pages split) -- a metadata bug on either side shifts the causal window
829 # or the block->page mapping and breaks the comparison.
830 @pytest.mark.skipif(
831 not current_platform.is_device_capability_family(100),
832 reason="fmha_sm100 indexer requires SM100 (Blackwell).",
833 )
834 @pytest.mark.parametrize("topk", [16])
835 @pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])
836 def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch):
837 import vllm.models.minimax_m3.common.indexer as indexer_mod
838 from tests.v1.attention.utils import (
839 BatchSpec,
840 create_common_attn_metadata,
841 create_vllm_config,
842 )
843 from vllm.config import set_current_vllm_config
844 from vllm.forward_context import set_forward_context
845 from vllm.models.minimax_m3.common.indexer import (
846 MiniMaxM3IndexerTritonImpl,
847 MiniMaxM3IndexerTritonMetadataBuilder,
848 )
849 from vllm.models.minimax_m3.nvidia.indexer_msa import (
850 MiniMaxM3IndexerMSAImpl,
851 MiniMaxM3IndexerMSAMetadataBuilder,
852 )
853
854 torch.manual_seed(0)
855 device = torch.device("cuda")
856 num_idx_heads, head_dim = 4, HEAD_DIM
857 # TP=1: avoid requiring an initialized distributed group in a unit test.
858 monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1)
859
860 vllm_config = create_vllm_config(
861 block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192
862 )
863 vllm_config.model_config.hf_config.sparse_attention_config = {
864 "sparse_num_index_heads": num_idx_heads
865 }
866
867 # Decode-first mixed batch: 2 decode reqs (q_len 1) then 2 prefill reqs. Long
868 # prefixes so every token sees > TOPK causal blocks (non-trivial selection).
869 batch = BatchSpec(seq_lens=[2305, 2561, 2624, 2720], query_lens=[1, 1, 64, 96])
870 common = create_common_attn_metadata(
871 batch, BLOCK_SIZE, device, arange_block_indices=True
872 )
873 num_tokens = batch.compute_num_tokens()
874 # Absolute token positions; the MSA builder derives per-token causal page
875 # counts from them.
876 common.positions = torch.cat(
877 [
878 torch.arange(s - q, s, device=device, dtype=torch.int64)
879 for s, q in zip(batch.seq_lens, batch.query_lens)
880 ]
881 )
882
883 # Deterministic index cache: distinct, monotonic per-logical-block values so
884 # the top-k is unambiguous (both kernels pick the same blocks, no fp ties).
885 block_table = common.block_table_tensor
886 num_pages = int(block_table.max().item()) + 1
887 index_cache = torch.zeros(
888 num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype
889 )
890 for r, seq_len in enumerate(batch.seq_lens):
891 for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE):
892 index_cache[block_table[r, b]] = float(b + 1)
893 index_q = torch.ones(
894 num_tokens, num_idx_heads * head_dim, device=device, dtype=index_dtype
895 )
896
897 spec = MLAAttentionSpec(
898 block_size=BLOCK_SIZE, num_kv_heads=1, head_size=head_dim, dtype=DTYPE
899 )
900 impl_kwargs = dict(
901 num_kv_heads=num_idx_heads,
902 scale=head_dim**-0.5,
903 topk_blocks=topk,
904 sparse_block_size=BLOCK_SIZE,
905 num_index_heads=num_idx_heads,
906 index_head_dim=head_dim,
907 init_blocks=0,
908 local_blocks=0,
909 )
910
911 with set_current_vllm_config(vllm_config):
912 msa_impl = MiniMaxM3IndexerMSAImpl(prefix="idx_msa", **impl_kwargs)
913 triton_impl = MiniMaxM3IndexerTritonImpl(prefix="idx_triton", **impl_kwargs)
914 msa_builder = MiniMaxM3IndexerMSAMetadataBuilder(
915 spec, [msa_impl.index_cache.prefix], vllm_config, device
916 )
917 triton_builder = MiniMaxM3IndexerTritonMetadataBuilder(
918 spec, [triton_impl.index_cache.prefix], vllm_config, device
919 )
920
921 # Both impls score against the same index keys.
922 msa_impl.index_cache.kv_cache = index_cache
923 triton_impl.index_cache.kv_cache = index_cache
924
925 # Exercise the shared persistent top-k buffer for BOTH impls: each must write
926 # decode ([:nd]) and prefill ([nd:]) into its token-major buffer and (Triton
927 # only) return views. Separate buffers so the two forwards don't clobber.
928 nd = sum(q for q in batch.query_lens if q <= 1)
929 msa_impl.topk_indices_buffer = torch.full(
930 (num_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device
931 )
932 triton_impl.topk_indices_buffer = torch.full(
933 (num_tokens, num_idx_heads, topk), -2, dtype=torch.int32, device=device
934 )
935
936 attn_metadata = {
937 msa_impl.index_cache.prefix: msa_builder.build(0, common),
938 triton_impl.index_cache.prefix: triton_builder.build(0, common),
939 }
940 with set_forward_context(attn_metadata, vllm_config):
941 msa_decode, msa_prefill = msa_impl(index_q)
942 tri_decode, tri_prefill = triton_impl(index_q)
943
944 # MSA's return is vestigial; the attend reads its buffer directly.
945 assert msa_decode is None and msa_prefill is None
946 assert tri_decode is not None and tri_prefill is not None
947 _assert_topk_indices_equal_unordered(
948 msa_impl.topk_indices_buffer[:num_tokens],
949 triton_impl.topk_indices_buffer[:num_tokens],
950 )
951 # Triton's decode/prefill outputs are views into its persistent buffer.
952 buf_htk = triton_impl.topk_indices_buffer.transpose(0, 1)
953 assert tri_decode.data_ptr() == buf_htk[:, :nd, :].data_ptr()
954 assert tri_prefill.data_ptr() == buf_htk[:, nd:, :].data_ptr()
955
956
957 @pytest.mark.parametrize(
958 ("decode_query_len", "max_decode_query_len"),
959 [
960 (1, 1),
961 (1, 4),
962 (4, 4),
963 ],
964 )
965 @pytest.mark.parametrize("num_padded_reqs", [0, 2])
966 def test_decode_index_topk_correctness(
967 decode_query_len: int,
968 max_decode_query_len: int,
969 num_padded_reqs: int,
970 ):
971 topk = 6
972 init_blocks = 0
973 local_blocks = 1
974 num_idx_heads = 2
975 head_dim = 16
976 active_seq_lens = torch.tensor((7, 129, 1025), device="cuda", dtype=torch.int32)
977 q_lens = torch.full_like(active_seq_lens, decode_query_len)
978 prefix_lens = active_seq_lens - decode_query_len
979 active_batch = active_seq_lens.numel()
980 batch = active_batch + num_padded_reqs
981 seq_lens = torch.cat(
982 [
983 active_seq_lens,
984 torch.zeros(num_padded_reqs, device="cuda", dtype=torch.int32),
985 ]
986 )
987 max_seq_len = active_seq_lens.max().item()
988 max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
989 num_pages = active_batch * max_blocks
990
991 active_block_table = torch.randperm(
992 num_pages, device="cuda", dtype=torch.int32
993 ).reshape(active_batch, max_blocks)
994 block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
995 block_table[:active_batch] = active_block_table
996 idx_q = torch.randn(
997 batch * decode_query_len, num_idx_heads, head_dim, device="cuda"
998 )
999 index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda")
1000
1001 actual = minimax_m3_index_decode(
1002 idx_q,
1003 index_kv_cache,
1004 block_table,
1005 seq_lens,
1006 max_seq_len=max_seq_len,
1007 topk=topk,
1008 init_blocks=init_blocks,
1009 local_blocks=local_blocks,
1010 num_kv_heads=num_idx_heads,
1011 decode_query_len=decode_query_len,
1012 max_decode_query_len=max_decode_query_len,
1013 )
1014 expected = torch.full_like(actual, -1)
1015 active_tokens = active_batch * decode_query_len
1016 expected[:, :active_tokens] = _reference_index_topk(
1017 idx_q[:active_tokens],
1018 index_kv_cache,
1019 block_table[:active_batch],
1020 q_lens,
1021 active_seq_lens,
1022 prefix_lens,
1023 topk,
1024 init_blocks,
1025 local_blocks,
1026 )
1027 _assert_topk_indices_equal_unordered(actual, expected)
1028
1029
1030 @pytest.mark.parametrize(
1031 (
1032 "max_block",
1033 "total_q",
1034 "num_idx_heads",
1035 "topk",
1036 "is_gfx950",
1037 "expected",
1038 ),
1039 [
1040 (0, 24, 1, 16, True, (2, 512, 8, 2, False, False)),
1041 (1, 1, 1, 16, True, (1, 128, 2, 1, True, True)),
1042 (74, 4, 1, 16, True, (1, 128, 2, 1, True, True)),
1043 (128, 4, 1, 16, True, (1, 128, 2, 1, True, True)),
1044 (129, 4, 1, 16, True, (16, 512, 4, 2, True, True)),
1045 (512, 4, 1, 16, True, (16, 512, 4, 2, True, True)),
1046 (513, 24, 4, 16, True, (16, 512, 4, 2, True, True)),
1047 (2760, 5, 2, 16, True, (16, 512, 4, 2, True, True)),
1048 (8192, 64, 4, 16, True, (16, 512, 4, 2, True, True)),
1049 (8193, 4, 1, 16, True, (16, 512, 8, 2, False, False)),
1050 (8193, 16, 1, 16, True, (4, 512, 8, 2, False, False)),
1051 (65534, 64, 1, 16, True, (1, 512, 8, 2, False, False)),
1052 (513, 16, 1, 13, True, (4, 512, 8, 2, False, False)),
1053 (513, 16, 4, 16, False, (1, 512, 8, 2, False, False)),
1054 ],
1055 )
1056 def test_amd_decode_topk_launch_policy(
1057 max_block: int,
1058 total_q: int,
1059 num_idx_heads: int,
1060 topk: int,
1061 is_gfx950: bool,
1062 expected: tuple[int, int, int, int, bool, bool],
1063 ):
1064 """The short-context specialization has strict dispatch boundaries."""
1065 from vllm.models.minimax_m3.amd.ops.index_topk import (
1066 _decode_topk_launch_policy,
1067 )
1068
1069 assert (
1070 _decode_topk_launch_policy(
1071 max_block,
1072 total_q,
1073 num_idx_heads,
1074 topk,
1075 is_gfx950=is_gfx950,
1076 )
1077 == expected
1078 )
1079
1080
1081 @pytest.mark.skipif(
1082 not current_platform.is_rocm(),
1083 reason="MiniMax-M3 fused decode top-k is ROCm-only",
1084 )
1085 @pytest.mark.parametrize(
1086 (
1087 "num_topk_chunks",
1088 "single_tile_guaranteed",
1089 "adaptive_final_merge",
1090 "max_blocks",
1091 "block_size_k",
1092 "num_warps",
1093 "num_stages",
1094 "topk",
1095 "num_idx_heads",
1096 ),
1097 [
1098 (1, False, False, 65534, 512, 8, 2, 16, 1),
1099 (2, False, False, 8193, 512, 8, 2, 16, 1),
1100 (4, False, False, 8193, 512, 8, 2, 16, 1),
1101 (8, False, False, 8193, 512, 8, 2, 16, 1),
1102 (16, False, False, 8193, 512, 8, 2, 16, 1),
1103 (16, False, False, 65534, 512, 8, 2, 16, 1),
1104 (2, True, False, 513, 512, 8, 2, 13, 1),
1105 (4, False, True, 8192, 512, 4, 2, 16, 2),
1106 (16, True, True, 8192, 512, 4, 2, 16, 4),
1107 (1, True, True, 74, 128, 2, 1, 16, 1),
1108 ],
1109 )
1110 def test_amd_decode_fused_topk_total_order_and_replay(
1111 num_topk_chunks: int,
1112 single_tile_guaranteed: bool,
1113 adaptive_final_merge: bool,
1114 max_blocks: int,
1115 block_size_k: int,
1116 num_warps: int,
1117 num_stages: int,
1118 topk: int,
1119 num_idx_heads: int,
1120 ):
1121 """Selector specializations preserve ordering and graph state."""
1122 from vllm.models.minimax_m3.amd.ops.index_topk import (
1123 SPARSE_BLOCK_SIZE,
1124 _decode_topk_fused_kernel,
1125 )
1126
1127 set_random_seed(0)
1128 block_size_t = 1 << (topk - 1).bit_length()
1129 block_counts = tuple(
1130 count
1131 for count in (
1132 15,
1133 16,
1134 17,
1135 block_size_k // 2 - 1,
1136 block_size_k // 2,
1137 block_size_k // 2 + 1,
1138 block_size_k - 1,
1139 block_size_k,
1140 block_size_k + 1,
1141 )
1142 if count <= max_blocks
1143 )
1144 long_seq_lens_list = [
1145 0,
1146 1,
1147 *(blocks * SPARSE_BLOCK_SIZE for blocks in block_counts),
1148 max_blocks * SPARSE_BLOCK_SIZE,
1149 ]
1150 short_seq_lens_list = [0, 1, 127, 128, 129, 0, 1, 128, 129, 0, 1, 129]
1151 boundary_limit = 0
1152 if num_topk_chunks == 16 and max_blocks >= 8192:
1153 boundary_limit = 15
1154 elif num_topk_chunks == 4 and adaptive_final_merge:
1155 boundary_limit = 4
1156 if boundary_limit:
1157 boundary_blocks = [
1158 (active_chunks - 1) * block_size_k + 1
1159 for active_chunks in range(3, boundary_limit + 1)
1160 ]
1161 assert [
1162 (num_blocks + block_size_k - 1) // block_size_k
1163 for num_blocks in boundary_blocks
1164 ] == list(range(3, boundary_limit + 1))
1165 long_seq_lens_list.extend(
1166 num_blocks * SPARSE_BLOCK_SIZE for num_blocks in boundary_blocks
1167 )
1168 short_seq_lens_list.extend(
1169 num_blocks * SPARSE_BLOCK_SIZE for num_blocks in reversed(boundary_blocks)
1170 )
1171 batch = len(long_seq_lens_list)
1172 repeats = (batch + len(short_seq_lens_list) - 1) // len(short_seq_lens_list)
1173 short_seq_lens_list = (short_seq_lens_list * repeats)[:batch]
1174 score = torch.randn((num_idx_heads, batch, max_blocks), device="cuda")
1175 score[:, :, 0] = float("nan")
1176 score[:, :, 1] = float("inf")
1177 score[:, :, 2] = float("-inf")
1178 score[:, :, 3] = 0.0
1179 score[:, :, 4] = -0.0
1180 score[:, :, 5] = 7.0
1181 score[:, :, 6] = -1e30
1182 for index in (block_size_k // 2, block_size_k):
1183 if index < max_blocks:
1184 score[:, :, index] = 7.0
1185
1186 seq_lens = torch.empty(batch, device="cuda", dtype=torch.int32)
1187 partial = torch.full(
1188 (num_topk_chunks, num_idx_heads, batch, block_size_t),
1189 -1,
1190 device="cuda",
1191 dtype=torch.int64,
1192 )
1193 counter = torch.zeros((num_idx_heads, batch), device="cuda", dtype=torch.int32)
1194 actual = torch.full(
1195 (num_idx_heads, batch, topk), -2, device="cuda", dtype=torch.int32
1196 )
1197 dummy_block_table = torch.zeros(
1198 (batch, max_blocks), device="cuda", dtype=torch.int32
1199 )
1200 addresses = (partial.data_ptr(), counter.data_ptr(), actual.data_ptr())
1201 long_seq_lens = torch.tensor(long_seq_lens_list, device="cuda", dtype=torch.int32)
1202 short_seq_lens = torch.tensor(short_seq_lens_list, device="cuda", dtype=torch.int32)
1203
1204 def launch() -> None:
1205 _decode_topk_fused_kernel[(batch, num_idx_heads, num_topk_chunks)](
1206 score,
1207 partial,
1208 counter,
1209 actual,
1210 seq_lens,
1211 dummy_block_table,
1212 actual,
1213 seq_lens,
1214 1,
1215 score.stride(0),
1216 score.stride(1),
1217 score.stride(2),
1218 partial.stride(0),
1219 partial.stride(1),
1220 partial.stride(2),
1221 partial.stride(3),
1222 counter.stride(0),
1223 counter.stride(1),
1224 actual.stride(0),
1225 actual.stride(1),
1226 actual.stride(2),
1227 dummy_block_table.stride(0),
1228 actual.stride(1),
1229 topk=topk,
1230 block_size=SPARSE_BLOCK_SIZE,
1231 pages_per_sparse_block=8,
1232 block_page_stride=8,
1233 NUM_TOPK_CHUNKS=num_topk_chunks,
1234 BLOCK_SIZE_K=block_size_k,
1235 BLOCK_SIZE_T=block_size_t,
1236 EMIT_SPARSE_TABLE=False,
1237 SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
1238 ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
1239 num_warps=num_warps,
1240 num_stages=num_stages,
1241 )
1242
1243 score_cpu = score.cpu()
1244
1245 def reference(seq_lens_list: list[int]) -> torch.Tensor:
1246 expected = torch.full((num_idx_heads, batch, topk), -1, dtype=torch.int32)
1247 for head in range(num_idx_heads):
1248 for row, seq_len in enumerate(seq_lens_list):
1249 num_blocks = (seq_len + SPARSE_BLOCK_SIZE - 1) // SPARSE_BLOCK_SIZE
1250 values = score_cpu[head, row, :num_blocks].tolist()
1251 values = [-1e30 if value != value else value for value in values]
1252 selected = sorted(
1253 range(num_blocks),
1254 key=lambda index: (-values[index], index),
1255 )[:topk]
1256 expected[head, row, : len(selected)] = torch.tensor(selected)
1257 return expected
1258
1259 seq_lens.copy_(long_seq_lens)
1260 launch()
1261 torch.accelerator.synchronize()
1262 assert torch.equal(actual.cpu(), reference(long_seq_lens_list))
1263 assert torch.count_nonzero(counter).item() == 0
1264
1265 seq_lens.copy_(short_seq_lens)
1266 launch()
1267 torch.accelerator.synchronize()
1268 assert torch.equal(actual.cpu(), reference(short_seq_lens_list))
1269 assert torch.count_nonzero(counter).item() == 0
1270
1271 replay_count = 8 if max_blocks == 65534 else 128
1272 for replay in range(replay_count):
1273 seq_lens.copy_(long_seq_lens if replay % 2 == 0 else short_seq_lens)
1274 launch()
1275 seq_lens.copy_(long_seq_lens)
1276 launch()
1277 torch.accelerator.synchronize()
1278 assert torch.equal(actual.cpu(), reference(long_seq_lens_list))
1279 assert torch.count_nonzero(counter).item() == 0
1280 assert addresses == (partial.data_ptr(), counter.data_ptr(), actual.data_ptr())
1281
1282 if (num_topk_chunks, single_tile_guaranteed, adaptive_final_merge, max_blocks) in (
1283 (4, False, False, 8193),
1284 (16, False, False, 8193),
1285 (4, False, True, 8192),
1286 (16, True, True, 8192),
1287 (1, True, True, 74),
1288 ):
1289 graph = torch.cuda.CUDAGraph()
1290 torch.accelerator.synchronize()
1291 with torch.cuda.graph(graph):
1292 launch()
1293 seq_lens.copy_(short_seq_lens)
1294 graph.replay()
1295 torch.accelerator.synchronize()
1296 assert torch.equal(actual.cpu(), reference(short_seq_lens_list))
1297 assert torch.count_nonzero(counter).item() == 0
1298 seq_lens.copy_(long_seq_lens)
1299 graph.replay()
1300 torch.accelerator.synchronize()
1301 assert torch.equal(actual.cpu(), reference(long_seq_lens_list))
1302 assert torch.count_nonzero(counter).item() == 0
1303 assert addresses == (partial.data_ptr(), counter.data_ptr(), actual.data_ptr())
1304
1305
1306 @pytest.mark.skipif(
1307 not current_platform.is_rocm(),
1308 reason="MiniMax-M3 fused decode top-k is ROCm-only",
1309 )
1310 @pytest.mark.parametrize(
1311 ("num_idx_heads", "decode_query_len", "max_decode_query_len"),
1312 [(1, 1, 1), (2, 1, 4), (4, 4, 4), (4, 3, 8)],
1313 )
1314 def test_amd_decode_index_topk_end_to_end(
1315 num_idx_heads: int,
1316 decode_query_len: int,
1317 max_decode_query_len: int,
1318 ):
1319 """ROCm BF16 score-to-selector parity crosses the atomic boundary."""
1320 from vllm.models.minimax_m3.amd.ops.index_topk import (
1321 minimax_m3_index_decode as amd_minimax_m3_index_decode,
1322 )
1323
1324 torch.manual_seed(0)
1325 topk = 16
1326 head_dim = 16
1327 seq_lens = torch.tensor([65537], device="cuda", dtype=torch.int32)
1328 q_lens = torch.tensor([decode_query_len], device="cuda", dtype=torch.int32)
1329 prefix_lens = seq_lens - q_lens
1330 max_seq_len = int(seq_lens.max())
1331 max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
1332 block_table = torch.randperm(
1333 max_blocks, device="cuda", dtype=torch.int32
1334 ).unsqueeze(0)
1335 idx_q = torch.randn(
1336 decode_query_len,
1337 num_idx_heads,
1338 head_dim,
1339 device="cuda",
1340 dtype=torch.bfloat16,
1341 )
1342 index_kv_cache = torch.randn(
1343 max_blocks,
1344 BLOCK_SIZE,
1345 head_dim,
1346 device="cuda",
1347 dtype=torch.bfloat16,
1348 )
1349 completion_counter = torch.zeros(
1350 (num_idx_heads, decode_query_len), device="cuda", dtype=torch.int32
1351 )
1352
1353 actual = amd_minimax_m3_index_decode(
1354 idx_q,
1355 index_kv_cache,
1356 block_table,
1357 seq_lens,
1358 max_seq_len=max_seq_len,
1359 topk=topk,
1360 init_blocks=0,
1361 local_blocks=1,
1362 num_kv_heads=num_idx_heads,
1363 decode_query_len=decode_query_len,
1364 max_decode_query_len=max_decode_query_len,
1365 completion_counter=completion_counter,
1366 )
1367 expected = _reference_index_topk(
1368 idx_q,
1369 index_kv_cache,
1370 block_table,
1371 q_lens,
1372 seq_lens,
1373 prefix_lens,
1374 topk,
1375 0,
1376 1,
1377 )
1378 torch.accelerator.synchronize()
1379
1380 _assert_topk_indices_equal_unordered(actual, expected)
1381 assert torch.count_nonzero(completion_counter).item() == 0
1382
1383
1384 @pytest.mark.skipif(
1385 not current_platform.is_device_capability_family(100),
1386 reason="fp8 e4m3 indexer cache is the SM100 (MSA) path.",
1387 )
1388 @pytest.mark.parametrize("num_idx_heads", [1, 4])
1389 def test_decode_index_topk_fp8(num_idx_heads: int):
1390 """The standalone Triton path must score FP8 inputs in FP32 so its top-k
1391 matches a reference computed from the dequantized FP8 values."""
1392 torch.manual_seed(0)
1393 topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128
1394 decode_query_len = 1
1395 active_seq_lens = torch.tensor((129, 1025, 4097), device="cuda", dtype=torch.int32)
1396 q_lens = torch.full_like(active_seq_lens, decode_query_len)
1397 prefix_lens = active_seq_lens - decode_query_len
1398 batch = active_seq_lens.numel()
1399 max_seq_len = int(active_seq_lens.max())
1400 max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
1401 num_pages = batch * max_blocks
1402 block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
1403 batch, max_blocks
1404 )
1405 idx_q = torch.randn(
1406 batch * decode_query_len, num_idx_heads, head_dim, device="cuda"
1407 ).to(torch.float8_e4m3fn)
1408 index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda").to(
1409 torch.float8_e4m3fn
1410 )
1411
1412 actual = minimax_m3_index_decode(
1413 idx_q,
1414 index_kv_cache,
1415 block_table,
1416 active_seq_lens,
1417 max_seq_len=max_seq_len,
1418 topk=topk,
1419 init_blocks=init_blocks,
1420 local_blocks=local_blocks,
1421 num_kv_heads=num_idx_heads,
1422 decode_query_len=decode_query_len,
1423 max_decode_query_len=decode_query_len,
1424 )
1425 # Reference from the DEQUANTIZED fp8 values (the kernel computes the fp8 QK
1426 # in fp32 with no scaling, so it must match an unscaled fp32 matmul of the
1427 # same e4m3 values).
1428 expected = _reference_index_topk(
1429 idx_q.float(),
1430 index_kv_cache.float(),
1431 block_table,
1432 q_lens,
1433 active_seq_lens,
1434 prefix_lens,
1435 topk,
1436 init_blocks,
1437 local_blocks,
1438 )
1439 _assert_topk_indices_equal_unordered(actual, expected)
1440
1441
1442 @pytest.mark.skipif(
1443 not current_platform.is_device_capability_family(100),
1444 reason="CuteDSL index decode score requires Blackwell.",
1445 )
1446 @pytest.mark.parametrize(
1447 ("dtype", "decode_query_len", "max_decode_query_len"),
1448 [
1449 (torch.bfloat16, 1, 1),
1450 (torch.bfloat16, 3, 8),
1451 (torch.float8_e4m3fn, 1, 1),
1452 (torch.float8_e4m3fn, 3, 8),
1453 (torch.float8_e4m3fn, 8, 8),
1454 ],
1455 )
1456 def test_decode_index_score_cutedsl_correctness(
1457 dtype: torch.dtype,
1458 decode_query_len: int,
1459 max_decode_query_len: int,
1460 ):
1461 pytest.importorskip("cutlass")
1462 from vllm.models.minimax_m3.nvidia.ops import (
1463 minimax_m3_index_decode_score_cutedsl,
1464 )
1465
1466 torch.manual_seed(0)
1467 init_blocks, local_blocks = 0, 0
1468 num_idx_heads, head_dim = 4, 128
1469 active_seq_lens = torch.tensor((1025, 4097), device="cuda", dtype=torch.int32)
1470 batch = active_seq_lens.numel()
1471 total_q = batch * decode_query_len
1472 max_seq_len = int(active_seq_lens.max())
1473 max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
1474 score_block_stride = ((max_blocks + 15) // 16) * 16
1475 num_pages = batch * max_blocks
1476 block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
1477 batch, max_blocks
1478 )
1479 idx_q = torch.randn(total_q, num_idx_heads, head_dim, device="cuda").to(dtype)
1480 index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda").to(
1481 dtype
1482 )
1483 unified_score = torch.full(
1484 (total_q, num_idx_heads, score_block_stride),
1485 -float("inf"),
1486 device="cuda",
1487 dtype=torch.float32,
1488 )
1489 score = unified_score.transpose(0, 1)
1490
1491 minimax_m3_index_decode_score_cutedsl(
1492 idx_q,
1493 index_kv_cache,
1494 block_table,
1495 active_seq_lens,
1496 max_seq_len=max_seq_len,
1497 init_blocks=init_blocks,
1498 local_blocks=local_blocks,
1499 num_kv_heads=num_idx_heads,
1500 decode_query_len=decode_query_len,
1501 max_decode_query_len=max_decode_query_len,
1502 score_out=score,
1503 )
1504 expected = _reference_decode_index_score(
1505 idx_q,
1506 index_kv_cache,
1507 block_table,
1508 active_seq_lens,
1509 decode_query_len,
1510 score_block_stride,
1511 )
1512 torch.testing.assert_close(score, expected)
1513
1514
1515 # Sparse attention kernels.
1516 def _reference_sparse_attn(
1517 q: torch.Tensor,
1518 kv_cache: torch.Tensor,
1519 topk_idx: torch.Tensor,
1520 block_table: torch.Tensor,
1521 q_lens: torch.Tensor,
1522 seq_lens: torch.Tensor,
1523 prefix_lens: torch.Tensor,
1524 ) -> torch.Tensor:
1525 out = torch.empty_like(q, dtype=torch.float32)
1526 gqa_group_size = NUM_Q_HEADS // NUM_KV_HEADS
1527 q_start = 0
1528 for req_id, (q_len, seq_len, prefix_len) in enumerate(
1529 zip(q_lens.tolist(), seq_lens.tolist(), prefix_lens.tolist())
1530 ):
1531 q_end = q_start + q_len
1532 q_req = q[q_start:q_end]
1533 positions = torch.arange(seq_len, device="cuda")
1534 pages = block_table[req_id, positions // BLOCK_SIZE]
1535 rows = positions % BLOCK_SIZE
1536 kv_req = kv_cache[pages, :, rows]
1537 k_req = kv_req[..., :HEAD_DIM]
1538 v_req = kv_req[..., HEAD_DIM:].float()
1539
1540 q_pos = prefix_len + torch.arange(q_len, device="cuda")
1541 key_blocks = positions // BLOCK_SIZE
1542 causal_mask = positions.unsqueeze(0) <= q_pos.unsqueeze(1)
1543
1544 for kv_head in range(NUM_KV_HEADS):
1545 selected = topk_idx[kv_head, q_start:q_end]
1546 selected_mask = (key_blocks[None, :, None] == selected[:, None, :]).any(-1)
1547 mask = causal_mask & selected_mask
1548 head_start = kv_head * gqa_group_size
1549 head_end = head_start + gqa_group_size
1550
1551 q_heads = q_req[:, head_start:head_end].transpose(0, 1)
1552 k_head = k_req[:, kv_head].T.expand(gqa_group_size, -1, -1)
1553 scores = torch.bmm(q_heads, k_head, out_dtype=torch.float32)
1554 scores = scores.transpose(0, 1) * SM_SCALE
1555 probs = torch.softmax(
1556 scores.masked_fill(~mask[:, None, :], -float("inf")), -1
1557 )
1558 out[q_start:q_end, head_start:head_end] = torch.einsum(
1559 "qhk,kd->qhd", probs, v_req[:, kv_head]
1560 )
1561 q_start += q_len
1562 return out.to(q.dtype)
1563
1564
1565 @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
1566 @pytest.mark.parametrize(
1567 ("q_lens", "kv_lens"),
1568 [
1569 ((129, 257), (129, 257)),
1570 ((65, 129, 257), (129, 257, 385)),
1571 ],
1572 )
1573 def test_prefill_sparse_attention_correctness(
1574 kv_layout: KVCacheLayout,
1575 q_lens: tuple[int, ...],
1576 kv_lens: tuple[int, ...],
1577 ):
1578 assert len(q_lens) == len(kv_lens)
1579 assert all(kv_len >= q_len for q_len, kv_len in zip(q_lens, kv_lens))
1580
1581 # Build paged-KV metadata, including a non-identity page order.
1582 batch = len(q_lens)
1583 pages_per_req = [(kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE for kv_len in kv_lens]
1584 max_blocks = max(pages_per_req)
1585 num_pages = sum(pages_per_req)
1586 physical_pages = torch.randperm(num_pages, device="cuda", dtype=torch.int32)
1587 block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
1588 base_page = 0
1589 for req_id, num_req_pages in enumerate(pages_per_req):
1590 block_table[req_id, :num_req_pages] = physical_pages[
1591 base_page : base_page + num_req_pages
1592 ]
1593 base_page += num_req_pages
1594
1595 q_lens_t = torch.tensor(q_lens, device="cuda", dtype=torch.int32)
1596 seq_lens = torch.tensor(kv_lens, device="cuda", dtype=torch.int32)
1597 prefix_lens = seq_lens - q_lens_t
1598 cu_seqlens = torch.zeros(batch + 1, device="cuda", dtype=torch.int32)
1599 cu_seqlens[1:] = q_lens_t.cumsum(0)
1600 total_q = sum(q_lens)
1601 max_seqlen_q = max(q_lens)
1602
1603 q_shape = (total_q, NUM_Q_HEADS, HEAD_DIM)
1604 q = torch.randn(q_shape, device="cuda", dtype=DTYPE)
1605 # Allocate the main KV cache through the backend layout contract so the
1606 # physical storage matches the active layout (contiguous NHD or strided
1607 # HND), while the kernels and reference see the logical-NHD view.
1608 kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout)
1609
1610 # Build sparse block indices with the same contract as the real M3 indexer:
1611 # one forced local block, then score-selected older causal blocks.
1612 topk_shape = (NUM_KV_HEADS, total_q, TOPK)
1613 topk_idx = torch.full(topk_shape, -1, device="cuda", dtype=torch.int32)
1614 q_start = 0
1615 for q_len, prefix_len in zip(q_lens_t.tolist(), prefix_lens.tolist()):
1616 for local_q in range(q_len):
1617 current_block = (prefix_len + local_q) // BLOCK_SIZE
1618 older_blocks = torch.randperm(
1619 current_block, device="cuda", dtype=torch.int32
1620 )
1621 selected = torch.cat(
1622 [
1623 torch.tensor([current_block], device="cuda", dtype=torch.int32),
1624 older_blocks[: TOPK - 1],
1625 ]
1626 )
1627 topk_idx[:, q_start + local_q, : selected.numel()] = selected
1628 q_start += q_len
1629
1630 actual = torch.empty_like(q)
1631 minimax_m3_sparse_attn(
1632 q,
1633 kv_cache,
1634 topk_idx,
1635 block_table,
1636 cu_seqlens,
1637 seq_lens,
1638 prefix_lens,
1639 max_seqlen_q,
1640 NUM_KV_HEADS,
1641 SM_SCALE,
1642 actual,
1643 )
1644
1645 expected = _reference_sparse_attn(
1646 q,
1647 kv_cache,
1648 topk_idx,
1649 block_table,
1650 q_lens_t,
1651 seq_lens,
1652 prefix_lens,
1653 )
1654 torch.accelerator.synchronize()
1655
1656 error = (actual.float() - expected.float()).abs()
1657 assert error.mean().item() < 2.5e-4
1658 assert error.max().item() < 1.7e-2
1659
1660
1661 def test_main_cache_layout_contract():
1662 """The standardized per-layer logical shape is [B, H, N, C] with packed
1663 K/V content, and the legacy layout aliases resolve to the expected
1664 per-layer stride orders."""
1665 nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
1666 logical = _main_kv_logical_shape(nb)
1667 assert logical == (nb, h, bs, 2 * d)
1668 # The old separate K/V-axis shape is no longer the logical shape.
1669 assert logical != (nb, 2, bs, h, d)
1670
1671 assert KVCacheLayout.LBHNC.layer_view_order == (0, 1, 2, 3)
1672 assert KVCacheLayout.LBNHC.layer_view_order == (0, 2, 1, 3)
1673
1674 cache_config = CacheConfig()
1675 record_kv_cache_layout(cache_config, "HND")
1676 assert cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBHNC
1677 cache_config = CacheConfig()
1678 record_kv_cache_layout(cache_config, "NHD")
1679 assert cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBNHC
1680
1681 for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC):
1682 order = layout.layer_view_order
1683 # Valid permutation: no duplicates, covers every axis.
1684 assert set(order) == set(range(len(order)))
1685
1686
1687 @pytest.mark.parametrize("layout", [KVCacheLayout.LHBNC, KVCacheLayout.LBHNC])
1688 def test_aiter_sparse_pa_layout_contract(monkeypatch, layout):
1689 """Build the exact separate and packed page-16 views used by the model."""
1690 import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
1691
1692 monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
1693 monkeypatch.setattr(
1694 sparse_attn_mod.rocm_aiter_ops,
1695 "is_shuffle_kv_cache_enabled",
1696 lambda: True,
1697 )
1698
1699 backend = sparse_attn_mod.MiniMaxM3SparseBackend
1700 layouts = backend.supported_kv_cache_layouts()
1701 assert layouts is not None
1702 assert KVCacheLayout.LHBNC in layouts
1703 assert KVCacheLayout.LBHNC in layouts
1704
1705 nb, h = 3, 1
1706 cache_dtype = current_platform.fp8_dtype()
1707 spec = FullAttentionSpec(
1708 block_size=BLOCK_SIZE,
1709 num_kv_heads=h,
1710 head_size=HEAD_DIM,
1711 dtype=cache_dtype,
1712 kv_quant_mode=get_kv_quant_mode("fp8"),
1713 num_head_slots=2,
1714 state_content_bytes=h * HEAD_DIM * cache_dtype.itemsize,
1715 )
1716 raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
1717 kv_cache = dense_kv_cache_views(raw, spec, nb, 1, layout)[0].view(cache_dtype)
1718 key_cache, value_cache = kv_cache.unbind(1)
1719 x = 16 // cache_dtype.itemsize
1720 pages_per_side = BLOCK_SIZE // 16
1721 if key_cache.is_contiguous() and value_cache.is_contiguous():
1722 k_src, v_src = key_cache, value_cache
1723 block_pages, v_page_offset = pages_per_side, 0
1724 else:
1725 assert kv_cache.is_contiguous()
1726 k_src = v_src = kv_cache
1727 block_pages, v_page_offset = 2 * pages_per_side, pages_per_side
1728
1729 num_phys16 = nb * block_pages
1730 k_view = k_src.view(num_phys16, h, HEAD_DIM // x, 16, x)
1731 v_view = v_src.view(num_phys16, h, 16 // x, HEAD_DIM, x)[v_page_offset:]
1732 assert k_view.is_contiguous()
1733 assert v_view.is_contiguous()
1734 if layout is KVCacheLayout.LHBNC:
1735 assert k_view.shape == (24, 1, 8, 16, 16)
1736 assert v_view.shape == (24, 1, 1, 128, 16)
1737 else:
1738 assert k_view.shape == (48, 1, 8, 16, 16)
1739 assert v_view.shape == (40, 1, 1, 128, 16)
1740
1741
1742 def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch):
1743 """Do not pair AITER's separated cache layout with the Triton fallback."""
1744 import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
1745
1746 monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
1747 monkeypatch.setattr(
1748 sparse_attn_mod.rocm_aiter_ops,
1749 "is_shuffle_kv_cache_enabled",
1750 lambda: True,
1751 )
1752
1753 with pytest.raises(ValueError, match="num_kv_heads == 1"):
1754 sparse_attn_mod.minimax_m3_use_aiter_sparse_pa(2)
1755
1756 # An indexer that emits the attend's page table addresses every head, so
1757 # the pairing the check guards against does not arise.
1758 assert sparse_attn_mod.minimax_m3_use_aiter_sparse_pa(
1759 2, emits_sparse_block_table=True
1760 )
1761 # The AITER path carries its own backend: its builder is what rebases the
1762 # block table the indexer's top-k resolves through, so selecting the impl
1763 # without it would leave that field unset.
1764 backend, impl_cls = sparse_attn_mod.select_main_backend_and_impl_cls(
1765 topk_blocks=TOPK,
1766 kv_cache_dtype="fp8_e4m3",
1767 num_kv_heads=2,
1768 emits_sparse_block_table=True,
1769 )
1770 assert impl_cls.__name__ == "MiniMaxM3SparseAiterPAImpl"
1771 assert backend.__name__ == "MiniMaxM3SparseAiterPABackend"
1772 assert backend.get_builder_cls().__name__ == "MiniMaxM3SparseAiterPAMetadataBuilder"
1773
1774
1775 def test_aiter_indexer_requires_the_aiter_attend(monkeypatch):
1776 """The top-k emits the attend's page table, so it needs that attend.
1777
1778 Selected without it, the indexer would reach its top-k with no table to
1779 resolve through; falling back instead surfaces the fp8 index cache as the
1780 configuration error it is.
1781 """
1782 import vllm.models.minimax_m3.amd.indexer_aiter as indexer_aiter_mod
1783
1784 monkeypatch.setattr(indexer_aiter_mod.current_platform, "is_rocm", lambda: True)
1785
1786 def unsupported_reason() -> str | None:
1787 return indexer_aiter_mod.aiter_indexer_unsupported_reason(
1788 topk_blocks=TOPK,
1789 sparse_block_size=BLOCK_SIZE,
1790 num_index_heads=1,
1791 index_head_dim=HEAD_DIM,
1792 # The implementation also accepts this legacy fp8 alias.
1793 indexer_kv_dtype="fp8_e4m3",
1794 max_model_len=8192,
1795 )
1796
1797 monkeypatch.setattr(
1798 indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: False
1799 )
1800 reason = unsupported_reason()
1801 assert reason is not None and "AITER sparse PA attend" in reason
1802
1803 # With the attend asked for, the gate moves on to the kernel-contract
1804 # checks rather than stopping here.
1805 monkeypatch.setattr(
1806 indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: True
1807 )
1808 reason = unsupported_reason()
1809 assert reason is None or "AITER sparse PA attend" not in reason
1810
1811
1812 def test_aiter_consolidated_qknorm_enabled(monkeypatch):
1813 from vllm._aiter_ops import rocm_aiter_ops
1814
1815 monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True)
1816 monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
1817
1818 scale = torch.tensor([0.25])
1819 assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto")
1820 assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8", scale, scale)
1821 assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e4m3", scale, scale)
1822 assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8")
1823 assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8", scale, None)
1824 assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e5m2", scale, scale)
1825
1826 monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", False)
1827 assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto")
1828
1829 monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
1830 monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: False)
1831 assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto")
1832
1833
1834 def _fake_aiter_qknorm_args(kv_cache_dtype: str = "auto") -> dict:
1835 tensor = torch.empty(1)
1836 return {
1837 "qkv": tensor,
1838 "q_norm_weight": tensor,
1839 "k_norm_weight": tensor,
1840 "cos_sin_cache": tensor,
1841 "positions": tensor,
1842 "num_heads": 16,
1843 "num_kv_heads": 1,
1844 "rotary_dim": 64,
1845 "eps": 1e-6,
1846 "slot_mapping": tensor,
1847 "kv_cache_k": tensor,
1848 "kv_cache_v": tensor,
1849 "q_out": tensor,
1850 "kv_cache_dtype": kv_cache_dtype,
1851 }
1852
1853
1854 def _install_fake_aiter_qknorm(monkeypatch: pytest.MonkeyPatch) -> list[dict]:
1855 """Stub ``aiter.fused_qknorm_idxrqknorm`` so mapping tests run without AITER."""
1856 import sys
1857 import types
1858 from typing import Any
1859
1860 calls: list[dict] = []
1861
1862 def fake_op(*args, **kwargs):
1863 calls.append(kwargs)
1864
1865 existing = sys.modules.get("aiter")
1866 if existing is not None:
1867 monkeypatch.setattr(existing, "fused_qknorm_idxrqknorm", fake_op, raising=False)
1868 else:
1869 fake_mod: Any = types.ModuleType("aiter")
1870 fake_mod.fused_qknorm_idxrqknorm = fake_op
1871 monkeypatch.setitem(sys.modules, "aiter", fake_mod)
1872 return calls
1873
1874
1875 def test_aiter_consolidated_qknorm_packed_shuffle_unequal_spans(monkeypatch):
1876 from vllm._aiter_ops import rocm_aiter_ops
1877
1878 calls = _install_fake_aiter_qknorm(monkeypatch)
1879 call_args = _fake_aiter_qknorm_args()
1880 call_args["kv_cache_k"] = torch.empty(2, 1)
1881 call_args["kv_cache_v"] = torch.empty(1, 1)
1882
1883 rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
1884 assert len(calls) == 1
1885
1886
1887 @pytest.mark.parametrize(
1888 ("kv_cache_dtype", "aiter_dtype", "passes_scales"),
1889 [
1890 ("auto", "auto", False),
1891 ("fp8", "fp8_e4m3_static", True),
1892 ("fp8_e4m3", "fp8_e4m3_static", True),
1893 ],
1894 )
1895 def test_aiter_consolidated_qknorm_dtype_and_layout_mapping(
1896 monkeypatch, kv_cache_dtype, aiter_dtype, passes_scales
1897 ):
1898 from vllm._aiter_ops import rocm_aiter_ops
1899
1900 calls = _install_fake_aiter_qknorm(monkeypatch)
1901 call_args = _fake_aiter_qknorm_args(kv_cache_dtype)
1902 k_scale = torch.tensor([0.25])
1903 v_scale = torch.tensor([0.5])
1904 call_args.update(k_scale=k_scale, v_scale=v_scale)
1905
1906 rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
1907 assert len(calls) == 1
1908 kwargs = calls[0]
1909 assert kwargs["kv_cache_dtype"] == aiter_dtype
1910 assert kwargs["index_cache_dtype"] == "auto"
1911 assert kwargs["block_size"] == 16
1912 assert kwargs["asm_layout"] is True
1913 if passes_scales:
1914 assert kwargs["k_scale"] is k_scale
1915 assert kwargs["v_scale"] is v_scale
1916 else:
1917 assert kwargs["k_scale"] is None
1918 assert kwargs["v_scale"] is None
1919
1920
1921 def test_aiter_consolidated_qknorm_fp8_index_dtype_mapping(monkeypatch):
1922 from vllm._aiter_ops import rocm_aiter_ops
1923
1924 calls = _install_fake_aiter_qknorm(monkeypatch)
1925 fp8 = current_platform.fp8_dtype()
1926 call_args = _fake_aiter_qknorm_args("fp8")
1927 call_args.update(
1928 k_scale=torch.tensor([0.25]),
1929 v_scale=torch.tensor([0.5]),
1930 index_cache=torch.empty(1, dtype=fp8),
1931 index_q_out=torch.empty(1, dtype=fp8),
1932 num_index_heads=1,
1933 )
1934
1935 rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
1936 assert calls[0]["index_cache_dtype"] == "fp8"
1937 assert calls[0]["kv_cache_dtype"] == "fp8_e4m3_static"
1938
1939
1940 @pytest.mark.parametrize("skip_index_branch", [False, True])
1941 def test_aiter_consolidated_qknorm_full_and_skip_index_args(
1942 monkeypatch, skip_index_branch
1943 ):
1944 from vllm._aiter_ops import rocm_aiter_ops
1945
1946 calls = _install_fake_aiter_qknorm(monkeypatch)
1947 call_args = _fake_aiter_qknorm_args()
1948 if skip_index_branch:
1949 call_args.update(skip_index_branch=True)
1950 else:
1951 tensor = torch.empty(1)
1952 call_args.update(
1953 index_q_norm_weight=tensor,
1954 index_k_norm_weight=tensor,
1955 num_index_heads=1,
1956 index_cache=tensor,
1957 index_q_out=tensor,
1958 index_slot_mapping=tensor,
1959 )
1960
1961 rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
1962 kwargs = calls[0]
1963 assert kwargs["skip_index_branch"] is skip_index_branch
1964 if skip_index_branch:
1965 assert kwargs["index_q_norm_weight"] is None
1966 assert kwargs["index_k_norm_weight"] is None
1967 assert kwargs["index_cache"] is None
1968 assert kwargs["index_q_out"] is None
1969 assert kwargs["index_slot_mapping"] is None
1970 else:
1971 assert kwargs["index_q_norm_weight"] is not None
1972 assert kwargs["index_k_norm_weight"] is not None
1973 assert kwargs["index_cache"] is not None
1974 assert kwargs["index_q_out"] is not None
1975 assert kwargs["index_slot_mapping"] is not None
1976
1977
1978 def test_aiter_consolidated_qknorm_rejects_fp8_e5m2(monkeypatch):
1979 from vllm._aiter_ops import rocm_aiter_ops
1980
1981 monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True)
1982 monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
1983 scale = torch.tensor([0.25])
1984 assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e5m2", scale, scale)
1985
1986
1987 def test_aiter_sparse_pa_slot_rebase_preserves_page_offsets():
1988 slots = torch.tensor([0, 15, 16, 127, 128, 129, -1])
1989 actual = minimax_m3_rebase_slots_to_page16(slots, block_size=128)
1990 expected = torch.tensor([0, 15, 16, 127, 256, 257, -1])
1991 assert torch.equal(actual, expected)
1992
1993
1994 @pytest.mark.skipif(
1995 not current_platform.is_rocm(),
1996 reason="The consolidated packed-SHUFFLE AITER op is ROCm-only",
1997 )
1998 @pytest.mark.parametrize("skip_index_branch", [False, True])
1999 @pytest.mark.parametrize("fp8_index", [False, True])
2000 def test_aiter_consolidated_qknorm_real_packed_shuffle(skip_index_branch, fp8_index):
2001 from vllm._aiter_ops import rocm_aiter_ops
2002
2003 torch.manual_seed(20260831)
2004 num_tokens = 4
2005 num_heads = 16
2006 num_kv_heads = 1
2007 num_index_heads = 1
2008 rotary_dim = 64
2009 head_dim = 128
2010 logical_block_size = 128
2011 num_blocks = 3
2012 fp8_dtype = current_platform.fp8_dtype()
2013
2014 qkv = torch.randn(
2015 num_tokens,
2016 (num_heads + 2 * num_kv_heads + num_index_heads + 1) * head_dim,
2017 dtype=torch.bfloat16,
2018 device="cuda",
2019 )
2020 weights = [
2021 torch.randn(head_dim, dtype=torch.bfloat16, device="cuda") for _ in range(4)
2022 ]
2023 cos_sin_cache = torch.randn(16, rotary_dim, dtype=torch.bfloat16, device="cuda")
2024 positions = torch.arange(num_tokens, dtype=torch.int64, device="cuda")
2025 logical_slots = torch.tensor([0, 127, 128, -1], dtype=torch.int64, device="cuda")
2026 slot_mapping = minimax_m3_rebase_slots_to_page16(logical_slots, logical_block_size)
2027 index_slot_mapping = logical_slots.clone()
2028
2029 packed = torch.zeros(
2030 num_blocks,
2031 2,
2032 logical_block_size,
2033 head_dim,
2034 dtype=fp8_dtype,
2035 device="cuda",
2036 )
2037 x = 16 // packed.element_size()
2038 pages_per_side = logical_block_size // 16
2039 num_phys16 = num_blocks * 2 * pages_per_side
2040 kv_cache_k = packed.view(num_phys16, num_kv_heads, head_dim // x, 16, x)
2041 kv_cache_v = packed.view(num_phys16, num_kv_heads, 16 // x, head_dim, x)[
2042 pages_per_side:
2043 ]
2044 assert kv_cache_k.shape == (48, 1, 8, 16, 16)
2045 assert kv_cache_v.shape == (40, 1, 1, 128, 16)
2046
2047 q_out = torch.zeros(
2048 num_tokens, num_heads * head_dim, dtype=torch.bfloat16, device="cuda"
2049 )
2050 index_dtype = fp8_dtype if fp8_index else torch.bfloat16
2051 if fp8_index:
2052 index_q_out = torch.empty(
2053 (num_tokens, num_index_heads * head_dim),
2054 dtype=index_dtype,
2055 device="cuda",
2056 )
2057 index_cache = torch.empty(
2058 (num_blocks, logical_block_size, head_dim),
2059 dtype=index_dtype,
2060 device="cuda",
2061 )
2062 index_q_out.view(torch.uint8).fill_(7)
2063 index_cache.view(torch.uint8).fill_(7)
2064 else:
2065 index_q_out = torch.full(
2066 (num_tokens, num_index_heads * head_dim),
2067 7,
2068 dtype=torch.bfloat16,
2069 device="cuda",
2070 )
2071 index_cache = torch.full(
2072 (num_blocks, logical_block_size, head_dim),
2073 7,
2074 dtype=torch.bfloat16,
2075 device="cuda",
2076 )
2077 scale = torch.tensor([0.02], dtype=torch.float32, device="cuda")
2078
2079 kwargs = {}
2080 if not skip_index_branch:
2081 kwargs = {
2082 "index_q_norm_weight": weights[2],
2083 "index_k_norm_weight": weights[3],
2084 "index_cache": index_cache,
2085 "index_q_out": index_q_out,
2086 "index_slot_mapping": index_slot_mapping,
2087 }
2088 rocm_aiter_ops.fused_qknorm_idxrqknorm(
2089 qkv,
2090 weights[0],
2091 weights[1],
2092 cos_sin_cache,
2093 positions,
2094 num_heads,
2095 num_kv_heads,
2096 rotary_dim,
2097 1e-6,
2098 slot_mapping,
2099 kv_cache_k,
2100 kv_cache_v,
2101 q_out,
2102 "fp8",
2103 scale,
2104 scale,
2105 num_index_heads=num_index_heads,
2106 skip_index_branch=skip_index_branch,
2107 **kwargs,
2108 )
2109 torch.accelerator.synchronize()
2110
2111 assert torch.isfinite(q_out).all()
2112 assert torch.count_nonzero(q_out) > 0
2113 assert torch.count_nonzero(packed) > 0
2114 if fp8_index:
2115 index_q_untouched = torch.all(index_q_out.view(torch.uint8) == 7)
2116 index_cache_untouched = torch.all(index_cache.view(torch.uint8) == 7)
2117 index_cache_t0 = torch.all(index_cache[0, 0].view(torch.uint8) == 7)
2118 index_cache_t127 = torch.all(index_cache[0, 127].view(torch.uint8) == 7)
2119 index_cache_t128 = torch.all(index_cache[1, 0].view(torch.uint8) == 7)
2120 else:
2121 index_q_untouched = torch.all(index_q_out == 7)
2122 index_cache_untouched = torch.all(index_cache == 7)
2123 index_cache_t0 = torch.all(index_cache[0, 0] == 7)
2124 index_cache_t127 = torch.all(index_cache[0, 127] == 7)
2125 index_cache_t128 = torch.all(index_cache[1, 0] == 7)
2126 if skip_index_branch:
2127 assert index_q_untouched
2128 assert index_cache_untouched
2129 else:
2130 assert not index_q_untouched
2131 assert not index_cache_t0
2132 assert not index_cache_t127
2133 assert not index_cache_t128
2134
2135
2136 def test_indexer_cache_squeezes_to_contiguous_3d():
2137 """The indexer side cache is standardized 4D with H=1: under both layouts
2138 the allocator's logical view stays contiguous and squeezes (as
2139 `MiniMaxM3IndexerCache.bind_kv_cache` does) to the 3-dim
2140 [num_blocks, block_size, head_dim] cache the kernels consume."""
2141 nb = 5
2142 ispec = MLAAttentionSpec(
2143 block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
2144 )
2145 shape_bytes = compute_layer_kv_cache_shape_bytes(ispec, nb)
2146 assert shape_bytes == (nb, 1, BLOCK_SIZE, HEAD_DIM * DTYPE.itemsize)
2147 assert _layer_stride_order(KVCacheLayout.LBNHC, 3) == (0, 1, 2)
2148
2149 for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC):
2150 iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
2151 view = dense_kv_cache_views(iraw, ispec, nb, 1, layout, BLOCK_SIZE)[0]
2152 assert tuple(view.shape) == (nb, 1, BLOCK_SIZE, HEAD_DIM)
2153 indexer_cache = view.squeeze(1)
2154 assert tuple(indexer_cache.shape) == (nb, BLOCK_SIZE, HEAD_DIM)
2155 assert indexer_cache.is_contiguous()
2156
2157
2158 def test_hnd_allocation_is_packed_head_major():
2159 """Under HND the backend-visible logical view is the packed head-major
2160 physical allocation."""
2161 nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
2162 logical = _main_kv_logical_shape(nb)
2163 stride_order = KVCacheLayout.LBHNC.layer_view_order
2164
2165 physical_shape = tuple(logical[i] for i in stride_order)
2166 assert physical_shape == (nb, h, bs, 2 * d)
2167
2168 inv_order = [stride_order.index(i) for i in range(len(stride_order))]
2169 raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
2170 view = raw.permute(*inv_order)
2171 expected = raw.view((nb, h, bs, 2 * d))
2172
2173 assert view.shape == expected.shape
2174 assert view.stride() == expected.stride()
2175 assert view.storage_offset() == expected.storage_offset()
2176
2177 # Negative: the NHD stride order under HND does not reproduce the
2178 # head-major view.
2179 wrong_order = (0, 2, 1, 3)
2180 wrong_inv = [wrong_order.index(i) for i in range(len(wrong_order))]
2181 wrong_view = raw.view(tuple(logical[i] for i in wrong_order)).permute(*wrong_inv)
2182 assert wrong_view.stride() != expected.stride()
2183
2184
2185 def test_main_cache_is_block_first_and_unpadded():
2186 """The allocator's contiguous-view branch (not the padded-strided branch)
2187 is used for the main GQA cache: its spec is unpadded and the physical
2188 layout keeps num_blocks as the first dimension under both layouts."""
2189 spec = _main_spec()
2190 # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided().
2191 assert spec.page_size_padded is None
2192
2193 logical = _main_kv_logical_shape(4)
2194 for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC):
2195 order = layout.layer_view_order
2196 inv_order = [order.index(i) for i in range(len(order))]
2197 # Physical first dim is num_blocks (block-first); required by the
2198 # padded-strided branch's block-first assumption if it were ever taken.
2199 assert inv_order[0] == 0
2200 assert logical[order[0]] == logical[0]
2201
2202
2203 def _build_decode_inputs(
2204 seq_lens_list: tuple[int, ...],
2205 decode_query_len: int = 1,
2206 num_padded_reqs: int = 0,
2207 ):
2208 """Shared decode setup: uniform query tokens per request, a non-identity
2209 block table, and topk indices selecting the current block plus older causal
2210 blocks for each query token."""
2211 active_batch = len(seq_lens_list)
2212 batch = active_batch + num_padded_reqs
2213 pages_per_req = [(s + BLOCK_SIZE - 1) // BLOCK_SIZE for s in seq_lens_list]
2214 max_blocks = max(pages_per_req)
2215 num_pages = sum(pages_per_req)
2216 physical_pages = torch.randperm(num_pages, device="cuda", dtype=torch.int32)
2217 block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
2218 base_page = 0
2219 for req_id, num_req_pages in enumerate(pages_per_req):
2220 block_table[req_id, :num_req_pages] = physical_pages[
2221 base_page : base_page + num_req_pages
2222 ]
2223 base_page += num_req_pages
2224
2225 seq_lens = torch.tensor(
2226 (*seq_lens_list, *([0] * num_padded_reqs)),
2227 device="cuda",
2228 dtype=torch.int32,
2229 )
2230 q = torch.randn(
2231 batch * decode_query_len, NUM_Q_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE
2232 )
2233
2234 topk_idx = torch.full(
2235 (NUM_KV_HEADS, batch * decode_query_len, TOPK),
2236 -1,
2237 device="cuda",
2238 dtype=torch.int32,
2239 )
2240 token_id = 0
2241 for req_id, seq_len in enumerate(seq_lens_list):
2242 for local_q in range(decode_query_len):
2243 query_pos = seq_len - decode_query_len + local_q
2244 current_block = query_pos // BLOCK_SIZE
2245 older_blocks = torch.randperm(
2246 current_block, device="cuda", dtype=torch.int32
2247 )
2248 selected = torch.cat(
2249 [
2250 torch.tensor([current_block], device="cuda", dtype=torch.int32),
2251 older_blocks[: TOPK - 1],
2252 ]
2253 )
2254 topk_idx[:, token_id, : selected.numel()] = selected
2255 token_id += 1
2256
2257 return q, block_table, seq_lens, topk_idx, num_pages
2258
2259
2260 @pytest.mark.skipif(
2261 not current_platform.is_rocm(),
2262 reason="MiniMax-M3 AITER sparse PA is ROCm-only",
2263 )
2264 def test_aiter_sparse_block_table_handles_padded_decode_rows():
2265 """Zero-length padded decode rows must produce empty sparse attention."""
2266 from vllm.models.minimax_m3.amd.ops.sparse_pa import (
2267 PAGES_PER_SPARSE_BLOCK,
2268 minimax_m3_build_sparse_block_table,
2269 )
2270
2271 topk_idx = torch.tensor(
2272 [[[1, 0, -1], [0, 1, -1]]],
2273 device="cuda",
2274 dtype=torch.int32,
2275 )
2276 block_table = torch.tensor(
2277 [[5, 7], [11, 13]],
2278 device="cuda",
2279 dtype=torch.int32,
2280 )
2281 seq_lens = torch.tensor([129, 0], device="cuda", dtype=torch.int32)
2282
2283 sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table(
2284 topk_idx, block_table, seq_lens
2285 )
2286 torch.accelerator.synchronize()
2287
2288 expected_bt = torch.zeros_like(sparse_bt)
2289 expected_ctx = torch.tensor([129, 0], device="cuda", dtype=torch.int32)
2290
2291 for slot, block_id in enumerate((0, 1)):
2292 base_phys = int(block_table[0, block_id]) * PAGES_PER_SPARSE_BLOCK
2293 start = slot * PAGES_PER_SPARSE_BLOCK
2294 expected_bt[0, start : start + PAGES_PER_SPARSE_BLOCK] = torch.arange(
2295 base_phys,
2296 base_phys + PAGES_PER_SPARSE_BLOCK,
2297 device="cuda",
2298 dtype=torch.int32,
2299 )
2300
2301 assert torch.equal(sparse_ctx, expected_ctx)
2302 assert torch.equal(sparse_bt, expected_bt)
2303
2304
2305 @pytest.mark.skipif(
2306 not current_platform.is_rocm(),
2307 reason="MiniMax-M3 AITER sparse PA is ROCm-only",
2308 )
2309 def test_aiter_decode_sparse_block_table_supports_spec_decode():
2310 """AITER sparse PA needs one compact page-16 table per speculative query.
2311
2312 The decode indexer flattens speculative rows as
2313 ``req_id * decode_query_len + local_q``. This verifies that the ROCm AITER
2314 block-table adapter uses the same mapping before handing rows to Gluon.
2315 """
2316 from vllm.models.minimax_m3.amd.ops.sparse_pa import (
2317 PAGES_PER_SPARSE_BLOCK,
2318 minimax_m3_build_sparse_block_table_decode,
2319 )
2320
2321 decode_query_len = 4
2322 seq_lens = torch.tensor([260, 132, 0], device="cuda", dtype=torch.int32)
2323 block_table = torch.tensor(
2324 [
2325 [5, 7, 11],
2326 [13, 17, 19],
2327 [0, 0, 0],
2328 ],
2329 device="cuda",
2330 dtype=torch.int32,
2331 )
2332 topk_idx = torch.tensor(
2333 [
2334 [
2335 [2, 0, 3, -1],
2336 [0, 2, 3, -1],
2337 [2, 1, 0, -1],
2338 [1, 2, 0, -1],
2339 [1, 0, 2, -1],
2340 [0, 1, 2, -1],
2341 [1, -1, 0, -1],
2342 [0, 2, 1, -1],
2343 [0, 1, -1, -1],
2344 [0, 1, -1, -1],
2345 [0, 1, -1, -1],
2346 [0, 1, -1, -1],
2347 ]
2348 ],
2349 device="cuda",
2350 dtype=torch.int32,
2351 )
2352
2353 sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_decode(
2354 topk_idx, block_table, seq_lens, decode_query_len
2355 )
2356 torch.accelerator.synchronize()
2357
2358 expected_bt = torch.zeros_like(sparse_bt)
2359 expected_ctx = torch.empty_like(sparse_ctx)
2360 topk_cpu = topk_idx.cpu()[0].tolist()
2361 block_table_cpu = block_table.cpu()
2362 seq_lens_cpu = seq_lens.cpu().tolist()
2363
2364 for row, selected_blocks in enumerate(topk_cpu):
2365 req_id = row // decode_query_len
2366 local_q = row - req_id * decode_query_len
2367 query_pos = seq_lens_cpu[req_id] - decode_query_len + local_q
2368 causal_len = max(query_pos + 1, 0)
2369 if causal_len == 0:
2370 expected_ctx[row] = 0
2371 continue
2372 self_block = query_pos // BLOCK_SIZE
2373
2374 valid = [blk for blk in selected_blocks if 0 <= blk <= self_block]
2375 full_blocks = [blk for blk in valid if blk < self_block]
2376 tail_blocks = [blk for blk in valid if blk == self_block]
2377 ordered_blocks = full_blocks + tail_blocks
2378
2379 for slot, block_id in enumerate(ordered_blocks):
2380 base_phys = int(block_table_cpu[req_id, block_id]) * PAGES_PER_SPARSE_BLOCK
2381 start = slot * PAGES_PER_SPARSE_BLOCK
2382 expected_bt[row, start : start + PAGES_PER_SPARSE_BLOCK] = torch.arange(
2383 base_phys,
2384 base_phys + PAGES_PER_SPARSE_BLOCK,
2385 device="cuda",
2386 dtype=torch.int32,
2387 )
2388
2389 if tail_blocks:
2390 expected_ctx[row] = (
2391 len(full_blocks) * BLOCK_SIZE + causal_len - self_block * BLOCK_SIZE
2392 )
2393 else:
2394 expected_ctx[row] = min(len(valid) * BLOCK_SIZE, causal_len)
2395
2396 assert torch.equal(sparse_ctx, expected_ctx)
2397 assert torch.equal(sparse_bt, expected_bt)
2398
2399
2400 @pytest.mark.skipif(
2401 not current_platform.is_rocm(),
2402 reason="MiniMax-M3 fused decode top-k is ROCm-only",
2403 )
2404 @pytest.mark.parametrize("block_page_stride", [8, 16])
2405 @pytest.mark.parametrize(
2406 "seq_lens_list",
2407 [
2408 [4097],
2409 [65537],
2410 [1, 129, 4097, 0],
2411 [127, 128, 129, 256, 4097, 0],
2412 [1, 127, 128, 129, 255, 256, 257, 4097, 0],
2413 ],
2414 )
2415 def test_amd_decode_fused_topk_emits_sparse_table(
2416 block_page_stride: int, seq_lens_list: list[int]
2417 ):
2418 """The fused selector must match top-k followed by table construction."""
2419 from vllm.models.minimax_m3.amd.ops.index_topk import (
2420 minimax_m3_index_decode as amd_minimax_m3_index_decode,
2421 )
2422 from vllm.models.minimax_m3.amd.ops.sparse_pa import (
2423 PAGES_PER_SPARSE_BLOCK,
2424 minimax_m3_build_sparse_block_table_decode,
2425 )
2426
2427 torch.manual_seed(0)
2428 decode_query_len = 4
2429 seq_lens = torch.tensor(seq_lens_list, device="cuda", dtype=torch.int32)
2430 batch = seq_lens.numel()
2431 total_q = batch * decode_query_len
2432 max_seq_len = int(seq_lens.max())
2433 max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
2434
2435 index_block_table = torch.zeros(
2436 (batch, max_blocks), device="cuda", dtype=torch.int32
2437 )
2438 for req_id, seq_len in enumerate(seq_lens.tolist()):
2439 num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
2440 pages = torch.arange(max_blocks, device="cuda", dtype=torch.int32)
2441 index_block_table[req_id, :num_blocks] = torch.roll(pages, req_id)[:num_blocks]
2442
2443 # The index cache and main-attention cache have independent physical page
2444 # mappings in production. Disjoint values make using the wrong table fail.
2445 attention_block_table = torch.full_like(index_block_table, -31337)
2446 for req_id, seq_len in enumerate(seq_lens.tolist()):
2447 num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
2448 first_page = 10_000 + req_id * max_blocks
2449 attention_block_table[req_id, :num_blocks] = torch.arange(
2450 first_page + num_blocks - 1,
2451 first_page - 1,
2452 -1,
2453 device="cuda",
2454 dtype=torch.int32,
2455 )
2456 assert not torch.equal(index_block_table, attention_block_table)
2457
2458 idx_q = torch.randn(total_q, 1, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
2459 index_kv_cache = torch.randn(
2460 max_blocks,
2461 BLOCK_SIZE,
2462 HEAD_DIM,
2463 device="cuda",
2464 dtype=torch.bfloat16,
2465 )
2466 decode_kwargs = {
2467 "max_seq_len": max_seq_len,
2468 "topk": TOPK,
2469 "init_blocks": 0,
2470 "local_blocks": 1,
2471 "num_kv_heads": 1,
2472 "decode_query_len": decode_query_len,
2473 "max_decode_query_len": decode_query_len,
2474 }
2475 completion_counter = torch.zeros((1, total_q), device="cuda", dtype=torch.int32)
2476
2477 expected_topk = amd_minimax_m3_index_decode(
2478 idx_q,
2479 index_kv_cache,
2480 index_block_table,
2481 seq_lens,
2482 completion_counter=completion_counter,
2483 **decode_kwargs,
2484 )
2485 expected_sparse_bt, expected_sparse_ctx = (
2486 minimax_m3_build_sparse_block_table_decode(
2487 expected_topk,
2488 attention_block_table,
2489 seq_lens,
2490 decode_query_len,
2491 block_page_stride,
2492 )
2493 )
2494
2495 topk_out = torch.full_like(expected_topk, -27182)
2496 sparse_bt_out = torch.full(
2497 (total_q, TOPK * PAGES_PER_SPARSE_BLOCK),
2498 -27182,
2499 device="cuda",
2500 dtype=torch.int32,
2501 )
2502 sparse_ctx_out = torch.full((total_q,), -27182, device="cuda", dtype=torch.int32)
2503 actual_topk = amd_minimax_m3_index_decode(
2504 idx_q,
2505 index_kv_cache,
2506 index_block_table,
2507 seq_lens,
2508 out=topk_out,
2509 attention_block_table=attention_block_table,
2510 sparse_block_table_out=sparse_bt_out,
2511 sparse_context_lens_out=sparse_ctx_out,
2512 block_page_stride=block_page_stride,
2513 completion_counter=completion_counter,
2514 **decode_kwargs,
2515 )
2516 torch.accelerator.synchronize()
2517
2518 assert actual_topk.data_ptr() == topk_out.data_ptr()
2519 assert torch.equal(actual_topk, expected_topk)
2520 assert torch.equal(sparse_bt_out, expected_sparse_bt)
2521 assert torch.equal(sparse_ctx_out, expected_sparse_ctx)
2522 assert torch.count_nonzero(completion_counter).item() == 0
2523
2524 if seq_lens_list[-1] == 0:
2525 padded_rows = slice(total_q - decode_query_len, total_q)
2526 assert torch.all(actual_topk[:, padded_rows] == -1)
2527 assert torch.count_nonzero(sparse_bt_out[padded_rows]) == 0
2528 assert torch.count_nonzero(sparse_ctx_out[padded_rows]) == 0
2529
2530
2531 @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
2532 @pytest.mark.parametrize(
2533 "seq_lens_list",
2534 [(130, 257), (129, 200, 384)],
2535 )
2536 @pytest.mark.parametrize("decode_query_len", [1, 4])
2537 @pytest.mark.parametrize("num_padded_reqs", [0, 2])
2538 def test_decode_sparse_attention_correctness(
2539 kv_layout: KVCacheLayout,
2540 seq_lens_list: tuple[int, ...],
2541 decode_query_len: int,
2542 num_padded_reqs: int,
2543 ):
2544 """Decode (split-K) parity under both layouts: this is the only coverage of
2545 the decode-site cache feed, and the strided HND case fails if the kernel
2546 ignores the cache strides."""
2547 torch.manual_seed(0)
2548 q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(
2549 seq_lens_list, decode_query_len, num_padded_reqs
2550 )
2551 kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout)
2552
2553 actual = torch.empty_like(q)
2554 minimax_m3_sparse_attn_decode(
2555 q,
2556 kv_cache,
2557 topk_idx,
2558 block_table,
2559 seq_lens,
2560 NUM_KV_HEADS,
2561 SM_SCALE,
2562 actual,
2563 decode_query_len,
2564 )
2565
2566 # Reuse the prefill reference: decode is a uniform query chunk ending at
2567 # seq_len - 1 for each request.
2568 active_batch = len(seq_lens_list)
2569 active_tokens = active_batch * decode_query_len
2570 q_lens_t = torch.full(
2571 (len(seq_lens_list),), decode_query_len, device="cuda", dtype=torch.int32
2572 )
2573 active_seq_lens = seq_lens[:active_batch]
2574 prefix_lens = active_seq_lens - q_lens_t
2575 expected = _reference_sparse_attn(
2576 q[:active_tokens],
2577 kv_cache,
2578 topk_idx[:, :active_tokens],
2579 block_table[:active_batch],
2580 q_lens_t,
2581 active_seq_lens,
2582 prefix_lens,
2583 )
2584 torch.accelerator.synchronize()
2585
2586 error = (actual[:active_tokens].float() - expected.float()).abs()
2587 assert error.mean().item() < 2.5e-4
2588 assert error.max().item() < 1.7e-2
2589
2590
2591 def test_decode_wrong_layout_breaks_parity():
2592 """Negative (AC-3/AC-5): consuming the physical HND buffer as if it were
2593 already contiguous-NHD (i.e. skipping the allocator's inverse permute)
2594 reorders the K/V content, so the decode output no longer matches the
2595 reference computed on the correct logical view. The mislabeled tensor keeps
2596 the same shape as the correct view, so the kernel stays in bounds."""
2597 torch.manual_seed(0)
2598 seq_lens_list = (130, 257)
2599 q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list)
2600
2601 # Physical HND storage [blocks, heads, block, packed_kv_dim].
2602 phys = torch.randn(
2603 (num_pages, NUM_KV_HEADS, BLOCK_SIZE, 2 * HEAD_DIM),
2604 device="cuda",
2605 dtype=DTYPE,
2606 )
2607 # Correct logical packed-HND view vs. the same bytes mislabeled as NHD
2608 # physical storage and then exposed as a logical cache.
2609 correct = phys
2610 wrong = phys.reshape(num_pages, BLOCK_SIZE, NUM_KV_HEADS, 2 * HEAD_DIM).permute(
2611 0, 2, 1, 3
2612 )
2613
2614 q_lens_t = torch.ones(len(seq_lens_list), device="cuda", dtype=torch.int32)
2615 prefix_lens = seq_lens - q_lens_t
2616 expected = _reference_sparse_attn(
2617 q, correct, topk_idx, block_table, q_lens_t, seq_lens, prefix_lens
2618 )
2619
2620 actual = torch.empty_like(q)
2621 minimax_m3_sparse_attn_decode(
2622 q, wrong, topk_idx, block_table, seq_lens, NUM_KV_HEADS, SM_SCALE, actual, 1
2623 )
2624 torch.accelerator.synchronize()
2625 assert (actual.float() - expected.float()).abs().max().item() > 1.7e-2
2626
2627
2628 def test_main_cache_byte_identical_through_production_allocator():
2629 """AC-2: drive the real allocator (`create_kv_cache_views`) for the M3 main
2630 `FullAttentionSpec` under HND and assert the backend-visible view has the
2631 same shape, stride, and storage offset as the packed-HND allocation."""
2632 nb = 4
2633 spec = _main_spec()
2634 raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
2635 view = dense_kv_cache_views(raw, spec, nb, 1, KVCacheLayout.LBHNC, BLOCK_SIZE)[0]
2636
2637 oracle = raw.view(DTYPE).view((nb, NUM_KV_HEADS, BLOCK_SIZE, 2 * HEAD_DIM))
2638 assert tuple(view.shape) == tuple(oracle.shape)
2639 assert view.stride() == oracle.stride()
2640 assert view.storage_offset() == oracle.storage_offset()
2641
2642
2643 def test_padded_main_cache_is_flagged():
2644 """AC-2.1 negative: the M3 main cache relies on the allocator's
2645 contiguous-view branch (`page_size_padded is None`). A spec that sets
2646 `page_size_padded` is explicitly flagged rather than silently wrong-strided."""
2647
2648 def _require_unpadded_block_first(spec, stride_order):
2649 inv_order = [stride_order.index(i) for i in range(len(stride_order))]
2650 assert spec.page_size_padded is None, (
2651 "main GQA cache must be unpadded to use the contiguous-view "
2652 "allocator branch"
2653 )
2654 assert inv_order[0] == 0, "main GQA cache must remain block-first"
2655
2656 stride_order = KVCacheLayout.LBHNC.layer_view_order
2657
2658 good = FullAttentionSpec(
2659 block_size=BLOCK_SIZE,
2660 num_kv_heads=NUM_KV_HEADS,
2661 head_size=HEAD_DIM,
2662 head_size_v=HEAD_DIM,
2663 dtype=DTYPE,
2664 )
2665 _require_unpadded_block_first(good, stride_order) # passes
2666
2667 padded = FullAttentionSpec(
2668 block_size=BLOCK_SIZE,
2669 num_kv_heads=NUM_KV_HEADS,
2670 head_size=HEAD_DIM,
2671 head_size_v=HEAD_DIM,
2672 dtype=DTYPE,
2673 page_size_padded=good.page_size_bytes + 128,
2674 )
2675 with pytest.raises(AssertionError):
2676 _require_unpadded_block_first(padded, stride_order)
2677
2678
2679 @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
2680 def test_reshape_and_cache_flash_write_persists(kv_layout: KVCacheLayout):
2681 """AC-5 write path: the `reshape_and_cache_flash` write site now consumes
2682 packed-content K/V split views. Writing through those views must persist
2683 into the bound storage under both layouts."""
2684 torch.manual_seed(0)
2685 num_pages = 4
2686 kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout)
2687 with torch.no_grad():
2688 kv_cache.zero_()
2689
2690 # Exactly the production write-site code under test.
2691 key_cache, value_cache = kv_cache.transpose(1, 2).split(HEAD_DIM, dim=-1)
2692
2693 num_tokens = 12
2694 slot_mapping = torch.randperm(num_pages * BLOCK_SIZE, device="cuda")[
2695 :num_tokens
2696 ].to(torch.int64)
2697 key = torch.randn(num_tokens, NUM_KV_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE)
2698 value = torch.randn(num_tokens, NUM_KV_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE)
2699 scale = torch.ones((), device="cuda")
2700 ops.reshape_and_cache_flash(
2701 key, value, key_cache, value_cache, slot_mapping, "auto", scale, scale
2702 )
2703 torch.accelerator.synchronize()
2704
2705 # Read back through the independent logical view; proves the writes landed
2706 # in the engine-bound storage, not a detached copy.
2707 for t in range(num_tokens):
2708 slot = int(slot_mapping[t].item())
2709 blk, intra = divmod(slot, BLOCK_SIZE)
2710 torch.testing.assert_close(kv_cache[blk, :, intra, :HEAD_DIM], key[t])
2711 torch.testing.assert_close(kv_cache[blk, :, intra, HEAD_DIM:], value[t])