iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Perf] Tune MiniMax-M3 decode top-k for short contexts (#55235)

Browse Source 
Signed-off-by: fai <fangzhouai@gmail.com>
This commit is contained in:
Fangzhou Ai
2026-09-13 21:36:27 -05:00
committed by GitHub
parent 58d45fd767
commit dfa1984e58
2 changed files with 90 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_minimax_m3.py
+63 -41
View File
@@ -1025,17 +1025,20 @@ def test_decode_index_topk_correctness(
"expected",
),
[
(0, 24, 1, 16, True, (2, False, False)),
(1, 1, 1, 16, True, (16, True, True)),
(512, 4, 1, 16, True, (16, True, True)),
(513, 24, 4, 16, True, (16, True, True)),
(2760, 5, 2, 16, True, (16, True, True)),
(8192, 64, 4, 16, True, (16, True, True)),
(8193, 4, 1, 16, True, (16, False, False)),
(8193, 16, 1, 16, True, (4, False, False)),
(65534, 64, 1, 16, True, (1, False, False)),
(513, 16, 1, 13, True, (4, False, False)),
(513, 16, 4, 16, False, (1, False, False)),
(0, 24, 1, 16, True, (2, 512, 8, 2, False, False)),
(1, 1, 1, 16, True, (1, 128, 2, 1, True, True)),
(74, 4, 1, 16, True, (1, 128, 2, 1, True, True)),
(128, 4, 1, 16, True, (1, 128, 2, 1, True, True)),
(129, 4, 1, 16, True, (16, 512, 4, 2, True, True)),
(512, 4, 1, 16, True, (16, 512, 4, 2, True, True)),
(513, 24, 4, 16, True, (16, 512, 4, 2, True, True)),
(2760, 5, 2, 16, True, (16, 512, 4, 2, True, True)),
(8192, 64, 4, 16, True, (16, 512, 4, 2, True, True)),
(8193, 4, 1, 16, True, (16, 512, 8, 2, False, False)),
(8193, 16, 1, 16, True, (4, 512, 8, 2, False, False)),
(65534, 64, 1, 16, True, (1, 512, 8, 2, False, False)),
(513, 16, 1, 13, True, (4, 512, 8, 2, False, False)),
(513, 16, 4, 16, False, (1, 512, 8, 2, False, False)),
],
)
def test_amd_decode_topk_launch_policy(
@@ -1044,9 +1047,9 @@ def test_amd_decode_topk_launch_policy(
num_idx_heads: int,
topk: int,
is_gfx950: bool,
expected: tuple[int, bool, bool],
expected: tuple[int, int, int, int, bool, bool],
):
"""All supported head and query layouts use adaptive selection."""
"""The short-context specialization has strict dispatch boundaries."""
from vllm.models.minimax_m3.amd.ops.index_topk import (
_decode_topk_launch_policy,
)
@@ -1073,19 +1076,23 @@ def test_amd_decode_topk_launch_policy(
"single_tile_guaranteed",
"adaptive_final_merge",
"max_blocks",
"block_size_k",
"num_warps",
"num_stages",
"topk",
"num_idx_heads",
),
[
(1, False, False, 65534, 16, 1),
(2, False, False, 8193, 16, 1),
(4, False, False, 8193, 16, 1),
(8, False, False, 8193, 16, 1),
(16, False, False, 8193, 16, 1),
(16, False, False, 65534, 16, 1),
(2, True, False, 513, 13, 1),
(4, False, True, 8192, 16, 2),
(16, True, True, 8192, 16, 4),
(1, False, False, 65534, 512, 8, 2, 16, 1),
(2, False, False, 8193, 512, 8, 2, 16, 1),
(4, False, False, 8193, 512, 8, 2, 16, 1),
(8, False, False, 8193, 512, 8, 2, 16, 1),
(16, False, False, 8193, 512, 8, 2, 16, 1),
(16, False, False, 65534, 512, 8, 2, 16, 1),
(2, True, False, 513, 512, 8, 2, 13, 1),
(4, False, True, 8192, 512, 4, 2, 16, 2),
(16, True, True, 8192, 512, 4, 2, 16, 4),
(1, True, True, 74, 128, 2, 1, 16, 1),
],
)
def test_amd_decode_fused_topk_total_order_and_replay(
@@ -1093,10 +1100,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
single_tile_guaranteed: bool,
adaptive_final_merge: bool,
max_blocks: int,
block_size_k: int,
num_warps: int,
num_stages: int,
topk: int,
num_idx_heads: int,
):
"""Atomic multi-chunk selection is ordered and resets graph state."""
"""Selector specializations preserve ordering and graph state."""
from vllm.models.minimax_m3.amd.ops.index_topk import (
SPARSE_BLOCK_SIZE,
_decode_topk_fused_kernel,
@@ -1104,18 +1114,25 @@ def test_amd_decode_fused_topk_total_order_and_replay(
set_random_seed(0)
block_size_t = 1 << (topk - 1).bit_length()
block_counts = tuple(
count
for count in (
15,
16,
17,
block_size_k // 2 - 1,
block_size_k // 2,
block_size_k // 2 + 1,
block_size_k - 1,
block_size_k,
block_size_k + 1,
)
if count <= max_blocks
)
long_seq_lens_list = [
0,
1,
15 * SPARSE_BLOCK_SIZE,
16 * SPARSE_BLOCK_SIZE,
17 * SPARSE_BLOCK_SIZE,
255 * SPARSE_BLOCK_SIZE,
256 * SPARSE_BLOCK_SIZE,
257 * SPARSE_BLOCK_SIZE,
511 * SPARSE_BLOCK_SIZE,
512 * SPARSE_BLOCK_SIZE,
513 * SPARSE_BLOCK_SIZE,
*(blocks * SPARSE_BLOCK_SIZE for blocks in block_counts),
max_blocks * SPARSE_BLOCK_SIZE,
]
short_seq_lens_list = [0, 1, 127, 128, 129, 0, 1, 128, 129, 0, 1, 129]
@@ -1126,12 +1143,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
boundary_limit = 4
if boundary_limit:
boundary_blocks = [
(active_chunks - 1) * 512 + 1
(active_chunks - 1) * block_size_k + 1
for active_chunks in range(3, boundary_limit + 1)
]
assert [(num_blocks + 511) // 512 for num_blocks in boundary_blocks] == list(
range(3, boundary_limit + 1)
)
assert [
(num_blocks + block_size_k - 1) // block_size_k
for num_blocks in boundary_blocks
] == list(range(3, boundary_limit + 1))
long_seq_lens_list.extend(
num_blocks * SPARSE_BLOCK_SIZE for num_blocks in boundary_blocks
)
@@ -1139,6 +1157,8 @@ def test_amd_decode_fused_topk_total_order_and_replay(
num_blocks * SPARSE_BLOCK_SIZE for num_blocks in reversed(boundary_blocks)
)
batch = len(long_seq_lens_list)
repeats = (batch + len(short_seq_lens_list) - 1) // len(short_seq_lens_list)
short_seq_lens_list = (short_seq_lens_list * repeats)[:batch]
score = torch.randn((num_idx_heads, batch, max_blocks), device="cuda")
score[:, :, 0] = float("nan")
score[:, :, 1] = float("inf")
@@ -1147,8 +1167,9 @@ def test_amd_decode_fused_topk_total_order_and_replay(
score[:, :, 4] = -0.0
score[:, :, 5] = 7.0
score[:, :, 6] = -1e30
score[:, :, 256] = 7.0
score[:, :, 512] = 7.0
for index in (block_size_k // 2, block_size_k):
if index < max_blocks:
score[:, :, index] = 7.0
seq_lens = torch.empty(batch, device="cuda", dtype=torch.int32)
partial = torch.full(
@@ -1198,13 +1219,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
pages_per_sparse_block=8,
block_page_stride=8,
NUM_TOPK_CHUNKS=num_topk_chunks,
BLOCK_SIZE_K=512,
BLOCK_SIZE_K=block_size_k,
BLOCK_SIZE_T=block_size_t,
EMIT_SPARSE_TABLE=False,
SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
num_warps=4 if adaptive_final_merge else 8,
num_stages=2,
num_warps=num_warps,
num_stages=num_stages,
)
score_cpu = score.cpu()
@@ -1251,6 +1272,7 @@ def test_amd_decode_fused_topk_total_order_and_replay(
(16, False, False, 8193),
(4, False, True, 8192),
(16, True, True, 8192),
(1, True, True, 74),
):
graph = torch.cuda.CUDAGraph()
torch.accelerator.synchronize()
vllm/models/minimax_m3/amd/ops/index_topk.py
+27 -7
View File
@@ -37,6 +37,7 @@ MIN_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS = 20
MAX_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS = 64
MAX_DECODE_SCORE_KV_CHUNKS = 256
DECODE_TOPK_BLOCKS_PER_CHUNK = 512
DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK = 128
MAX_DECODE_TOPK_FAST_CHUNKS = 16
DECODE_TOPK_TARGET_GRID = 64
@@ -103,14 +104,23 @@ def _decode_topk_launch_policy(
topk: int,
*,
is_gfx950: bool,
) -> tuple[int, bool, bool]:
"""Choose the graph grid and bounded-context selector specialization."""
) -> tuple[int, int, int, int, bool, bool]:
"""Choose the selector grid and compile-time launch configuration."""
if (
is_gfx950
and topk == 16
and 0 < max_block <= MAX_DECODE_TOPK_FAST_CHUNKS * DECODE_TOPK_BLOCKS_PER_CHUNK
):
return MAX_DECODE_TOPK_FAST_CHUNKS, True, True
if max_block <= DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK:
return 1, DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK, 2, 1, True, True
return (
MAX_DECODE_TOPK_FAST_CHUNKS,
DECODE_TOPK_BLOCKS_PER_CHUNK,
4,
2,
True,
True,
)
target = max(
1,
@@ -119,7 +129,14 @@ def _decode_topk_launch_policy(
DECODE_TOPK_TARGET_GRID // max(1, total_q * num_idx_heads),
),
)
return 1 << (target.bit_length() - 1), False, False
return (
1 << (target.bit_length() - 1),
DECODE_TOPK_BLOCKS_PER_CHUNK,
8,
2,
False,
False,
)
# ---------------------------------------------------------------------------
@@ -1525,6 +1542,9 @@ def minimax_m3_index_decode(
# only the number of chunks needed for its live context.
(
num_topk_chunks,
topk_block_size,
selector_num_warps,
selector_num_stages,
single_tile_guaranteed,
adaptive_final_merge,
) = _decode_topk_launch_policy(
@@ -1608,12 +1628,12 @@ def minimax_m3_index_decode(
pages_per_sparse_block=PAGES_PER_SPARSE_BLOCK,
block_page_stride=selector_block_page_stride,
NUM_TOPK_CHUNKS=num_topk_chunks,
BLOCK_SIZE_K=512,
BLOCK_SIZE_K=topk_block_size,
BLOCK_SIZE_T=block_size_t,
EMIT_SPARSE_TABLE=emit_sparse_table,
SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
num_warps=4 if adaptive_final_merge else 8,
num_stages=2,
num_warps=selector_num_warps,
num_stages=selector_num_stages,
)
return topk_idx