diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 8fe2342808fc..01974badab1c 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -124,6 +124,46 @@ def test_sparse_kernels_recognize_fp8_dtypes(dtype: torch.dtype): # Index top-k kernels. +def _assert_prefill_index_scores( + actual: torch.Tensor, + idx_q: torch.Tensor, + index_kv_cache: torch.Tensor, + block_table: torch.Tensor, + q_lens: torch.Tensor, + seq_lens: torch.Tensor, + prefix_lens: torch.Tensor, + block_size_q: int, +) -> None: + q_start = 0 + for req_id, (q_len, seq_len, prefix_len) in enumerate( + zip(q_lens.tolist(), seq_lens.tolist(), prefix_lens.tolist()) + ): + q = idx_q[q_start : q_start + q_len] + num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE + pages = block_table[req_id, :num_blocks] + k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1) + expected = torch.einsum("qhd,kd->hqk", q.float(), k.float()) + + q_pos = prefix_len + torch.arange(q_len, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + expected.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + expected = ( + expected.reshape(idx_q.shape[1], q_len, num_blocks, BLOCK_SIZE) + .max(dim=3) + .values + ) + + for local_q in range(q_len): + q_block_end = min(q_len, (local_q // block_size_q + 1) * block_size_q) + hi = min(seq_len, prefix_len + q_block_end) + written_blocks = (hi + BLOCK_SIZE - 1) // BLOCK_SIZE + torch.testing.assert_close( + actual[:, q_start + local_q, :written_blocks], + expected[:, local_q, :written_blocks], + ) + q_start += q_len + + def _reference_index_topk( idx_q: torch.Tensor, index_kv_cache: torch.Tensor, @@ -222,14 +262,34 @@ def _reference_decode_index_score( return out -def test_prefill_index_topk_correctness(): +@pytest.mark.parametrize("long_context", [False, True]) +def test_prefill_index_topk_correctness(long_context: bool): + if current_platform.is_rocm(): + from vllm.models.minimax_m3.amd.ops.index_topk import ( + minimax_m3_index_score as amd_index_score, + ) + + index_score = amd_index_score + else: + index_score = minimax_m3_index_score + + if long_context: + if not current_platform.is_rocm(): + pytest.skip("The split-K index-score path is ROCm-specific.") + from vllm.platforms.rocm import on_gfx942 + + if not on_gfx942(): + pytest.skip("The split-K index-score path is enabled on gfx942.") + topk = 6 init_blocks = 0 local_blocks = 1 num_idx_heads = 2 head_dim = 16 - q_lens = torch.tensor((4, 3), device="cuda", dtype=torch.int32) - prefix_lens = torch.tensor((0, 1024), device="cuda", dtype=torch.int32) + q_lens_values = (128, 129) if long_context else (4, 3) + prefix_lens_values = (8192, 16384) if long_context else (0, 1024) + q_lens = torch.tensor(q_lens_values, device="cuda", dtype=torch.int32) + prefix_lens = torch.tensor(prefix_lens_values, device="cuda", dtype=torch.int32) seq_lens = prefix_lens + q_lens batch = q_lens.numel() max_seq_len = seq_lens.max().item() @@ -242,13 +302,15 @@ def test_prefill_index_topk_correctness(): batch, max_blocks ) idx_q = torch.ones(q_lens.sum().item(), num_idx_heads, head_dim, device="cuda") - index_kv_cache = torch.empty(num_pages, BLOCK_SIZE, head_dim, device="cuda") - for req_id in range(batch): - for block_id in range(max_blocks): - page = block_table[req_id, block_id] - index_kv_cache[page].fill_(block_id + 1) + block_values = torch.empty(num_pages, device="cuda") + block_values[block_table] = torch.arange( + 1, max_blocks + 1, device="cuda", dtype=torch.float32 + ).expand(batch, -1) + index_kv_cache = ( + block_values[:, None, None].expand(-1, BLOCK_SIZE, head_dim).contiguous() + ) - score = minimax_m3_index_score( + score = index_score( idx_q, index_kv_cache, block_table, @@ -259,6 +321,16 @@ def test_prefill_index_topk_correctness(): max_seq_len=max_seq_len, num_kv_heads=num_idx_heads, ) + _assert_prefill_index_scores( + score, + idx_q, + index_kv_cache, + block_table, + q_lens, + seq_lens, + prefix_lens, + block_size_q=128 if long_context else 64, + ) actual = minimax_m3_index_topk( score, cu_seqlens, diff --git a/vllm/models/minimax_m3/amd/ops/index_topk.py b/vllm/models/minimax_m3/amd/ops/index_topk.py index 2b076a38b89d..1885e4384ebd 100644 --- a/vllm/models/minimax_m3/amd/ops/index_topk.py +++ b/vllm/models/minimax_m3/amd/ops/index_topk.py @@ -17,12 +17,20 @@ import torch from vllm.platforms import current_platform +from vllm.platforms.rocm import on_gfx942 from vllm.triton_utils import tl, triton from vllm.utils.math_utils import round_up +from vllm.utils.platform_utils import num_compute_units # One sparse block == one KV page. SPARSE_BLOCK_SIZE = 128 +# gfx942 split-K prefill tuning. Target bounded K-loop lengths while retaining +# enough work per split to amortize launch overhead. +_INDEX_SCORE_TARGET_K_TILES_PER_SPLIT = 64 +_INDEX_SCORE_MAX_SPLITS = 16 +_INDEX_SCORE_MIN_SEQ_LEN = _INDEX_SCORE_TARGET_K_TILES_PER_SPLIT * SPARSE_BLOCK_SIZE + # --------------------------------------------------------------------------- # Bitonic top-k helpers (layout-agnostic). @@ -101,8 +109,10 @@ def _index_block_score_kernel( stride_bt_b, BLOCK_SIZE_Q: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_K_SPLITS: tl.constexpr, ): - pid_q = tl.program_id(0) + pid_q = tl.program_id(0) // NUM_K_SPLITS + pid_split = tl.program_id(0) % NUM_K_SPLITS pid_bh = tl.program_id(1) pid_b = pid_bh // num_idx_heads pid_h = pid_bh % num_idx_heads @@ -114,6 +124,16 @@ def _index_block_score_kernel( if BLOCK_SIZE_Q * pid_q >= q_len: return + # Split the state-free K loop across programs. Every score element still + # has exactly one writer, so no reduction or atomic operation is needed. + hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q) + num_k_tiles = tl.cdiv(hi, BLOCK_SIZE_K) + tiles_per_split = tl.cdiv(num_k_tiles, NUM_K_SPLITS) + tile_lo = pid_split * tiles_per_split + tile_hi = min(num_k_tiles, tile_lo + tiles_per_split) + if tile_lo >= tile_hi: + return + q_ptrs = tl.make_block_ptr( base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h, shape=(q_len, head_dim), @@ -130,10 +150,8 @@ def _index_block_score_kernel( off_d = tl.arange(0, head_dim) # Block table row for this request. bt_row = block_table_ptr + pid_b * stride_bt_b - # Causal window: only blocks up to the last query token's position. - hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q) - for i in tl.range(0, hi, BLOCK_SIZE_K): - blk = i // BLOCK_SIZE_K + for blk in tl.range(tile_lo, tile_hi): + i = blk * BLOCK_SIZE_K page = tl.load(bt_row + blk).to(tl.int64) pos = i + off_k # index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed) @@ -651,6 +669,39 @@ def _topk_index_merge_kernel( # --------------------------------------------------------------------------- # Python wrappers # --------------------------------------------------------------------------- +def _index_score_launch_config( + max_query_len: int, + max_seq_len: int, + batch: int, + num_idx_heads: int, +) -> tuple[int, int, dict]: + """Select a split-K launch only for shapes with useful parallel work.""" + if not on_gfx942() or max_query_len < 128 or max_seq_len < _INDEX_SCORE_MIN_SEQ_LEN: + return 64, 1, {} + + block_size_q = 128 + num_k_tiles = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) + q_programs = triton.cdiv(max_query_len, block_size_q) * batch * num_idx_heads + + # Aim for two programs per CU when the split cap and K work permit, and + # avoid leaving very long K loops. Power-of-two splits limit compilations. + occupancy_splits = triton.cdiv(2 * num_compute_units(), q_programs) + work_splits = triton.cdiv(num_k_tiles, _INDEX_SCORE_TARGET_K_TILES_PER_SPLIT) + num_splits = triton.next_power_of_2(max(occupancy_splits, work_splits)) + num_splits = min(_INDEX_SCORE_MAX_SPLITS, num_k_tiles, num_splits) + if num_splits == 1: + return 64, 1, {} + + return ( + block_size_q, + num_splits, + { + "num_warps": 4, + "num_stages": 1, + }, + ) + + @torch.no_grad() def minimax_m3_index_score( idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim] @@ -682,8 +733,16 @@ def minimax_m3_index_score( dtype=torch.float32, device=idx_q.device, ) - BLOCK_SIZE_Q = 64 - grid_score = (triton.cdiv(max_query_len, BLOCK_SIZE_Q), batch * num_idx_heads) + block_size_q, num_k_splits, launch_kwargs = _index_score_launch_config( + max_query_len, + max_seq_len, + batch, + num_idx_heads, + ) + grid_score = ( + triton.cdiv(max_query_len, block_size_q) * num_k_splits, + batch * num_idx_heads, + ) _index_block_score_kernel[grid_score]( idx_q, index_kv_cache, @@ -704,8 +763,10 @@ def minimax_m3_index_score( score.stride(1), score.stride(2), block_table.stride(0), - BLOCK_SIZE_Q=BLOCK_SIZE_Q, + BLOCK_SIZE_Q=block_size_q, BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_K_SPLITS=num_k_splits, + **launch_kwargs, ) return score