128 lines
3.7 KiB
Python
128 lines
3.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.platforms import current_platform
|
|
from vllm.triton_utils import tl, triton
|
|
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
|
|
_use_rocm_sparse_triton,
|
|
fit_kpool_indices_to_aiter,
|
|
)
|
|
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
|
_sparse_kv_row_offset,
|
|
_validate_dsv4_sparse_dims,
|
|
_validate_sparse_dims,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def _store_sparse_kv_row_offset_kernel(slot_ptr, output_ptr, stride: tl.constexpr):
|
|
slot = tl.load(slot_ptr)
|
|
tl.store(output_ptr, _sparse_kv_row_offset(slot, stride))
|
|
|
|
|
|
def test_fit_kpool_indices_preserves_tail_and_best_history():
|
|
token_indices = torch.tensor(
|
|
[
|
|
[10, 9, 8, 7, 6, 5, 100, 101],
|
|
[10, 9, 8, -1, -1, -1, 100, -1],
|
|
[-1, -1, -1, -1, -1, -1, -1, -1],
|
|
],
|
|
dtype=torch.int32,
|
|
)
|
|
|
|
fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=6)
|
|
|
|
assert fitted.tolist() == [
|
|
[10, 9, 8, 7, 100, 101],
|
|
[10, 9, 8, 100, -1, -1],
|
|
[-1, -1, -1, -1, -1, -1],
|
|
]
|
|
|
|
|
|
def test_fit_kpool_indices_exact_width_is_noop():
|
|
token_indices = torch.tensor([[3, 2, 1, -1]], dtype=torch.int32)
|
|
|
|
fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=4)
|
|
|
|
assert fitted.data_ptr() == token_indices.data_ptr()
|
|
|
|
|
|
def test_fit_kpool_indices_rejects_narrow_input():
|
|
with pytest.raises(ValueError, match="at least topk_tokens"):
|
|
fit_kpool_indices_to_aiter(
|
|
torch.zeros((1, 3), dtype=torch.int32), topk_tokens=4
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
(
|
|
"kv_cache_dtype",
|
|
"head_size",
|
|
"num_prefills",
|
|
"num_decodes",
|
|
"num_decode_tokens",
|
|
"max_query_len",
|
|
"expected",
|
|
),
|
|
[
|
|
("auto", 512, 1, 0, 0, 32, True),
|
|
("auto", 512, 1, 2, 2, 32, True),
|
|
("auto", 512, 0, 2, 2, 1, True),
|
|
("fp8", 512, 1, 0, 0, 32, False),
|
|
("auto", 576, 1, 0, 0, 32, False),
|
|
("auto", 512, 0, 2, 4, 2, False),
|
|
],
|
|
)
|
|
def test_rocm_sparse_triton_route(
|
|
kv_cache_dtype,
|
|
head_size,
|
|
num_prefills,
|
|
num_decodes,
|
|
num_decode_tokens,
|
|
max_query_len,
|
|
expected,
|
|
):
|
|
assert (
|
|
_use_rocm_sparse_triton(
|
|
kv_cache_dtype=kv_cache_dtype,
|
|
head_size=head_size,
|
|
kv_lora_rank=512,
|
|
num_prefills=num_prefills,
|
|
num_decodes=num_decodes,
|
|
num_decode_tokens=num_decode_tokens,
|
|
max_query_len=max_query_len,
|
|
)
|
|
is expected
|
|
)
|
|
|
|
|
|
def test_rocm_sparse_attention_accepts_glm_nope_dimensions():
|
|
_validate_sparse_dims(512, 512, 0, "test")
|
|
|
|
|
|
def test_rocm_sparse_attention_rejects_inconsistent_dimensions():
|
|
with pytest.raises(AssertionError, match="expected head_dim"):
|
|
_validate_sparse_dims(511, 512, 0, "test")
|
|
|
|
|
|
def test_dsv4_sparse_attention_keeps_layout_constraint():
|
|
_validate_dsv4_sparse_dims(512, 448, 64, "test")
|
|
with pytest.raises(AssertionError, match="expects 448 NoPE dims"):
|
|
_validate_dsv4_sparse_dims(512, 512, 0, "test")
|
|
|
|
|
|
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
|
|
def test_sparse_prefill_kv_row_offset_does_not_overflow_int32():
|
|
# GLM's 640-token pages cross the signed-int32 address boundary at block
|
|
# 6554 for a 512-element KV row. The production kernel must promote the
|
|
# slot before multiplying by the row stride.
|
|
slot = torch.tensor([6554 * 640], dtype=torch.int32, device="cuda")
|
|
output = torch.empty(1, dtype=torch.int64, device="cuda")
|
|
|
|
_store_sparse_kv_row_offset_kernel[(1,)](slot, output, stride=512)
|
|
|
|
assert output.item() == 6554 * 640 * 512
|