[ROCm][Perf] Insert MiniMax-M3 sparse-PA K/V without a contiguous copy (#56849)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96 committed
Sep 17, 2026 at 17:44 UTC
438434b5b5f8936a4dc7d651c708928c90f085d6
2 files changed
+104
-2
tests/kernels/test_minimax_m3_amd_ops.py
+76
@@ -10,6 +10,7 @@ tests assert the two agree within tolerance:
10
* Fused MXFP8 activation quant (Triton) -> _mxfp8_e4m3_quantize_torch
11
* Native MXFP8 linear (dot_scaled) -> dequant-to-bf16 @ matmul
12
* Native MXFP8 MoE (dot_scaled grouped GEMM) -> dequant-to-bf16 MoE math
13
+ * Strided K/V sparse-PA cache insert -> contiguous-copy insert
14
15
The native MXFP8 GEMMs also guard the ``dot_scaled`` rhs-scale orientation: the
16
scale is loaded ``[N, K//32]`` and passed WITHOUT transpose; a stray ``.T``
@@ -39,6 +40,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( # noqa:
40
_mxfp8_e4m3_quantize_triton,
41
dequant_mxfp8_to_bf16,
42
)
43
+from vllm.models.minimax_m3.amd.model import _kv_insert_operand # noqa: E402
44
from vllm.models.minimax_m3.amd.ops import ( # noqa: E402
45
gemma_fused_add_rmsnorm,
46
gemma_rmsnorm,
@@ -608,3 +610,77 @@ def test_routed_experts_expert_map_delegates_to_kernel():
610
assert torch.equal(resolve(True), mask) # AITER kernel -> 0/1 mask
611
assert torch.equal(resolve(False), canonical) # non-AITER -> canonical map
612
assert torch.equal(resolve(False, has_moe_kernel=False), canonical) # non-modular
613
+
614
+
615
+# --------------------------------------------------------------------------- #
616
+# Strided K/V operand for the AITER sparse-PA cache insert
617
+# --------------------------------------------------------------------------- #
618
+def _fused_qkv_kv_slices(num_tokens, num_kv_heads, head_dim, num_q_heads=16):
619
+ """K/V as the sparse-PA insert sees them: column slices of a fused qkv."""
620
+ kv = num_kv_heads * head_dim
621
+ row = num_q_heads * head_dim + 2 * kv + 2 * head_dim
622
+ qkv = torch.randn((num_tokens, row), dtype=torch.bfloat16, device=DEVICE)
623
+ k_start = num_q_heads * head_dim
624
+ k = qkv[:, k_start : k_start + kv].view(num_tokens, num_kv_heads, head_dim)
625
+ v = qkv[:, k_start + kv : k_start + 2 * kv].view(num_tokens, num_kv_heads, head_dim)
626
+ return k, v
627
+
628
+
629
+def _asm_kv_cache(pages, num_kv_heads, head_dim):
630
+ """Page-16 K/V cache views matching ``_ensure_aiter_sparse_pa_kv_cache``."""
631
+ x = 16 // torch.tensor([], dtype=torch.bfloat16).element_size()
632
+ key = torch.zeros(
633
+ (pages, num_kv_heads, head_dim // x, 16, x),
634
+ dtype=torch.bfloat16,
635
+ device=DEVICE,
636
+ )
637
+ value = torch.zeros(
638
+ (pages, num_kv_heads, 16 // x, head_dim, x),
639
+ dtype=torch.bfloat16,
640
+ device=DEVICE,
641
+ )
642
+ return key, value
643
+
644
+
645
+@pytest.mark.parametrize("num_tokens", [1, 8, 129, 1024])
646
+@pytest.mark.parametrize("num_kv_heads", [1, 2])
647
+def test_kv_insert_operand_matches_contiguous(num_tokens, num_kv_heads):
648
+ """A strided K/V slice must insert bit-identically to a contiguous copy."""
649
+ reshape_and_cache = pytest.importorskip("aiter").reshape_and_cache
650
+ head_dim = 128
651
+ pages = num_tokens // 16 + 4
652
+ slots = torch.randperm(pages * 16, device=DEVICE)[:num_tokens].to(torch.int32)
653
+
654
+ torch.manual_seed(num_tokens)
655
+ k, v = _fused_qkv_kv_slices(num_tokens, num_kv_heads, head_dim)
656
+ assert num_tokens == 1 or not k.is_contiguous()
657
+
658
+ caches = []
659
+ for operand in (lambda t: t.contiguous(), _kv_insert_operand):
660
+ key_cache, value_cache = _asm_kv_cache(pages, num_kv_heads, head_dim)
661
+ reshape_and_cache(
662
+ operand(k),
663
+ operand(v),
664
+ key_cache,
665
+ value_cache,
666
+ slots,
667
+ kv_cache_dtype="auto",
668
+ asm_layout=True,
669
+ )
670
+ caches.append((key_cache, value_cache))
671
+
672
+ (ref_k, ref_v), (got_k, got_v) = caches
673
+ assert torch.equal(ref_k, got_k)
674
+ assert torch.equal(ref_v, got_v)
675
+
676
+
677
+def test_kv_insert_operand_copies_unsupported_layouts():
678
+ """Only a row-strided slice passes through; anything else is copied."""
679
+ k, _ = _fused_qkv_kv_slices(32, 2, 128)
680
+ assert _kv_insert_operand(k) is k
681
+
682
+ inner_strided = k[:, :, ::2]
683
+ assert _kv_insert_operand(inner_strided).is_contiguous()
684
+
685
+ flat = k.reshape(32, -1)
686
+ assert _kv_insert_operand(flat).is_contiguous()
vllm/models/minimax_m3/amd/model.py
+28
-2
@@ -579,6 +579,32 @@ class MiniMaxM3Attention(nn.Module):
579
return output
580
581
582
+# Widest row pitch aiter's reshape_and_cache can carry in its `int` stride args.
583
+_MAX_KV_INSERT_ROW_STRIDE = 2**31 - 1
584
+
585
+
586
+def _kv_insert_operand(t: torch.Tensor) -> torch.Tensor:
587
+ """K or V operand for ``aiter.reshape_and_cache``, without a needless copy.
588
+
589
+ The AITER sparse-PA insert takes K and V as column slices of the fused
590
+ ``[q | k | v | index_q | index_k]`` projection, so they are row-strided
591
+ views and never contiguous. ``reshape_and_cache`` does not need them to be:
592
+ it passes ``key.stride(0)`` to the kernel as the row pitch and reads
593
+ ``key[token_idx * key_stride + i]`` for ``i`` over ``num_heads *
594
+ head_size``, so the only requirement is that the trailing
595
+ ``(num_heads, head_size)`` block be contiguous within a row, which such a
596
+ slice always satisfies. Any other layout falls back to a copy.
597
+ """
598
+ if (
599
+ t.dim() == 3
600
+ and t.stride(2) == 1
601
+ and t.stride(1) == t.shape[2]
602
+ and t.stride(0) <= _MAX_KV_INSERT_ROW_STRIDE
603
+ ):
604
+ return t
605
+ return t.contiguous()
606
+
607
+
608
class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
609
"""Block-sparse attention layer with the lightning-indexer branch.
610
@@ -902,8 +928,8 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
928
else "auto"
929
)
930
reshape_and_cache(
905
- k.contiguous(),
906
- v.contiguous(),
931
+ _kv_insert_operand(k),
932
+ _kv_insert_operand(v),
933
key_cache,
934
value_cache,
935
slot_mapping,