diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 745d0fd587..d88540a3b5 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -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", diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 235b76d74c..71babc80ec 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -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) diff --git a/vllm/models/minimax_m3/amd/ops/index_topk.py b/vllm/models/minimax_m3/amd/ops/index_topk.py index 2b076a38b8..b5c75371f0 100644 --- a/vllm/models/minimax_m3/amd/ops/index_topk.py +++ b/vllm/models/minimax_m3/amd/ops/index_topk.py @@ -7,21 +7,119 @@ then the top-k blocks (plus forced init/local blocks) are selected per query token. Adapted to vLLM's paged KV cache: the KV page size is forced to equal the sparse block size (128), so one sparse block maps to exactly one page. -Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (single head). +Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (one shared +key vector per token). Only the paths MiniMax M3 uses are implemented: score_type="max", index value -disabled (score-only indexer), single shared index head. The selected block ids -feed the block-sparse attention kernels in ``sparse_attn``. +disabled (score-only indexer), and shared index keys. Each local index-query +head selects its own block ids for the block-sparse attention kernels in +``sparse_attn``. """ import torch +from vllm.models.minimax_m3.amd.ops.sparse_pa import ( + PAGES_PER_SPARSE_BLOCK, + _write_sparse_block_table_row_from_values, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.math_utils import round_up # One sparse block == one KV page. SPARSE_BLOCK_SIZE = 128 +DECODE_SCORE_BALANCED_PROGRAM_BUDGET = 1024 +DECODE_SCORE_HIGH_BATCH_PROGRAM_BUDGET = 768 +MAX_DECODE_SCORE_BALANCED_REQUESTS = 11 +DECODE_SCORE_DEFAULT_TARGET_GRID = 512 +DECODE_SCORE_GFX950_HIGH_BATCH_TARGET_GRID = 1024 +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 +MAX_DECODE_TOPK_FAST_CHUNKS = 16 +DECODE_TOPK_TARGET_GRID = 64 + + +def _decode_score_program_budget( + num_reqs: int, + head_dim: int, + query_dtype: torch.dtype, + cache_dtype: torch.dtype, + *, + is_gfx950: bool, +) -> int | None: + if not ( + is_gfx950 + and head_dim == 128 + and query_dtype == torch.bfloat16 + and cache_dtype == torch.bfloat16 + ): + return None + if 1 <= num_reqs <= 8: + return DECODE_SCORE_BALANCED_PROGRAM_BUDGET + if 9 <= num_reqs <= MAX_DECODE_SCORE_BALANCED_REQUESTS: + return DECODE_SCORE_HIGH_BATCH_PROGRAM_BUDGET + return None + + +def _decode_score_split_launch_policy( + num_reqs: int, + head_dim: int, + query_dtype: torch.dtype, + cache_dtype: torch.dtype, + *, + is_gfx950: bool, +) -> tuple[int, bool]: + """Choose the generic split-K launch and high-batch specialization.""" + use_high_batch_config = ( + is_gfx950 + and MIN_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS + <= num_reqs + <= MAX_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS + and head_dim == 128 + and query_dtype == torch.bfloat16 + and cache_dtype == torch.bfloat16 + ) + target_grid = ( + DECODE_SCORE_GFX950_HIGH_BATCH_TARGET_GRID + if use_high_batch_config + else DECODE_SCORE_DEFAULT_TARGET_GRID + ) + target = max( + 1, + min( + MAX_DECODE_SCORE_KV_CHUNKS, + target_grid // max(1, num_reqs), + ), + ) + return 1 << (target.bit_length() - 1), use_high_batch_config + + +def _decode_topk_launch_policy( + max_block: int, + total_q: int, + num_idx_heads: int, + topk: int, + *, + is_gfx950: bool, +) -> tuple[int, bool, bool]: + """Choose the graph grid and bounded-context selector specialization.""" + 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 + + target = max( + 1, + min( + MAX_DECODE_TOPK_FAST_CHUNKS, + DECODE_TOPK_TARGET_GRID // max(1, total_q * num_idx_heads), + ), + ) + return 1 << (target.bit_length() - 1), False, False # --------------------------------------------------------------------------- @@ -397,255 +495,670 @@ def _decode_index_score_kernel( ) +@triton.jit +def _ceil_div_nonnegative(value, divisor): + return value // divisor + tl.where(value % divisor != 0, 1, 0) + + +@triton.jit +def _decode_index_score_mapped_range( + q_ptr, + ik_cache_ptr, + score_ptr, + block_table_ptr, + pid_r, + seq_len, + chunk_start_block, + blocks_per_program, + num_idx_heads: tl.constexpr, + head_dim: tl.constexpr, + init_blocks, + local_blocks, + decode_query_len, + stride_q_n, + stride_q_h, + stride_q_d, + stride_ik_blk, + stride_ik_pos, + stride_ik_d, + stride_s_h, + stride_s_n, + stride_s_k, + stride_bt_b, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_Q: tl.constexpr, +): + BLOCK_SIZE_HQ: tl.constexpr = num_idx_heads * BLOCK_SIZE_Q + 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 + + query_pos = seq_len - decode_query_len + q_offsets + kv_len = tl.maximum(query_pos + 1, 0) + 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 + chunk_end_block = tl.minimum( + chunk_start_block + blocks_per_program, + num_blocks, + ) + + if chunk_start_block < chunk_end_block: + off_k = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, head_dim) + bt_row = block_table_ptr + pid_r * stride_bt_b + local_start = tl.maximum(0, num_blocks_q - local_blocks) + q = tl.load( + q_ptr + + q_ids[None, :] * stride_q_n + + h_offsets[None, :] * stride_q_h + + off_d[:, None] * stride_q_d, + mask=q_mask[None, :], + other=0.0, + ) + 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[:, None] < kv_len[None, :] + k = tl.load( + ik_cache_ptr + + page * stride_ik_blk + + off_k[:, None] * stride_ik_pos + + off_d * stride_ik_d, + ) + if BLOCK_SIZE_HQ == 1: + q_vec = tl.sum(q, axis=1).to(tl.float32) + kq = tl.sum( + k.to(tl.float32) * q_vec[None, :], + axis=1, + )[:, None] + else: + kq = tl.dot(k, q, out_dtype=tl.float32) + kq = tl.where( + pos_mask & q_mask[None, :], + kq, + float("-inf"), + ) + score = tl.max(kq, axis=0) + 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 + + h_offsets * stride_s_h + + q_ids * stride_s_n + + blk * stride_s_k, + score, + mask=q_mask, + ) + + +@triton.jit(do_not_specialize=["score_program_budget", "decode_query_len"]) +def _decode_index_score_balanced_kernel( + q_ptr, + ik_cache_ptr, + score_ptr, + block_table_ptr, + seq_lens, + num_idx_heads: tl.constexpr, + head_dim: tl.constexpr, + init_blocks, + local_blocks, + NUM_REQUESTS: tl.constexpr, + score_program_budget, + decode_query_len, + stride_q_n, + stride_q_h, + stride_q_d, + stride_ik_blk, + stride_ik_pos, + stride_ik_d, + stride_s_h, + stride_s_n, + stride_s_k, + stride_bt_b, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_Q: tl.constexpr, +): + pid = tl.program_id(0) + total_blocks = tl.zeros((), dtype=tl.int32) + for req in tl.static_range(0, NUM_REQUESTS): + req_seq_len = tl.maximum(tl.load(seq_lens + req), 0) + total_blocks += _ceil_div_nonnegative( + req_seq_len, + BLOCK_SIZE_K, + ) + blocks_per_program = tl.maximum( + 1, + _ceil_div_nonnegative( + total_blocks, + score_program_budget, + ), + ) + + program_offset = tl.zeros((), dtype=tl.int32) + pid_r = tl.zeros((), dtype=tl.int32) + chunk_start_block = tl.zeros((), dtype=tl.int32) + selected_seq_len = tl.zeros((), dtype=tl.int32) + owns_work = tl.zeros((), dtype=tl.int32) + for req in tl.static_range(0, NUM_REQUESTS): + req_seq_len = tl.maximum(tl.load(seq_lens + req), 0) + req_num_blocks = _ceil_div_nonnegative( + req_seq_len, + BLOCK_SIZE_K, + ) + req_programs = _ceil_div_nonnegative( + req_num_blocks, + blocks_per_program, + ) + owns_request = (pid >= program_offset) & (pid < program_offset + req_programs) + pid_r = tl.where(owns_request, req, pid_r) + local_program = tl.where( + owns_request, + pid - program_offset, + 0, + ) + chunk_start_block = tl.where( + owns_request, + local_program * blocks_per_program, + chunk_start_block, + ) + selected_seq_len = tl.where( + owns_request, + req_seq_len, + selected_seq_len, + ) + owns_work += owns_request.to(tl.int32) + program_offset += req_programs + if owns_work == 0: + return + + _decode_index_score_mapped_range( + q_ptr, + ik_cache_ptr, + score_ptr, + block_table_ptr, + pid_r, + selected_seq_len, + chunk_start_block, + blocks_per_program, + num_idx_heads, + head_dim, + init_blocks, + local_blocks, + decode_query_len, + stride_q_n, + stride_q_h, + stride_q_d, + stride_ik_blk, + stride_ik_pos, + stride_ik_d, + stride_s_h, + stride_s_n, + stride_s_k, + stride_bt_b, + BLOCK_SIZE_K, + BLOCK_SIZE_Q, + ) + + # --------------------------------------------------------------------------- -# Decode top-k (split-K): per-chunk partial top-k + merge. Forced init/local -# blocks are already encoded in the scores. +# Decode top-k. Each CTA keeps a packed score/id top-k. Short rows emit +# directly; long rows synchronize their live CTAs and the last arrival merges +# the partial lists and optionally builds the sparse page table. # --------------------------------------------------------------------------- -@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])}) -@triton.autotune( - configs=[ - triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2), - triton.Config({"BLOCK_SIZE_K": 256}, num_warps=4, num_stages=2), - triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2), - triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=3), - triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2), - ], - key=["topk"], -) -@triton.jit(do_not_specialize=["chunk_blocks", "decode_query_len"]) -def _topk_index_partial_kernel( - s_ptr, # score: [num_idx_heads, total_q, max_block] - ts_partial_ptr, # partial scores out: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T] - ti_partial_ptr, # partial idx out (1-indexed global, 0=invalid): same shape - seq_lens, # [num_reqs] - block_size: tl.constexpr, # sparse block size (128) +@triton.jit +def _decode_topk_key(score, index, valid): + score = tl.where(score != score, -1e30, score) + score = tl.where(score == 0.0, 0.0, score) + bits = score.to(tl.uint32, bitcast=True) + ordered = bits ^ tl.where(bits >> 31 != 0, 0xFFFFFFFF, 0x80000000) + tie = (0xFFFF - index).to(tl.int64) + key = (1 << 48) | (ordered.to(tl.int64) << 16) | tie + return tl.where(valid, key, 0) + + +@triton.jit +def _store_decode_topk( + winners, + topk_ptr, + attention_block_table_ptr, + sparse_bt_ptr, + sparse_ctx_ptr, + pid_b, + pid_h, + req_id, + query_pos, + num_blocks, topk: tl.constexpr, - chunk_blocks, # how many score-blocks each chunk owns + block_size: tl.constexpr, + pages_per_sparse_block: tl.constexpr, + block_page_stride: tl.constexpr, + stride_topk_h, + stride_topk_b, + stride_topk_t, + stride_attention_bt_b, + stride_sparse_bt_b, + BLOCK_SIZE_T: tl.constexpr, + EMIT_SPARSE_TABLE: tl.constexpr, +): + off_t = tl.arange(0, BLOCK_SIZE_T) + topk_idx = (0xFFFF - (winners & 0xFFFF)).to(tl.int32) - 1 + topk_idx = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx, -1) + tl.store( + topk_ptr + + pid_h * stride_topk_h + + pid_b * stride_topk_b + + off_t * stride_topk_t, + topk_idx, + mask=off_t < topk, + ) + if EMIT_SPARSE_TABLE: + _write_sparse_block_table_row_from_values( + topk_idx, + attention_block_table_ptr + req_id * stride_attention_bt_b, + sparse_bt_ptr + pid_b * stride_sparse_bt_b, + sparse_ctx_ptr + pid_b, + query_pos, + topk, + block_size, + pages_per_sparse_block, + block_page_stride, + BLOCK_SIZE_T, + ) + + +@triton.jit +def _merge_store_decode_topk( + partial_ptr, + counter, + topk_ptr, + attention_block_table_ptr, + sparse_bt_ptr, + sparse_ctx_ptr, + active_chunks, + pid_b, + pid_h, + req_id, + query_pos, + num_blocks, + stride_partial_c, + stride_partial_h, + stride_partial_b, + stride_partial_t, + stride_topk_h, + stride_topk_b, + stride_topk_t, + stride_attention_bt_b, + stride_sparse_bt_b, + topk: tl.constexpr, + block_size: tl.constexpr, + pages_per_sparse_block: tl.constexpr, + block_page_stride: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + MERGE_CHUNKS: tl.constexpr, + EMIT_SPARSE_TABLE: tl.constexpr, +): + candidate_width: tl.constexpr = MERGE_CHUNKS * BLOCK_SIZE_T + off_candidate = tl.arange(0, candidate_width) + candidate_chunk = off_candidate // BLOCK_SIZE_T + candidate_slot = off_candidate % BLOCK_SIZE_T + candidates = tl.load( + partial_ptr + + candidate_chunk * stride_partial_c + + pid_h * stride_partial_h + + pid_b * stride_partial_b + + candidate_slot * stride_partial_t, + mask=candidate_chunk < active_chunks, + other=0, + volatile=True, + ) + winners = tl.topk(candidates, BLOCK_SIZE_T) + _store_decode_topk( + winners, + topk_ptr, + attention_block_table_ptr, + sparse_bt_ptr, + sparse_ctx_ptr, + pid_b, + pid_h, + req_id, + query_pos, + num_blocks, + topk, + block_size, + pages_per_sparse_block, + block_page_stride, + stride_topk_h, + stride_topk_b, + stride_topk_t, + stride_attention_bt_b, + stride_sparse_bt_b, + BLOCK_SIZE_T, + EMIT_SPARSE_TABLE, + ) + tl.atomic_xchg(counter, 0, sem="acq_rel", scope="gpu") + + +@triton.jit(do_not_specialize=["decode_query_len"]) +def _decode_topk_fused_kernel( + score_ptr, # [num_idx_heads, total_q, max_block] + partial_ptr, # [NUM_TOPK_CHUNKS, num_idx_heads, total_q, BLOCK_SIZE_T] + counter_ptr, # [num_idx_heads, total_q] + topk_ptr, # [num_idx_heads, total_q, topk] + seq_lens_ptr, # [num_reqs] + attention_block_table_ptr, # [num_reqs, max_blocks] + sparse_bt_ptr, # [total_q, topk * PAGES_PER_SPARSE_BLOCK] + sparse_ctx_ptr, # [total_q] decode_query_len, stride_s_h, stride_s_b, stride_s_k, - stride_ts_c, - stride_ts_h, - stride_ts_b, - stride_ts_t, - stride_ti_c, - stride_ti_h, - stride_ti_b, - stride_ti_t, + stride_partial_c, + stride_partial_h, + stride_partial_b, + stride_partial_t, + stride_counter_h, + stride_counter_b, + stride_topk_h, + stride_topk_b, + stride_topk_t, + stride_attention_bt_b, + stride_sparse_bt_b, + topk: tl.constexpr, + block_size: tl.constexpr, + pages_per_sparse_block: tl.constexpr, + block_page_stride: tl.constexpr, + NUM_TOPK_CHUNKS: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, - USE_PDL: tl.constexpr, + EMIT_SPARSE_TABLE: tl.constexpr, + SINGLE_TILE_GUARANTEED: tl.constexpr, + ADAPTIVE_FINAL_MERGE: tl.constexpr, ): - tl.static_assert(topk < BLOCK_SIZE_K) - pid_b = tl.program_id(0) # flattened query-token id + tl.static_assert(topk <= BLOCK_SIZE_T) + tl.static_assert(BLOCK_SIZE_T <= BLOCK_SIZE_K) + pid_b = tl.program_id(0) pid_h = tl.program_id(1) pid_chunk = tl.program_id(2) req_id = pid_b // decode_query_len q_offset = pid_b - req_id * decode_query_len - if USE_PDL: - tl.extra.cuda.gdc_wait() - - seq_len = tl.load(seq_lens + req_id) + seq_len = tl.load(seq_lens_ptr + req_id) query_pos = seq_len - decode_query_len + q_offset - # 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 - 1) // block_size - # Slice this chunk owns within [0, num_blocks). - chunk_start = pid_chunk * chunk_blocks - chunk_end = tl.minimum(chunk_start + chunk_blocks, num_blocks) - chunk_actual = tl.maximum(chunk_end - chunk_start, 0) + desired_chunks = (num_blocks + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + active_chunks = tl.maximum(1, tl.minimum(NUM_TOPK_CHUNKS, desired_chunks)) + if pid_chunk >= active_chunks: + return + + blocks_per_chunk = num_blocks // active_chunks + extra_chunks = num_blocks - blocks_per_chunk * active_chunks + chunk_start = pid_chunk * blocks_per_chunk + tl.minimum(pid_chunk, extra_chunks) + chunk_blocks = blocks_per_chunk + (pid_chunk < extra_chunks) off_k = tl.arange(0, BLOCK_SIZE_K) off_t = tl.arange(0, BLOCK_SIZE_T) - - s_ptrs = ( - s_ptr - + pid_b * stride_s_b - + pid_h * stride_s_h - + (chunk_start + off_k) * stride_s_k - ) - - topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32) - topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32) - left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2 - - # Streaming top-K within this chunk. tl.range(0, 0) is a no-op so empty - # chunks (chunk_actual == 0) skip the body and store sentinel -1e30 / 0. - for i in tl.range(0, chunk_actual, BLOCK_SIZE_K): - mask = off_k < chunk_actual - i - score = tl.load(s_ptrs, mask=mask, other=-1e30).to(tl.float32) - score = tl.where(score != score, -1e30, score) - s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K - topk_score, last_topk_score = score, topk_score - topk_idx, last_topk_idx = ( - tl.where(mask, chunk_start + i + off_k + 1, 0), # 1-indexed global - topk_idx, + if SINGLE_TILE_GUARANTEED: + valid = off_k < chunk_blocks + score = tl.load( + score_ptr + + pid_h * stride_s_h + + pid_b * stride_s_b + + (chunk_start + off_k) * stride_s_k, + mask=valid, + other=-1e30, + ).to(tl.float32) + index = (chunk_start + off_k + 1).to(tl.int32) + local_topk = tl.topk( + _decode_topk_key(score, index, valid), + BLOCK_SIZE_T, ) - n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) - for j in tl.static_range(1, n_dims): - topk_score, topk_idx = _bitonic_merge( - topk_score, topk_idx.to(tl.int32), j, 2, n_dims + else: + score_ptrs = ( + score_ptr + + pid_h * stride_s_h + + pid_b * stride_s_b + + (chunk_start + off_k) * stride_s_k + ) + valid = off_k < chunk_blocks + score = tl.load(score_ptrs, mask=valid, other=-1e30).to(tl.float32) + index = (chunk_start + off_k + 1).to(tl.int32) + local_topk = tl.topk( + _decode_topk_key(score, index, valid), + BLOCK_SIZE_T, + ) + score_ptrs += BLOCK_SIZE_K * stride_s_k + for offset in tl.range(BLOCK_SIZE_K, chunk_blocks, BLOCK_SIZE_K): + valid = off_k < chunk_blocks - offset + score = tl.load(score_ptrs, mask=valid, other=-1e30).to(tl.float32) + index = (chunk_start + offset + off_k + 1).to(tl.int32) + tile_topk = tl.topk( + _decode_topk_key(score, index, valid), + BLOCK_SIZE_T, ) - if i != 0: - topk_score, topk_idx = _bitonic_merge( - topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims - ) - topk_score_new = last_topk_score * left_half_mask + topk_score * ( - 1 - left_half_mask - ) - topk_idx_new = last_topk_idx * left_half_mask + topk_idx * ( - 1 - left_half_mask - ) - topk_score, topk_idx = _bitonic_merge( - topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims - ) - else: - topk_score, topk_idx = _bitonic_merge( - topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims + local_topk = tl.topk( + tl.cat(local_topk, tile_topk, can_reorder=True), + BLOCK_SIZE_T, ) + score_ptrs += BLOCK_SIZE_K * stride_s_k - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() + if active_chunks == 1: + _store_decode_topk( + local_topk, + topk_ptr, + attention_block_table_ptr, + sparse_bt_ptr, + sparse_ctx_ptr, + pid_b, + pid_h, + req_id, + query_pos, + num_blocks, + topk, + block_size, + pages_per_sparse_block, + block_page_stride, + stride_topk_h, + stride_topk_b, + stride_topk_t, + stride_attention_bt_b, + stride_sparse_bt_b, + BLOCK_SIZE_T, + EMIT_SPARSE_TABLE, + ) + return - # Extract first BLOCK_SIZE_T entries (top-K of this chunk after the sort). - topk_mask_extract = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 - final_score = tl.sum( - topk_mask_extract[:, None] - * tl.reshape(topk_score, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), - axis=0, - ) - final_idx = tl.sum( - topk_mask_extract[:, None] - * tl.reshape(topk_idx, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), - axis=0, - ) - - # Always write all BLOCK_SIZE_T slots — invalid slots carry -1e30 / 0 - # sentinels and lose to real scores in the merge stage. - ts_ptrs = ( - ts_partial_ptr - + pid_chunk * stride_ts_c - + pid_b * stride_ts_b - + pid_h * stride_ts_h - + off_t * stride_ts_t - ) - ti_ptrs = ( - ti_partial_ptr - + pid_chunk * stride_ti_c - + pid_b * stride_ti_b - + pid_h * stride_ti_h - + off_t * stride_ti_t - ) - tl.store(ts_ptrs, final_score) - tl.store(ti_ptrs, final_idx) - - -@triton.heuristics( - { - "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]), - "BLOCK_SIZE_K": lambda args: triton.next_power_of_2( - args["num_topk_chunks"] * triton.next_power_of_2(args["topk"]) - ), - } -) -@triton.jit(do_not_specialize=["num_topk_chunks", "decode_query_len"]) -def _topk_index_merge_kernel( - ts_partial_ptr, # partial scores: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T] - ti_partial_ptr, # partial idx (1-indexed global, 0=invalid): same shape - ti_final_ptr, # final idx (0-indexed, -1=invalid): [num_idx_heads, total_q, topk] - seq_lens, # [num_reqs] - block_size: tl.constexpr, # sparse block size (128) - topk: tl.constexpr, - decode_query_len, - stride_ts_c, - stride_ts_h, - stride_ts_b, - stride_ts_t, - stride_ti_c, - stride_ti_h, - stride_ti_b, - stride_ti_t, - stride_tif_h, - stride_tif_b, - stride_tif_t, - num_topk_chunks, - BLOCK_SIZE_K: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - USE_PDL: tl.constexpr, -): - pid_b = tl.program_id(0) # flattened query-token id - pid_h = tl.program_id(1) - req_id = pid_b // decode_query_len - q_offset = pid_b - req_id * decode_query_len - - 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 - # 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 - 1) // block_size - - # Load NUM_TOPK_CHUNKS * BLOCK_SIZE_T candidates, padded to BLOCK_SIZE_K. - # Candidate at flat position p comes from chunk = p // BLOCK_SIZE_T, - # in_chunk = p % BLOCK_SIZE_T. - off = tl.arange(0, BLOCK_SIZE_K) - chunk_idx = off // BLOCK_SIZE_T - in_chunk_idx = off % BLOCK_SIZE_T - valid = chunk_idx < num_topk_chunks - - score_offset = ( - chunk_idx * stride_ts_c - + pid_h * stride_ts_h - + pid_b * stride_ts_b - + in_chunk_idx * stride_ts_t - ) - idx_offset = ( - chunk_idx * stride_ti_c - + pid_h * stride_ti_h - + pid_b * stride_ti_b - + in_chunk_idx * stride_ti_t - ) - - score = tl.load(ts_partial_ptr + score_offset, mask=valid, other=-1e30).to( - tl.float32 - ) - score = tl.where(score != score, -1e30, score) - idx = tl.load(ti_partial_ptr + idx_offset, mask=valid, other=0).to(tl.int32) - - # Full bitonic descending sort of BLOCK_SIZE_K items. - n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) - for j in tl.static_range(1, n_dims): - score, idx = _bitonic_merge(score, idx.to(tl.int32), j, 2, n_dims) - score, idx = _bitonic_merge(score, idx.to(tl.int32), n_dims, True, n_dims) - - # Extract first BLOCK_SIZE_T positions — these are the global top-K. - extract_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 - topk_idx_final = tl.sum( - extract_mask[:, None] - * tl.reshape(idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), - axis=0, - ) - - off_t = tl.arange(0, BLOCK_SIZE_T) - tif_ptrs = ( - ti_final_ptr - + pid_h * stride_tif_h - + pid_b * stride_tif_b - + off_t * stride_tif_t - ) - store_mask = off_t < topk - topk_idx_final = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx_final, -1) tl.store( - tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty), mask=store_mask + partial_ptr + + pid_chunk * stride_partial_c + + pid_h * stride_partial_h + + pid_b * stride_partial_b + + off_t * stride_partial_t, + local_topk, ) + counter = counter_ptr + pid_h * stride_counter_h + pid_b * stride_counter_b + arrival = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") + if arrival != active_chunks - 1: + return + + if ADAPTIVE_FINAL_MERGE: + if active_chunks <= 2: + _merge_store_decode_topk( + partial_ptr, + counter, + topk_ptr, + attention_block_table_ptr, + sparse_bt_ptr, + sparse_ctx_ptr, + active_chunks, + pid_b, + pid_h, + req_id, + query_pos, + num_blocks, + stride_partial_c, + stride_partial_h, + stride_partial_b, + stride_partial_t, + stride_topk_h, + stride_topk_b, + stride_topk_t, + stride_attention_bt_b, + stride_sparse_bt_b, + topk, + block_size, + pages_per_sparse_block, + block_page_stride, + BLOCK_SIZE_T, + 2, + EMIT_SPARSE_TABLE, + ) + return + if active_chunks <= 4: + _merge_store_decode_topk( + partial_ptr, + counter, + topk_ptr, + attention_block_table_ptr, + sparse_bt_ptr, + sparse_ctx_ptr, + active_chunks, + pid_b, + pid_h, + req_id, + query_pos, + num_blocks, + stride_partial_c, + stride_partial_h, + stride_partial_b, + stride_partial_t, + stride_topk_h, + stride_topk_b, + stride_topk_t, + stride_attention_bt_b, + stride_sparse_bt_b, + topk, + block_size, + pages_per_sparse_block, + block_page_stride, + BLOCK_SIZE_T, + 4, + EMIT_SPARSE_TABLE, + ) + return + if active_chunks <= 8: + _merge_store_decode_topk( + partial_ptr, + counter, + topk_ptr, + attention_block_table_ptr, + sparse_bt_ptr, + sparse_ctx_ptr, + active_chunks, + pid_b, + pid_h, + req_id, + query_pos, + num_blocks, + stride_partial_c, + stride_partial_h, + stride_partial_b, + stride_partial_t, + stride_topk_h, + stride_topk_b, + stride_topk_t, + stride_attention_bt_b, + stride_sparse_bt_b, + topk, + block_size, + pages_per_sparse_block, + block_page_stride, + BLOCK_SIZE_T, + 8, + EMIT_SPARSE_TABLE, + ) + return + _merge_store_decode_topk( + partial_ptr, + counter, + topk_ptr, + attention_block_table_ptr, + sparse_bt_ptr, + sparse_ctx_ptr, + active_chunks, + pid_b, + pid_h, + req_id, + query_pos, + num_blocks, + stride_partial_c, + stride_partial_h, + stride_partial_b, + stride_partial_t, + stride_topk_h, + stride_topk_b, + stride_topk_t, + stride_attention_bt_b, + stride_sparse_bt_b, + topk, + block_size, + pages_per_sparse_block, + block_page_stride, + BLOCK_SIZE_T, + 16, + EMIT_SPARSE_TABLE, + ) + return + + candidate_width: tl.constexpr = NUM_TOPK_CHUNKS * BLOCK_SIZE_T + off_candidate = tl.arange(0, candidate_width) + candidate_chunk = off_candidate // BLOCK_SIZE_T + candidate_slot = off_candidate % BLOCK_SIZE_T + candidates = tl.load( + partial_ptr + + candidate_chunk * stride_partial_c + + pid_h * stride_partial_h + + pid_b * stride_partial_b + + candidate_slot * stride_partial_t, + mask=candidate_chunk < active_chunks, + other=0, + volatile=True, + ) + winners = tl.topk(candidates, BLOCK_SIZE_T) + _store_decode_topk( + winners, + topk_ptr, + attention_block_table_ptr, + sparse_bt_ptr, + sparse_ctx_ptr, + pid_b, + pid_h, + req_id, + query_pos, + num_blocks, + topk, + block_size, + pages_per_sparse_block, + block_page_stride, + stride_topk_h, + stride_topk_b, + stride_topk_t, + stride_attention_bt_b, + stride_sparse_bt_b, + BLOCK_SIZE_T, + EMIT_SPARSE_TABLE, + ) + tl.atomic_xchg(counter, 0, sem="acq_rel", scope="gpu") # --------------------------------------------------------------------------- @@ -777,21 +1290,118 @@ def minimax_m3_index_decode( decode_query_len: int, max_decode_query_len: int, out: torch.Tensor | None = None, + *, + 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, + completion_counter: torch.Tensor | None = None, ) -> torch.Tensor: - """Decode index block-score + top-k, both split-K (cudagraph-safe). + """Decode index block-score followed by fused adaptive top-k selection. Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad). When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating. + The optional sparse-table arguments fuse current-layer table construction + into the selector. They must be provided together. ``completion_counter`` + provides stable per-query synchronization storage for CUDA graphs. It must + be zero before its first launch and must not be shared by overlapping + selector invocations; every completed launch resets its active entries. """ total_q, num_idx_heads, head_dim = idx_q.shape 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 0 < decode_query_len <= max_decode_query_len assert total_q == seq_lens.shape[0] * decode_query_len batch = total_q + emit_sparse_table = attention_block_table is not None + if emit_sparse_table and ( + sparse_block_table_out is None + or sparse_context_lens_out is None + or block_page_stride is None + ): + raise ValueError( + "MiniMax-M3 fused decode sparse-table arguments must be provided together" + ) + if not emit_sparse_table and ( + sparse_block_table_out is not None + or sparse_context_lens_out is not None + or block_page_stride is not None + ): + raise ValueError( + "MiniMax-M3 fused decode sparse-table arguments must be provided together" + ) + if emit_sparse_table: + assert attention_block_table is not None + assert sparse_block_table_out is not None + assert sparse_context_lens_out is not None + assert block_page_stride is not None + if num_idx_heads != 1: + raise ValueError( + "MiniMax-M3 fused decode sparse-table construction requires " + f"one index head, got {num_idx_heads}" + ) + if block_page_stride not in ( + PAGES_PER_SPARSE_BLOCK, + 2 * PAGES_PER_SPARSE_BLOCK, + ): + raise ValueError( + "MiniMax-M3 fused decode sparse-table page stride must be " + f"{PAGES_PER_SPARSE_BLOCK} or {2 * PAGES_PER_SPARSE_BLOCK}, " + f"got {block_page_stride}" + ) + expected_sparse_shape = (total_q, topk * PAGES_PER_SPARSE_BLOCK) + if sparse_block_table_out.shape != expected_sparse_shape: + raise ValueError( + "MiniMax-M3 fused decode sparse block table has shape " + f"{tuple(sparse_block_table_out.shape)}, expected " + f"{expected_sparse_shape}" + ) + if sparse_context_lens_out.shape != (total_q,): + raise ValueError( + "MiniMax-M3 fused decode sparse context lengths have shape " + f"{tuple(sparse_context_lens_out.shape)}, expected {(total_q,)}" + ) + if ( + attention_block_table.dim() != 2 + or attention_block_table.shape[0] != seq_lens.shape[0] + ): + raise ValueError( + "MiniMax-M3 fused decode attention block table must be " + "[num_requests, max_blocks]" + ) + int32_tensors = ( + attention_block_table, + sparse_block_table_out, + sparse_context_lens_out, + ) + if any(tensor.dtype != torch.int32 for tensor in int32_tensors): + raise ValueError("MiniMax-M3 fused decode sparse tables require int32") + if any(tensor.device != idx_q.device for tensor in int32_tensors): + raise ValueError( + "MiniMax-M3 fused decode sparse tables must share the query device" + ) + if ( + attention_block_table.stride(1) != 1 + or sparse_block_table_out.stride(1) != 1 + or sparse_context_lens_out.stride(0) != 1 + ): + raise ValueError( + "MiniMax-M3 fused decode sparse tables require contiguous rows" + ) max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) + if max_block >= 0xFFFF: + raise ValueError( + "MiniMax-M3 decode top-k supports fewer than 65535 sparse blocks" + ) + if emit_sparse_table: + assert attention_block_table is not None + if attention_block_table.shape[1] < max_block: + raise ValueError( + "MiniMax-M3 fused decode attention block table is shorter " + f"than the required {max_block} sparse blocks" + ) use_pdl = current_platform.is_arch_support_pdl() # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA # SM9+); this ROCm Triton rejects it even when False ("Keyword argument @@ -800,9 +1410,13 @@ def minimax_m3_index_decode( 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. + is_gfx950 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx950 + + is_gfx950 = on_gfx950() + # Multi-head spec decode scores a wider head-position tile per K block; + # reduce stages to ease memory/register pressure on the fallback path. 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}) @@ -814,47 +1428,90 @@ def minimax_m3_index_decode( dtype=torch.float32, device=idx_q.device, ) - # split-K over seq blocks; chunk count depends only on shape constants so - # the grid is fixed within a cuda graph. - 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, score_ctas_per_chunk)), - ) - num_kv_chunks = 1 << (target.bit_length() - 1) - grid_score = (seq_lens.shape[0], num_kv_chunks) - _decode_index_score_kernel[grid_score]( - idx_q, - index_kv_cache, - score, - block_table, - seq_lens, - num_idx_heads, + num_reqs = seq_lens.shape[0] + score_program_budget = _decode_score_program_budget( + num_reqs, 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), - score.stride(0), - score.stride(1), - 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, - **score_kwargs, + idx_q.dtype, + index_kv_cache.dtype, + is_gfx950=is_gfx950, ) + grid_score: tuple[int, ...] + if score_program_budget is not None: + grid_score = (score_program_budget + num_reqs - 1,) + _decode_index_score_balanced_kernel[grid_score]( + idx_q, + index_kv_cache, + score, + block_table, + seq_lens, + num_idx_heads, + head_dim, + init_blocks, + local_blocks, + num_reqs, + score_program_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), + score.stride(0), + score.stride(1), + score.stride(2), + block_table.stride(0), + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + BLOCK_SIZE_Q=BLOCK_SIZE_Q, + num_warps=2, + num_stages=1, + ) + else: + # Increase independent work for the measured high-batch gfx950 decode + # range while preserving the deployed split for every other shape. + num_kv_chunks, use_high_batch_config = _decode_score_split_launch_policy( + num_reqs, + head_dim, + idx_q.dtype, + index_kv_cache.dtype, + is_gfx950=is_gfx950, + ) + if use_high_batch_config and not ( + num_idx_heads > 1 and max_decode_query_len > 1 + ): + score_kwargs.update({"num_warps": 2, "num_stages": 1}) + grid_score = (num_reqs, num_kv_chunks) + _decode_index_score_kernel[grid_score]( + idx_q, + index_kv_cache, + 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), + score.stride(0), + score.stride(1), + 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, + **score_kwargs, + ) if out is not None: topk_idx = out[:, :total_q, :] @@ -864,76 +1521,99 @@ def minimax_m3_index_decode( dtype=torch.int32, device=idx_q.device, ) - # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts - # pow2(num_topk_chunks * pow2(topk)) candidates. - TOPK_TARGET_GRID = 64 - MAX_NUM_TOPK_CHUNKS = 16 - topk_target = max( - 1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads)) - ) - num_topk_chunks = 1 << (topk_target.bit_length() - 1) - block_size_t = triton.next_power_of_2(topk) - chunk_blocks = (max_block + num_topk_chunks - 1) // num_topk_chunks - topk_score_partial = torch.empty( + # The launch grid remains shape-constant for CUDA graphs. Each query uses + # only the number of chunks needed for its live context. + ( num_topk_chunks, - num_idx_heads, + single_tile_guaranteed, + adaptive_final_merge, + ) = _decode_topk_launch_policy( + max_block, batch, - block_size_t, - dtype=torch.float32, - device=idx_q.device, - ) - topk_idx_partial = torch.empty( - num_topk_chunks, num_idx_heads, - batch, - block_size_t, - dtype=torch.int32, - device=idx_q.device, - ) - _topk_index_partial_kernel[(batch, num_idx_heads, num_topk_chunks)]( - score, - topk_score_partial, - topk_idx_partial, - seq_lens, - SPARSE_BLOCK_SIZE, topk, - chunk_blocks, + is_gfx950=is_gfx950, + ) + block_size_t = triton.next_power_of_2(topk) + topk_partial = torch.empty( + num_topk_chunks, + num_idx_heads, + batch, + block_size_t, + dtype=torch.int64, + device=idx_q.device, + ) + if completion_counter is None: + active_counter = torch.zeros( + (num_idx_heads, batch), + dtype=torch.int32, + device=idx_q.device, + ) + else: + if ( + completion_counter.dim() != 2 + or completion_counter.shape[0] != num_idx_heads + or completion_counter.shape[1] < batch + or completion_counter.dtype != torch.int32 + or completion_counter.device != idx_q.device + or not completion_counter.is_contiguous() + ): + raise ValueError( + "MiniMax-M3 completion counter must be contiguous int32 " + "[num_idx_heads, >=total_q] on the query device" + ) + active_counter = completion_counter[:, :batch] + + selector_attention_block_table = block_table + selector_sparse_block_table = topk_idx + selector_sparse_context_lens = seq_lens + selector_block_page_stride = PAGES_PER_SPARSE_BLOCK + selector_sparse_block_stride = topk_idx.stride(1) + if emit_sparse_table: + assert attention_block_table is not None + assert sparse_block_table_out is not None + assert sparse_context_lens_out is not None + assert block_page_stride is not None + selector_attention_block_table = attention_block_table + selector_sparse_block_table = sparse_block_table_out + selector_sparse_context_lens = sparse_context_lens_out + selector_block_page_stride = block_page_stride + selector_sparse_block_stride = sparse_block_table_out.stride(0) + _decode_topk_fused_kernel[(batch, num_idx_heads, num_topk_chunks)]( + score, + topk_partial, + active_counter, + topk_idx, + seq_lens, + selector_attention_block_table, + selector_sparse_block_table, + selector_sparse_context_lens, decode_query_len, score.stride(0), score.stride(1), score.stride(2), - topk_score_partial.stride(0), - topk_score_partial.stride(1), - topk_score_partial.stride(2), - topk_score_partial.stride(3), - topk_idx_partial.stride(0), - topk_idx_partial.stride(1), - topk_idx_partial.stride(2), - topk_idx_partial.stride(3), - USE_PDL=use_pdl, - **pdl_kwargs, - ) - _topk_index_merge_kernel[(batch, num_idx_heads)]( - topk_score_partial, - topk_idx_partial, - topk_idx, - seq_lens, - SPARSE_BLOCK_SIZE, - topk, - decode_query_len, - topk_score_partial.stride(0), - topk_score_partial.stride(1), - topk_score_partial.stride(2), - topk_score_partial.stride(3), - topk_idx_partial.stride(0), - topk_idx_partial.stride(1), - topk_idx_partial.stride(2), - topk_idx_partial.stride(3), + topk_partial.stride(0), + topk_partial.stride(1), + topk_partial.stride(2), + topk_partial.stride(3), + active_counter.stride(0), + active_counter.stride(1), topk_idx.stride(0), topk_idx.stride(1), topk_idx.stride(2), - num_topk_chunks=num_topk_chunks, - USE_PDL=use_pdl, - **pdl_kwargs, + selector_attention_block_table.stride(0), + selector_sparse_block_stride, + topk=topk, + block_size=SPARSE_BLOCK_SIZE, + 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_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, ) return topk_idx diff --git a/vllm/models/minimax_m3/amd/ops/sparse_pa.py b/vllm/models/minimax_m3/amd/ops/sparse_pa.py index f7010cac48..1431bc579c 100644 --- a/vllm/models/minimax_m3/amd/ops/sparse_pa.py +++ b/vllm/models/minimax_m3/amd/ops/sparse_pa.py @@ -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, diff --git a/vllm/models/minimax_m3/amd/sparse_attention_msa.py b/vllm/models/minimax_m3/amd/sparse_attention_msa.py index 272a075aa0..c5e0d31681 100644 --- a/vllm/models/minimax_m3/amd/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/amd/sparse_attention_msa.py @@ -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: diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index 49c1fc83b8..f5a8db0c24 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -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, + )