index_topk.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 1,639 lines 55.2 KB

Raw

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)`` (one shared
11	key vector per token).
12
13	Only the paths MiniMax M3 uses are implemented: score_type="max", index value
14	disabled (score-only indexer), and shared index keys. Each local index-query
15	head selects its own block ids for the block-sparse attention kernels in
16	``sparse_attn``.
17	"""
18
19	import torch
20
21	from vllm.models.minimax_m3.amd.ops.sparse_pa import (
22	PAGES_PER_SPARSE_BLOCK,
23	_write_sparse_block_table_row_from_values,
24	)
25	from vllm.platforms import current_platform
26	from vllm.triton_utils import tl, triton
27	from vllm.utils.math_utils import round_up
28
29	# One sparse block == one KV page.
30	SPARSE_BLOCK_SIZE = 128
31	DECODE_SCORE_BALANCED_PROGRAM_BUDGET = 1024
32	DECODE_SCORE_HIGH_BATCH_PROGRAM_BUDGET = 768
33	MAX_DECODE_SCORE_BALANCED_REQUESTS = 11
34	DECODE_SCORE_DEFAULT_TARGET_GRID = 512
35	DECODE_SCORE_GFX950_HIGH_BATCH_TARGET_GRID = 1024
36	MIN_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS = 20
37	MAX_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS = 64
38	MAX_DECODE_SCORE_KV_CHUNKS = 256
39	DECODE_TOPK_BLOCKS_PER_CHUNK = 512
40	DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK = 128
41	MAX_DECODE_TOPK_FAST_CHUNKS = 16
42	DECODE_TOPK_TARGET_GRID = 64
43
44
45	def _decode_score_program_budget(
46	num_reqs: int,
47	head_dim: int,
48	query_dtype: torch.dtype,
49	cache_dtype: torch.dtype,
50	*,
51	is_gfx950: bool,
52	) -> int \| None:
53	if not (
54	is_gfx950
55	and head_dim == 128
56	and query_dtype == torch.bfloat16
57	and cache_dtype == torch.bfloat16
58	):
59	return None
60	if 1 <= num_reqs <= 8:
61	return DECODE_SCORE_BALANCED_PROGRAM_BUDGET
62	if 9 <= num_reqs <= MAX_DECODE_SCORE_BALANCED_REQUESTS:
63	return DECODE_SCORE_HIGH_BATCH_PROGRAM_BUDGET
64	return None
65
66
67	def _decode_score_split_launch_policy(
68	num_reqs: int,
69	head_dim: int,
70	query_dtype: torch.dtype,
71	cache_dtype: torch.dtype,
72	*,
73	is_gfx950: bool,
74	) -> tuple[int, bool]:
75	"""Choose the generic split-K launch and high-batch specialization."""
76	use_high_batch_config = (
77	is_gfx950
78	and MIN_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS
79	<= num_reqs
80	<= MAX_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS
81	and head_dim == 128
82	and query_dtype == torch.bfloat16
83	and cache_dtype == torch.bfloat16
84	)
85	target_grid = (
86	DECODE_SCORE_GFX950_HIGH_BATCH_TARGET_GRID
87	if use_high_batch_config
88	else DECODE_SCORE_DEFAULT_TARGET_GRID
89	)
90	target = max(
91	1,
92	min(
93	MAX_DECODE_SCORE_KV_CHUNKS,
94	target_grid // max(1, num_reqs),
95	),
96	)
97	return 1 << (target.bit_length() - 1), use_high_batch_config
98
99
100	def _decode_topk_launch_policy(
101	max_block: int,
102	total_q: int,
103	num_idx_heads: int,
104	topk: int,
105	*,
106	is_gfx950: bool,
107	) -> tuple[int, int, int, int, bool, bool]:
108	"""Choose the selector grid and compile-time launch configuration."""
109	if (
110	is_gfx950
111	and topk == 16
112	and 0 < max_block <= MAX_DECODE_TOPK_FAST_CHUNKS * DECODE_TOPK_BLOCKS_PER_CHUNK
113	):
114	if max_block <= DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK:
115	return 1, DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK, 2, 1, True, True
116	return (
117	MAX_DECODE_TOPK_FAST_CHUNKS,
118	DECODE_TOPK_BLOCKS_PER_CHUNK,
119	4,
120	2,
121	True,
122	True,
123	)
124
125	target = max(
126	1,
127	min(
128	MAX_DECODE_TOPK_FAST_CHUNKS,
129	DECODE_TOPK_TARGET_GRID // max(1, total_q * num_idx_heads),
130	),
131	)
132	return (
133	1 << (target.bit_length() - 1),
134	DECODE_TOPK_BLOCKS_PER_CHUNK,
135	8,
136	2,
137	False,
138	False,
139	)
140
141
142	# ---------------------------------------------------------------------------
143	# Bitonic top-k helpers (layout-agnostic).
144	# ---------------------------------------------------------------------------
145	@triton.jit
146	def _compare_and_swap(x, ids, flip, i: tl.constexpr, n_dims: tl.constexpr):
147	n_outer: tl.constexpr = x.numel >> n_dims
148	shape: tl.constexpr = [n_outer * 2i, 2, 2 (n_dims - i - 1)]
149	y = tl.reshape(x, shape)
150	mask = tl.arange(0, 2)[None, :, None]
151	left = tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype)
152	right = tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape).to(y.dtype)
153	left = tl.reshape(left, x.shape)
154	right = tl.reshape(right, x.shape)
155	y_idx = tl.reshape(ids, shape)
156	left_idx = tl.broadcast_to(tl.sum(y_idx * (1 - mask), 1)[:, None, :], shape)
157	right_idx = tl.broadcast_to(tl.sum(y_idx * mask, 1)[:, None, :], shape)
158	left_idx = tl.reshape(left_idx, x.shape).to(y_idx.dtype)
159	right_idx = tl.reshape(right_idx, x.shape).to(y_idx.dtype)
160	idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True)
161	ileft = left.to(idtype, bitcast=True)
162	iright = right.to(idtype, bitcast=True)
163	ix = x.to(idtype, bitcast=True)
164	cond = (left > right) != flip
165	ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix))
166	new_ids = ids ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(ids))
167	return ret.to(x.dtype, bitcast=True), new_ids
168
169
170	@triton.jit
171	def _bitonic_merge(
172	x, ids, stage: tl.constexpr, order: tl.constexpr, n_dims: tl.constexpr
173	):
174	n_outer: tl.constexpr = x.numel >> n_dims
175	tl.static_assert(stage <= n_dims)
176	if order == 2:
177	shape: tl.constexpr = [n_outer * 2 (n_dims - 1 - stage), 2, 2stage]
178	flip = tl.reshape(
179	tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape
180	)
181	else:
182	flip = order
183	for i in tl.static_range(stage):
184	x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims)
185	return x, ids
186
187
188	# ---------------------------------------------------------------------------
189	# Index block-score kernel (paged). score[h, token, block] = max over the
190	# 128-token block of (idx_q . index_k), causal-masked. BLOCK_SIZE_K == 128 so
191	# each K-tile is exactly one page (BLOCKS_PER_K_BLOCK == 1).
192	# ---------------------------------------------------------------------------
193	# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens
194	# might lose pointer alignment, which trigger Triton recompiles. we don't actually
195	# need pointer alignment for those tensors anyway because we do scalar load.
196	@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
197	def _index_block_score_kernel(
198	q_ptr, # idx_q: [total_q, num_idx_heads, head_dim]
199	ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim]
200	score_ptr, # [num_idx_heads, total_q, max_block]
201	block_table_ptr, # [num_reqs, max_blocks]
202	cu_seqlens, # [batch+1] query start offsets
203	seq_lens, # [batch] total K length
204	prefix_lens, # [batch] context length before this chunk's queries
205	num_idx_heads,
206	head_dim: tl.constexpr,
207	stride_q_n,
208	stride_q_h,
209	stride_q_d,
210	stride_ik_blk,
211	stride_ik_pos,
212	stride_ik_d,
213	stride_s_h,
214	stride_s_n,
215	stride_s_k,
216	stride_bt_b,
217	BLOCK_SIZE_Q: tl.constexpr,
218	BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
219	):
220	pid_q = tl.program_id(0)
221	pid_bh = tl.program_id(1)
222	pid_b = pid_bh // num_idx_heads
223	pid_h = pid_bh % num_idx_heads
224
225	seq_start = tl.load(cu_seqlens + pid_b)
226	q_len = tl.load(cu_seqlens + pid_b + 1) - seq_start
227	seq_len = tl.load(seq_lens + pid_b)
228	prefix_len = tl.load(prefix_lens + pid_b)
229	if BLOCK_SIZE_Q * pid_q >= q_len:
230	return
231
232	q_ptrs = tl.make_block_ptr(
233	base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h,
234	shape=(q_len, head_dim),
235	strides=(stride_q_n, stride_q_d),
236	offsets=(pid_q * BLOCK_SIZE_Q, 0),
237	block_shape=(BLOCK_SIZE_Q, head_dim),
238	order=(1, 0),
239	)
240	q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero")
241	q_start = prefix_len + pid_q * BLOCK_SIZE_Q
242
243	off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len
244	off_k = tl.arange(0, BLOCK_SIZE_K)
245	off_d = tl.arange(0, head_dim)
246	# Block table row for this request.
247	bt_row = block_table_ptr + pid_b * stride_bt_b
248	# Causal window: only blocks up to the last query token's position.
249	hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q)
250	for i in tl.range(0, hi, BLOCK_SIZE_K):
251	blk = i // BLOCK_SIZE_K
252	page = tl.load(bt_row + blk).to(tl.int64)
253	pos = i + off_k
254	# index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed)
255	# we don't need masked load for K, because KV cache ensures
256	# allocation is multiple of BLOCK_SIZE_K.
257	# for tokens beyond seqlen, they will be masked in qk later.
258	k = tl.load(
259	ik_cache_ptr
260	+ page * stride_ik_blk
261	+ off_k[None, :] * stride_ik_pos
262	+ off_d[:, None] * stride_ik_d,
263	)
264	qk = tl.dot(q, k)
265	# apply causal mask as needed
266	if q_start < i + BLOCK_SIZE_K:
267	qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf"))
268	# one sparse block per K-tile -> max over the 128 positions
269	score = tl.max(qk, axis=1) # [BLOCK_SIZE_Q]
270	s_ptrs = (
271	score_ptr
272	+ pid_h * stride_s_h
273	+ (seq_start + pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q))
274	* stride_s_n
275	+ blk * stride_s_k
276	)
277	q_store_mask = (pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) < q_len
278	tl.store(s_ptrs, score, mask=q_store_mask)
279
280
281	# ---------------------------------------------------------------------------
282	# Top-k selection over per-token block scores (layout-agnostic). block_size_q
283	# is 1 for M3, so top-k is computed per query token.
284	# ---------------------------------------------------------------------------
285	# since prefill metadata is sliced from mixed batch metadata, prefix_lens
286	# might lose pointer alignment, which trigger Triton recompiles. we don't actually
287	# need pointer alignment for those tensors anyway because we do scalar load.
288	@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])})
289	@triton.autotune(
290	configs=[
291	triton.Config({"BLOCK_SIZE_K": 2048}, num_warps=8, num_stages=2),
292	triton.Config({"BLOCK_SIZE_K": 1024}, num_warps=8, num_stages=2),
293	triton.Config({"BLOCK_SIZE_K": 512}, num_warps=8, num_stages=2),
294	triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2),
295	triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2),
296	triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2),
297	],
298	key=["BLOCK_SIZE_T"],
299	)
300	@triton.jit(do_not_specialize_on_alignment=["prefix_lens"])
301	def _topk_index_kernel(
302	s_ptr, # [num_heads, total_q, max_block]
303	ti_ptr, # [num_heads, total_q, topk]
304	sample_interval: tl.constexpr, # block_size_q (1 for M3)
305	block_size: tl.constexpr, # sparse block size (128)
306	cu_seqlens,
307	cu_seqblocks_q,
308	prefix_lens,
309	topk,
310	init_blocks: tl.constexpr,
311	local_blocks: tl.constexpr,
312	stride_s_h,
313	stride_s_n,
314	stride_s_k,
315	stride_ti_h,
316	stride_ti_n,
317	stride_ti_t,
318	BLOCK_SIZE_K: tl.constexpr,
319	BLOCK_SIZE_T: tl.constexpr,
320	MASK_INIT: tl.constexpr,
321	MASK_LOCAL: tl.constexpr,
322	):
323	tl.static_assert(BLOCK_SIZE_K > BLOCK_SIZE_T)
324	pid_q = tl.program_id(0)
325	pid_b = tl.program_id(1)
326	pid_h = tl.program_id(2)
327	seq_start = tl.load(cu_seqlens + pid_b)
328	block_start = tl.load(cu_seqblocks_q + pid_b)
329	block_num = tl.load(cu_seqblocks_q + pid_b + 1) - block_start
330	prefix_len = tl.load(prefix_lens + pid_b)
331	if pid_q >= block_num:
332	return
333	off_k = tl.arange(0, BLOCK_SIZE_K)
334	off_t = tl.arange(0, BLOCK_SIZE_T)
335	s_ptrs = (
336	s_ptr
337	+ (seq_start + pid_q * sample_interval) * stride_s_n
338	+ pid_h * stride_s_h
339	+ off_k * stride_s_k
340	)
341	topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32)
342	topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32)
343	left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2
344	valid_blocks = (prefix_len + pid_q * sample_interval + block_size) // block_size
345	for i in tl.range(0, valid_blocks, BLOCK_SIZE_K):
346	causal_mask = i + off_k < valid_blocks
347	local_mask = i + off_k >= max(0, valid_blocks - local_blocks)
348	init_mask = i + off_k < init_blocks
349	score = tl.load(s_ptrs, mask=causal_mask, other=-1e30).to(tl.float32)
350	score = tl.where(score != score, -1e30, score)
351	s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K
352	if MASK_INIT:
353	score = tl.where(causal_mask & init_mask, score - 1e29, score)
354	else:
355	score = tl.where(causal_mask & init_mask, 1e30, score)
356	if MASK_LOCAL:
357	score = tl.where(causal_mask & local_mask, score - 1e28, score)
358	else:
359	score = tl.where(causal_mask & local_mask, 1e29, score)
360	topk_score, last_topk_score = score, topk_score
361	topk_idx, last_topk_idx = (tl.where(causal_mask, i + off_k + 1, 0), topk_idx)
362	n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
363	for j in tl.static_range(1, n_dims):
364	topk_score, topk_idx = _bitonic_merge(
365	topk_score, topk_idx.to(tl.int32), j, 2, n_dims
366	)
367	if i != 0:
368	topk_score, topk_idx = _bitonic_merge(
369	topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims
370	)
371	topk_score_new = last_topk_score * left_half_mask + topk_score * (
372	1 - left_half_mask
373	)
374	topk_idx_new = last_topk_idx * left_half_mask + topk_idx * (
375	1 - left_half_mask
376	)
377	topk_score, topk_idx = _bitonic_merge(
378	topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims
379	)
380	else:
381	topk_score, topk_idx = _bitonic_merge(
382	topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims
383	)
384	topk_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
385	topk_idx = tl.sum(
386	topk_mask[:, None]
387	* tl.reshape(topk_idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
388	axis=0,
389	)
390	ti_ptrs = (
391	ti_ptr
392	+ (block_start + pid_q) * stride_ti_n
393	+ pid_h * stride_ti_h
394	+ off_t * stride_ti_t
395	)
396	store_mask = off_t < topk
397	valid_mask = off_t < valid_blocks
398	topk_idx = tl.where(store_mask & valid_mask, topk_idx, -1)
399	tl.store(ti_ptrs, topk_idx.to(ti_ptrs.dtype.element_ty), mask=store_mask)
400
401
402	# ---------------------------------------------------------------------------
403	# Decode index-score kernel (split-K over seq blocks). Decode batches are
404	# flattened request-major, with a runtime query length used to map each query
405	# token back to its request metadata. Chunk counts depend only on shape
406	# constants so the grid is fixed within a cuda graph. The score scale is omitted
407	# because decode only consumes block ordering.
408	# ---------------------------------------------------------------------------
409	@triton.jit(do_not_specialize=["num_kv_chunks", "decode_query_len"])
410	def _decode_index_score_kernel(
411	q_ptr, # idx_q: [total_q, num_idx_heads, head_dim]
412	ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim]
413	score_ptr, # [num_idx_heads, total_q, max_block]
414	block_table_ptr, # [num_reqs, max_blocks]
415	seq_lens, # [num_reqs]
416	num_idx_heads: tl.constexpr,
417	head_dim: tl.constexpr,
418	init_blocks,
419	local_blocks,
420	decode_query_len,
421	stride_q_n,
422	stride_q_h,
423	stride_q_d,
424	stride_ik_blk,
425	stride_ik_pos,
426	stride_ik_d,
427	stride_s_h,
428	stride_s_n,
429	stride_s_k,
430	stride_bt_b,
431	BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
432	BLOCK_SIZE_Q: tl.constexpr,
433	num_kv_chunks,
434	USE_PDL: tl.constexpr,
435	):
436	BLOCK_SIZE_HQ: tl.constexpr = num_idx_heads * BLOCK_SIZE_Q
437	pid_r = tl.program_id(0)
438	pid_c = tl.program_id(1)
439	hq_offsets = tl.arange(0, BLOCK_SIZE_HQ)
440	h_offsets = hq_offsets // BLOCK_SIZE_Q
441	q_offsets = hq_offsets % BLOCK_SIZE_Q
442	q_mask = q_offsets < decode_query_len
443	q_ids = pid_r * decode_query_len + q_offsets
444
445	if USE_PDL:
446	tl.extra.cuda.gdc_wait()
447	tl.extra.cuda.gdc_launch_dependents()
448
449	seq_len = tl.load(seq_lens + pid_r)
450	query_pos = seq_len - decode_query_len + q_offsets
451	# Full-CG padding uses zero-length request rows. Clamp to an empty
452	# attention range instead of letting padded rows produce negative lengths.
453	kv_len = tl.maximum(query_pos + 1, 0)
454	num_blocks_q = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
455	kv_len_max = tl.max(tl.where(q_mask, kv_len, 0), axis=0)
456	num_blocks = (kv_len_max + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
457
458	# block-aligned fixed-count split: grid independent of seq_len (cuda graph).
459	chunk_size_blocks = (num_blocks + num_kv_chunks - 1) // num_kv_chunks
460	chunk_start_block = pid_c * chunk_size_blocks
461	chunk_end_block = tl.minimum(chunk_start_block + chunk_size_blocks, num_blocks)
462	if chunk_start_block >= chunk_end_block:
463	return
464	off_k = tl.arange(0, BLOCK_SIZE_K) # positions within a 128-block
465	off_d = tl.arange(0, head_dim)
466	bt_row = block_table_ptr + pid_r * stride_bt_b
467	# Force-select init (1e30) and local (1e29, higher priority) blocks.
468	local_start = tl.maximum(0, num_blocks_q - local_blocks)
469	# Query vectors for all index heads in a small spec-decode block.
470	q = tl.load(
471	q_ptr
472	+ q_ids[None, :] * stride_q_n
473	+ h_offsets[None, :] * stride_q_h
474	+ off_d[:, None] * stride_q_d,
475	mask=q_mask[None, :],
476	other=0.0,
477	) # [D,HQ]
478	for blk in tl.range(chunk_start_block, chunk_end_block):
479	page = tl.load(bt_row + blk).to(tl.int64)
480	pos = blk * BLOCK_SIZE_K + off_k
481	pos_mask = pos[:, None] < kv_len[None, :]
482	# we don't need masked load for K, because KV cache ensures
483	# allocation is multiple of BLOCK_SIZE_K.
484	# for tokens beyond seqlen, they will be masked in qk later.
485	k = tl.load(
486	ik_cache_ptr
487	+ page * stride_ik_blk
488	+ off_k[:, None] * stride_ik_pos
489	+ off_d * stride_ik_d,
490	) # [N,D]
491	if BLOCK_SIZE_HQ == 1:
492	# Degenerate GEMV (q is [D,1]): vectorized fp32 multiply + reduce
493	# instead of an MFMA tile. Numerically equivalent to tl.dot.
494	q_vec = tl.sum(q, axis=1).to(tl.float32) # [D]
495	kq = tl.sum(k.to(tl.float32) * q_vec[None, :], axis=1)[:, None] # [N,1]
496	else:
497	# fp32 accumulation is required for the fp8 (e4m3) index cache: q/k
498	# are loaded in their stored dtype (bf16 or e4m3) and the MMA
499	# accumulates in fp32 so the per-block max score is exact for the
500	# fp8 indexer too.
501	kq = tl.dot(k, q, out_dtype=tl.float32) # [N,HQ]
502	kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf"))
503	score = tl.max(kq, axis=0) # [HQ]
504	is_visible_block = blk < num_blocks_q
505	is_init = (blk < init_blocks) & is_visible_block
506	is_local = (blk >= local_start) & is_visible_block
507	score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score))
508	tl.store(
509	score_ptr + h_offsets * stride_s_h + q_ids * stride_s_n + blk * stride_s_k,
510	score,
511	mask=q_mask,
512	)
513
514
515	@triton.jit
516	def _ceil_div_nonnegative(value, divisor):
517	return value // divisor + tl.where(value % divisor != 0, 1, 0)
518
519
520	@triton.jit
521	def _decode_index_score_mapped_range(
522	q_ptr,
523	ik_cache_ptr,
524	score_ptr,
525	block_table_ptr,
526	pid_r,
527	seq_len,
528	chunk_start_block,
529	blocks_per_program,
530	num_idx_heads: tl.constexpr,
531	head_dim: tl.constexpr,
532	init_blocks,
533	local_blocks,
534	decode_query_len,
535	stride_q_n,
536	stride_q_h,
537	stride_q_d,
538	stride_ik_blk,
539	stride_ik_pos,
540	stride_ik_d,
541	stride_s_h,
542	stride_s_n,
543	stride_s_k,
544	stride_bt_b,
545	BLOCK_SIZE_K: tl.constexpr,
546	BLOCK_SIZE_Q: tl.constexpr,
547	):
548	BLOCK_SIZE_HQ: tl.constexpr = num_idx_heads * BLOCK_SIZE_Q
549	hq_offsets = tl.arange(0, BLOCK_SIZE_HQ)
550	h_offsets = hq_offsets // BLOCK_SIZE_Q
551	q_offsets = hq_offsets % BLOCK_SIZE_Q
552	q_mask = q_offsets < decode_query_len
553	q_ids = pid_r * decode_query_len + q_offsets
554
555	query_pos = seq_len - decode_query_len + q_offsets
556	kv_len = tl.maximum(query_pos + 1, 0)
557	num_blocks_q = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
558	kv_len_max = tl.max(tl.where(q_mask, kv_len, 0), axis=0)
559	num_blocks = (kv_len_max + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
560	chunk_end_block = tl.minimum(
561	chunk_start_block + blocks_per_program,
562	num_blocks,
563	)
564
565	if chunk_start_block < chunk_end_block:
566	off_k = tl.arange(0, BLOCK_SIZE_K)
567	off_d = tl.arange(0, head_dim)
568	bt_row = block_table_ptr + pid_r * stride_bt_b
569	local_start = tl.maximum(0, num_blocks_q - local_blocks)
570	q = tl.load(
571	q_ptr
572	+ q_ids[None, :] * stride_q_n
573	+ h_offsets[None, :] * stride_q_h
574	+ off_d[:, None] * stride_q_d,
575	mask=q_mask[None, :],
576	other=0.0,
577	)
578	for blk in tl.range(chunk_start_block, chunk_end_block):
579	page = tl.load(bt_row + blk).to(tl.int64)
580	pos = blk * BLOCK_SIZE_K + off_k
581	pos_mask = pos[:, None] < kv_len[None, :]
582	k = tl.load(
583	ik_cache_ptr
584	+ page * stride_ik_blk
585	+ off_k[:, None] * stride_ik_pos
586	+ off_d * stride_ik_d,
587	)
588	if BLOCK_SIZE_HQ == 1:
589	q_vec = tl.sum(q, axis=1).to(tl.float32)
590	kq = tl.sum(
591	k.to(tl.float32) * q_vec[None, :],
592	axis=1,
593	)[:, None]
594	else:
595	kq = tl.dot(k, q, out_dtype=tl.float32)
596	kq = tl.where(
597	pos_mask & q_mask[None, :],
598	kq,
599	float("-inf"),
600	)
601	score = tl.max(kq, axis=0)
602	is_visible_block = blk < num_blocks_q
603	is_init = (blk < init_blocks) & is_visible_block
604	is_local = (blk >= local_start) & is_visible_block
605	score = tl.where(
606	is_local,
607	1e29,
608	tl.where(is_init, 1e30, score),
609	)
610	tl.store(
611	score_ptr
612	+ h_offsets * stride_s_h
613	+ q_ids * stride_s_n
614	+ blk * stride_s_k,
615	score,
616	mask=q_mask,
617	)
618
619
620	@triton.jit(do_not_specialize=["score_program_budget", "decode_query_len"])
621	def _decode_index_score_balanced_kernel(
622	q_ptr,
623	ik_cache_ptr,
624	score_ptr,
625	block_table_ptr,
626	seq_lens,
627	num_idx_heads: tl.constexpr,
628	head_dim: tl.constexpr,
629	init_blocks,
630	local_blocks,
631	NUM_REQUESTS: tl.constexpr,
632	score_program_budget,
633	decode_query_len,
634	stride_q_n,
635	stride_q_h,
636	stride_q_d,
637	stride_ik_blk,
638	stride_ik_pos,
639	stride_ik_d,
640	stride_s_h,
641	stride_s_n,
642	stride_s_k,
643	stride_bt_b,
644	BLOCK_SIZE_K: tl.constexpr,
645	BLOCK_SIZE_Q: tl.constexpr,
646	):
647	pid = tl.program_id(0)
648	total_blocks = tl.zeros((), dtype=tl.int32)
649	for req in tl.static_range(0, NUM_REQUESTS):
650	req_seq_len = tl.maximum(tl.load(seq_lens + req), 0)
651	total_blocks += _ceil_div_nonnegative(
652	req_seq_len,
653	BLOCK_SIZE_K,
654	)
655	blocks_per_program = tl.maximum(
656	1,
657	_ceil_div_nonnegative(
658	total_blocks,
659	score_program_budget,
660	),
661	)
662
663	program_offset = tl.zeros((), dtype=tl.int32)
664	pid_r = tl.zeros((), dtype=tl.int32)
665	chunk_start_block = tl.zeros((), dtype=tl.int32)
666	selected_seq_len = tl.zeros((), dtype=tl.int32)
667	owns_work = tl.zeros((), dtype=tl.int32)
668	for req in tl.static_range(0, NUM_REQUESTS):
669	req_seq_len = tl.maximum(tl.load(seq_lens + req), 0)
670	req_num_blocks = _ceil_div_nonnegative(
671	req_seq_len,
672	BLOCK_SIZE_K,
673	)
674	req_programs = _ceil_div_nonnegative(
675	req_num_blocks,
676	blocks_per_program,
677	)
678	owns_request = (pid >= program_offset) & (pid < program_offset + req_programs)
679	pid_r = tl.where(owns_request, req, pid_r)
680	local_program = tl.where(
681	owns_request,
682	pid - program_offset,
683	0,
684	)
685	chunk_start_block = tl.where(
686	owns_request,
687	local_program * blocks_per_program,
688	chunk_start_block,
689	)
690	selected_seq_len = tl.where(
691	owns_request,
692	req_seq_len,
693	selected_seq_len,
694	)
695	owns_work += owns_request.to(tl.int32)
696	program_offset += req_programs
697	if owns_work == 0:
698	return
699
700	_decode_index_score_mapped_range(
701	q_ptr,
702	ik_cache_ptr,
703	score_ptr,
704	block_table_ptr,
705	pid_r,
706	selected_seq_len,
707	chunk_start_block,
708	blocks_per_program,
709	num_idx_heads,
710	head_dim,
711	init_blocks,
712	local_blocks,
713	decode_query_len,
714	stride_q_n,
715	stride_q_h,
716	stride_q_d,
717	stride_ik_blk,
718	stride_ik_pos,
719	stride_ik_d,
720	stride_s_h,
721	stride_s_n,
722	stride_s_k,
723	stride_bt_b,
724	BLOCK_SIZE_K,
725	BLOCK_SIZE_Q,
726	)
727
728
729	# ---------------------------------------------------------------------------
730	# Decode top-k. Each CTA keeps a packed score/id top-k. Short rows emit
731	# directly; long rows synchronize their live CTAs and the last arrival merges
732	# the partial lists and optionally builds the sparse page table.
733	# ---------------------------------------------------------------------------
734	@triton.jit
735	def _decode_topk_key(score, index, valid):
736	score = tl.where(score != score, -1e30, score)
737	score = tl.where(score == 0.0, 0.0, score)
738	bits = score.to(tl.uint32, bitcast=True)
739	ordered = bits ^ tl.where(bits >> 31 != 0, 0xFFFFFFFF, 0x80000000)
740	tie = (0xFFFF - index).to(tl.int64)
741	key = (1 << 48) \| (ordered.to(tl.int64) << 16) \| tie
742	return tl.where(valid, key, 0)
743
744
745	@triton.jit
746	def _store_decode_topk(
747	winners,
748	topk_ptr,
749	attention_block_table_ptr,
750	sparse_bt_ptr,
751	sparse_ctx_ptr,
752	pid_b,
753	pid_h,
754	req_id,
755	query_pos,
756	num_blocks,
757	topk: tl.constexpr,
758	block_size: tl.constexpr,
759	pages_per_sparse_block: tl.constexpr,
760	block_page_stride: tl.constexpr,
761	stride_topk_h,
762	stride_topk_b,
763	stride_topk_t,
764	stride_attention_bt_b,
765	stride_sparse_bt_b,
766	BLOCK_SIZE_T: tl.constexpr,
767	EMIT_SPARSE_TABLE: tl.constexpr,
768	):
769	off_t = tl.arange(0, BLOCK_SIZE_T)
770	topk_idx = (0xFFFF - (winners & 0xFFFF)).to(tl.int32) - 1
771	topk_idx = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx, -1)
772	tl.store(
773	topk_ptr
774	+ pid_h * stride_topk_h
775	+ pid_b * stride_topk_b
776	+ off_t * stride_topk_t,
777	topk_idx,
778	mask=off_t < topk,
779	)
780	if EMIT_SPARSE_TABLE:
781	_write_sparse_block_table_row_from_values(
782	topk_idx,
783	attention_block_table_ptr + req_id * stride_attention_bt_b,
784	sparse_bt_ptr + pid_b * stride_sparse_bt_b,
785	sparse_ctx_ptr + pid_b,
786	query_pos,
787	topk,
788	block_size,
789	pages_per_sparse_block,
790	block_page_stride,
791	BLOCK_SIZE_T,
792	)
793
794
795	@triton.jit
796	def _merge_store_decode_topk(
797	partial_ptr,
798	counter,
799	topk_ptr,
800	attention_block_table_ptr,
801	sparse_bt_ptr,
802	sparse_ctx_ptr,
803	active_chunks,
804	pid_b,
805	pid_h,
806	req_id,
807	query_pos,
808	num_blocks,
809	stride_partial_c,
810	stride_partial_h,
811	stride_partial_b,
812	stride_partial_t,
813	stride_topk_h,
814	stride_topk_b,
815	stride_topk_t,
816	stride_attention_bt_b,
817	stride_sparse_bt_b,
818	topk: tl.constexpr,
819	block_size: tl.constexpr,
820	pages_per_sparse_block: tl.constexpr,
821	block_page_stride: tl.constexpr,
822	BLOCK_SIZE_T: tl.constexpr,
823	MERGE_CHUNKS: tl.constexpr,
824	EMIT_SPARSE_TABLE: tl.constexpr,
825	):
826	candidate_width: tl.constexpr = MERGE_CHUNKS * BLOCK_SIZE_T
827	off_candidate = tl.arange(0, candidate_width)
828	candidate_chunk = off_candidate // BLOCK_SIZE_T
829	candidate_slot = off_candidate % BLOCK_SIZE_T
830	candidates = tl.load(
831	partial_ptr
832	+ candidate_chunk * stride_partial_c
833	+ pid_h * stride_partial_h
834	+ pid_b * stride_partial_b
835	+ candidate_slot * stride_partial_t,
836	mask=candidate_chunk < active_chunks,
837	other=0,
838	volatile=True,
839	)
840	winners = tl.topk(candidates, BLOCK_SIZE_T)
841	_store_decode_topk(
842	winners,
843	topk_ptr,
844	attention_block_table_ptr,
845	sparse_bt_ptr,
846	sparse_ctx_ptr,
847	pid_b,
848	pid_h,
849	req_id,
850	query_pos,
851	num_blocks,
852	topk,
853	block_size,
854	pages_per_sparse_block,
855	block_page_stride,
856	stride_topk_h,
857	stride_topk_b,
858	stride_topk_t,
859	stride_attention_bt_b,
860	stride_sparse_bt_b,
861	BLOCK_SIZE_T,
862	EMIT_SPARSE_TABLE,
863	)
864	tl.atomic_xchg(counter, 0, sem="acq_rel", scope="gpu")
865
866
867	@triton.jit(do_not_specialize=["decode_query_len"])
868	def _decode_topk_fused_kernel(
869	score_ptr, # [num_idx_heads, total_q, max_block]
870	partial_ptr, # [NUM_TOPK_CHUNKS, num_idx_heads, total_q, BLOCK_SIZE_T]
871	counter_ptr, # [num_idx_heads, total_q]
872	topk_ptr, # [num_idx_heads, total_q, topk]
873	seq_lens_ptr, # [num_reqs]
874	attention_block_table_ptr, # [num_reqs, max_blocks]
875	sparse_bt_ptr, # [total_q, topk * PAGES_PER_SPARSE_BLOCK]
876	sparse_ctx_ptr, # [total_q]
877	decode_query_len,
878	stride_s_h,
879	stride_s_b,
880	stride_s_k,
881	stride_partial_c,
882	stride_partial_h,
883	stride_partial_b,
884	stride_partial_t,
885	stride_counter_h,
886	stride_counter_b,
887	stride_topk_h,
888	stride_topk_b,
889	stride_topk_t,
890	stride_attention_bt_b,
891	stride_sparse_bt_b,
892	topk: tl.constexpr,
893	block_size: tl.constexpr,
894	pages_per_sparse_block: tl.constexpr,
895	block_page_stride: tl.constexpr,
896	NUM_TOPK_CHUNKS: tl.constexpr,
897	BLOCK_SIZE_K: tl.constexpr,
898	BLOCK_SIZE_T: tl.constexpr,
899	EMIT_SPARSE_TABLE: tl.constexpr,
900	SINGLE_TILE_GUARANTEED: tl.constexpr,
901	ADAPTIVE_FINAL_MERGE: tl.constexpr,
902	):
903	tl.static_assert(topk <= BLOCK_SIZE_T)
904	tl.static_assert(BLOCK_SIZE_T <= BLOCK_SIZE_K)
905	pid_b = tl.program_id(0)
906	pid_h = tl.program_id(1)
907	pid_chunk = tl.program_id(2)
908	req_id = pid_b // decode_query_len
909	q_offset = pid_b - req_id * decode_query_len
910
911	seq_len = tl.load(seq_lens_ptr + req_id)
912	query_pos = seq_len - decode_query_len + q_offset
913	kv_len = tl.maximum(query_pos + 1, 0)
914	num_blocks = (kv_len + block_size - 1) // block_size
915
916	desired_chunks = (num_blocks + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
917	active_chunks = tl.maximum(1, tl.minimum(NUM_TOPK_CHUNKS, desired_chunks))
918	if pid_chunk >= active_chunks:
919	return
920
921	blocks_per_chunk = num_blocks // active_chunks
922	extra_chunks = num_blocks - blocks_per_chunk * active_chunks
923	chunk_start = pid_chunk * blocks_per_chunk + tl.minimum(pid_chunk, extra_chunks)
924	chunk_blocks = blocks_per_chunk + (pid_chunk < extra_chunks)
925
926	off_k = tl.arange(0, BLOCK_SIZE_K)
927	off_t = tl.arange(0, BLOCK_SIZE_T)
928	if SINGLE_TILE_GUARANTEED:
929	valid = off_k < chunk_blocks
930	score = tl.load(
931	score_ptr
932	+ pid_h * stride_s_h
933	+ pid_b * stride_s_b
934	+ (chunk_start + off_k) * stride_s_k,
935	mask=valid,
936	other=-1e30,
937	).to(tl.float32)
938	index = (chunk_start + off_k + 1).to(tl.int32)
939	local_topk = tl.topk(
940	_decode_topk_key(score, index, valid),
941	BLOCK_SIZE_T,
942	)
943	else:
944	score_ptrs = (
945	score_ptr
946	+ pid_h * stride_s_h
947	+ pid_b * stride_s_b
948	+ (chunk_start + off_k) * stride_s_k
949	)
950	valid = off_k < chunk_blocks
951	score = tl.load(score_ptrs, mask=valid, other=-1e30).to(tl.float32)
952	index = (chunk_start + off_k + 1).to(tl.int32)
953	local_topk = tl.topk(
954	_decode_topk_key(score, index, valid),
955	BLOCK_SIZE_T,
956	)
957	score_ptrs += BLOCK_SIZE_K * stride_s_k
958	for offset in tl.range(BLOCK_SIZE_K, chunk_blocks, BLOCK_SIZE_K):
959	valid = off_k < chunk_blocks - offset
960	score = tl.load(score_ptrs, mask=valid, other=-1e30).to(tl.float32)
961	index = (chunk_start + offset + off_k + 1).to(tl.int32)
962	tile_topk = tl.topk(
963	_decode_topk_key(score, index, valid),
964	BLOCK_SIZE_T,
965	)
966	local_topk = tl.topk(
967	tl.cat(local_topk, tile_topk, can_reorder=True),
968	BLOCK_SIZE_T,
969	)
970	score_ptrs += BLOCK_SIZE_K * stride_s_k
971
972	if active_chunks == 1:
973	_store_decode_topk(
974	local_topk,
975	topk_ptr,
976	attention_block_table_ptr,
977	sparse_bt_ptr,
978	sparse_ctx_ptr,
979	pid_b,
980	pid_h,
981	req_id,
982	query_pos,
983	num_blocks,
984	topk,
985	block_size,
986	pages_per_sparse_block,
987	block_page_stride,
988	stride_topk_h,
989	stride_topk_b,
990	stride_topk_t,
991	stride_attention_bt_b,
992	stride_sparse_bt_b,
993	BLOCK_SIZE_T,
994	EMIT_SPARSE_TABLE,
995	)
996	return
997
998	tl.store(
999	partial_ptr
1000	+ pid_chunk * stride_partial_c
1001	+ pid_h * stride_partial_h
1002	+ pid_b * stride_partial_b
1003	+ off_t * stride_partial_t,
1004	local_topk,
1005	)
1006	counter = counter_ptr + pid_h * stride_counter_h + pid_b * stride_counter_b
1007	arrival = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu")
1008	if arrival != active_chunks - 1:
1009	return
1010
1011	if ADAPTIVE_FINAL_MERGE:
1012	if active_chunks <= 2:
1013	_merge_store_decode_topk(
1014	partial_ptr,
1015	counter,
1016	topk_ptr,
1017	attention_block_table_ptr,
1018	sparse_bt_ptr,
1019	sparse_ctx_ptr,
1020	active_chunks,
1021	pid_b,
1022	pid_h,
1023	req_id,
1024	query_pos,
1025	num_blocks,
1026	stride_partial_c,
1027	stride_partial_h,
1028	stride_partial_b,
1029	stride_partial_t,
1030	stride_topk_h,
1031	stride_topk_b,
1032	stride_topk_t,
1033	stride_attention_bt_b,
1034	stride_sparse_bt_b,
1035	topk,
1036	block_size,
1037	pages_per_sparse_block,
1038	block_page_stride,
1039	BLOCK_SIZE_T,
1040	2,
1041	EMIT_SPARSE_TABLE,
1042	)
1043	return
1044	if active_chunks <= 4:
1045	_merge_store_decode_topk(
1046	partial_ptr,
1047	counter,
1048	topk_ptr,
1049	attention_block_table_ptr,
1050	sparse_bt_ptr,
1051	sparse_ctx_ptr,
1052	active_chunks,
1053	pid_b,
1054	pid_h,
1055	req_id,
1056	query_pos,
1057	num_blocks,
1058	stride_partial_c,
1059	stride_partial_h,
1060	stride_partial_b,
1061	stride_partial_t,
1062	stride_topk_h,
1063	stride_topk_b,
1064	stride_topk_t,
1065	stride_attention_bt_b,
1066	stride_sparse_bt_b,
1067	topk,
1068	block_size,
1069	pages_per_sparse_block,
1070	block_page_stride,
1071	BLOCK_SIZE_T,
1072	4,
1073	EMIT_SPARSE_TABLE,
1074	)
1075	return
1076	if active_chunks <= 8:
1077	_merge_store_decode_topk(
1078	partial_ptr,
1079	counter,
1080	topk_ptr,
1081	attention_block_table_ptr,
1082	sparse_bt_ptr,
1083	sparse_ctx_ptr,
1084	active_chunks,
1085	pid_b,
1086	pid_h,
1087	req_id,
1088	query_pos,
1089	num_blocks,
1090	stride_partial_c,
1091	stride_partial_h,
1092	stride_partial_b,
1093	stride_partial_t,
1094	stride_topk_h,
1095	stride_topk_b,
1096	stride_topk_t,
1097	stride_attention_bt_b,
1098	stride_sparse_bt_b,
1099	topk,
1100	block_size,
1101	pages_per_sparse_block,
1102	block_page_stride,
1103	BLOCK_SIZE_T,
1104	8,
1105	EMIT_SPARSE_TABLE,
1106	)
1107	return
1108	_merge_store_decode_topk(
1109	partial_ptr,
1110	counter,
1111	topk_ptr,
1112	attention_block_table_ptr,
1113	sparse_bt_ptr,
1114	sparse_ctx_ptr,
1115	active_chunks,
1116	pid_b,
1117	pid_h,
1118	req_id,
1119	query_pos,
1120	num_blocks,
1121	stride_partial_c,
1122	stride_partial_h,
1123	stride_partial_b,
1124	stride_partial_t,
1125	stride_topk_h,
1126	stride_topk_b,
1127	stride_topk_t,
1128	stride_attention_bt_b,
1129	stride_sparse_bt_b,
1130	topk,
1131	block_size,
1132	pages_per_sparse_block,
1133	block_page_stride,
1134	BLOCK_SIZE_T,
1135	16,
1136	EMIT_SPARSE_TABLE,
1137	)
1138	return
1139
1140	candidate_width: tl.constexpr = NUM_TOPK_CHUNKS * BLOCK_SIZE_T
1141	off_candidate = tl.arange(0, candidate_width)
1142	candidate_chunk = off_candidate // BLOCK_SIZE_T
1143	candidate_slot = off_candidate % BLOCK_SIZE_T
1144	candidates = tl.load(
1145	partial_ptr
1146	+ candidate_chunk * stride_partial_c
1147	+ pid_h * stride_partial_h
1148	+ pid_b * stride_partial_b
1149	+ candidate_slot * stride_partial_t,
1150	mask=candidate_chunk < active_chunks,
1151	other=0,
1152	volatile=True,
1153	)
1154	winners = tl.topk(candidates, BLOCK_SIZE_T)
1155	_store_decode_topk(
1156	winners,
1157	topk_ptr,
1158	attention_block_table_ptr,
1159	sparse_bt_ptr,
1160	sparse_ctx_ptr,
1161	pid_b,
1162	pid_h,
1163	req_id,
1164	query_pos,
1165	num_blocks,
1166	topk,
1167	block_size,
1168	pages_per_sparse_block,
1169	block_page_stride,
1170	stride_topk_h,
1171	stride_topk_b,
1172	stride_topk_t,
1173	stride_attention_bt_b,
1174	stride_sparse_bt_b,
1175	BLOCK_SIZE_T,
1176	EMIT_SPARSE_TABLE,
1177	)
1178	tl.atomic_xchg(counter, 0, sem="acq_rel", scope="gpu")
1179
1180
1181	# ---------------------------------------------------------------------------
1182	# Python wrappers
1183	# ---------------------------------------------------------------------------
1184	@torch.no_grad()
1185	def minimax_m3_index_score(
1186	idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim]
1187	index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim]
1188	block_table: torch.Tensor, # [batch, max_blocks]
1189	cu_seqlens_q: torch.Tensor, # [batch+1] int32
1190	seq_lens: torch.Tensor, # [batch] int32
1191	prefix_lens: torch.Tensor, # [batch] int32
1192	max_query_len: int,
1193	max_seq_len: int,
1194	num_kv_heads: int,
1195	) -> torch.Tensor:
1196	"""Compute per-token index scores for each visible sparse block.
1197
1198	Returns score [num_kv_heads, total_q, max_block], where each score is the
1199	max over a 128-token index-K block. M3 has num_idx_heads == num_kv_heads.
1200	"""
1201	total_q, num_idx_heads, head_dim = idx_q.shape
1202	assert num_idx_heads == num_kv_heads, (
1203	"M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
1204	)
1205	batch = cu_seqlens_q.shape[0] - 1
1206	max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
1207
1208	# Keep score strides 16-divisible to avoid Triton recompiles.
1209	score_block_stride = round_up(max_block, 16)
1210	score = torch.empty(
1211	(num_idx_heads, total_q, score_block_stride),
1212	dtype=torch.float32,
1213	device=idx_q.device,
1214	)
1215	BLOCK_SIZE_Q = 64
1216	grid_score = (triton.cdiv(max_query_len, BLOCK_SIZE_Q), batch * num_idx_heads)
1217	_index_block_score_kernel[grid_score](
1218	idx_q,
1219	index_kv_cache,
1220	score,
1221	block_table,
1222	cu_seqlens_q,
1223	seq_lens,
1224	prefix_lens,
1225	num_idx_heads,
1226	head_dim,
1227	idx_q.stride(0),
1228	idx_q.stride(1),
1229	idx_q.stride(2),
1230	index_kv_cache.stride(0),
1231	index_kv_cache.stride(1),
1232	index_kv_cache.stride(2),
1233	score.stride(0),
1234	score.stride(1),
1235	score.stride(2),
1236	block_table.stride(0),
1237	BLOCK_SIZE_Q=BLOCK_SIZE_Q,
1238	BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
1239	)
1240	return score
1241
1242
1243	@torch.no_grad()
1244	def minimax_m3_index_topk(
1245	score: torch.Tensor, # [num_idx_heads, total_q, max_block]
1246	cu_seqlens_q: torch.Tensor, # [batch+1] int32
1247	prefix_lens: torch.Tensor, # [batch] int32
1248	max_query_len: int,
1249	topk: int,
1250	init_blocks: int,
1251	local_blocks: int,
1252	out: torch.Tensor \| None = None,
1253	) -> torch.Tensor:
1254	"""Select index top-k from a precomputed score tensor.
1255
1256	When ``out`` is provided (a ``[num_idx_heads, >=total_q, topk]`` buffer), the
1257	result is written into ``out[:, :total_q, :]`` instead of a fresh tensor --
1258	used to keep the top-k output at a stable address for cudagraph capture.
1259	"""
1260	num_idx_heads = score.shape[0]
1261	batch = cu_seqlens_q.shape[0] - 1
1262	total_q = score.shape[1]
1263	if out is not None:
1264	topk_idx = out[:, :total_q, :]
1265	else:
1266	topk_idx = torch.empty(
1267	(num_idx_heads, total_q, topk),
1268	dtype=torch.int32,
1269	device=score.device,
1270	)
1271	# block_size_q == 1 -> query blocks coincide with query tokens.
1272	grid_topk = (max_query_len, batch, num_idx_heads)
1273	_topk_index_kernel[grid_topk](
1274	score,
1275	topk_idx,
1276	1, # sample_interval (block_size_q)
1277	SPARSE_BLOCK_SIZE,
1278	cu_seqlens_q,
1279	cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
1280	prefix_lens,
1281	topk,
1282	init_blocks,
1283	local_blocks,
1284	score.stride(0),
1285	score.stride(1),
1286	score.stride(2),
1287	topk_idx.stride(0),
1288	topk_idx.stride(1),
1289	topk_idx.stride(2),
1290	MASK_INIT=False,
1291	MASK_LOCAL=False,
1292	)
1293	return topk_idx
1294
1295
1296	@torch.no_grad()
1297	def minimax_m3_index_decode(
1298	idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim]
1299	index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim]
1300	block_table: torch.Tensor, # [num_reqs, max_blocks]
1301	seq_lens: torch.Tensor, # [num_reqs] int32
1302	max_seq_len: int,
1303	topk: int,
1304	init_blocks: int,
1305	local_blocks: int,
1306	num_kv_heads: int,
1307	decode_query_len: int,
1308	max_decode_query_len: int,
1309	out: torch.Tensor \| None = None,
1310	*,
1311	attention_block_table: torch.Tensor \| None = None,
1312	sparse_block_table_out: torch.Tensor \| None = None,
1313	sparse_context_lens_out: torch.Tensor \| None = None,
1314	block_page_stride: int \| None = None,
1315	completion_counter: torch.Tensor \| None = None,
1316	) -> torch.Tensor:
1317	"""Decode index block-score followed by fused adaptive top-k selection.
1318
1319	Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).
1320	When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into
1321	``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating.
1322	The optional sparse-table arguments fuse current-layer table construction
1323	into the selector. They must be provided together. ``completion_counter``
1324	provides stable per-query synchronization storage for CUDA graphs. It must
1325	be zero before its first launch and must not be shared by overlapping
1326	selector invocations; every completed launch resets its active entries.
1327	"""
1328	total_q, num_idx_heads, head_dim = idx_q.shape
1329	assert num_idx_heads == num_kv_heads, (
1330	"M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
1331	)
1332	assert 0 < decode_query_len <= max_decode_query_len
1333	assert total_q == seq_lens.shape[0] * decode_query_len
1334	batch = total_q
1335	emit_sparse_table = attention_block_table is not None
1336	if emit_sparse_table and (
1337	sparse_block_table_out is None
1338	or sparse_context_lens_out is None
1339	or block_page_stride is None
1340	):
1341	raise ValueError(
1342	"MiniMax-M3 fused decode sparse-table arguments must be provided together"
1343	)
1344	if not emit_sparse_table and (
1345	sparse_block_table_out is not None
1346	or sparse_context_lens_out is not None
1347	or block_page_stride is not None
1348	):
1349	raise ValueError(
1350	"MiniMax-M3 fused decode sparse-table arguments must be provided together"
1351	)
1352	if emit_sparse_table:
1353	assert attention_block_table is not None
1354	assert sparse_block_table_out is not None
1355	assert sparse_context_lens_out is not None
1356	assert block_page_stride is not None
1357	if num_idx_heads != 1:
1358	raise ValueError(
1359	"MiniMax-M3 fused decode sparse-table construction requires "
1360	f"one index head, got {num_idx_heads}"
1361	)
1362	if block_page_stride not in (
1363	PAGES_PER_SPARSE_BLOCK,
1364	2 * PAGES_PER_SPARSE_BLOCK,
1365	):
1366	raise ValueError(
1367	"MiniMax-M3 fused decode sparse-table page stride must be "
1368	f"{PAGES_PER_SPARSE_BLOCK} or {2 * PAGES_PER_SPARSE_BLOCK}, "
1369	f"got {block_page_stride}"
1370	)
1371	expected_sparse_shape = (total_q, topk * PAGES_PER_SPARSE_BLOCK)
1372	if sparse_block_table_out.shape != expected_sparse_shape:
1373	raise ValueError(
1374	"MiniMax-M3 fused decode sparse block table has shape "
1375	f"{tuple(sparse_block_table_out.shape)}, expected "
1376	f"{expected_sparse_shape}"
1377	)
1378	if sparse_context_lens_out.shape != (total_q,):
1379	raise ValueError(
1380	"MiniMax-M3 fused decode sparse context lengths have shape "
1381	f"{tuple(sparse_context_lens_out.shape)}, expected {(total_q,)}"
1382	)
1383	if (
1384	attention_block_table.dim() != 2
1385	or attention_block_table.shape[0] != seq_lens.shape[0]
1386	):
1387	raise ValueError(
1388	"MiniMax-M3 fused decode attention block table must be "
1389	"[num_requests, max_blocks]"
1390	)
1391	int32_tensors = (
1392	attention_block_table,
1393	sparse_block_table_out,
1394	sparse_context_lens_out,
1395	)
1396	if any(tensor.dtype != torch.int32 for tensor in int32_tensors):
1397	raise ValueError("MiniMax-M3 fused decode sparse tables require int32")
1398	if any(tensor.device != idx_q.device for tensor in int32_tensors):
1399	raise ValueError(
1400	"MiniMax-M3 fused decode sparse tables must share the query device"
1401	)
1402	if (
1403	attention_block_table.stride(1) != 1
1404	or sparse_block_table_out.stride(1) != 1
1405	or sparse_context_lens_out.stride(0) != 1
1406	):
1407	raise ValueError(
1408	"MiniMax-M3 fused decode sparse tables require contiguous rows"
1409	)
1410	max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
1411	if max_block >= 0xFFFF:
1412	raise ValueError(
1413	"MiniMax-M3 decode top-k supports fewer than 65535 sparse blocks"
1414	)
1415	if emit_sparse_table:
1416	assert attention_block_table is not None
1417	if attention_block_table.shape[1] < max_block:
1418	raise ValueError(
1419	"MiniMax-M3 fused decode attention block table is shorter "
1420	f"than the required {max_block} sparse blocks"
1421	)
1422	use_pdl = current_platform.is_arch_support_pdl()
1423	# `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
1424	# SM9+); this ROCm Triton rejects it even when False ("Keyword argument
1425	# launch_pdl was specified but unrecognised"). Only pass it when PDL is
1426	# actually supported -- on ROCm use_pdl is always False, so it's omitted.
1427	pdl_kwargs: dict[str, bool \| int] = {}
1428	if use_pdl:
1429	pdl_kwargs.update({"launch_pdl": True})
1430	is_gfx950 = False
1431	if current_platform.is_rocm():
1432	from vllm.platforms.rocm import on_gfx950
1433
1434	is_gfx950 = on_gfx950()
1435	# Multi-head spec decode scores a wider head-position tile per K block;
1436	# reduce stages to ease memory/register pressure on the fallback path.
1437	score_kwargs = pdl_kwargs.copy()
1438	if num_idx_heads > 1 and max_decode_query_len > 1:
1439	score_kwargs.update({"num_warps": 4, "num_stages": 2})
1440
1441	# Keep score strides 16-divisible to avoid Triton recompiles.
1442	score_block_stride = round_up(max_block, 16)
1443	score = torch.empty(
1444	(num_idx_heads, total_q, score_block_stride),
1445	dtype=torch.float32,
1446	device=idx_q.device,
1447	)
1448	# Use the configured max decode length to avoid Triton recompiles when
1449	# switching between qlen=1 and spec-decode verification batches.
1450	BLOCK_SIZE_Q = triton.next_power_of_2(max_decode_query_len)
1451	num_reqs = seq_lens.shape[0]
1452	score_program_budget = _decode_score_program_budget(
1453	num_reqs,
1454	head_dim,
1455	idx_q.dtype,
1456	index_kv_cache.dtype,
1457	is_gfx950=is_gfx950,
1458	)
1459	grid_score: tuple[int, ...]
1460	if score_program_budget is not None:
1461	grid_score = (score_program_budget + num_reqs - 1,)
1462	_decode_index_score_balanced_kernel[grid_score](
1463	idx_q,
1464	index_kv_cache,
1465	score,
1466	block_table,
1467	seq_lens,
1468	num_idx_heads,
1469	head_dim,
1470	init_blocks,
1471	local_blocks,
1472	num_reqs,
1473	score_program_budget,
1474	decode_query_len,
1475	idx_q.stride(0),
1476	idx_q.stride(1),
1477	idx_q.stride(2),
1478	index_kv_cache.stride(0),
1479	index_kv_cache.stride(1),
1480	index_kv_cache.stride(2),
1481	score.stride(0),
1482	score.stride(1),
1483	score.stride(2),
1484	block_table.stride(0),
1485	BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
1486	BLOCK_SIZE_Q=BLOCK_SIZE_Q,
1487	num_warps=2,
1488	num_stages=1,
1489	)
1490	else:
1491	# Increase independent work for the measured high-batch gfx950 decode
1492	# range while preserving the deployed split for every other shape.
1493	num_kv_chunks, use_high_batch_config = _decode_score_split_launch_policy(
1494	num_reqs,
1495	head_dim,
1496	idx_q.dtype,
1497	index_kv_cache.dtype,
1498	is_gfx950=is_gfx950,
1499	)
1500	if use_high_batch_config and not (
1501	num_idx_heads > 1 and max_decode_query_len > 1
1502	):
1503	score_kwargs.update({"num_warps": 2, "num_stages": 1})
1504	grid_score = (num_reqs, num_kv_chunks)
1505	_decode_index_score_kernel[grid_score](
1506	idx_q,
1507	index_kv_cache,
1508	score,
1509	block_table,
1510	seq_lens,
1511	num_idx_heads,
1512	head_dim,
1513	init_blocks,
1514	local_blocks,
1515	decode_query_len,
1516	idx_q.stride(0),
1517	idx_q.stride(1),
1518	idx_q.stride(2),
1519	index_kv_cache.stride(0),
1520	index_kv_cache.stride(1),
1521	index_kv_cache.stride(2),
1522	score.stride(0),
1523	score.stride(1),
1524	score.stride(2),
1525	block_table.stride(0),
1526	BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
1527	BLOCK_SIZE_Q=BLOCK_SIZE_Q,
1528	num_kv_chunks=num_kv_chunks,
1529	USE_PDL=use_pdl,
1530	**score_kwargs,
1531	)
1532
1533	if out is not None:
1534	topk_idx = out[:, :total_q, :]
1535	else:
1536	topk_idx = torch.empty(
1537	(num_idx_heads, total_q, topk),
1538	dtype=torch.int32,
1539	device=idx_q.device,
1540	)
1541	# The launch grid remains shape-constant for CUDA graphs. Each query uses
1542	# only the number of chunks needed for its live context.
1543	(
1544	num_topk_chunks,
1545	topk_block_size,
1546	selector_num_warps,
1547	selector_num_stages,
1548	single_tile_guaranteed,
1549	adaptive_final_merge,
1550	) = _decode_topk_launch_policy(
1551	max_block,
1552	batch,
1553	num_idx_heads,
1554	topk,
1555	is_gfx950=is_gfx950,
1556	)
1557	block_size_t = triton.next_power_of_2(topk)
1558	topk_partial = torch.empty(
1559	num_topk_chunks,
1560	num_idx_heads,
1561	batch,
1562	block_size_t,
1563	dtype=torch.int64,
1564	device=idx_q.device,
1565	)
1566	if completion_counter is None:
1567	active_counter = torch.zeros(
1568	(num_idx_heads, batch),
1569	dtype=torch.int32,
1570	device=idx_q.device,
1571	)
1572	else:
1573	if (
1574	completion_counter.dim() != 2
1575	or completion_counter.shape[0] != num_idx_heads
1576	or completion_counter.shape[1] < batch
1577	or completion_counter.dtype != torch.int32
1578	or completion_counter.device != idx_q.device
1579	or not completion_counter.is_contiguous()
1580	):
1581	raise ValueError(
1582	"MiniMax-M3 completion counter must be contiguous int32 "
1583	"[num_idx_heads, >=total_q] on the query device"
1584	)
1585	active_counter = completion_counter[:, :batch]
1586
1587	selector_attention_block_table = block_table
1588	selector_sparse_block_table = topk_idx
1589	selector_sparse_context_lens = seq_lens
1590	selector_block_page_stride = PAGES_PER_SPARSE_BLOCK
1591	selector_sparse_block_stride = topk_idx.stride(1)
1592	if emit_sparse_table:
1593	assert attention_block_table is not None
1594	assert sparse_block_table_out is not None
1595	assert sparse_context_lens_out is not None
1596	assert block_page_stride is not None
1597	selector_attention_block_table = attention_block_table
1598	selector_sparse_block_table = sparse_block_table_out
1599	selector_sparse_context_lens = sparse_context_lens_out
1600	selector_block_page_stride = block_page_stride
1601	selector_sparse_block_stride = sparse_block_table_out.stride(0)
1602	_decode_topk_fused_kernel[(batch, num_idx_heads, num_topk_chunks)](
1603	score,
1604	topk_partial,
1605	active_counter,
1606	topk_idx,
1607	seq_lens,
1608	selector_attention_block_table,
1609	selector_sparse_block_table,
1610	selector_sparse_context_lens,
1611	decode_query_len,
1612	score.stride(0),
1613	score.stride(1),
1614	score.stride(2),
1615	topk_partial.stride(0),
1616	topk_partial.stride(1),
1617	topk_partial.stride(2),
1618	topk_partial.stride(3),
1619	active_counter.stride(0),
1620	active_counter.stride(1),
1621	topk_idx.stride(0),
1622	topk_idx.stride(1),
1623	topk_idx.stride(2),
1624	selector_attention_block_table.stride(0),
1625	selector_sparse_block_stride,
1626	topk=topk,
1627	block_size=SPARSE_BLOCK_SIZE,
1628	pages_per_sparse_block=PAGES_PER_SPARSE_BLOCK,
1629	block_page_stride=selector_block_page_stride,
1630	NUM_TOPK_CHUNKS=num_topk_chunks,
1631	BLOCK_SIZE_K=topk_block_size,
1632	BLOCK_SIZE_T=block_size_t,
1633	EMIT_SPARSE_TABLE=emit_sparse_table,
1634	SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
1635	ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
1636	num_warps=selector_num_warps,
1637	num_stages=selector_num_stages,
1638	)
1639	return topk_idx