@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 95aab66e95

[ROCm][MiniMax-M3][Spec Decode] Support speculative decode with AITER sparse PA (#47984)

Signed-off-by: Tan Pin Siang <tanpinsiang@gmail.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: Jun Kang Chow <junkangchow@gmail.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

Tan Pin Siang committed Jul 14, 2026 at 12:12 UTC 95aab66e9585e0f7889f2adc5ef6067a4a0e0f25
3 files changed +182 -15
tests/kernels/attention/test_minimax_m3.py
+140
@@ -1010,6 +1010,146 @@ def _build_decode_inputs(
1010 return q, block_table, seq_lens, topk_idx, num_pages
1011
1012
1013 +@pytest.mark.skipif(
1014 + not current_platform.is_rocm(),
1015 + reason="MiniMax-M3 AITER sparse PA is ROCm-only",
1016 +)
1017 +def test_aiter_sparse_block_table_handles_padded_decode_rows():
1018 + """Zero-length padded decode rows must produce empty sparse attention."""
1019 + from vllm.models.minimax_m3.amd.ops.sparse_pa import (
1020 + PAGES_PER_SPARSE_BLOCK,
1021 + minimax_m3_build_sparse_block_table,
1022 + )
1023 +
1024 + topk_idx = torch.tensor(
1025 + [[[1, 0, -1], [0, 1, -1]]],
1026 + device="cuda",
1027 + dtype=torch.int32,
1028 + )
1029 + block_table = torch.tensor(
1030 + [[5, 7], [11, 13]],
1031 + device="cuda",
1032 + dtype=torch.int32,
1033 + )
1034 + seq_lens = torch.tensor([129, 0], device="cuda", dtype=torch.int32)
1035 +
1036 + sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table(
1037 + topk_idx, block_table, seq_lens
1038 + )
1039 + torch.accelerator.synchronize()
1040 +
1041 + expected_bt = torch.zeros_like(sparse_bt)
1042 + expected_ctx = torch.tensor([129, 0], device="cuda", dtype=torch.int32)
1043 +
1044 + for slot, block_id in enumerate((0, 1)):
1045 + base_phys = int(block_table[0, block_id]) * PAGES_PER_SPARSE_BLOCK
1046 + start = slot * PAGES_PER_SPARSE_BLOCK
1047 + expected_bt[0, start : start + PAGES_PER_SPARSE_BLOCK] = torch.arange(
1048 + base_phys,
1049 + base_phys + PAGES_PER_SPARSE_BLOCK,
1050 + device="cuda",
1051 + dtype=torch.int32,
1052 + )
1053 +
1054 + assert torch.equal(sparse_ctx, expected_ctx)
1055 + assert torch.equal(sparse_bt, expected_bt)
1056 +
1057 +
1058 +@pytest.mark.skipif(
1059 + not current_platform.is_rocm(),
1060 + reason="MiniMax-M3 AITER sparse PA is ROCm-only",
1061 +)
1062 +def test_aiter_decode_sparse_block_table_supports_spec_decode():
1063 + """AITER sparse PA needs one compact page-16 table per speculative query.
1064 +
1065 + The decode indexer flattens speculative rows as
1066 + ``req_id * decode_query_len + local_q``. This verifies that the ROCm AITER
1067 + block-table adapter uses the same mapping before handing rows to Gluon.
1068 + """
1069 + from vllm.models.minimax_m3.amd.ops.sparse_pa import (
1070 + PAGES_PER_SPARSE_BLOCK,
1071 + minimax_m3_build_sparse_block_table_decode,
1072 + )
1073 +
1074 + decode_query_len = 4
1075 + seq_lens = torch.tensor([260, 132, 0], device="cuda", dtype=torch.int32)
1076 + block_table = torch.tensor(
1077 + [
1078 + [5, 7, 11],
1079 + [13, 17, 19],
1080 + [0, 0, 0],
1081 + ],
1082 + device="cuda",
1083 + dtype=torch.int32,
1084 + )
1085 + topk_idx = torch.tensor(
1086 + [
1087 + [
1088 + [2, 0, 3, -1],
1089 + [0, 2, 3, -1],
1090 + [2, 1, 0, -1],
1091 + [1, 2, 0, -1],
1092 + [1, 0, 2, -1],
1093 + [0, 1, 2, -1],
1094 + [1, -1, 0, -1],
1095 + [0, 2, 1, -1],
1096 + [0, 1, -1, -1],
1097 + [0, 1, -1, -1],
1098 + [0, 1, -1, -1],
1099 + [0, 1, -1, -1],
1100 + ]
1101 + ],
1102 + device="cuda",
1103 + dtype=torch.int32,
1104 + )
1105 +
1106 + sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_decode(
1107 + topk_idx, block_table, seq_lens, decode_query_len
1108 + )
1109 + torch.accelerator.synchronize()
1110 +
1111 + expected_bt = torch.zeros_like(sparse_bt)
1112 + expected_ctx = torch.empty_like(sparse_ctx)
1113 + topk_cpu = topk_idx.cpu()[0].tolist()
1114 + block_table_cpu = block_table.cpu()
1115 + seq_lens_cpu = seq_lens.cpu().tolist()
1116 +
1117 + for row, selected_blocks in enumerate(topk_cpu):
1118 + req_id = row // decode_query_len
1119 + local_q = row - req_id * decode_query_len
1120 + query_pos = seq_lens_cpu[req_id] - decode_query_len + local_q
1121 + causal_len = max(query_pos + 1, 0)
1122 + if causal_len == 0:
1123 + expected_ctx[row] = 0
1124 + continue
1125 + self_block = query_pos // BLOCK_SIZE
1126 +
1127 + valid = [blk for blk in selected_blocks if 0 <= blk <= self_block]
1128 + full_blocks = [blk for blk in valid if blk < self_block]
1129 + tail_blocks = [blk for blk in valid if blk == self_block]
1130 + ordered_blocks = full_blocks + tail_blocks
1131 +
1132 + for slot, block_id in enumerate(ordered_blocks):
1133 + base_phys = int(block_table_cpu[req_id, block_id]) * PAGES_PER_SPARSE_BLOCK
1134 + start = slot * PAGES_PER_SPARSE_BLOCK
1135 + expected_bt[row, start : start + PAGES_PER_SPARSE_BLOCK] = torch.arange(
1136 + base_phys,
1137 + base_phys + PAGES_PER_SPARSE_BLOCK,
1138 + device="cuda",
1139 + dtype=torch.int32,
1140 + )
1141 +
1142 + if tail_blocks:
1143 + expected_ctx[row] = (
1144 + len(full_blocks) * BLOCK_SIZE + causal_len - self_block * BLOCK_SIZE
1145 + )
1146 + else:
1147 + expected_ctx[row] = min(len(valid) * BLOCK_SIZE, causal_len)
1148 +
1149 + assert torch.equal(sparse_ctx, expected_ctx)
1150 + assert torch.equal(sparse_bt, expected_bt)
1151 +
1152 +
1153 @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
1154 @pytest.mark.parametrize(
1155 "seq_lens_list",
vllm/models/minimax_m3/amd/ops/sparse_pa.py
+41 -10
@@ -49,7 +49,8 @@ def _build_sparse_block_table_kernel(
49 ):
50 pid_b = tl.program_id(0)
51 seq_len = tl.load(seq_lens_ptr + pid_b)
52 - last_blk = (seq_len - 1) // SPARSE_BLOCK_SIZE_C
52 + has_tokens = seq_len > 0
53 + last_blk = tl.maximum((seq_len - 1) // SPARSE_BLOCK_SIZE_C, 0)
54
55 topk_row = topk_ptr + pid_b * stride_topk_n
56 bt_row = block_table_ptr + pid_b * stride_bt_b
@@ -57,7 +58,7 @@ def _build_sparse_block_table_kernel(
58
59 off_t = tl.arange(0, BLOCK_SIZE_T)
60 blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
60 - valid = blk >= 0
61 + valid = has_tokens & (blk >= 0)
62 is_tail = valid & (blk == last_blk)
63 is_full = valid & (blk != last_blk)
64
@@ -82,7 +83,7 @@ def _build_sparse_block_table_kernel(
83 mask=(off_w >= n_used) & (off_w < row_width),
84 )
85
85 - tail_tokens = seq_len - last_blk * SPARSE_BLOCK_SIZE_C
86 + tail_tokens = tl.where(has_tokens, seq_len - last_blk * SPARSE_BLOCK_SIZE_C, 0)
87 has_tail = tl.sum(is_tail.to(tl.int32), axis=0) > 0
88 ctx = n_full * SPARSE_BLOCK_SIZE_C + tl.where(has_tail, tail_tokens, 0)
89 ctx = tl.where(has_tail, ctx, tl.minimum(n_valid * SPARSE_BLOCK_SIZE_C, seq_len))
@@ -140,7 +141,8 @@ def _build_sparse_block_table_prefill_kernel(
141 pid_n = tl.program_id(0)
142 req_id = tl.load(req_id_ptr + pid_n)
143 abs_pos = tl.load(abs_pos_ptr + pid_n)
143 - causal_len = abs_pos + 1
144 + # Padded speculative rows can have negative positions; clamp to empty range.
145 + causal_len = tl.maximum(abs_pos + 1, 0)
146 self_blk = abs_pos // SPARSE_BLOCK_SIZE_C
147
148 topk_row = topk_ptr + pid_n * stride_topk_n
@@ -149,7 +151,7 @@ def _build_sparse_block_table_prefill_kernel(
151
152 off_t = tl.arange(0, BLOCK_SIZE_T)
153 blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
152 - valid = (blk >= 0) & (blk <= self_blk)
154 + valid = (causal_len > 0) & (blk >= 0) & (blk <= self_blk)
155 is_tail = valid & (blk == self_blk)
156 is_full = valid & (blk < self_blk)
157
@@ -214,6 +216,34 @@ def minimax_m3_build_sparse_block_table_prefill(
216 return sparse_bt, sparse_ctx
217
218
219 +@torch.no_grad()
220 +def minimax_m3_build_sparse_block_table_decode(
221 + topk_idx: torch.Tensor, # [1, batch * decode_query_len, topk]
222 + block_table: torch.Tensor, # [batch, max_blocks]
223 + seq_lens: torch.Tensor, # [batch]
224 + decode_query_len: int,
225 +) -> tuple[torch.Tensor, torch.Tensor]:
226 + """Build one page-16 sparse block table row per decode query token."""
227 + if decode_query_len == 1:
228 + return minimax_m3_build_sparse_block_table(topk_idx, block_table, seq_lens)
229 +
230 + total_q = topk_idx.shape[1]
231 + expected_q = block_table.shape[0] * decode_query_len
232 + assert total_q == expected_q, (
233 + "MiniMax-M3 decode top-k rows must equal batch * decode_query_len: "
234 + f"{total_q} != {expected_q}"
235 + )
236 + pos = torch.arange(total_q, dtype=torch.int32, device=topk_idx.device)
237 + query_req_id = torch.div(pos, decode_query_len, rounding_mode="floor")
238 + q_offset = pos - query_req_id * decode_query_len
239 + query_abs_pos = (seq_lens[query_req_id] - decode_query_len + q_offset).to(
240 + torch.int32
241 + )
242 + return minimax_m3_build_sparse_block_table_prefill(
243 + topk_idx, block_table, query_req_id, query_abs_pos
244 + )
245 +
246 +
247 @triton.jit
248 def _insert_index_cache_kernel(
249 index_k_ptr, # [num_tokens, head_dim]
@@ -317,20 +347,21 @@ def _gluon_scale_arg(
347
348 @torch.no_grad()
349 def minimax_m3_sparse_attn_decode_aiter(
320 - q: torch.Tensor, # [batch, num_heads, head_dim]
350 + q: torch.Tensor, # [batch * decode_query_len, num_heads, head_dim]
351 k_cache: torch.Tensor, # [phys16, num_kv_heads, head_dim // x, 16, x]
352 v_cache: torch.Tensor, # [phys16, num_kv_heads, 16 // x, head_dim, x]
323 - topk_idx: torch.Tensor, # [1, batch, topk]
353 + topk_idx: torch.Tensor, # [1, batch * decode_query_len, topk]
354 block_table: torch.Tensor, # [batch, max_blocks]
355 seq_lens: torch.Tensor, # [batch]
356 num_kv_heads: int,
357 sm_scale: float,
328 - output: torch.Tensor, # [batch, num_heads, head_dim]
358 + output: torch.Tensor, # [batch * decode_query_len, num_heads, head_dim]
359 k_scale: torch.Tensor | None = None,
360 v_scale: torch.Tensor | None = None,
361 + decode_query_len: int = 1,
362 ) -> None:
332 - sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table(
333 - topk_idx, block_table, seq_lens
363 + sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_decode(
364 + topk_idx, block_table, seq_lens, decode_query_len
365 )
366 _run_gluon_decode(
367 q,
vllm/models/minimax_m3/amd/sparse_attention_msa.py
+1 -5
@@ -55,11 +55,6 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
55 if main_md.num_decodes > 0:
56 d = main_md.decode
57 assert d is not None
58 - if d.decode_query_len != 1:
59 - raise NotImplementedError(
60 - "MiniMax-M3 AITER sparse PA does not support speculative "
61 - f"decode_query_len={d.decode_query_len}"
62 - )
58 minimax_m3_sparse_attn_decode_aiter(
59 q[:nd],
60 k_cache,
@@ -72,6 +67,7 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
67 out[:nd],
68 k_scale=k_scale,
69 v_scale=v_scale,
70 + decode_query_len=d.decode_query_len,
71 )
72
73 if main_md.num_prefills > 0: