diff --git a/tests/v1/attention/test_rocm_glm5next_sparse.py b/tests/v1/attention/test_rocm_glm5next_sparse.py index ccc9e5c617..e69686f500 100644 --- a/tests/v1/attention/test_rocm_glm5next_sparse.py +++ b/tests/v1/attention/test_rocm_glm5next_sparse.py @@ -8,6 +8,7 @@ import torch from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( _use_rocm_sparse_triton, @@ -26,6 +27,26 @@ def _store_sparse_kv_row_offset_kernel(slot_ptr, output_ptr, stride: tl.constexp tl.store(output_ptr, _sparse_kv_row_offset(slot, stride)) +def _fit_kpool_indices_reference( + token_indices: torch.Tensor, topk_tokens: int +) -> torch.Tensor: + history = token_indices[:, :topk_tokens] + tail = token_indices[:, topk_tokens:] + valid_history = (history >= 0).sum(dim=1) + valid_tail = (tail >= 0).sum(dim=1) + keep_history = torch.minimum(valid_history, topk_tokens - valid_tail) + + columns = torch.arange(topk_tokens, device=token_indices.device).unsqueeze(0) + tail_offsets = columns - keep_history.unsqueeze(1) + tail_values = torch.gather( + tail, 1, tail_offsets.clamp(min=0, max=tail.shape[1] - 1) + ) + output = torch.where(columns < keep_history.unsqueeze(1), history, tail_values) + valid_output = columns < (keep_history + valid_tail).unsqueeze(1) + return torch.where(valid_output, output, -1) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required") def test_fit_kpool_indices_preserves_tail_and_best_history(): token_indices = torch.tensor( [ @@ -34,6 +55,7 @@ def test_fit_kpool_indices_preserves_tail_and_best_history(): [-1, -1, -1, -1, -1, -1, -1, -1], ], dtype=torch.int32, + device="cuda", ) fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=6) @@ -45,6 +67,28 @@ def test_fit_kpool_indices_preserves_tail_and_best_history(): ] +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required") +@pytest.mark.parametrize( + ("num_tokens", "topk_tokens", "tail_width"), + [(1, 4, 1), (7, 6, 2), (13, 128, 128), (3, 2048, 128)], +) +def test_fit_kpool_indices_matches_eager_reference(num_tokens, topk_tokens, tail_width): + """The Triton fit must reproduce the eager packing rule exactly.""" + set_random_seed(topk_tokens + num_tokens) + width = topk_tokens + tail_width + token_indices = torch.randint(0, 4096, (num_tokens, width), dtype=torch.int32) + valid = torch.rand((num_tokens, width)) < 0.5 + valid[0] = False + valid[-1] = True + token_indices = torch.where(valid, token_indices, -1).cuda() + + fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=topk_tokens) + + torch.testing.assert_close( + fitted, _fit_kpool_indices_reference(token_indices, topk_tokens) + ) + + def test_fit_kpool_indices_exact_width_is_noop(): token_indices = torch.tensor([[3, 2, 1, -1]], dtype=torch.int32) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index c5a672e5d4..e473f85948 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -70,6 +70,50 @@ def _use_rocm_sparse_triton( ) +@triton.jit +def _fit_kpool_indices_kernel( + token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS + TAIL_WIDTH] + out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + ti_stride0, + ti_stride1, + out_stride0, + out_stride1, + NUM_TOPK_TOKENS: tl.constexpr, + TAIL_WIDTH: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_TAIL: tl.constexpr, +): + row_ptr = token_indices_ptr + tl.program_id(0) * ti_stride0 + columns = tl.arange(0, BLOCK_T) + column_mask = columns < NUM_TOPK_TOKENS + history = tl.load(row_ptr + columns * ti_stride1, mask=column_mask, other=-1) + + tail_columns = tl.arange(0, BLOCK_TAIL) + tail = tl.load( + row_ptr + (NUM_TOPK_TOKENS + tail_columns) * ti_stride1, + mask=tail_columns < TAIL_WIDTH, + other=-1, + ) + + valid_history = tl.sum((history >= 0).to(tl.int32)) + valid_tail = tl.sum((tail >= 0).to(tl.int32)) + keep_history = tl.minimum(valid_history, NUM_TOPK_TOKENS - valid_tail) + + tail_offsets = columns - keep_history + tail_values = tl.load( + row_ptr + (NUM_TOPK_TOKENS + tail_offsets) * ti_stride1, + mask=column_mask & (tail_offsets >= 0) & (tail_offsets < TAIL_WIDTH), + other=-1, + ) + fitted = tl.where(columns < keep_history, history, tail_values) + fitted = tl.where(columns < keep_history + valid_tail, fitted, -1) + tl.store( + out_ptr + tl.program_id(0) * out_stride0 + columns * out_stride1, + fitted, + mask=column_mask, + ) + + def fit_kpool_indices_to_aiter( token_indices: torch.Tensor, topk_tokens: int ) -> torch.Tensor: @@ -79,20 +123,29 @@ def fit_kpool_indices_to_aiter( if token_indices.shape[1] == topk_tokens: return token_indices - history = token_indices[:, :topk_tokens] - tail = token_indices[:, topk_tokens:] - valid_history = (history >= 0).sum(dim=1) - valid_tail = (tail >= 0).sum(dim=1) - keep_history = torch.minimum(valid_history, topk_tokens - valid_tail) - - columns = torch.arange(topk_tokens, device=token_indices.device).unsqueeze(0) - tail_offsets = columns - keep_history.unsqueeze(1) - tail_values = torch.gather( - tail, 1, tail_offsets.clamp(min=0, max=tail.shape[1] - 1) + num_tokens, width = token_indices.shape + tail_width = width - topk_tokens + output = torch.empty( + (num_tokens, topk_tokens), + dtype=token_indices.dtype, + device=token_indices.device, ) - output = torch.where(columns < keep_history.unsqueeze(1), history, tail_values) - valid_output = columns < (keep_history + valid_tail).unsqueeze(1) - return torch.where(valid_output, output, -1) + if num_tokens == 0: + return output + + _fit_kpool_indices_kernel[(num_tokens,)]( + token_indices, + output, + token_indices.stride(0), + token_indices.stride(1), + output.stride(0), + output.stride(1), + NUM_TOPK_TOKENS=topk_tokens, + TAIL_WIDTH=tail_width, + BLOCK_T=triton.next_power_of_2(topk_tokens), + BLOCK_TAIL=triton.next_power_of_2(tail_width), + ) + return output @triton.jit