02e55d2af3
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Signed-off-by: Dakai An <dakaian108@gmail.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Validate the registry reference kernel against runtime dispatch."""
|
|
|
|
import pytest
|
|
|
|
from vllm.platforms import current_platform
|
|
|
|
if not current_platform.is_cuda_alike():
|
|
pytest.skip("NVIDIA dispatch tests require CUDA", allow_module_level=True)
|
|
|
|
from vllm.v1.worker.block_table import ComputeSlotMappingKernel
|
|
|
|
|
|
@pytest.mark.parametrize("kv_cache_block_size", [256, 64, 8, 4])
|
|
def test_compute_slot_mapping_warmup_matches_runtime_specializations(
|
|
kv_cache_block_size: int,
|
|
) -> None:
|
|
kernel = ComputeSlotMappingKernel()
|
|
kwargs = dict(
|
|
kv_cache_block_size=kv_cache_block_size,
|
|
total_cp_world_size=1,
|
|
total_cp_rank=0,
|
|
cp_kv_cache_interleave_size=1,
|
|
block_table_stride=32768,
|
|
)
|
|
expected = kernel.CompileKey(
|
|
kv_cache_block_size=kv_cache_block_size,
|
|
total_cp_world_size=1,
|
|
total_cp_rank=0,
|
|
cp_kv_cache_interleave_size=1,
|
|
block_table_stride=16,
|
|
)
|
|
|
|
assert kernel.dispatch(**kwargs) == expected
|
|
assert kernel.get_warmup_keys(**kwargs) == [expected]
|