@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5bdc01bcc3

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[M3] Tune Triton indexer score decode for spec-decode (#45743)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Thien Tran committed Jun 17, 2026 at 12:07 UTC 5bdc01bcc35ff75a4686d30dac2c8e86b2e6e195

3 files changed +72 -49

tests/kernels/attention/test_minimax_m3.py

+11 -7

@@ -134,7 +134,6 @@ def _reference_index_topk(

134 topk: int,

135 init_blocks: int,

136 local_blocks: int,

137 - sm_scale: float,

137 ) -> torch.Tensor:

138 total_q, num_idx_heads, _ = idx_q.shape

139 out = torch.full(

@@ -150,7 +149,7 @@ def _reference_index_topk(

149 num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE

150 pages = block_table[req_id, :num_blocks]

151 k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)

153 - score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale

152 + score = torch.einsum("qhd,kd->hqk", q.float(), k.float())

153

154 q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)

155 k_pos = torch.arange(k.shape[0], device=idx_q.device)

@@ -221,7 +220,6 @@ def test_prefill_index_topk_correctness():

220 max_query_len=q_lens.max().item(),

221 max_seq_len=max_seq_len,

222 num_kv_heads=num_idx_heads,

224 - sm_scale=head_dim**-0.5,

223 )

224 actual = minimax_m3_index_topk(

225 score,

@@ -242,15 +240,22 @@ def test_prefill_index_topk_correctness():

240 topk,

241 init_blocks,

242 local_blocks,

245 - head_dim**-0.5,

243 )

244 _assert_topk_indices_equal_unordered(actual, expected)

245

246

250 -@pytest.mark.parametrize("decode_query_len", [1, 4])

247 +@pytest.mark.parametrize(

248 + ("decode_query_len", "max_decode_query_len"),

249 + [

250 + (1, 1),

251 + (1, 4),

252 + (4, 4),

253 + ],

254 +)

255 @pytest.mark.parametrize("num_padded_reqs", [0, 2])

256 def test_decode_index_topk_correctness(

257 decode_query_len: int,

258 + max_decode_query_len: int,

259 num_padded_reqs: int,

260 ):

261 topk = 6

@@ -293,8 +298,8 @@ def test_decode_index_topk_correctness(

298 init_blocks=init_blocks,

299 local_blocks=local_blocks,

300 num_kv_heads=num_idx_heads,

296 - sm_scale=head_dim**-0.5,

301 decode_query_len=decode_query_len,

302 + max_decode_query_len=max_decode_query_len,

303 )

304 expected = torch.full_like(actual, -1)

305 active_tokens = active_batch * decode_query_len

@@ -308,7 +313,6 @@ def test_decode_index_topk_correctness(

313 topk,

314 init_blocks,

315 local_blocks,

311 - head_dim**-0.5,

316 )

317 _assert_topk_indices_equal_unordered(actual, expected)

318

vllm/models/minimax_m3/common/indexer.py

+5 -2

@@ -175,6 +175,7 @@ class MiniMaxM3IndexerDecodeMetadata:

175 block_table: torch.Tensor

176 max_seq_len: int

177 decode_query_len: int

178 + max_decode_query_len: int

179

180

181 @dataclass

@@ -229,6 +230,8 @@ class MiniMaxM3IndexerMetadataBuilder(

230 assert tp_size % total_index_heads == 0

231 self.num_index_heads = max(1, total_index_heads // tp_size)

232 self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)

233 + assert self.reorder_batch_threshold is not None

234 + self.max_decode_query_len = self.reorder_batch_threshold

235

236 # Stable context-length buffer for decode cudagraph replays.

237 self.context_len_buffer = torch.empty(

@@ -297,6 +300,7 @@ class MiniMaxM3IndexerTritonMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):

300 block_table=block_table[:num_decodes],

301 max_seq_len=common_attn_metadata.max_seq_len,

302 decode_query_len=decode_query_len,

303 + max_decode_query_len=self.max_decode_query_len,

304 )

305

306 return MiniMaxM3IndexerMetadata(

@@ -403,8 +407,8 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):

407 self.init_blocks,

408 self.local_blocks,

409 self.num_kv_heads,

406 - self.scale,

410 d.decode_query_len,

411 + d.max_decode_query_len,

412 )

413 if index_md.num_prefills > 0:

414 p = index_md.prefill

@@ -419,7 +423,6 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):

423 p.max_query_len,

424 p.max_seq_len,

425 self.num_kv_heads,

422 - self.scale,

426 )

427 prefill_topk = minimax_m3_index_topk(

428 score,

vllm/models/minimax_m3/common/ops/index_topk.py

+56 -40

@@ -89,7 +89,6 @@ def _index_block_score_kernel(

89 prefix_lens, # [batch] context length before this chunk's queries

90 num_idx_heads,

91 head_dim: tl.constexpr,

92 - sm_scale,

92 stride_q_n,

93 stride_q_h,

94 stride_q_d,

@@ -103,7 +102,6 @@ def _index_block_score_kernel(

102 BLOCK_SIZE_Q: tl.constexpr,

103 BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)

104 ):

106 - sm_scale_log2e = sm_scale * 1.4426950409

105 pid_q = tl.program_id(0)

106 pid_bh = tl.program_id(1)

107 pid_b = pid_bh // num_idx_heads

@@ -148,7 +146,7 @@ def _index_block_score_kernel(

146 + off_k[None, :] * stride_ik_pos

147 + off_d[:, None] * stride_ik_d,

148 )

151 - qk = tl.dot(q, k) * sm_scale_log2e

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"))

@@ -290,8 +288,8 @@ def _topk_index_kernel(

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

293 -# constants so the grid is fixed within a cuda graph. Base-2 (exp2/log2)

294 -# softmax matches prefill.

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(

@@ -304,7 +302,6 @@ def _decode_index_score_kernel(

302 head_dim: tl.constexpr,

303 init_blocks,

304 local_blocks,

307 - sm_scale,

305 decode_query_len,

306 stride_q_n,

307 stride_q_h,

@@ -317,25 +314,31 @@ def _decode_index_score_kernel(

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 ):

323 - sm_scale_log2e = sm_scale * 1.4426950409

324 - pid_b = tl.program_id(0) # flattened query-token id

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)

326 - req_id = pid_b // decode_query_len

327 - q_offset = pid_b - req_id * decode_query_len

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

333 - seq_len = tl.load(seq_lens + req_id)

334 - query_pos = seq_len - decode_query_len + q_offset

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)

338 - num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K

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,20 +348,22 @@ def _decode_index_score_kernel(

348 return

349 off_k = tl.arange(0, BLOCK_SIZE_K) # positions within a 128-block

350 off_d = tl.arange(0, head_dim)

348 - bt_row = block_table_ptr + req_id * stride_bt_b

351 + bt_row = block_table_ptr + pid_r * stride_bt_b

352 # Force-select init (1e30) and local (1e29, higher priority) blocks.

350 - local_start = tl.maximum(0, num_blocks - local_blocks)

351 - # query vectors across all heads

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

354 - + pid_b * stride_q_n

355 - + tl.arange(0, num_idx_heads) * stride_q_h

357 + + q_ids[None, :] * stride_q_n

358 + + h_offsets[None, :] * stride_q_h

359 + off_d[:, None] * stride_q_d,

357 - ) # [D,H]

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

361 - pos_mask = pos < kv_len

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.

@@ -368,18 +373,17 @@ def _decode_index_score_kernel(

373 + off_k[:, None] * stride_ik_pos

374 + off_d * stride_ik_d,

375 ) # [N,D]

371 - kq = tl.dot(k, q) * sm_scale_log2e # [N,H]

372 - kq = tl.where(pos_mask[:, None], kq, float("-inf"))

373 - score = tl.max(kq, axis=0) # [H]

374 - is_init = blk < init_blocks

375 - is_local = (blk >= local_start) & (blk < num_blocks)

376 + kq = tl.dot(k, q) # [N,HQ]

377 + kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf"))

378 + score = tl.max(kq, axis=0) # [HQ]

379 + is_visible_block = blk < num_blocks_q

380 + is_init = (blk < init_blocks) & is_visible_block

381 + is_local = (blk >= local_start) & is_visible_block

382 score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score))

383 tl.store(

378 - score_ptr

379 - + tl.arange(0, num_idx_heads) * stride_s_h

380 - + pid_b * stride_s_n

381 - + blk * stride_s_k,

384 + score_ptr + h_offsets * stride_s_h + q_ids * stride_s_n + blk * stride_s_k,

385 score,

386 + mask=q_mask,

387 )

388

389

@@ -648,7 +652,6 @@ def minimax_m3_index_score(

652 max_query_len: int,

653 max_seq_len: int,

654 num_kv_heads: int,

651 - sm_scale: float,

655 ) -> torch.Tensor:

656 """Compute per-token index scores for each visible sparse block.

657

@@ -681,7 +684,6 @@ def minimax_m3_index_score(

684 prefix_lens,

685 num_idx_heads,

686 head_dim,

684 - sm_scale,

687 idx_q.stride(0),

688 idx_q.stride(1),

689 idx_q.stride(2),

@@ -753,8 +755,8 @@ def minimax_m3_index_decode(

755 init_blocks: int,

756 local_blocks: int,

757 num_kv_heads: int,

756 - sm_scale: float,

758 decode_query_len: int,

759 + max_decode_query_len: int,

760 ) -> torch.Tensor:

761 """Decode index block-score + top-k, both split-K (cudagraph-safe).

762

@@ -764,6 +766,7 @@ def minimax_m3_index_decode(

766 assert num_idx_heads == num_kv_heads, (

767 "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"

768 )

769 + assert decode_query_len <= max_decode_query_len

770 assert total_q == seq_lens.shape[0] * decode_query_len

771 batch = total_q

772 max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)

@@ -772,7 +775,15 @@ def minimax_m3_index_decode(

775 # SM9+); this ROCm Triton rejects it even when False ("Keyword argument

776 # launch_pdl was specified but unrecognised"). Only pass it when PDL is

777 # actually supported -- on ROCm use_pdl is always False, so it's omitted.

775 - pdl_launch = {"launch_pdl": True} if use_pdl else {}

778 + pdl_kwargs: dict[str, bool | int] = {}

779 + if use_pdl:

780 + pdl_kwargs.update({"launch_pdl": True})

781 + # TP=1 spec decode scores a wide 4-head x 4-position query tile per K block;

782 + # reduce stages to ease memory/register pressure. Keep no-spec and TP=4

783 + # single-head codegen unchanged.

784 + score_kwargs = pdl_kwargs.copy()

785 + if num_idx_heads > 1 and max_decode_query_len > 1:

786 + score_kwargs.update({"num_warps": 4, "num_stages": 2})

787

788 # Keep score strides 16-divisible to avoid Triton recompiles.

789 score_block_stride = round_up(max_block, 16)

@@ -783,13 +794,18 @@ def minimax_m3_index_decode(

794 )

795 # split-K over seq blocks; chunk count depends only on shape constants so

796 # the grid is fixed within a cuda graph.

786 - TARGET_GRID = 4096

797 + TARGET_GRID = 512

798 MAX_NUM_KV_CHUNKS = 256

799 + # Use the configured max decode length to avoid Triton recompiles when

800 + # switching between qlen=1 and spec-decode verification batches.

801 + BLOCK_SIZE_Q = triton.next_power_of_2(max_decode_query_len)

802 + score_ctas_per_chunk = seq_lens.shape[0]

803 target = max(

789 - 1, min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, batch * num_idx_heads))

804 + 1,

805 + min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, score_ctas_per_chunk)),

806 )

807 num_kv_chunks = 1 << (target.bit_length() - 1)

792 - grid_score = (batch, num_kv_chunks)

808 + grid_score = (seq_lens.shape[0], num_kv_chunks)

809 _decode_index_score_kernel[grid_score](

810 idx_q,

811 index_kv_cache,

@@ -800,7 +816,6 @@ def minimax_m3_index_decode(

816 head_dim,

817 init_blocks,

818 local_blocks,

803 - sm_scale,

819 decode_query_len,

820 idx_q.stride(0),

821 idx_q.stride(1),

@@ -813,9 +828,10 @@ def minimax_m3_index_decode(

828 score.stride(2),

829 block_table.stride(0),

830 BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,

831 + BLOCK_SIZE_Q=BLOCK_SIZE_Q,

832 num_kv_chunks=num_kv_chunks,

833 USE_PDL=use_pdl,

818 - **pdl_launch,

834 + **score_kwargs,

835 )

836

837 topk_idx = torch.empty(

@@ -870,7 +886,7 @@ def minimax_m3_index_decode(

886 topk_idx_partial.stride(2),

887 topk_idx_partial.stride(3),

888 USE_PDL=use_pdl,

873 - **pdl_launch,

889 + **pdl_kwargs,

890 )

891 _topk_index_merge_kernel[(batch, num_idx_heads)](

892 topk_score_partial,

@@ -893,6 +909,6 @@ def minimax_m3_index_decode(

909 topk_idx.stride(2),

910 num_topk_chunks=num_topk_chunks,

911 USE_PDL=use_pdl,

896 - **pdl_launch,

912 + **pdl_kwargs,

913 )

914 return topk_idx