From 5bdc01bcc35ff75a4686d30dac2c8e86b2e6e195 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Wed, 17 Jun 2026 12:07:34 +0800 Subject: [PATCH] [M3] Tune Triton indexer score decode for spec-decode (#45743) Signed-off-by: Thien Tran Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/kernels/attention/test_minimax_m3.py | 18 ++-- vllm/models/minimax_m3/common/indexer.py | 7 +- .../minimax_m3/common/ops/index_topk.py | 96 +++++++++++-------- 3 files changed, 72 insertions(+), 49 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index a69f85f7e08..1246f1721c2 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -134,7 +134,6 @@ def _reference_index_topk( topk: int, init_blocks: int, local_blocks: int, - sm_scale: float, ) -> torch.Tensor: total_q, num_idx_heads, _ = idx_q.shape out = torch.full( @@ -150,7 +149,7 @@ def _reference_index_topk( num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE pages = block_table[req_id, :num_blocks] k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1) - score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale + score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) q_pos = prefix_len + torch.arange(q_len, device=idx_q.device) k_pos = torch.arange(k.shape[0], device=idx_q.device) @@ -221,7 +220,6 @@ def test_prefill_index_topk_correctness(): max_query_len=q_lens.max().item(), max_seq_len=max_seq_len, num_kv_heads=num_idx_heads, - sm_scale=head_dim**-0.5, ) actual = minimax_m3_index_topk( score, @@ -242,15 +240,22 @@ def test_prefill_index_topk_correctness(): topk, init_blocks, local_blocks, - head_dim**-0.5, ) _assert_topk_indices_equal_unordered(actual, expected) -@pytest.mark.parametrize("decode_query_len", [1, 4]) +@pytest.mark.parametrize( + ("decode_query_len", "max_decode_query_len"), + [ + (1, 1), + (1, 4), + (4, 4), + ], +) @pytest.mark.parametrize("num_padded_reqs", [0, 2]) def test_decode_index_topk_correctness( decode_query_len: int, + max_decode_query_len: int, num_padded_reqs: int, ): topk = 6 @@ -293,8 +298,8 @@ def test_decode_index_topk_correctness( init_blocks=init_blocks, local_blocks=local_blocks, num_kv_heads=num_idx_heads, - sm_scale=head_dim**-0.5, decode_query_len=decode_query_len, + max_decode_query_len=max_decode_query_len, ) expected = torch.full_like(actual, -1) active_tokens = active_batch * decode_query_len @@ -308,7 +313,6 @@ def test_decode_index_topk_correctness( topk, init_blocks, local_blocks, - head_dim**-0.5, ) _assert_topk_indices_equal_unordered(actual, expected) diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index e43ad60914f..4da52805604 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -175,6 +175,7 @@ class MiniMaxM3IndexerDecodeMetadata: block_table: torch.Tensor max_seq_len: int decode_query_len: int + max_decode_query_len: int @dataclass @@ -229,6 +230,8 @@ class MiniMaxM3IndexerMetadataBuilder( assert tp_size % total_index_heads == 0 self.num_index_heads = max(1, total_index_heads // tp_size) self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) + assert self.reorder_batch_threshold is not None + self.max_decode_query_len = self.reorder_batch_threshold # Stable context-length buffer for decode cudagraph replays. self.context_len_buffer = torch.empty( @@ -297,6 +300,7 @@ class MiniMaxM3IndexerTritonMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): block_table=block_table[:num_decodes], max_seq_len=common_attn_metadata.max_seq_len, decode_query_len=decode_query_len, + max_decode_query_len=self.max_decode_query_len, ) return MiniMaxM3IndexerMetadata( @@ -403,8 +407,8 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): self.init_blocks, self.local_blocks, self.num_kv_heads, - self.scale, d.decode_query_len, + d.max_decode_query_len, ) if index_md.num_prefills > 0: p = index_md.prefill @@ -419,7 +423,6 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): p.max_query_len, p.max_seq_len, self.num_kv_heads, - self.scale, ) prefill_topk = minimax_m3_index_topk( score, diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py index c32ff38d998..208c2d69006 100644 --- a/vllm/models/minimax_m3/common/ops/index_topk.py +++ b/vllm/models/minimax_m3/common/ops/index_topk.py @@ -89,7 +89,6 @@ def _index_block_score_kernel( prefix_lens, # [batch] context length before this chunk's queries num_idx_heads, head_dim: tl.constexpr, - sm_scale, stride_q_n, stride_q_h, stride_q_d, @@ -103,7 +102,6 @@ def _index_block_score_kernel( BLOCK_SIZE_Q: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) ): - sm_scale_log2e = sm_scale * 1.4426950409 pid_q = tl.program_id(0) pid_bh = tl.program_id(1) pid_b = pid_bh // num_idx_heads @@ -148,7 +146,7 @@ def _index_block_score_kernel( + off_k[None, :] * stride_ik_pos + off_d[:, None] * stride_ik_d, ) - qk = tl.dot(q, k) * sm_scale_log2e + qk = tl.dot(q, k) # apply causal mask as needed if q_start < i + BLOCK_SIZE_K: qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf")) @@ -290,8 +288,8 @@ def _topk_index_kernel( # Decode index-score kernel (split-K over seq blocks). Decode batches are # flattened request-major, with a runtime query length used to map each query # token back to its request metadata. Chunk counts depend only on shape -# constants so the grid is fixed within a cuda graph. Base-2 (exp2/log2) -# softmax matches prefill. +# constants so the grid is fixed within a cuda graph. The score scale is omitted +# because decode only consumes block ordering. # --------------------------------------------------------------------------- @triton.jit(do_not_specialize=["num_kv_chunks", "decode_query_len"]) def _decode_index_score_kernel( @@ -304,7 +302,6 @@ def _decode_index_score_kernel( head_dim: tl.constexpr, init_blocks, local_blocks, - sm_scale, decode_query_len, stride_q_n, stride_q_h, @@ -317,25 +314,31 @@ def _decode_index_score_kernel( stride_s_k, stride_bt_b, BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + BLOCK_SIZE_Q: tl.constexpr, num_kv_chunks, USE_PDL: tl.constexpr, ): - sm_scale_log2e = sm_scale * 1.4426950409 - pid_b = tl.program_id(0) # flattened query-token id + BLOCK_SIZE_HQ: tl.constexpr = num_idx_heads * BLOCK_SIZE_Q + pid_r = tl.program_id(0) pid_c = tl.program_id(1) - req_id = pid_b // decode_query_len - q_offset = pid_b - req_id * decode_query_len + hq_offsets = tl.arange(0, BLOCK_SIZE_HQ) + h_offsets = hq_offsets // BLOCK_SIZE_Q + q_offsets = hq_offsets % BLOCK_SIZE_Q + q_mask = q_offsets < decode_query_len + q_ids = pid_r * decode_query_len + q_offsets if USE_PDL: tl.extra.cuda.gdc_wait() tl.extra.cuda.gdc_launch_dependents() - seq_len = tl.load(seq_lens + req_id) - query_pos = seq_len - decode_query_len + q_offset + seq_len = tl.load(seq_lens + pid_r) + query_pos = seq_len - decode_query_len + q_offsets # Full-CG padding uses zero-length request rows. Clamp to an empty # attention range instead of letting padded rows produce negative lengths. kv_len = tl.maximum(query_pos + 1, 0) - num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + num_blocks_q = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + kv_len_max = tl.max(tl.where(q_mask, kv_len, 0), axis=0) + num_blocks = (kv_len_max + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K # block-aligned fixed-count split: grid independent of seq_len (cuda graph). chunk_size_blocks = (num_blocks + num_kv_chunks - 1) // num_kv_chunks @@ -345,20 +348,22 @@ def _decode_index_score_kernel( return off_k = tl.arange(0, BLOCK_SIZE_K) # positions within a 128-block off_d = tl.arange(0, head_dim) - bt_row = block_table_ptr + req_id * stride_bt_b + bt_row = block_table_ptr + pid_r * stride_bt_b # Force-select init (1e30) and local (1e29, higher priority) blocks. - local_start = tl.maximum(0, num_blocks - local_blocks) - # query vectors across all heads + local_start = tl.maximum(0, num_blocks_q - local_blocks) + # Query vectors for all index heads in a small spec-decode block. q = tl.load( q_ptr - + pid_b * stride_q_n - + tl.arange(0, num_idx_heads) * stride_q_h + + q_ids[None, :] * stride_q_n + + h_offsets[None, :] * stride_q_h + off_d[:, None] * stride_q_d, - ) # [D,H] + mask=q_mask[None, :], + other=0.0, + ) # [D,HQ] for blk in tl.range(chunk_start_block, chunk_end_block): page = tl.load(bt_row + blk).to(tl.int64) pos = blk * BLOCK_SIZE_K + off_k - pos_mask = pos < kv_len + pos_mask = pos[:, None] < kv_len[None, :] # we don't need masked load for K, because KV cache ensures # allocation is multiple of BLOCK_SIZE_K. # for tokens beyond seqlen, they will be masked in qk later. @@ -368,18 +373,17 @@ def _decode_index_score_kernel( + off_k[:, None] * stride_ik_pos + off_d * stride_ik_d, ) # [N,D] - kq = tl.dot(k, q) * sm_scale_log2e # [N,H] - kq = tl.where(pos_mask[:, None], kq, float("-inf")) - score = tl.max(kq, axis=0) # [H] - is_init = blk < init_blocks - is_local = (blk >= local_start) & (blk < num_blocks) + kq = tl.dot(k, q) # [N,HQ] + kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf")) + score = tl.max(kq, axis=0) # [HQ] + is_visible_block = blk < num_blocks_q + is_init = (blk < init_blocks) & is_visible_block + is_local = (blk >= local_start) & is_visible_block score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score)) tl.store( - score_ptr - + tl.arange(0, num_idx_heads) * stride_s_h - + pid_b * stride_s_n - + blk * stride_s_k, + score_ptr + h_offsets * stride_s_h + q_ids * stride_s_n + blk * stride_s_k, score, + mask=q_mask, ) @@ -648,7 +652,6 @@ def minimax_m3_index_score( max_query_len: int, max_seq_len: int, num_kv_heads: int, - sm_scale: float, ) -> torch.Tensor: """Compute per-token index scores for each visible sparse block. @@ -681,7 +684,6 @@ def minimax_m3_index_score( prefix_lens, num_idx_heads, head_dim, - sm_scale, idx_q.stride(0), idx_q.stride(1), idx_q.stride(2), @@ -753,8 +755,8 @@ def minimax_m3_index_decode( init_blocks: int, local_blocks: int, num_kv_heads: int, - sm_scale: float, decode_query_len: int, + max_decode_query_len: int, ) -> torch.Tensor: """Decode index block-score + top-k, both split-K (cudagraph-safe). @@ -764,6 +766,7 @@ def minimax_m3_index_decode( assert num_idx_heads == num_kv_heads, ( "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)" ) + assert decode_query_len <= max_decode_query_len assert total_q == seq_lens.shape[0] * decode_query_len batch = total_q max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) @@ -772,7 +775,15 @@ def minimax_m3_index_decode( # SM9+); this ROCm Triton rejects it even when False ("Keyword argument # launch_pdl was specified but unrecognised"). Only pass it when PDL is # actually supported -- on ROCm use_pdl is always False, so it's omitted. - pdl_launch = {"launch_pdl": True} if use_pdl else {} + pdl_kwargs: dict[str, bool | int] = {} + if use_pdl: + pdl_kwargs.update({"launch_pdl": True}) + # TP=1 spec decode scores a wide 4-head x 4-position query tile per K block; + # reduce stages to ease memory/register pressure. Keep no-spec and TP=4 + # single-head codegen unchanged. + score_kwargs = pdl_kwargs.copy() + if num_idx_heads > 1 and max_decode_query_len > 1: + score_kwargs.update({"num_warps": 4, "num_stages": 2}) # Keep score strides 16-divisible to avoid Triton recompiles. score_block_stride = round_up(max_block, 16) @@ -783,13 +794,18 @@ def minimax_m3_index_decode( ) # split-K over seq blocks; chunk count depends only on shape constants so # the grid is fixed within a cuda graph. - TARGET_GRID = 4096 + TARGET_GRID = 512 MAX_NUM_KV_CHUNKS = 256 + # Use the configured max decode length to avoid Triton recompiles when + # switching between qlen=1 and spec-decode verification batches. + BLOCK_SIZE_Q = triton.next_power_of_2(max_decode_query_len) + score_ctas_per_chunk = seq_lens.shape[0] target = max( - 1, min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, batch * num_idx_heads)) + 1, + min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, score_ctas_per_chunk)), ) num_kv_chunks = 1 << (target.bit_length() - 1) - grid_score = (batch, num_kv_chunks) + grid_score = (seq_lens.shape[0], num_kv_chunks) _decode_index_score_kernel[grid_score]( idx_q, index_kv_cache, @@ -800,7 +816,6 @@ def minimax_m3_index_decode( head_dim, init_blocks, local_blocks, - sm_scale, decode_query_len, idx_q.stride(0), idx_q.stride(1), @@ -813,9 +828,10 @@ def minimax_m3_index_decode( score.stride(2), block_table.stride(0), BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + BLOCK_SIZE_Q=BLOCK_SIZE_Q, num_kv_chunks=num_kv_chunks, USE_PDL=use_pdl, - **pdl_launch, + **score_kwargs, ) topk_idx = torch.empty( @@ -870,7 +886,7 @@ def minimax_m3_index_decode( topk_idx_partial.stride(2), topk_idx_partial.stride(3), USE_PDL=use_pdl, - **pdl_launch, + **pdl_kwargs, ) _topk_index_merge_kernel[(batch, num_idx_heads)]( topk_score_partial, @@ -893,6 +909,6 @@ def minimax_m3_index_decode( topk_idx.stride(2), num_topk_chunks=num_topk_chunks, USE_PDL=use_pdl, - **pdl_launch, + **pdl_kwargs, ) return topk_idx