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

[ROCm][Perf] Optimize MiniMax-M3 decode indexer and top-k (#54682)

Browse Source 
Signed-off-by: fai <fangzhouai@gmail.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Fangzhou Ai
2026-09-03 13:20:18 -07:00
committed by GitHub
co-authored by Andreas Karatzas OpenAI Codex
parent 2a336d8239
commit cee0f92c02
6 changed files with 2028 additions and 345 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_minimax_m3.py
+840
View File
@@ -21,6 +21,7 @@ from vllm.models.minimax_m3.common.ops.sparse_attn import (
)
from vllm.models.minimax_m3.common.sparse_attention import MiniMaxM3SparseTritonImpl
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.attention.backends.utils import record_kv_cache_layout
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
@@ -92,6 +93,382 @@ SM_SCALE = HEAD_DIM**-0.5
TOPK = 16
@pytest.mark.parametrize(
(
"num_reqs",
"head_dim",
"query_dtype",
"cache_dtype",
"is_gfx950",
"expected",
),
[
(19, 128, torch.bfloat16, torch.bfloat16, True, (16, False)),
(20, 128, torch.bfloat16, torch.bfloat16, True, (32, True)),
(32, 128, torch.bfloat16, torch.bfloat16, True, (32, True)),
(33, 128, torch.bfloat16, torch.bfloat16, True, (16, True)),
(64, 128, torch.bfloat16, torch.bfloat16, True, (16, True)),
(65, 128, torch.bfloat16, torch.bfloat16, True, (4, False)),
(20, 128, torch.bfloat16, torch.bfloat16, False, (16, False)),
(20, 64, torch.bfloat16, torch.bfloat16, True, (16, False)),
(20, 128, torch.float16, torch.bfloat16, True, (16, False)),
(20, 128, torch.bfloat16, torch.float16, True, (16, False)),
],
)
def test_amd_decode_score_split_launch_policy(
num_reqs: int,
head_dim: int,
query_dtype: torch.dtype,
cache_dtype: torch.dtype,
is_gfx950: bool,
expected: tuple[int, bool],
):
from vllm.models.minimax_m3.amd.ops.index_topk import (
_decode_score_split_launch_policy,
)
assert (
_decode_score_split_launch_policy(
num_reqs,
head_dim,
query_dtype,
cache_dtype,
is_gfx950=is_gfx950,
)
== expected
)
@pytest.mark.parametrize("num_reqs", range(1, 12))
def test_amd_balanced_decode_score_mapping_exact_coverage(num_reqs: int):
from vllm.models.minimax_m3.amd.ops.index_topk import (
_decode_score_program_budget,
)
budget = _decode_score_program_budget(
num_reqs,
128,
torch.bfloat16,
torch.bfloat16,
is_gfx950=True,
)
assert budget is not None
patterns = (
(0,) * num_reqs,
tuple((request * 7919 + 17) % 65535 for request in range(num_reqs)),
(65534,) * num_reqs,
)
for request_blocks in patterns:
width = max(1, (sum(request_blocks) + budget - 1) // budget)
grid = budget + len(request_blocks) - 1
intervals = []
for request, blocks in enumerate(request_blocks):
programs = (blocks + width - 1) // width
for local_program in range(programs):
begin = local_program * width
intervals.append((request, begin, min(begin + width, blocks)))
assert len(intervals) <= grid
by_request: list[list[tuple[int, int]]] = [[] for _ in request_blocks]
for request, begin, end in intervals:
assert 0 <= request < len(request_blocks)
assert 0 <= begin < end <= request_blocks[request]
by_request[request].append((begin, end))
for blocks, request_intervals in zip(
request_blocks,
by_request,
strict=True,
):
cursor = 0
for begin, end in request_intervals:
assert begin == cursor
cursor = end
assert cursor == blocks
@pytest.mark.skipif(
not current_platform.is_rocm(),
reason="MiniMax-M3 balanced decode score is ROCm-only",
)
@pytest.mark.parametrize(
(
"num_reqs",
"max_block",
"num_idx_heads",
"decode_query_len",
"max_decode_query_len",
"init_blocks",
"local_blocks",
),
[
(1, 37, 1, 1, 1, 1, 0),
(8, 37, 1, 1, 4, 2, 1),
(9, 37, 1, 4, 4, 0, 2),
(11, 37, 1, 3, 5, 1, 3),
(1, 37, 2, 1, 1, 2, 0),
(8, 37, 2, 1, 4, 0, 1),
(9, 37, 2, 4, 4, 1, 2),
(11, 37, 2, 3, 5, 2, 3),
(1, 37, 4, 1, 1, 0, 0),
(8, 37, 4, 1, 4, 1, 1),
(9, 37, 4, 4, 4, 2, 2),
(11, 37, 4, 3, 5, 0, 3),
(20, 65, 1, 4, 4, 1, 3),
(36, 65, 1, 4, 4, 2, 2),
(64, 65, 1, 4, 4, 0, 1),
],
)
def test_amd_optimized_decode_score_bitwise_and_graph_replay(
num_reqs: int,
max_block: int,
num_idx_heads: int,
decode_query_len: int,
max_decode_query_len: int,
init_blocks: int,
local_blocks: int,
):
from vllm.models.minimax_m3.amd.ops.index_topk import (
_decode_index_score_balanced_kernel,
_decode_index_score_kernel,
_decode_score_program_budget,
_decode_score_split_launch_policy,
)
set_random_seed(0)
head_dim = 128
block_size_q = 1 << (max_decode_query_len - 1).bit_length()
score_stride = (max_block + 15) // 16 * 16
num_pages = num_reqs * max_block
block_table = torch.arange(
num_pages,
device="cuda",
dtype=torch.int32,
).reshape(num_reqs, max_block)
idx_q = torch.randn(
(num_reqs * decode_query_len, num_idx_heads, head_dim),
device="cuda",
dtype=torch.bfloat16,
)
index_kv_cache = torch.randn(
(num_pages, BLOCK_SIZE, head_dim),
device="cuda",
dtype=torch.bfloat16,
)
seq_lens = torch.empty(num_reqs, device="cuda", dtype=torch.int32)
deployed_score = torch.empty(
(num_idx_heads, num_reqs * decode_query_len, score_stride),
device="cuda",
dtype=torch.float32,
)
optimized_score = torch.empty_like(deployed_score)
budget = _decode_score_program_budget(
num_reqs,
head_dim,
idx_q.dtype,
index_kv_cache.dtype,
is_gfx950=True,
)
optimized_chunks = None
if budget is None:
optimized_chunks, use_high_batch_config = _decode_score_split_launch_policy(
num_reqs,
head_dim,
idx_q.dtype,
index_kv_cache.dtype,
is_gfx950=True,
)
assert use_high_batch_config
target = max(1, min(256, 512 // num_reqs))
deployed_chunks = 1 << (target.bit_length() - 1)
def launch_deployed() -> None:
_decode_index_score_kernel[(num_reqs, deployed_chunks)](
idx_q,
index_kv_cache,
deployed_score,
block_table,
seq_lens,
num_idx_heads,
head_dim,
init_blocks,
local_blocks,
decode_query_len,
idx_q.stride(0),
idx_q.stride(1),
idx_q.stride(2),
index_kv_cache.stride(0),
index_kv_cache.stride(1),
index_kv_cache.stride(2),
deployed_score.stride(0),
deployed_score.stride(1),
deployed_score.stride(2),
block_table.stride(0),
BLOCK_SIZE_K=BLOCK_SIZE,
BLOCK_SIZE_Q=block_size_q,
num_kv_chunks=deployed_chunks,
USE_PDL=False,
)
def launch_optimized() -> None:
if budget is not None:
_decode_index_score_balanced_kernel[(budget + num_reqs - 1,)](
idx_q,
index_kv_cache,
optimized_score,
block_table,
seq_lens,
num_idx_heads,
head_dim,
init_blocks,
local_blocks,
num_reqs,
budget,
decode_query_len,
idx_q.stride(0),
idx_q.stride(1),
idx_q.stride(2),
index_kv_cache.stride(0),
index_kv_cache.stride(1),
index_kv_cache.stride(2),
optimized_score.stride(0),
optimized_score.stride(1),
optimized_score.stride(2),
block_table.stride(0),
BLOCK_SIZE_K=BLOCK_SIZE,
BLOCK_SIZE_Q=block_size_q,
num_warps=2,
num_stages=1,
)
else:
assert optimized_chunks is not None
_decode_index_score_kernel[(num_reqs, optimized_chunks)](
idx_q,
index_kv_cache,
optimized_score,
block_table,
seq_lens,
num_idx_heads,
head_dim,
init_blocks,
local_blocks,
decode_query_len,
idx_q.stride(0),
idx_q.stride(1),
idx_q.stride(2),
index_kv_cache.stride(0),
index_kv_cache.stride(1),
index_kv_cache.stride(2),
optimized_score.stride(0),
optimized_score.stride(1),
optimized_score.stride(2),
block_table.stride(0),
BLOCK_SIZE_K=BLOCK_SIZE,
BLOCK_SIZE_Q=block_size_q,
num_kv_chunks=optimized_chunks,
USE_PDL=False,
num_warps=2,
num_stages=1,
)
def score_masks(
expected_seq_lens: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
query_offsets = torch.arange(decode_query_len, device="cuda")
kv_lens = torch.clamp(
expected_seq_lens[:, None] - decode_query_len + query_offsets + 1,
min=0,
)
query_blocks = (kv_lens + BLOCK_SIZE - 1) // BLOCK_SIZE
block_ids = torch.arange(score_stride, device="cuda")
topk_mask = block_ids[None, :] < query_blocks.reshape(-1, 1)
request_blocks = query_blocks.max(dim=1).values
write_mask = (
block_ids[None, :]
< request_blocks.repeat_interleave(decode_query_len)[:, None]
)
return write_mask, topk_mask
def packed_keys(score: torch.Tensor, live: torch.Tensor) -> torch.Tensor:
normalized = torch.where(torch.isnan(score), -1e30, score)
normalized = torch.where(normalized == 0.0, 0.0, normalized)
bits = normalized.contiguous().view(torch.int32).to(torch.int64)
bits &= 0xFFFFFFFF
ordered = bits ^ torch.where(
bits >> 31 != 0,
0xFFFFFFFF,
0x80000000,
)
index = torch.arange(score_stride, device="cuda") + 1
keys = (1 << 48) | (ordered << 16) | (0xFFFF - index)
return torch.where(live, keys, 0)
def assert_exact(expected_seq_lens: torch.Tensor) -> None:
write_mask, topk_mask = score_masks(expected_seq_lens)
write_mask = write_mask.unsqueeze(0).expand_as(deployed_score)
topk_mask = topk_mask.unsqueeze(0)
deployed = deployed_score
optimized = optimized_score
assert torch.equal(
deployed[write_mask].contiguous().view(torch.int32),
optimized[write_mask].contiguous().view(torch.int32),
)
assert torch.isnan(deployed[~write_mask]).all()
assert torch.isnan(optimized[~write_mask]).all()
deployed_keys = packed_keys(deployed, topk_mask)
optimized_keys = packed_keys(optimized, topk_mask)
assert torch.equal(deployed_keys, optimized_keys)
deployed_topk = torch.topk(deployed_keys, 16, dim=-1, sorted=True)
optimized_topk = torch.topk(optimized_keys, 16, dim=-1, sorted=True)
assert torch.equal(deployed_topk.values, optimized_topk.values)
assert torch.equal(deployed_topk.indices, optimized_topk.indices)
uniform = torch.full(
(num_reqs,),
max_block * BLOCK_SIZE,
device="cuda",
dtype=torch.int32,
)
one_active = torch.zeros(num_reqs, device="cuda", dtype=torch.int32)
one_active[0] = max(1, max_block // 2) * BLOCK_SIZE
patterns = [uniform, one_active]
non_aligned = torch.tensor(
[(max_block - 1) * BLOCK_SIZE + request % 3 + 1 for request in range(num_reqs)],
device="cuda",
dtype=torch.int32,
)
patterns.append(non_aligned)
seq_lens.copy_(uniform)
deployed_score.fill_(float("nan"))
optimized_score.fill_(float("nan"))
launch_deployed()
launch_optimized()
assert_exact(uniform)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
launch_optimized()
addresses = (
seq_lens.data_ptr(),
optimized_score.data_ptr(),
block_table.data_ptr(),
index_kv_cache.data_ptr(),
)
for pattern in patterns:
seq_lens.copy_(pattern)
deployed_score.fill_(float("nan"))
optimized_score.fill_(float("nan"))
launch_deployed()
graph.replay()
assert_exact(pattern)
assert addresses == (
seq_lens.data_ptr(),
optimized_score.data_ptr(),
block_table.data_ptr(),
index_kv_cache.data_ptr(),
)
@pytest.mark.parametrize(
("kv_cache_dtype", "expected_dtype"),
[
@@ -638,6 +1015,338 @@ def test_decode_index_topk_correctness(
_assert_topk_indices_equal_unordered(actual, expected)
@pytest.mark.parametrize(
(
"max_block",
"total_q",
"num_idx_heads",
"topk",
"is_gfx950",
"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)),
],
)
def test_amd_decode_topk_launch_policy(
max_block: int,
total_q: int,
num_idx_heads: int,
topk: int,
is_gfx950: bool,
expected: tuple[int, bool, bool],
):
"""All supported head and query layouts use adaptive selection."""
from vllm.models.minimax_m3.amd.ops.index_topk import (
_decode_topk_launch_policy,
)
assert (
_decode_topk_launch_policy(
max_block,
total_q,
num_idx_heads,
topk,
is_gfx950=is_gfx950,
)
== expected
)
@pytest.mark.skipif(
not current_platform.is_rocm(),
reason="MiniMax-M3 fused decode top-k is ROCm-only",
)
@pytest.mark.parametrize(
(
"num_topk_chunks",
"single_tile_guaranteed",
"adaptive_final_merge",
"max_blocks",
"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),
],
)
def test_amd_decode_fused_topk_total_order_and_replay(
num_topk_chunks: int,
single_tile_guaranteed: bool,
adaptive_final_merge: bool,
max_blocks: int,
topk: int,
num_idx_heads: int,
):
"""Atomic multi-chunk selection is ordered and resets graph state."""
from vllm.models.minimax_m3.amd.ops.index_topk import (
SPARSE_BLOCK_SIZE,
_decode_topk_fused_kernel,
)
set_random_seed(0)
block_size_t = 1 << (topk - 1).bit_length()
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,
max_blocks * SPARSE_BLOCK_SIZE,
]
short_seq_lens_list = [0, 1, 127, 128, 129, 0, 1, 128, 129, 0, 1, 129]
boundary_limit = 0
if num_topk_chunks == 16 and max_blocks >= 8192:
boundary_limit = 15
elif num_topk_chunks == 4 and adaptive_final_merge:
boundary_limit = 4
if boundary_limit:
boundary_blocks = [
(active_chunks - 1) * 512 + 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)
)
long_seq_lens_list.extend(
num_blocks * SPARSE_BLOCK_SIZE for num_blocks in boundary_blocks
)
short_seq_lens_list.extend(
num_blocks * SPARSE_BLOCK_SIZE for num_blocks in reversed(boundary_blocks)
)
batch = len(long_seq_lens_list)
score = torch.randn((num_idx_heads, batch, max_blocks), device="cuda")
score[:, :, 0] = float("nan")
score[:, :, 1] = float("inf")
score[:, :, 2] = float("-inf")
score[:, :, 3] = 0.0
score[:, :, 4] = -0.0
score[:, :, 5] = 7.0
score[:, :, 6] = -1e30
score[:, :, 256] = 7.0
score[:, :, 512] = 7.0
seq_lens = torch.empty(batch, device="cuda", dtype=torch.int32)
partial = torch.full(
(num_topk_chunks, num_idx_heads, batch, block_size_t),
-1,
device="cuda",
dtype=torch.int64,
)
counter = torch.zeros((num_idx_heads, batch), device="cuda", dtype=torch.int32)
actual = torch.full(
(num_idx_heads, batch, topk), -2, device="cuda", dtype=torch.int32
)
dummy_block_table = torch.zeros(
(batch, max_blocks), device="cuda", dtype=torch.int32
)
addresses = (partial.data_ptr(), counter.data_ptr(), actual.data_ptr())
long_seq_lens = torch.tensor(long_seq_lens_list, device="cuda", dtype=torch.int32)
short_seq_lens = torch.tensor(short_seq_lens_list, device="cuda", dtype=torch.int32)
def launch() -> None:
_decode_topk_fused_kernel[(batch, num_idx_heads, num_topk_chunks)](
score,
partial,
counter,
actual,
seq_lens,
dummy_block_table,
actual,
seq_lens,
1,
score.stride(0),
score.stride(1),
score.stride(2),
partial.stride(0),
partial.stride(1),
partial.stride(2),
partial.stride(3),
counter.stride(0),
counter.stride(1),
actual.stride(0),
actual.stride(1),
actual.stride(2),
dummy_block_table.stride(0),
actual.stride(1),
topk=topk,
block_size=SPARSE_BLOCK_SIZE,
pages_per_sparse_block=8,
block_page_stride=8,
NUM_TOPK_CHUNKS=num_topk_chunks,
BLOCK_SIZE_K=512,
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,
)
score_cpu = score.cpu()
def reference(seq_lens_list: list[int]) -> torch.Tensor:
expected = torch.full((num_idx_heads, batch, topk), -1, dtype=torch.int32)
for head in range(num_idx_heads):
for row, seq_len in enumerate(seq_lens_list):
num_blocks = (seq_len + SPARSE_BLOCK_SIZE - 1) // SPARSE_BLOCK_SIZE
values = score_cpu[head, row, :num_blocks].tolist()
values = [-1e30 if value != value else value for value in values]
selected = sorted(
range(num_blocks),
key=lambda index: (-values[index], index),
)[:topk]
expected[head, row, : len(selected)] = torch.tensor(selected)
return expected
seq_lens.copy_(long_seq_lens)
launch()
torch.accelerator.synchronize()
assert torch.equal(actual.cpu(), reference(long_seq_lens_list))
assert torch.count_nonzero(counter).item() == 0
seq_lens.copy_(short_seq_lens)
launch()
torch.accelerator.synchronize()
assert torch.equal(actual.cpu(), reference(short_seq_lens_list))
assert torch.count_nonzero(counter).item() == 0
replay_count = 8 if max_blocks == 65534 else 128
for replay in range(replay_count):
seq_lens.copy_(long_seq_lens if replay % 2 == 0 else short_seq_lens)
launch()
seq_lens.copy_(long_seq_lens)
launch()
torch.accelerator.synchronize()
assert torch.equal(actual.cpu(), reference(long_seq_lens_list))
assert torch.count_nonzero(counter).item() == 0
assert addresses == (partial.data_ptr(), counter.data_ptr(), actual.data_ptr())
if (num_topk_chunks, single_tile_guaranteed, adaptive_final_merge, max_blocks) in (
(4, False, False, 8193),
(16, False, False, 8193),
(4, False, True, 8192),
(16, True, True, 8192),
):
graph = torch.cuda.CUDAGraph()
torch.accelerator.synchronize()
with torch.cuda.graph(graph):
launch()
seq_lens.copy_(short_seq_lens)
graph.replay()
torch.accelerator.synchronize()
assert torch.equal(actual.cpu(), reference(short_seq_lens_list))
assert torch.count_nonzero(counter).item() == 0
seq_lens.copy_(long_seq_lens)
graph.replay()
torch.accelerator.synchronize()
assert torch.equal(actual.cpu(), reference(long_seq_lens_list))
assert torch.count_nonzero(counter).item() == 0
assert addresses == (partial.data_ptr(), counter.data_ptr(), actual.data_ptr())
@pytest.mark.skipif(
not current_platform.is_rocm(),
reason="MiniMax-M3 fused decode top-k is ROCm-only",
)
@pytest.mark.parametrize(
("num_idx_heads", "decode_query_len", "max_decode_query_len"),
[(1, 1, 1), (2, 1, 4), (4, 4, 4), (4, 3, 8)],
)
def test_amd_decode_index_topk_end_to_end(
num_idx_heads: int,
decode_query_len: int,
max_decode_query_len: int,
):
"""ROCm BF16 score-to-selector parity crosses the atomic boundary."""
from vllm.models.minimax_m3.amd.ops.index_topk import (
minimax_m3_index_decode as amd_minimax_m3_index_decode,
)
torch.manual_seed(0)
topk = 16
head_dim = 16
seq_lens = torch.tensor([65537], device="cuda", dtype=torch.int32)
q_lens = torch.tensor([decode_query_len], device="cuda", dtype=torch.int32)
prefix_lens = seq_lens - q_lens
max_seq_len = int(seq_lens.max())
max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
block_table = torch.randperm(
max_blocks, device="cuda", dtype=torch.int32
).unsqueeze(0)
idx_q = torch.randn(
decode_query_len,
num_idx_heads,
head_dim,
device="cuda",
dtype=torch.bfloat16,
)
index_kv_cache = torch.randn(
max_blocks,
BLOCK_SIZE,
head_dim,
device="cuda",
dtype=torch.bfloat16,
)
completion_counter = torch.zeros(
(num_idx_heads, decode_query_len), device="cuda", dtype=torch.int32
)
actual = amd_minimax_m3_index_decode(
idx_q,
index_kv_cache,
block_table,
seq_lens,
max_seq_len=max_seq_len,
topk=topk,
init_blocks=0,
local_blocks=1,
num_kv_heads=num_idx_heads,
decode_query_len=decode_query_len,
max_decode_query_len=max_decode_query_len,
completion_counter=completion_counter,
)
expected = _reference_index_topk(
idx_q,
index_kv_cache,
block_table,
q_lens,
seq_lens,
prefix_lens,
topk,
0,
1,
)
torch.accelerator.synchronize()
_assert_topk_indices_equal_unordered(actual, expected)
assert torch.count_nonzero(completion_counter).item() == 0
@pytest.mark.skipif(
not current_platform.is_device_capability_family(100),
reason="fp8 e4m3 indexer cache is the SM100 (MSA) path.",
@@ -1250,6 +1959,137 @@ def test_aiter_decode_sparse_block_table_supports_spec_decode():
assert torch.equal(sparse_bt, expected_bt)
@pytest.mark.skipif(
not current_platform.is_rocm(),
reason="MiniMax-M3 fused decode top-k is ROCm-only",
)
@pytest.mark.parametrize("block_page_stride", [8, 16])
@pytest.mark.parametrize(
"seq_lens_list",
[
[4097],
[65537],
[1, 129, 4097, 0],
[127, 128, 129, 256, 4097, 0],
[1, 127, 128, 129, 255, 256, 257, 4097, 0],
],
)
def test_amd_decode_fused_topk_emits_sparse_table(
block_page_stride: int, seq_lens_list: list[int]
):
"""The fused selector must match top-k followed by table construction."""
from vllm.models.minimax_m3.amd.ops.index_topk import (
minimax_m3_index_decode as amd_minimax_m3_index_decode,
)
from vllm.models.minimax_m3.amd.ops.sparse_pa import (
PAGES_PER_SPARSE_BLOCK,
minimax_m3_build_sparse_block_table_decode,
)
torch.manual_seed(0)
decode_query_len = 4
seq_lens = torch.tensor(seq_lens_list, device="cuda", dtype=torch.int32)
batch = seq_lens.numel()
total_q = batch * decode_query_len
max_seq_len = int(seq_lens.max())
max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
index_block_table = torch.zeros(
(batch, max_blocks), device="cuda", dtype=torch.int32
)
for req_id, seq_len in enumerate(seq_lens.tolist()):
num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
pages = torch.arange(max_blocks, device="cuda", dtype=torch.int32)
index_block_table[req_id, :num_blocks] = torch.roll(pages, req_id)[:num_blocks]
# The index cache and main-attention cache have independent physical page
# mappings in production. Disjoint values make using the wrong table fail.
attention_block_table = torch.full_like(index_block_table, -31337)
for req_id, seq_len in enumerate(seq_lens.tolist()):
num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
first_page = 10_000 + req_id * max_blocks
attention_block_table[req_id, :num_blocks] = torch.arange(
first_page + num_blocks - 1,
first_page - 1,
-1,
device="cuda",
dtype=torch.int32,
)
assert not torch.equal(index_block_table, attention_block_table)
idx_q = torch.randn(total_q, 1, HEAD_DIM, device="cuda", dtype=torch.bfloat16)
index_kv_cache = torch.randn(
max_blocks,
BLOCK_SIZE,
HEAD_DIM,
device="cuda",
dtype=torch.bfloat16,
)
decode_kwargs = {
"max_seq_len": max_seq_len,
"topk": TOPK,
"init_blocks": 0,
"local_blocks": 1,
"num_kv_heads": 1,
"decode_query_len": decode_query_len,
"max_decode_query_len": decode_query_len,
}
completion_counter = torch.zeros((1, total_q), device="cuda", dtype=torch.int32)
expected_topk = amd_minimax_m3_index_decode(
idx_q,
index_kv_cache,
index_block_table,
seq_lens,
completion_counter=completion_counter,
**decode_kwargs,
)
expected_sparse_bt, expected_sparse_ctx = (
minimax_m3_build_sparse_block_table_decode(
expected_topk,
attention_block_table,
seq_lens,
decode_query_len,
block_page_stride,
)
)
topk_out = torch.full_like(expected_topk, -27182)
sparse_bt_out = torch.full(
(total_q, TOPK * PAGES_PER_SPARSE_BLOCK),
-27182,
device="cuda",
dtype=torch.int32,
)
sparse_ctx_out = torch.full((total_q,), -27182, device="cuda", dtype=torch.int32)
actual_topk = amd_minimax_m3_index_decode(
idx_q,
index_kv_cache,
index_block_table,
seq_lens,
out=topk_out,
attention_block_table=attention_block_table,
sparse_block_table_out=sparse_bt_out,
sparse_context_lens_out=sparse_ctx_out,
block_page_stride=block_page_stride,
completion_counter=completion_counter,
**decode_kwargs,
)
torch.accelerator.synchronize()
assert actual_topk.data_ptr() == topk_out.data_ptr()
assert torch.equal(actual_topk, expected_topk)
assert torch.equal(sparse_bt_out, expected_sparse_bt)
assert torch.equal(sparse_ctx_out, expected_sparse_ctx)
assert torch.count_nonzero(completion_counter).item() == 0
if seq_lens_list[-1] == 0:
padded_rows = slice(total_q - decode_query_len, total_q)
assert torch.all(actual_topk[:, padded_rows] == -1)
assert torch.count_nonzero(sparse_bt_out[padded_rows]) == 0
assert torch.count_nonzero(sparse_ctx_out[padded_rows]) == 0
@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
@pytest.mark.parametrize(
"seq_lens_list",
vllm/models/minimax_m3/amd/model.py
+46 -1
View File
@@ -92,6 +92,13 @@ from vllm.models.minimax_m3.amd.ops import (
gemma_rmsnorm,
swiglu_oai_split,
)
from vllm.models.minimax_m3.amd.ops.sparse_pa import (
minimax_m3_alloc_sparse_block_table,
minimax_m3_sparse_block_page_stride,
)
from vllm.models.minimax_m3.amd.sparse_attention_msa import (
MiniMaxM3SparseAiterPAImpl,
)
from vllm.models.minimax_m3.common.indexer import MiniMaxM3Indexer
from vllm.models.minimax_m3.common.mm_preprocess import (
MiniMaxM3VLDummyInputsBuilder,
@@ -701,6 +708,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
self.kv_cache_k = torch.tensor([])
self.kv_cache_v = torch.tensor([])
self._aiter_sparse_pa_cache_data_ptr = 0
self._aiter_sparse_pa_block_page_stride = 0
# Self-contained nn.Module: owns its side cache, selects its impl in init
# (Triton on ROCm, where the SM100 gate is always False).
self.indexer = MiniMaxM3Indexer(
@@ -796,6 +804,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
self.head_dim,
x,
)[v_page_offset:]
self._aiter_sparse_pa_block_page_stride = minimax_m3_sparse_block_page_stride(
self.kv_cache_k, self.kv_cache_v
)
self._aiter_sparse_pa_cache_data_ptr = self.kv_cache.data_ptr()
def get_aiter_sparse_pa_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]:
@@ -1020,9 +1031,43 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
# top-k into the shared ``topk_indices_buffer``; the attend reads it back.
# When skip_index_topk is set (ATOM index_topk_freq), reuse the selection
# the preceding compute layer wrote into the shared buffer this forward.
decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None
if not self.skip_index_topk:
assert index_query is not None
self.indexer(index_query)
attn_metadata = get_forward_context().attn_metadata
if self.use_aiter_sparse_pa and isinstance(attn_metadata, dict):
main_md = attn_metadata[self.layer_name]
assert isinstance(main_md, MiniMaxM3SparseMetadata)
if main_md.num_decodes > 0:
d = main_md.decode
assert d is not None
topk = self.topk_indices_buffer
assert topk is not None
decode_sparse_table = minimax_m3_alloc_sparse_block_table(
topk[:, : main_md.num_decode_tokens, :]
)
block_page_stride = self._aiter_sparse_pa_block_page_stride
assert block_page_stride > 0
self.indexer(
index_query,
attention_block_table=d.block_table,
sparse_block_table_out=decode_sparse_table[0],
sparse_context_lens_out=decode_sparse_table[1],
block_page_stride=block_page_stride,
)
else:
self.indexer(index_query)
else:
self.indexer(index_query)
if self.use_aiter_sparse_pa:
assert isinstance(self.impl, MiniMaxM3SparseAiterPAImpl)
return self.impl.forward(
self,
query,
self.kv_cache,
output,
decode_sparse_table=decode_sparse_table,
)
return self.impl.forward(self, query, self.kv_cache, output)
vllm/models/minimax_m3/amd/ops/index_topk.py
+1004 -324
View File
File diff suppressed because it is too large Load Diff
vllm/models/minimax_m3/amd/ops/sparse_pa.py
+60 -19
View File
@@ -32,14 +32,13 @@ def _is_fp8_kv_cache_tensor(kv_cache: torch.Tensor) -> bool:
@triton.jit
def _write_sparse_block_table_row(
topk_row, # [topk] int32, selected logical block ids for this query
def _write_sparse_block_table_row_from_values(
blk, # [BLOCK_SIZE_T] int32 selected logical block ids for this query
bt_row, # [max_blocks] int32, this request's logical page table
sbt_row, # [topk * PAGES_PER_BLOCK] int32, physical 16-page table
ctx_ptr, # int32, this query's attended context length
abs_pos, # absolute position of this query token, may be negative
max_topk,
stride_topk_k,
SPARSE_BLOCK_SIZE_C: tl.constexpr,
PAGES_PER_BLOCK: tl.constexpr,
BLOCK_PAGE_STRIDE: tl.constexpr,
@@ -56,8 +55,7 @@ def _write_sparse_block_table_row(
self_blk = abs_pos // SPARSE_BLOCK_SIZE_C
off_t = tl.arange(0, BLOCK_SIZE_T)
blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
valid = (causal_len > 0) & (blk >= 0) & (blk <= self_blk)
valid = (off_t < max_topk) & (causal_len > 0) & (blk >= 0) & (blk <= self_blk)
is_tail = valid & (blk == self_blk)
is_full = valid & (blk < self_blk)
@@ -89,6 +87,36 @@ def _write_sparse_block_table_row(
tl.store(ctx_ptr, ctx)
@triton.jit
def _write_sparse_block_table_row(
topk_row, # [topk] int32, selected logical block ids for this query
bt_row, # [max_blocks] int32, this request's logical page table
sbt_row, # [topk * PAGES_PER_BLOCK] int32, physical 16-page table
ctx_ptr, # int32, this query's attended context length
abs_pos, # absolute position of this query token, may be negative
max_topk,
stride_topk_k,
SPARSE_BLOCK_SIZE_C: tl.constexpr,
PAGES_PER_BLOCK: tl.constexpr,
BLOCK_PAGE_STRIDE: tl.constexpr,
BLOCK_SIZE_T: tl.constexpr,
):
off_t = tl.arange(0, BLOCK_SIZE_T)
blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
_write_sparse_block_table_row_from_values(
blk,
bt_row,
sbt_row,
ctx_ptr,
abs_pos,
max_topk,
SPARSE_BLOCK_SIZE_C,
PAGES_PER_BLOCK,
BLOCK_PAGE_STRIDE,
BLOCK_SIZE_T,
)
@triton.jit
def _build_sparse_block_table_kernel(
topk_ptr, # [1, batch * DECODE_QUERY_LEN, topk] int32, logical block ids
@@ -169,7 +197,7 @@ def _build_sparse_block_table_prefill_kernel(
)
def _alloc_sparse_block_table(
def minimax_m3_alloc_sparse_block_table(
topk_idx: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
assert topk_idx.shape[0] == 1, "AITER sparse PA requires num_kv_heads == 1"
@@ -190,7 +218,7 @@ def minimax_m3_build_sparse_block_table_prefill(
block_page_stride: int = PAGES_PER_SPARSE_BLOCK,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build one page-16 sparse block table row per prefill query token."""
sparse_bt, sparse_ctx = _alloc_sparse_block_table(topk_idx)
sparse_bt, sparse_ctx = minimax_m3_alloc_sparse_block_table(topk_idx)
topk = topk_idx.shape[-1]
_build_sparse_block_table_prefill_kernel[(topk_idx.shape[1],)](
topk_idx,
@@ -227,7 +255,7 @@ def minimax_m3_build_sparse_block_table_decode(
"MiniMax-M3 decode top-k rows must equal batch * decode_query_len: "
f"{total_q} != {expected_q}"
)
sparse_bt, sparse_ctx = _alloc_sparse_block_table(topk_idx)
sparse_bt, sparse_ctx = minimax_m3_alloc_sparse_block_table(topk_idx)
topk = topk_idx.shape[-1]
_build_sparse_block_table_kernel[(total_q,)](
topk_idx,
@@ -328,7 +356,9 @@ def _sides_are_packed(k_cache: torch.Tensor, v_cache: torch.Tensor) -> bool:
return v_cache.shape[0] != k_cache.shape[0]
def _block_page_stride(k_cache: torch.Tensor, v_cache: torch.Tensor) -> int:
def minimax_m3_sparse_block_page_stride(
k_cache: torch.Tensor, v_cache: torch.Tensor
) -> int:
"""How many page ids one sparse block spans.
Each side owns ``PAGES_PER_SPARSE_BLOCK`` pages. When the sides are dense
@@ -384,20 +414,31 @@ def minimax_m3_sparse_attn_decode_aiter(
k_scale: torch.Tensor | None = None,
v_scale: torch.Tensor | None = None,
decode_query_len: int = 1,
sparse_block_table: torch.Tensor | None = None,
sparse_context_lens: torch.Tensor | None = None,
) -> None:
sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_decode(
topk_idx,
block_table,
seq_lens,
decode_query_len,
_block_page_stride(k_cache, v_cache),
)
if (sparse_block_table is None) != (sparse_context_lens is None):
raise ValueError(
"MiniMax-M3 prepared sparse block table and context lengths "
"must be provided together"
)
if sparse_block_table is None:
sparse_block_table, sparse_context_lens = (
minimax_m3_build_sparse_block_table_decode(
topk_idx,
block_table,
seq_lens,
decode_query_len,
minimax_m3_sparse_block_page_stride(k_cache, v_cache),
)
)
assert sparse_context_lens is not None
_run_gluon_decode(
q,
k_cache,
v_cache,
sparse_bt,
sparse_ctx,
sparse_block_table,
sparse_context_lens,
num_kv_heads,
sm_scale,
output,
@@ -426,7 +467,7 @@ def minimax_m3_sparse_attn_prefill_aiter(
block_table,
query_req_id,
query_abs_pos,
_block_page_stride(k_cache, v_cache),
minimax_m3_sparse_block_page_stride(k_cache, v_cache),
)
_run_gluon_decode(
q,
vllm/models/minimax_m3/amd/sparse_attention_msa.py
+7
View File
@@ -25,6 +25,7 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
output: torch.Tensor,
*,
query_fp8: torch.Tensor | None = None,
decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
from vllm.models.minimax_m3.amd.ops.sparse_pa import (
minimax_m3_sparse_attn_decode_aiter,
@@ -70,6 +71,12 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
k_scale=k_scale,
v_scale=v_scale,
decode_query_len=d.decode_query_len,
sparse_block_table=(
decode_sparse_table[0] if decode_sparse_table is not None else None
),
sparse_context_lens=(
decode_sparse_table[1] if decode_sparse_table is not None else None
),
)
if main_md.num_prefills > 0:
vllm/models/minimax_m3/common/indexer.py
+71 -1
View File
@@ -369,6 +369,18 @@ class MiniMaxM3IndexerImpl(nn.Module):
# Shared, stable-address top-k output buffer (set by the model for the
# cudagraph-safe MSA impl); None -> impl allocates fresh (eager).
self.topk_indices_buffer = topk_indices_buffer
topk_completion_counter = None
if current_platform.is_rocm() and topk_indices_buffer is not None:
topk_completion_counter = torch.zeros(
topk_indices_buffer.shape[:2],
dtype=torch.int32,
device=topk_indices_buffer.device,
)
self.register_buffer(
"topk_completion_counter",
topk_completion_counter,
persistent=False,
)
# Owns the side cache (registers itself in the static forward context).
self.index_cache = MiniMaxM3IndexerCache(
head_dim=index_head_dim,
@@ -392,6 +404,11 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
def forward(
self,
index_query: torch.Tensor,
*,
attention_block_table: torch.Tensor | None = None,
sparse_block_table_out: torch.Tensor | None = None,
sparse_context_lens_out: torch.Tensor | None = None,
block_page_stride: int | None = None,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
attn_metadata = get_forward_context().attn_metadata
if not isinstance(attn_metadata, dict):
@@ -417,6 +434,19 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
if index_md.num_decodes > 0:
d = index_md.decode
assert d is not None
fused_sparse_kwargs = {}
if attention_block_table is not None:
fused_sparse_kwargs = {
"attention_block_table": attention_block_table,
"sparse_block_table_out": sparse_block_table_out,
"sparse_context_lens_out": sparse_context_lens_out,
"block_page_stride": block_page_stride,
}
decode_backend_kwargs = {}
if current_platform.is_rocm():
decode_backend_kwargs["completion_counter"] = (
self.topk_completion_counter
)
decode_topk = minimax_m3_index_decode(
iq[:nd],
kv,
@@ -430,6 +460,8 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
d.decode_query_len,
d.max_decode_query_len,
out=buf_htk,
**fused_sparse_kwargs,
**decode_backend_kwargs,
)
if index_md.num_prefills > 0:
p = index_md.prefill
@@ -567,5 +599,43 @@ class MiniMaxM3Indexer(nn.Module):
def forward(
self,
index_query: torch.Tensor,
*,
attention_block_table: torch.Tensor | None = None,
sparse_block_table_out: torch.Tensor | None = None,
sparse_context_lens_out: torch.Tensor | None = None,
block_page_stride: int | None = None,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
return self.impl(index_query)
if attention_block_table is None and all(
arg is None
for arg in (
sparse_block_table_out,
sparse_context_lens_out,
block_page_stride,
)
):
return self.impl(index_query)
if attention_block_table is None or any(
arg is None
for arg in (
sparse_block_table_out,
sparse_context_lens_out,
block_page_stride,
)
):
raise ValueError(
"MiniMax-M3 fused decode sparse-table arguments must be "
"provided together"
)
if not current_platform.is_rocm() or not isinstance(
self.impl, MiniMaxM3IndexerTritonImpl
):
raise ValueError(
"MiniMax-M3 fused decode sparse-table construction is ROCm-only"
)
return self.impl(
index_query,
attention_block_table=attention_block_table,
sparse_block_table_out=sparse_block_table_out,
sparse_context_lens_out=sparse_context_lens_out,
block_page_stride=block_page_stride,
)