karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

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

Browse Source 
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Thien Tran
2026-06-16 21:07:34 -07:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 20a5f8b43b
commit 5bdc01bcc3
3 changed files with 72 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_minimax_m3.py
+11 -7
View File
@@ -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)
vllm/models/minimax_m3/common/indexer.py
+5 -2
View File
@@ -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,
vllm/models/minimax_m3/common/ops/index_topk.py
+56 -40
View File
@@ -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