44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for the GLM kpool indexer's top-k backend wiring."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.platforms import current_platform
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
not current_platform.is_cuda(), reason="CUDA-only dispatch"
|
|
)
|
|
|
|
|
|
def _require_deep_gemm() -> None:
|
|
from vllm.utils.deep_gemm import has_deep_gemm
|
|
|
|
if not has_deep_gemm():
|
|
pytest.skip("kpool indexer requires DeepGEMM")
|
|
|
|
|
|
@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"])
|
|
def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None:
|
|
"""The kpool indexer must read kernel_config.sparse_indexer_topk_backend
|
|
and hand it to the shared SparseIndexerTopk dispatcher, rather than
|
|
hard-coding its own cooperative/persistent/per_row choice."""
|
|
_require_deep_gemm()
|
|
from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool
|
|
|
|
cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
|
|
with set_current_vllm_config(cfg):
|
|
op = SparseAttnIndexerKpool(
|
|
k_cache=None,
|
|
quant_block_size=128,
|
|
scale_fmt="ue8m0",
|
|
topk_tokens=2048,
|
|
head_dim=128,
|
|
max_pool_len=4096,
|
|
max_total_seq_len=8192,
|
|
topk_indices_buffer=torch.empty(8, 2176, dtype=torch.int32, device="cuda"),
|
|
)
|
|
assert op.topk_backend == backend
|