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

[ROCm][ [Perf] sparse attention optimization on minimax-m3 (#46546)

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: yueliu14 <yue.liu4@amd.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>

Hongxia Yang committed Jun 25, 2026 at 04:56 UTC c63cd4906c2a67f18b3714786cc036c1ad97a64f
5 files changed +1238 -36
vllm/models/minimax_m3/amd/ops/index_topk.py new
+939
@@ -0,0 +1,939 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton kernels for MiniMax M3 lightning-indexer block scoring + top-k.
4 +
5 +Index queries score each 128-token block of index keys (max over the block),
6 +then the top-k blocks (plus forced init/local blocks) are selected per query
7 +token. Adapted to vLLM's paged KV cache: the KV page size is forced to equal the
8 +sparse block size (128), so one sparse block maps to exactly one page.
9 +
10 +Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (single head).
11 +
12 +Only the paths MiniMax M3 uses are implemented: score_type="max", index value
13 +disabled (score-only indexer), single shared index head. The selected block ids
14 +feed the block-sparse attention kernels in ``sparse_attn``.
15 +"""
16 +
17 +import torch
18 +
19 +from vllm.platforms import current_platform
20 +from vllm.triton_utils import tl, triton
21 +from vllm.utils.math_utils import round_up
22 +
23 +# One sparse block == one KV page.
24 +SPARSE_BLOCK_SIZE = 128
25 +
26 +
27 +# ---------------------------------------------------------------------------
28 +# Bitonic top-k helpers (layout-agnostic).
29 +# ---------------------------------------------------------------------------
30 +@triton.jit
31 +def _compare_and_swap(x, ids, flip, i: tl.constexpr, n_dims: tl.constexpr):
32 + n_outer: tl.constexpr = x.numel >> n_dims
33 + shape: tl.constexpr = [n_outer * 2**i, 2, 2 ** (n_dims - i - 1)]
34 + y = tl.reshape(x, shape)
35 + mask = tl.arange(0, 2)[None, :, None]
36 + left = tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype)
37 + right = tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape).to(y.dtype)
38 + left = tl.reshape(left, x.shape)
39 + right = tl.reshape(right, x.shape)
40 + y_idx = tl.reshape(ids, shape)
41 + left_idx = tl.broadcast_to(tl.sum(y_idx * (1 - mask), 1)[:, None, :], shape)
42 + right_idx = tl.broadcast_to(tl.sum(y_idx * mask, 1)[:, None, :], shape)
43 + left_idx = tl.reshape(left_idx, x.shape).to(y_idx.dtype)
44 + right_idx = tl.reshape(right_idx, x.shape).to(y_idx.dtype)
45 + idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True)
46 + ileft = left.to(idtype, bitcast=True)
47 + iright = right.to(idtype, bitcast=True)
48 + ix = x.to(idtype, bitcast=True)
49 + cond = (left > right) != flip
50 + ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix))
51 + new_ids = ids ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(ids))
52 + return ret.to(x.dtype, bitcast=True), new_ids
53 +
54 +
55 +@triton.jit
56 +def _bitonic_merge(
57 + x, ids, stage: tl.constexpr, order: tl.constexpr, n_dims: tl.constexpr
58 +):
59 + n_outer: tl.constexpr = x.numel >> n_dims
60 + tl.static_assert(stage <= n_dims)
61 + if order == 2:
62 + shape: tl.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2**stage]
63 + flip = tl.reshape(
64 + tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape
65 + )
66 + else:
67 + flip = order
68 + for i in tl.static_range(stage):
69 + x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims)
70 + return x, ids
71 +
72 +
73 +# ---------------------------------------------------------------------------
74 +# Index block-score kernel (paged). score[h, token, block] = max over the
75 +# 128-token block of (idx_q . index_k), causal-masked. BLOCK_SIZE_K == 128 so
76 +# each K-tile is exactly one page (BLOCKS_PER_K_BLOCK == 1).
77 +# ---------------------------------------------------------------------------
78 +# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens
79 +# might lose pointer alignment, which trigger Triton recompiles. we don't actually
80 +# need pointer alignment for those tensors anyway because we do scalar load.
81 +@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
82 +def _index_block_score_kernel(
83 + q_ptr, # idx_q: [total_q, num_idx_heads, head_dim]
84 + ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim]
85 + score_ptr, # [num_idx_heads, total_q, max_block]
86 + block_table_ptr, # [num_reqs, max_blocks]
87 + cu_seqlens, # [batch+1] query start offsets
88 + seq_lens, # [batch] total K length
89 + prefix_lens, # [batch] context length before this chunk's queries
90 + num_idx_heads,
91 + head_dim: tl.constexpr,
92 + stride_q_n,
93 + stride_q_h,
94 + stride_q_d,
95 + stride_ik_blk,
96 + stride_ik_pos,
97 + stride_ik_d,
98 + stride_s_h,
99 + stride_s_n,
100 + stride_s_k,
101 + stride_bt_b,
102 + BLOCK_SIZE_Q: tl.constexpr,
103 + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
104 +):
105 + pid_q = tl.program_id(0)
106 + pid_bh = tl.program_id(1)
107 + pid_b = pid_bh // num_idx_heads
108 + pid_h = pid_bh % num_idx_heads
109 +
110 + seq_start = tl.load(cu_seqlens + pid_b)
111 + q_len = tl.load(cu_seqlens + pid_b + 1) - seq_start
112 + seq_len = tl.load(seq_lens + pid_b)
113 + prefix_len = tl.load(prefix_lens + pid_b)
114 + if BLOCK_SIZE_Q * pid_q >= q_len:
115 + return
116 +
117 + q_ptrs = tl.make_block_ptr(
118 + base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h,
119 + shape=(q_len, head_dim),
120 + strides=(stride_q_n, stride_q_d),
121 + offsets=(pid_q * BLOCK_SIZE_Q, 0),
122 + block_shape=(BLOCK_SIZE_Q, head_dim),
123 + order=(1, 0),
124 + )
125 + q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero")
126 + q_start = prefix_len + pid_q * BLOCK_SIZE_Q
127 +
128 + off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len
129 + off_k = tl.arange(0, BLOCK_SIZE_K)
130 + off_d = tl.arange(0, head_dim)
131 + # Block table row for this request.
132 + bt_row = block_table_ptr + pid_b * stride_bt_b
133 + # Causal window: only blocks up to the last query token's position.
134 + hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q)
135 + for i in tl.range(0, hi, BLOCK_SIZE_K):
136 + blk = i // BLOCK_SIZE_K
137 + page = tl.load(bt_row + blk).to(tl.int64)
138 + pos = i + off_k
139 + # index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed)
140 + # we don't need masked load for K, because KV cache ensures
141 + # allocation is multiple of BLOCK_SIZE_K.
142 + # for tokens beyond seqlen, they will be masked in qk later.
143 + k = tl.load(
144 + ik_cache_ptr
145 + + page * stride_ik_blk
146 + + off_k[None, :] * stride_ik_pos
147 + + off_d[:, None] * stride_ik_d,
148 + )
149 + qk = tl.dot(q, k)
150 + # apply causal mask as needed
151 + if q_start < i + BLOCK_SIZE_K:
152 + qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf"))
153 + # one sparse block per K-tile -> max over the 128 positions
154 + score = tl.max(qk, axis=1) # [BLOCK_SIZE_Q]
155 + s_ptrs = (
156 + score_ptr
157 + + pid_h * stride_s_h
158 + + (seq_start + pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q))
159 + * stride_s_n
160 + + blk * stride_s_k
161 + )
162 + q_store_mask = (pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) < q_len
163 + tl.store(s_ptrs, score, mask=q_store_mask)
164 +
165 +
166 +# ---------------------------------------------------------------------------
167 +# Top-k selection over per-token block scores (layout-agnostic). block_size_q
168 +# is 1 for M3, so top-k is computed per query token.
169 +# ---------------------------------------------------------------------------
170 +# since prefill metadata is sliced from mixed batch metadata, prefix_lens
171 +# might lose pointer alignment, which trigger Triton recompiles. we don't actually
172 +# need pointer alignment for those tensors anyway because we do scalar load.
173 +@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])})
174 +@triton.autotune(
175 + configs=[
176 + triton.Config({"BLOCK_SIZE_K": 2048}, num_warps=8, num_stages=2),
177 + triton.Config({"BLOCK_SIZE_K": 1024}, num_warps=8, num_stages=2),
178 + triton.Config({"BLOCK_SIZE_K": 512}, num_warps=8, num_stages=2),
179 + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2),
180 + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2),
181 + triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2),
182 + ],
183 + key=["BLOCK_SIZE_T"],
184 +)
185 +@triton.jit(do_not_specialize_on_alignment=["prefix_lens"])
186 +def _topk_index_kernel(
187 + s_ptr, # [num_heads, total_q, max_block]
188 + ti_ptr, # [num_heads, total_q, topk]
189 + sample_interval: tl.constexpr, # block_size_q (1 for M3)
190 + block_size: tl.constexpr, # sparse block size (128)
191 + cu_seqlens,
192 + cu_seqblocks_q,
193 + prefix_lens,
194 + topk,
195 + init_blocks: tl.constexpr,
196 + local_blocks: tl.constexpr,
197 + stride_s_h,
198 + stride_s_n,
199 + stride_s_k,
200 + stride_ti_h,
201 + stride_ti_n,
202 + stride_ti_t,
203 + BLOCK_SIZE_K: tl.constexpr,
204 + BLOCK_SIZE_T: tl.constexpr,
205 + MASK_INIT: tl.constexpr,
206 + MASK_LOCAL: tl.constexpr,
207 +):
208 + tl.static_assert(BLOCK_SIZE_K > BLOCK_SIZE_T)
209 + pid_q = tl.program_id(0)
210 + pid_b = tl.program_id(1)
211 + pid_h = tl.program_id(2)
212 + seq_start = tl.load(cu_seqlens + pid_b)
213 + block_start = tl.load(cu_seqblocks_q + pid_b)
214 + block_num = tl.load(cu_seqblocks_q + pid_b + 1) - block_start
215 + prefix_len = tl.load(prefix_lens + pid_b)
216 + if pid_q >= block_num:
217 + return
218 + off_k = tl.arange(0, BLOCK_SIZE_K)
219 + off_t = tl.arange(0, BLOCK_SIZE_T)
220 + s_ptrs = (
221 + s_ptr
222 + + (seq_start + pid_q * sample_interval) * stride_s_n
223 + + pid_h * stride_s_h
224 + + off_k * stride_s_k
225 + )
226 + topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32)
227 + topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32)
228 + left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2
229 + valid_blocks = (prefix_len + pid_q * sample_interval + block_size) // block_size
230 + for i in tl.range(0, valid_blocks, BLOCK_SIZE_K):
231 + causal_mask = i + off_k < valid_blocks
232 + local_mask = i + off_k >= max(0, valid_blocks - local_blocks)
233 + init_mask = i + off_k < init_blocks
234 + score = tl.load(s_ptrs, mask=causal_mask, other=-1e30).to(tl.float32)
235 + score = tl.where(score != score, -1e30, score)
236 + s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K
237 + if MASK_INIT:
238 + score = tl.where(causal_mask & init_mask, score - 1e29, score)
239 + else:
240 + score = tl.where(causal_mask & init_mask, 1e30, score)
241 + if MASK_LOCAL:
242 + score = tl.where(causal_mask & local_mask, score - 1e28, score)
243 + else:
244 + score = tl.where(causal_mask & local_mask, 1e29, score)
245 + topk_score, last_topk_score = score, topk_score
246 + topk_idx, last_topk_idx = (tl.where(causal_mask, i + off_k + 1, 0), topk_idx)
247 + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
248 + for j in tl.static_range(1, n_dims):
249 + topk_score, topk_idx = _bitonic_merge(
250 + topk_score, topk_idx.to(tl.int32), j, 2, n_dims
251 + )
252 + if i != 0:
253 + topk_score, topk_idx = _bitonic_merge(
254 + topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims
255 + )
256 + topk_score_new = last_topk_score * left_half_mask + topk_score * (
257 + 1 - left_half_mask
258 + )
259 + topk_idx_new = last_topk_idx * left_half_mask + topk_idx * (
260 + 1 - left_half_mask
261 + )
262 + topk_score, topk_idx = _bitonic_merge(
263 + topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims
264 + )
265 + else:
266 + topk_score, topk_idx = _bitonic_merge(
267 + topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims
268 + )
269 + topk_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
270 + topk_idx = tl.sum(
271 + topk_mask[:, None]
272 + * tl.reshape(topk_idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
273 + axis=0,
274 + )
275 + ti_ptrs = (
276 + ti_ptr
277 + + (block_start + pid_q) * stride_ti_n
278 + + pid_h * stride_ti_h
279 + + off_t * stride_ti_t
280 + )
281 + store_mask = off_t < topk
282 + valid_mask = off_t < valid_blocks
283 + topk_idx = tl.where(store_mask & valid_mask, topk_idx, -1)
284 + tl.store(ti_ptrs, topk_idx.to(ti_ptrs.dtype.element_ty), mask=store_mask)
285 +
286 +
287 +# ---------------------------------------------------------------------------
288 +# Decode index-score kernel (split-K over seq blocks). Decode batches are
289 +# flattened request-major, with a runtime query length used to map each query
290 +# token back to its request metadata. Chunk counts depend only on shape
291 +# constants so the grid is fixed within a cuda graph. The score scale is omitted
292 +# because decode only consumes block ordering.
293 +# ---------------------------------------------------------------------------
294 +@triton.jit(do_not_specialize=["num_kv_chunks", "decode_query_len"])
295 +def _decode_index_score_kernel(
296 + q_ptr, # idx_q: [total_q, num_idx_heads, head_dim]
297 + ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim]
298 + score_ptr, # [num_idx_heads, total_q, max_block]
299 + block_table_ptr, # [num_reqs, max_blocks]
300 + seq_lens, # [num_reqs]
301 + num_idx_heads: tl.constexpr,
302 + head_dim: tl.constexpr,
303 + init_blocks,
304 + local_blocks,
305 + decode_query_len,
306 + stride_q_n,
307 + stride_q_h,
308 + stride_q_d,
309 + stride_ik_blk,
310 + stride_ik_pos,
311 + stride_ik_d,
312 + stride_s_h,
313 + stride_s_n,
314 + stride_s_k,
315 + stride_bt_b,
316 + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
317 + BLOCK_SIZE_Q: tl.constexpr,
318 + num_kv_chunks,
319 + USE_PDL: tl.constexpr,
320 +):
321 + BLOCK_SIZE_HQ: tl.constexpr = num_idx_heads * BLOCK_SIZE_Q
322 + pid_r = tl.program_id(0)
323 + pid_c = tl.program_id(1)
324 + hq_offsets = tl.arange(0, BLOCK_SIZE_HQ)
325 + h_offsets = hq_offsets // BLOCK_SIZE_Q
326 + q_offsets = hq_offsets % BLOCK_SIZE_Q
327 + q_mask = q_offsets < decode_query_len
328 + q_ids = pid_r * decode_query_len + q_offsets
329 +
330 + if USE_PDL:
331 + tl.extra.cuda.gdc_wait()
332 + tl.extra.cuda.gdc_launch_dependents()
333 +
334 + seq_len = tl.load(seq_lens + pid_r)
335 + query_pos = seq_len - decode_query_len + q_offsets
336 + # Full-CG padding uses zero-length request rows. Clamp to an empty
337 + # attention range instead of letting padded rows produce negative lengths.
338 + kv_len = tl.maximum(query_pos + 1, 0)
339 + num_blocks_q = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
340 + kv_len_max = tl.max(tl.where(q_mask, kv_len, 0), axis=0)
341 + num_blocks = (kv_len_max + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
342 +
343 + # block-aligned fixed-count split: grid independent of seq_len (cuda graph).
344 + chunk_size_blocks = (num_blocks + num_kv_chunks - 1) // num_kv_chunks
345 + chunk_start_block = pid_c * chunk_size_blocks
346 + chunk_end_block = tl.minimum(chunk_start_block + chunk_size_blocks, num_blocks)
347 + if chunk_start_block >= chunk_end_block:
348 + return
349 + off_k = tl.arange(0, BLOCK_SIZE_K) # positions within a 128-block
350 + off_d = tl.arange(0, head_dim)
351 + bt_row = block_table_ptr + pid_r * stride_bt_b
352 + # Force-select init (1e30) and local (1e29, higher priority) blocks.
353 + local_start = tl.maximum(0, num_blocks_q - local_blocks)
354 + # Query vectors for all index heads in a small spec-decode block.
355 + q = tl.load(
356 + q_ptr
357 + + q_ids[None, :] * stride_q_n
358 + + h_offsets[None, :] * stride_q_h
359 + + off_d[:, None] * stride_q_d,
360 + mask=q_mask[None, :],
361 + other=0.0,
362 + ) # [D,HQ]
363 + for blk in tl.range(chunk_start_block, chunk_end_block):
364 + page = tl.load(bt_row + blk).to(tl.int64)
365 + pos = blk * BLOCK_SIZE_K + off_k
366 + pos_mask = pos[:, None] < kv_len[None, :]
367 + # we don't need masked load for K, because KV cache ensures
368 + # allocation is multiple of BLOCK_SIZE_K.
369 + # for tokens beyond seqlen, they will be masked in qk later.
370 + k = tl.load(
371 + ik_cache_ptr
372 + + page * stride_ik_blk
373 + + off_k[:, None] * stride_ik_pos
374 + + off_d * stride_ik_d,
375 + ) # [N,D]
376 + if BLOCK_SIZE_HQ == 1:
377 + # Degenerate GEMV (q is [D,1]): vectorized fp32 multiply + reduce
378 + # instead of an MFMA tile. Numerically equivalent to tl.dot.
379 + q_vec = tl.sum(q, axis=1).to(tl.float32) # [D]
380 + kq = tl.sum(k.to(tl.float32) * q_vec[None, :], axis=1)[:, None] # [N,1]
381 + else:
382 + # fp32 accumulation is required for the fp8 (e4m3) index cache: q/k
383 + # are loaded in their stored dtype (bf16 or e4m3) and the MMA
384 + # accumulates in fp32 so the per-block max score is exact for the
385 + # fp8 indexer too.
386 + kq = tl.dot(k, q, out_dtype=tl.float32) # [N,HQ]
387 + kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf"))
388 + score = tl.max(kq, axis=0) # [HQ]
389 + is_visible_block = blk < num_blocks_q
390 + is_init = (blk < init_blocks) & is_visible_block
391 + is_local = (blk >= local_start) & is_visible_block
392 + score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score))
393 + tl.store(
394 + score_ptr + h_offsets * stride_s_h + q_ids * stride_s_n + blk * stride_s_k,
395 + score,
396 + mask=q_mask,
397 + )
398 +
399 +
400 +# ---------------------------------------------------------------------------
401 +# Decode top-k (split-K): per-chunk partial top-k + merge. Forced init/local
402 +# blocks are already encoded in the scores.
403 +# ---------------------------------------------------------------------------
404 +@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])})
405 +@triton.autotune(
406 + configs=[
407 + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2),
408 + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=4, num_stages=2),
409 + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2),
410 + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=3),
411 + triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2),
412 + ],
413 + key=["topk"],
414 +)
415 +@triton.jit(do_not_specialize=["chunk_blocks", "decode_query_len"])
416 +def _topk_index_partial_kernel(
417 + s_ptr, # score: [num_idx_heads, total_q, max_block]
418 + ts_partial_ptr, # partial scores out: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T]
419 + ti_partial_ptr, # partial idx out (1-indexed global, 0=invalid): same shape
420 + seq_lens, # [num_reqs]
421 + block_size: tl.constexpr, # sparse block size (128)
422 + topk: tl.constexpr,
423 + chunk_blocks, # how many score-blocks each chunk owns
424 + decode_query_len,
425 + stride_s_h,
426 + stride_s_b,
427 + stride_s_k,
428 + stride_ts_c,
429 + stride_ts_h,
430 + stride_ts_b,
431 + stride_ts_t,
432 + stride_ti_c,
433 + stride_ti_h,
434 + stride_ti_b,
435 + stride_ti_t,
436 + BLOCK_SIZE_K: tl.constexpr,
437 + BLOCK_SIZE_T: tl.constexpr,
438 + USE_PDL: tl.constexpr,
439 +):
440 + tl.static_assert(topk < BLOCK_SIZE_K)
441 + pid_b = tl.program_id(0) # flattened query-token id
442 + pid_h = tl.program_id(1)
443 + pid_chunk = tl.program_id(2)
444 + req_id = pid_b // decode_query_len
445 + q_offset = pid_b - req_id * decode_query_len
446 +
447 + if USE_PDL:
448 + tl.extra.cuda.gdc_wait()
449 +
450 + seq_len = tl.load(seq_lens + req_id)
451 + query_pos = seq_len - decode_query_len + q_offset
452 + # Full-CG padding uses zero-length request rows. Clamp to an empty
453 + # attention range instead of letting padded rows produce negative lengths.
454 + kv_len = tl.maximum(query_pos + 1, 0)
455 + num_blocks = (kv_len + block_size - 1) // block_size
456 +
457 + # Slice this chunk owns within [0, num_blocks).
458 + chunk_start = pid_chunk * chunk_blocks
459 + chunk_end = tl.minimum(chunk_start + chunk_blocks, num_blocks)
460 + chunk_actual = tl.maximum(chunk_end - chunk_start, 0)
461 +
462 + off_k = tl.arange(0, BLOCK_SIZE_K)
463 + off_t = tl.arange(0, BLOCK_SIZE_T)
464 +
465 + s_ptrs = (
466 + s_ptr
467 + + pid_b * stride_s_b
468 + + pid_h * stride_s_h
469 + + (chunk_start + off_k) * stride_s_k
470 + )
471 +
472 + topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32)
473 + topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32)
474 + left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2
475 +
476 + # Streaming top-K within this chunk. tl.range(0, 0) is a no-op so empty
477 + # chunks (chunk_actual == 0) skip the body and store sentinel -1e30 / 0.
478 + for i in tl.range(0, chunk_actual, BLOCK_SIZE_K):
479 + mask = off_k < chunk_actual - i
480 + score = tl.load(s_ptrs, mask=mask, other=-1e30).to(tl.float32)
481 + score = tl.where(score != score, -1e30, score)
482 + s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K
483 + topk_score, last_topk_score = score, topk_score
484 + topk_idx, last_topk_idx = (
485 + tl.where(mask, chunk_start + i + off_k + 1, 0), # 1-indexed global
486 + topk_idx,
487 + )
488 + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
489 + for j in tl.static_range(1, n_dims):
490 + topk_score, topk_idx = _bitonic_merge(
491 + topk_score, topk_idx.to(tl.int32), j, 2, n_dims
492 + )
493 + if i != 0:
494 + topk_score, topk_idx = _bitonic_merge(
495 + topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims
496 + )
497 + topk_score_new = last_topk_score * left_half_mask + topk_score * (
498 + 1 - left_half_mask
499 + )
500 + topk_idx_new = last_topk_idx * left_half_mask + topk_idx * (
501 + 1 - left_half_mask
502 + )
503 + topk_score, topk_idx = _bitonic_merge(
504 + topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims
505 + )
506 + else:
507 + topk_score, topk_idx = _bitonic_merge(
508 + topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims
509 + )
510 +
511 + if USE_PDL:
512 + tl.extra.cuda.gdc_launch_dependents()
513 +
514 + # Extract first BLOCK_SIZE_T entries (top-K of this chunk after the sort).
515 + topk_mask_extract = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
516 + final_score = tl.sum(
517 + topk_mask_extract[:, None]
518 + * tl.reshape(topk_score, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
519 + axis=0,
520 + )
521 + final_idx = tl.sum(
522 + topk_mask_extract[:, None]
523 + * tl.reshape(topk_idx, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
524 + axis=0,
525 + )
526 +
527 + # Always write all BLOCK_SIZE_T slots — invalid slots carry -1e30 / 0
528 + # sentinels and lose to real scores in the merge stage.
529 + ts_ptrs = (
530 + ts_partial_ptr
531 + + pid_chunk * stride_ts_c
532 + + pid_b * stride_ts_b
533 + + pid_h * stride_ts_h
534 + + off_t * stride_ts_t
535 + )
536 + ti_ptrs = (
537 + ti_partial_ptr
538 + + pid_chunk * stride_ti_c
539 + + pid_b * stride_ti_b
540 + + pid_h * stride_ti_h
541 + + off_t * stride_ti_t
542 + )
543 + tl.store(ts_ptrs, final_score)
544 + tl.store(ti_ptrs, final_idx)
545 +
546 +
547 +@triton.heuristics(
548 + {
549 + "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]),
550 + "BLOCK_SIZE_K": lambda args: triton.next_power_of_2(
551 + args["num_topk_chunks"] * triton.next_power_of_2(args["topk"])
552 + ),
553 + }
554 +)
555 +@triton.jit(do_not_specialize=["num_topk_chunks", "decode_query_len"])
556 +def _topk_index_merge_kernel(
557 + ts_partial_ptr, # partial scores: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T]
558 + ti_partial_ptr, # partial idx (1-indexed global, 0=invalid): same shape
559 + ti_final_ptr, # final idx (0-indexed, -1=invalid): [num_idx_heads, total_q, topk]
560 + seq_lens, # [num_reqs]
561 + block_size: tl.constexpr, # sparse block size (128)
562 + topk: tl.constexpr,
563 + decode_query_len,
564 + stride_ts_c,
565 + stride_ts_h,
566 + stride_ts_b,
567 + stride_ts_t,
568 + stride_ti_c,
569 + stride_ti_h,
570 + stride_ti_b,
571 + stride_ti_t,
572 + stride_tif_h,
573 + stride_tif_b,
574 + stride_tif_t,
575 + num_topk_chunks,
576 + BLOCK_SIZE_K: tl.constexpr,
577 + BLOCK_SIZE_T: tl.constexpr,
578 + USE_PDL: tl.constexpr,
579 +):
580 + pid_b = tl.program_id(0) # flattened query-token id
581 + pid_h = tl.program_id(1)
582 + req_id = pid_b // decode_query_len
583 + q_offset = pid_b - req_id * decode_query_len
584 +
585 + if USE_PDL:
586 + tl.extra.cuda.gdc_wait()
587 + tl.extra.cuda.gdc_launch_dependents()
588 +
589 + seq_len = tl.load(seq_lens + req_id)
590 + query_pos = seq_len - decode_query_len + q_offset
591 + # Full-CG padding uses zero-length request rows. Clamp to an empty
592 + # attention range instead of letting padded rows produce negative lengths.
593 + kv_len = tl.maximum(query_pos + 1, 0)
594 + num_blocks = (kv_len + block_size - 1) // block_size
595 +
596 + # Load NUM_TOPK_CHUNKS * BLOCK_SIZE_T candidates, padded to BLOCK_SIZE_K.
597 + # Candidate at flat position p comes from chunk = p // BLOCK_SIZE_T,
598 + # in_chunk = p % BLOCK_SIZE_T.
599 + off = tl.arange(0, BLOCK_SIZE_K)
600 + chunk_idx = off // BLOCK_SIZE_T
601 + in_chunk_idx = off % BLOCK_SIZE_T
602 + valid = chunk_idx < num_topk_chunks
603 +
604 + score_offset = (
605 + chunk_idx * stride_ts_c
606 + + pid_h * stride_ts_h
607 + + pid_b * stride_ts_b
608 + + in_chunk_idx * stride_ts_t
609 + )
610 + idx_offset = (
611 + chunk_idx * stride_ti_c
612 + + pid_h * stride_ti_h
613 + + pid_b * stride_ti_b
614 + + in_chunk_idx * stride_ti_t
615 + )
616 +
617 + score = tl.load(ts_partial_ptr + score_offset, mask=valid, other=-1e30).to(
618 + tl.float32
619 + )
620 + score = tl.where(score != score, -1e30, score)
621 + idx = tl.load(ti_partial_ptr + idx_offset, mask=valid, other=0).to(tl.int32)
622 +
623 + # Full bitonic descending sort of BLOCK_SIZE_K items.
624 + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
625 + for j in tl.static_range(1, n_dims):
626 + score, idx = _bitonic_merge(score, idx.to(tl.int32), j, 2, n_dims)
627 + score, idx = _bitonic_merge(score, idx.to(tl.int32), n_dims, True, n_dims)
628 +
629 + # Extract first BLOCK_SIZE_T positions — these are the global top-K.
630 + extract_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
631 + topk_idx_final = tl.sum(
632 + extract_mask[:, None]
633 + * tl.reshape(idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
634 + axis=0,
635 + )
636 +
637 + off_t = tl.arange(0, BLOCK_SIZE_T)
638 + tif_ptrs = (
639 + ti_final_ptr
640 + + pid_h * stride_tif_h
641 + + pid_b * stride_tif_b
642 + + off_t * stride_tif_t
643 + )
644 + store_mask = off_t < topk
645 + topk_idx_final = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx_final, -1)
646 + tl.store(
647 + tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty), mask=store_mask
648 + )
649 +
650 +
651 +# ---------------------------------------------------------------------------
652 +# Python wrappers
653 +# ---------------------------------------------------------------------------
654 +@torch.no_grad()
655 +def minimax_m3_index_score(
656 + idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim]
657 + index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim]
658 + block_table: torch.Tensor, # [batch, max_blocks]
659 + cu_seqlens_q: torch.Tensor, # [batch+1] int32
660 + seq_lens: torch.Tensor, # [batch] int32
661 + prefix_lens: torch.Tensor, # [batch] int32
662 + max_query_len: int,
663 + max_seq_len: int,
664 + num_kv_heads: int,
665 +) -> torch.Tensor:
666 + """Compute per-token index scores for each visible sparse block.
667 +
668 + Returns score [num_kv_heads, total_q, max_block], where each score is the
669 + max over a 128-token index-K block. M3 has num_idx_heads == num_kv_heads.
670 + """
671 + total_q, num_idx_heads, head_dim = idx_q.shape
672 + assert num_idx_heads == num_kv_heads, (
673 + "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
674 + )
675 + batch = cu_seqlens_q.shape[0] - 1
676 + max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
677 +
678 + # Keep score strides 16-divisible to avoid Triton recompiles.
679 + score_block_stride = round_up(max_block, 16)
680 + score = torch.empty(
681 + (num_idx_heads, total_q, score_block_stride),
682 + dtype=torch.float32,
683 + device=idx_q.device,
684 + )
685 + BLOCK_SIZE_Q = 64
686 + grid_score = (triton.cdiv(max_query_len, BLOCK_SIZE_Q), batch * num_idx_heads)
687 + _index_block_score_kernel[grid_score](
688 + idx_q,
689 + index_kv_cache,
690 + score,
691 + block_table,
692 + cu_seqlens_q,
693 + seq_lens,
694 + prefix_lens,
695 + num_idx_heads,
696 + head_dim,
697 + idx_q.stride(0),
698 + idx_q.stride(1),
699 + idx_q.stride(2),
700 + index_kv_cache.stride(0),
701 + index_kv_cache.stride(1),
702 + index_kv_cache.stride(2),
703 + score.stride(0),
704 + score.stride(1),
705 + score.stride(2),
706 + block_table.stride(0),
707 + BLOCK_SIZE_Q=BLOCK_SIZE_Q,
708 + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
709 + )
710 + return score
711 +
712 +
713 +@torch.no_grad()
714 +def minimax_m3_index_topk(
715 + score: torch.Tensor, # [num_idx_heads, total_q, max_block]
716 + cu_seqlens_q: torch.Tensor, # [batch+1] int32
717 + prefix_lens: torch.Tensor, # [batch] int32
718 + max_query_len: int,
719 + topk: int,
720 + init_blocks: int,
721 + local_blocks: int,
722 + out: torch.Tensor | None = None,
723 +) -> torch.Tensor:
724 + """Select index top-k from a precomputed score tensor.
725 +
726 + When ``out`` is provided (a ``[num_idx_heads, >=total_q, topk]`` buffer), the
727 + result is written into ``out[:, :total_q, :]`` instead of a fresh tensor --
728 + used to keep the top-k output at a stable address for cudagraph capture.
729 + """
730 + num_idx_heads = score.shape[0]
731 + batch = cu_seqlens_q.shape[0] - 1
732 + total_q = score.shape[1]
733 + if out is not None:
734 + topk_idx = out[:, :total_q, :]
735 + else:
736 + topk_idx = torch.empty(
737 + (num_idx_heads, total_q, topk),
738 + dtype=torch.int32,
739 + device=score.device,
740 + )
741 + # block_size_q == 1 -> query blocks coincide with query tokens.
742 + grid_topk = (max_query_len, batch, num_idx_heads)
743 + _topk_index_kernel[grid_topk](
744 + score,
745 + topk_idx,
746 + 1, # sample_interval (block_size_q)
747 + SPARSE_BLOCK_SIZE,
748 + cu_seqlens_q,
749 + cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
750 + prefix_lens,
751 + topk,
752 + init_blocks,
753 + local_blocks,
754 + score.stride(0),
755 + score.stride(1),
756 + score.stride(2),
757 + topk_idx.stride(0),
758 + topk_idx.stride(1),
759 + topk_idx.stride(2),
760 + MASK_INIT=False,
761 + MASK_LOCAL=False,
762 + )
763 + return topk_idx
764 +
765 +
766 +@torch.no_grad()
767 +def minimax_m3_index_decode(
768 + idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim]
769 + index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim]
770 + block_table: torch.Tensor, # [num_reqs, max_blocks]
771 + seq_lens: torch.Tensor, # [num_reqs] int32
772 + max_seq_len: int,
773 + topk: int,
774 + init_blocks: int,
775 + local_blocks: int,
776 + num_kv_heads: int,
777 + decode_query_len: int,
778 + max_decode_query_len: int,
779 + out: torch.Tensor | None = None,
780 +) -> torch.Tensor:
781 + """Decode index block-score + top-k, both split-K (cudagraph-safe).
782 +
783 + Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).
784 + When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into
785 + ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating.
786 + """
787 + total_q, num_idx_heads, head_dim = idx_q.shape
788 + assert num_idx_heads == num_kv_heads, (
789 + "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
790 + )
791 + assert decode_query_len <= max_decode_query_len
792 + assert total_q == seq_lens.shape[0] * decode_query_len
793 + batch = total_q
794 + max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
795 + use_pdl = current_platform.is_arch_support_pdl()
796 + # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
797 + # SM9+); this ROCm Triton rejects it even when False ("Keyword argument
798 + # launch_pdl was specified but unrecognised"). Only pass it when PDL is
799 + # actually supported -- on ROCm use_pdl is always False, so it's omitted.
800 + pdl_kwargs: dict[str, bool | int] = {}
801 + if use_pdl:
802 + pdl_kwargs.update({"launch_pdl": True})
803 + # TP=1 spec decode scores a wide 4-head x 4-position query tile per K block;
804 + # reduce stages to ease memory/register pressure. Keep no-spec and TP=4
805 + # single-head codegen unchanged.
806 + score_kwargs = pdl_kwargs.copy()
807 + if num_idx_heads > 1 and max_decode_query_len > 1:
808 + score_kwargs.update({"num_warps": 4, "num_stages": 2})
809 +
810 + # Keep score strides 16-divisible to avoid Triton recompiles.
811 + score_block_stride = round_up(max_block, 16)
812 + score = torch.empty(
813 + (num_idx_heads, total_q, score_block_stride),
814 + dtype=torch.float32,
815 + device=idx_q.device,
816 + )
817 + # split-K over seq blocks; chunk count depends only on shape constants so
818 + # the grid is fixed within a cuda graph.
819 + TARGET_GRID = 512
820 + MAX_NUM_KV_CHUNKS = 256
821 + # Use the configured max decode length to avoid Triton recompiles when
822 + # switching between qlen=1 and spec-decode verification batches.
823 + BLOCK_SIZE_Q = triton.next_power_of_2(max_decode_query_len)
824 + score_ctas_per_chunk = seq_lens.shape[0]
825 + target = max(
826 + 1,
827 + min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, score_ctas_per_chunk)),
828 + )
829 + num_kv_chunks = 1 << (target.bit_length() - 1)
830 + grid_score = (seq_lens.shape[0], num_kv_chunks)
831 + _decode_index_score_kernel[grid_score](
832 + idx_q,
833 + index_kv_cache,
834 + score,
835 + block_table,
836 + seq_lens,
837 + num_idx_heads,
838 + head_dim,
839 + init_blocks,
840 + local_blocks,
841 + decode_query_len,
842 + idx_q.stride(0),
843 + idx_q.stride(1),
844 + idx_q.stride(2),
845 + index_kv_cache.stride(0),
846 + index_kv_cache.stride(1),
847 + index_kv_cache.stride(2),
848 + score.stride(0),
849 + score.stride(1),
850 + score.stride(2),
851 + block_table.stride(0),
852 + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
853 + BLOCK_SIZE_Q=BLOCK_SIZE_Q,
854 + num_kv_chunks=num_kv_chunks,
855 + USE_PDL=use_pdl,
856 + **score_kwargs,
857 + )
858 +
859 + if out is not None:
860 + topk_idx = out[:, :total_q, :]
861 + else:
862 + topk_idx = torch.empty(
863 + (num_idx_heads, total_q, topk),
864 + dtype=torch.int32,
865 + device=idx_q.device,
866 + )
867 + # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts
868 + # pow2(num_topk_chunks * pow2(topk)) candidates.
869 + TOPK_TARGET_GRID = 64
870 + MAX_NUM_TOPK_CHUNKS = 16
871 + topk_target = max(
872 + 1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads))
873 + )
874 + num_topk_chunks = 1 << (topk_target.bit_length() - 1)
875 + block_size_t = triton.next_power_of_2(topk)
876 + chunk_blocks = (max_block + num_topk_chunks - 1) // num_topk_chunks
877 + topk_score_partial = torch.empty(
878 + num_topk_chunks,
879 + num_idx_heads,
880 + batch,
881 + block_size_t,
882 + dtype=torch.float32,
883 + device=idx_q.device,
884 + )
885 + topk_idx_partial = torch.empty(
886 + num_topk_chunks,
887 + num_idx_heads,
888 + batch,
889 + block_size_t,
890 + dtype=torch.int32,
891 + device=idx_q.device,
892 + )
893 + _topk_index_partial_kernel[(batch, num_idx_heads, num_topk_chunks)](
894 + score,
895 + topk_score_partial,
896 + topk_idx_partial,
897 + seq_lens,
898 + SPARSE_BLOCK_SIZE,
899 + topk,
900 + chunk_blocks,
901 + decode_query_len,
902 + score.stride(0),
903 + score.stride(1),
904 + score.stride(2),
905 + topk_score_partial.stride(0),
906 + topk_score_partial.stride(1),
907 + topk_score_partial.stride(2),
908 + topk_score_partial.stride(3),
909 + topk_idx_partial.stride(0),
910 + topk_idx_partial.stride(1),
911 + topk_idx_partial.stride(2),
912 + topk_idx_partial.stride(3),
913 + USE_PDL=use_pdl,
914 + **pdl_kwargs,
915 + )
916 + _topk_index_merge_kernel[(batch, num_idx_heads)](
917 + topk_score_partial,
918 + topk_idx_partial,
919 + topk_idx,
920 + seq_lens,
921 + SPARSE_BLOCK_SIZE,
922 + topk,
923 + decode_query_len,
924 + topk_score_partial.stride(0),
925 + topk_score_partial.stride(1),
926 + topk_score_partial.stride(2),
927 + topk_score_partial.stride(3),
928 + topk_idx_partial.stride(0),
929 + topk_idx_partial.stride(1),
930 + topk_idx_partial.stride(2),
931 + topk_idx_partial.stride(3),
932 + topk_idx.stride(0),
933 + topk_idx.stride(1),
934 + topk_idx.stride(2),
935 + num_topk_chunks=num_topk_chunks,
936 + USE_PDL=use_pdl,
937 + **pdl_kwargs,
938 + )
939 + return topk_idx
vllm/models/minimax_m3/amd/ops/sparse_attn.py new
+271
@@ -0,0 +1,271 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""ROCm gfx942/gfx950 block-sparse GQA prefill kernel for MiniMax-M3.
4 +
5 +Only the prefill path is specialized on CDNA: each 128-token KV block is split
6 +into SUB_K-token sub-tiles to right-size the per-block QK/PV MFMAs. Everything
7 +else -- the decode split-K kernels, the FP8 dtype set, the sparse block size --
8 +is reused unchanged from ``common.ops.sparse_attn``.
9 +"""
10 +
11 +import torch
12 +
13 +from vllm.models.minimax_m3.common.ops.sparse_attn import (
14 + _FP8_DTYPES,
15 + SPARSE_BLOCK_SIZE,
16 + minimax_m3_sparse_attn_decode,
17 +)
18 +from vllm.platforms.rocm import on_gfx950, on_mi3xx
19 +from vllm.triton_utils import tl, triton
20 +
21 +__all__ = ["minimax_m3_sparse_attn", "minimax_m3_sparse_attn_decode"]
22 +
23 +
24 +# Sub-tile width for the prefill kernel's per-block QK/PV GEMMs. gfx950 -> 64,
25 +# gfx942 -> 32 (re-tune with tune_sparse_attn.py). Must divide SPARSE_BLOCK_SIZE.
26 +_SPARSE_ATTN_SUB_K = SPARSE_BLOCK_SIZE // 2 if on_gfx950() else SPARSE_BLOCK_SIZE // 4
27 +
28 +_SPARSE_ATTN_PREFILL_KWARG: dict | None = None
29 +
30 +
31 +def _sparse_attn_prefill_kwargs() -> dict:
32 + """MFMA + pipeline launch params for the sub-tiled prefill kernel.
33 +
34 + gfx942 and gfx950 share the same params: ``num_warps=1`` keeps one wave
35 + resident on the small per-sub-tile GEMM, ``matrix_instr_nonkdim=16`` /
36 + ``kpack=2`` select the MFMA_16x16 path, and ``num_stages=1`` fits LDS and is
37 + fastest in the sweep. Only the sub-tile width (``_SPARSE_ATTN_SUB_K``)
38 + differs by arch. Empty on other AMD archs. Cached: arch is fixed per process.
39 + """
40 + global _SPARSE_ATTN_PREFILL_KWARG
41 + if _SPARSE_ATTN_PREFILL_KWARG is None:
42 + kwarg: dict = {}
43 + if on_mi3xx():
44 + kwarg = {
45 + "num_warps": 1,
46 + "matrix_instr_nonkdim": 16,
47 + "kpack": 2,
48 + "num_stages": 1,
49 + }
50 + _SPARSE_ATTN_PREFILL_KWARG = kwarg
51 + return _SPARSE_ATTN_PREFILL_KWARG
52 +
53 +
54 +# ---------------------------------------------------------------------------
55 +# GQA block-sparse attention (paged). Main heads attend only to the selected
56 +# blocks. BLOCK_SIZE_K == 128 so each selected block is one page.
57 +# ---------------------------------------------------------------------------
58 +# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens
59 +# might lose pointer alignment, which trigger Triton recompiles. we don't actually
60 +# need pointer alignment for those tensors anyway because we do scalar load.
61 +@triton.heuristics(
62 + {
63 + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
64 + "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]),
65 + "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
66 + "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"]
67 + * triton.next_power_of_2(args["gqa_group_size"]),
68 + }
69 +)
70 +@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
71 +def _gqa_sparse_fwd_kernel(
72 + q_ptr, # [total_q, num_heads, head_dim]
73 + kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim]
74 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
75 + o_ptr, # [total_q, num_heads, head_dim]
76 + block_table_ptr, # [num_reqs, max_blocks]
77 + cu_seqlens_q,
78 + cu_seqblocks_q,
79 + seq_lens,
80 + prefix_lens,
81 + num_kv_heads,
82 + gqa_group_size,
83 + head_dim,
84 + max_topk,
85 + num_q_loop,
86 + sm_scale,
87 + stride_qn,
88 + stride_qh,
89 + stride_qd,
90 + stride_kv_blk,
91 + stride_kv_kv,
92 + stride_kv_pos,
93 + stride_kv_h,
94 + stride_kv_d,
95 + stride_th,
96 + stride_tn,
97 + stride_tk,
98 + stride_on,
99 + stride_oh,
100 + stride_od,
101 + stride_bt_b,
102 + BLOCK_SIZE_Q: tl.constexpr,
103 + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
104 + BLOCK_SIZE_D: tl.constexpr,
105 + BLOCK_SIZE_H: tl.constexpr,
106 + BLOCK_SIZE_T: tl.constexpr,
107 + BLOCK_SIZE_QH: tl.constexpr,
108 + USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
109 + SUB_K: tl.constexpr, # CDNA only: KV sub-tile width (see _IS_MI3XX)
110 +):
111 + sm_scale_log2e = sm_scale * 1.4426950409
112 + pid_q = tl.program_id(0)
113 + pid_kh = tl.program_id(1)
114 + pid_b = tl.program_id(2)
115 + pid_h = pid_kh * gqa_group_size
116 + q_start = tl.load(cu_seqlens_q + pid_b)
117 + q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start
118 + q_block_start = tl.load(cu_seqblocks_q + pid_b)
119 + q_block_len = tl.load(cu_seqblocks_q + pid_b + 1) - q_block_start
120 + seq_len = tl.load(seq_lens + pid_b)
121 + prefix_len = tl.load(prefix_lens + pid_b)
122 + if pid_q * num_q_loop >= q_block_len:
123 + return
124 + real_q_loop = min(num_q_loop, q_block_len - pid_q * num_q_loop)
125 + bt_row = block_table_ptr + pid_b * stride_bt_b
126 + off_d = tl.arange(0, BLOCK_SIZE_D)
127 + d_mask = off_d < head_dim
128 + for j in range(real_q_loop):
129 + pid_q_j = pid_q * num_q_loop + j
130 + t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th
131 + off_t = tl.arange(0, BLOCK_SIZE_T)
132 + topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1)
133 + real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
134 + q_ptrs = tl.make_block_ptr(
135 + base=q_ptr + q_start * stride_qn + pid_h * stride_qh,
136 + shape=(q_len, gqa_group_size, head_dim),
137 + strides=(stride_qn, stride_qh, stride_qd),
138 + offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
139 + block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
140 + order=(2, 1, 0),
141 + )
142 + q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero")
143 + m_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
144 + lse_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
145 + acc_o = tl.zeros((BLOCK_SIZE_QH, BLOCK_SIZE_D), dtype=tl.float32)
146 + q = tl.reshape(q, BLOCK_SIZE_QH, BLOCK_SIZE_D)
147 +
148 + # CDNA: process each 128-token KV block in SUB_K-token sub-tiles so
149 + # each QK/PV MFMA is right-sized. Numerically equivalent to the dense
150 + # path below (flash-softmax reassociation).
151 + NUM_SUB: tl.constexpr = BLOCK_SIZE_K // SUB_K
152 + for _ in tl.range(real_topk):
153 + blk = tl.load(t_ptr_j).to(tl.int32)
154 + t_ptr_j = t_ptr_j + stride_tk
155 + c = blk * BLOCK_SIZE_K
156 + page = tl.load(bt_row + blk).to(tl.int64)
157 + kv_base = kv_cache_ptr + page * stride_kv_blk + pid_kh * stride_kv_h
158 + for sub_i in range(NUM_SUB):
159 + off_sub = tl.arange(0, SUB_K) + sub_i * SUB_K
160 + pos_sub = c + off_sub
161 + pos_mask_sub = pos_sub < seq_len
162 + k_sub = tl.load(
163 + kv_base
164 + + 0 * stride_kv_kv
165 + + off_sub[None, :] * stride_kv_pos
166 + + off_d[:, None] * stride_kv_d,
167 + mask=d_mask[:, None] & pos_mask_sub[None, :],
168 + other=0.0,
169 + )
170 + if USE_FP8:
171 + k_sub = k_sub.to(q.dtype)
172 + off_q_sub = (
173 + tl.arange(0, BLOCK_SIZE_Q)[:, None]
174 + + pid_q_j * BLOCK_SIZE_Q
175 + + prefix_len
176 + - off_sub[None, :]
177 + )
178 + qk_sub = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_H, SUB_K), dtype=tl.float32)
179 + # causal: q_abs_pos - k_off >= block_start (c)
180 + qk_sub += tl.where(off_q_sub[:, None, :] >= c, 0, float("-inf"))
181 + qk_sub = tl.reshape(qk_sub, BLOCK_SIZE_QH, SUB_K)
182 + qk_sub += tl.dot(q, k_sub) * sm_scale_log2e
183 + qk_sub += tl.where(pos_mask_sub[None, :], 0, float("-inf"))
184 + m_ij = tl.maximum(m_i, tl.max(qk_sub, axis=1))
185 + p_sub = tl.exp2(qk_sub - m_ij[:, None])
186 + l_ij = tl.sum(p_sub, axis=1)
187 + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
188 + v_sub = tl.load(
189 + kv_base
190 + + 1 * stride_kv_kv
191 + + off_sub[:, None] * stride_kv_pos
192 + + off_d[None, :] * stride_kv_d,
193 + mask=pos_mask_sub[:, None] & d_mask[None, :],
194 + other=0.0,
195 + )
196 + if USE_FP8:
197 + v_sub = v_sub.to(q.dtype)
198 + acc_o += tl.dot(p_sub.to(v_sub.dtype), v_sub)
199 + m_i = m_ij
200 + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
201 + acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None]
202 + acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D)
203 + o_ptrs = tl.make_block_ptr(
204 + base=o_ptr + q_start * stride_on + pid_h * stride_oh,
205 + shape=(q_len, gqa_group_size, head_dim),
206 + strides=(stride_on, stride_oh, stride_od),
207 + offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
208 + block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
209 + order=(2, 1, 0),
210 + )
211 + tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2))
212 +
213 +
214 +@torch.no_grad()
215 +def minimax_m3_sparse_attn(
216 + q: torch.Tensor, # [total_q, num_heads, head_dim]
217 + kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim]
218 + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
219 + block_table: torch.Tensor, # [batch, max_blocks]
220 + cu_seqlens_q: torch.Tensor, # [batch+1] int32
221 + seq_lens: torch.Tensor, # [batch] int32
222 + prefix_lens: torch.Tensor, # [batch] int32
223 + max_query_len: int,
224 + num_kv_heads: int,
225 + sm_scale: float,
226 + output: torch.Tensor, # [total_q, num_heads, head_dim]
227 +) -> None:
228 + """GQA block-sparse attention over the selected blocks. block_size_q == 1."""
229 + total_q, num_heads, head_dim = q.shape
230 + batch = cu_seqlens_q.shape[0] - 1
231 + topk = topk_idx.shape[-1]
232 + gqa_group_size = num_heads // num_kv_heads
233 + use_fp8 = kv_cache.dtype in _FP8_DTYPES
234 + grid = (max_query_len, num_kv_heads, batch)
235 + _gqa_sparse_fwd_kernel[grid](
236 + q,
237 + kv_cache,
238 + topk_idx,
239 + output,
240 + block_table,
241 + cu_seqlens_q,
242 + cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
243 + seq_lens,
244 + prefix_lens,
245 + num_kv_heads,
246 + gqa_group_size,
247 + head_dim,
248 + topk,
249 + 1, # num_q_loop
250 + sm_scale,
251 + q.stride(0),
252 + q.stride(1),
253 + q.stride(2),
254 + kv_cache.stride(0),
255 + kv_cache.stride(1),
256 + kv_cache.stride(2),
257 + kv_cache.stride(3),
258 + kv_cache.stride(4),
259 + topk_idx.stride(0),
260 + topk_idx.stride(1),
261 + topk_idx.stride(2),
262 + output.stride(0),
263 + output.stride(1),
264 + output.stride(2),
265 + block_table.stride(0),
266 + BLOCK_SIZE_Q=1,
267 + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
268 + USE_FP8=use_fp8,
269 + SUB_K=_SPARSE_ATTN_SUB_K,
270 + **_sparse_attn_prefill_kwargs(),
271 + )
vllm/models/minimax_m3/common/indexer.py
+14 -5
@@ -27,12 +27,21 @@ from vllm.distributed import get_tensor_model_parallel_world_size
27 from vllm.forward_context import get_forward_context
28 from vllm.logger import init_logger
29 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
30 -from vllm.models.minimax_m3.common.ops.index_topk import (
31 - minimax_m3_index_decode,
32 - minimax_m3_index_score,
33 - minimax_m3_index_topk,
34 -)
30 from vllm.platforms import current_platform
31 +
32 +if current_platform.is_rocm():
33 + from vllm.models.minimax_m3.amd.ops.index_topk import (
34 + minimax_m3_index_decode,
35 + minimax_m3_index_score,
36 + minimax_m3_index_topk,
37 + )
38 +else:
39 + from vllm.models.minimax_m3.common.ops.index_topk import (
40 + minimax_m3_index_decode,
41 + minimax_m3_index_score,
42 + minimax_m3_index_topk,
43 + )
44 +
45 from vllm.v1.attention.backend import (
46 AttentionBackend,
47 AttentionCGSupport,
vllm/models/minimax_m3/common/ops/sparse_attn.py
-26
@@ -31,30 +31,6 @@ _FP8_DTYPES = (
31 torch.float8_e5m2fnuz,
32 )
33
34 -_SPARSE_ATTN_NUM_STAGES_KWARG: dict | None = None
35 -
36 -
37 -def _sparse_attn_num_stages_kwarg() -> dict:
38 - """Triton ``num_stages`` override for the sparse-attn GEMM kernels.
39 -
40 - Forced only where required: CDNA3 (gfx942) caps LDS at
41 - 64 KB, and the default 2-stage pipeline double-buffers the 128x128 K/V tiles
42 - to ~66 KB ("out of resource: shared memory"), so pin gfx942 to a single
43 - stage (~32 KB, which fits). Everywhere else (NVIDIA, CDNA4 gfx950) return an
44 - empty kwarg and let Triton keep its own default -- don't second-guess it.
45 - Cached: the arch is fixed per process.
46 - """
47 - global _SPARSE_ATTN_NUM_STAGES_KWARG
48 - if _SPARSE_ATTN_NUM_STAGES_KWARG is None:
49 - kwarg: dict = {}
50 - if current_platform.is_rocm():
51 - from vllm.platforms.rocm import on_gfx942
52 -
53 - if on_gfx942():
54 - kwarg = {"num_stages": 1}
55 - _SPARSE_ATTN_NUM_STAGES_KWARG = kwarg
56 - return _SPARSE_ATTN_NUM_STAGES_KWARG
57 -
34
35 # ---------------------------------------------------------------------------
36 # GQA block-sparse attention (paged). Main heads attend only to the selected
@@ -498,7 +474,6 @@ def minimax_m3_sparse_attn(
474 BLOCK_SIZE_Q=1,
475 BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
476 USE_FP8=use_fp8,
501 - **_sparse_attn_num_stages_kwarg(),
477 )
478
479
@@ -574,7 +549,6 @@ def minimax_m3_sparse_attn_decode(
549 NUM_TOPK_CHUNKS=num_topk_chunks,
550 USE_FP8=use_fp8,
551 USE_PDL=use_pdl,
577 - **_sparse_attn_num_stages_kwarg(),
552 **pdl_launch,
553 )
554 merge_grid = (total_q, num_heads)
vllm/models/minimax_m3/common/sparse_attention.py
+14 -5
@@ -24,12 +24,21 @@ from vllm.config import VllmConfig
24 from vllm.config.cache import CacheDType
25 from vllm.forward_context import get_forward_context
26 from vllm.logger import init_logger
27 -from vllm.models.minimax_m3.common.ops.sparse_attn import (
28 - SPARSE_BLOCK_SIZE,
29 - minimax_m3_sparse_attn,
30 - minimax_m3_sparse_attn_decode,
31 -)
27 +from vllm.models.minimax_m3.common.ops.sparse_attn import SPARSE_BLOCK_SIZE
28 from vllm.platforms import current_platform
29 +
30 +# AMD/ROCm uses the gfx942/gfx950-optimized block-sparse kernels in amd.ops;
31 +# every other platform uses the generic common.ops implementation.
32 +if current_platform.is_rocm():
33 + from vllm.models.minimax_m3.amd.ops.sparse_attn import (
34 + minimax_m3_sparse_attn,
35 + minimax_m3_sparse_attn_decode,
36 + )
37 +else:
38 + from vllm.models.minimax_m3.common.ops.sparse_attn import (
39 + minimax_m3_sparse_attn,
40 + minimax_m3_sparse_attn_decode,
41 + )
42 from vllm.v1.attention.backend import (
43 AttentionBackend,
44 AttentionCGSupport,