| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the GLM kpool indexer's top-k backend wiring.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.config import VllmConfig, set_current_vllm_config |
| 9 | from vllm.platforms import current_platform |
| 10 | |
| 11 | |
| 12 | def _require_deep_gemm() -> None: |
| 13 | from vllm.utils.deep_gemm import has_deep_gemm |
| 14 | |
| 15 | if not has_deep_gemm(): |
| 16 | pytest.skip("kpool indexer requires DeepGEMM") |
| 17 | |
| 18 | |
| 19 | def _build(indexer_cls, backend: str): |
| 20 | cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) |
| 21 | with set_current_vllm_config(cfg): |
| 22 | return indexer_cls( |
| 23 | k_cache=None, |
| 24 | quant_block_size=128, |
| 25 | scale_fmt="ue8m0", |
| 26 | topk_tokens=2048, |
| 27 | head_dim=128, |
| 28 | max_pool_len=4096, |
| 29 | max_total_seq_len=8192, |
| 30 | topk_indices_buffer=torch.empty(8, 2176, dtype=torch.int32, device="cuda"), |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only dispatch") |
| 35 | @pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"]) |
| 36 | def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None: |
| 37 | """The kpool indexer must read kernel_config.sparse_indexer_topk_backend |
| 38 | and hand it to the shared SparseIndexerTopk dispatcher, rather than |
| 39 | hard-coding its own cooperative/persistent/per_row choice.""" |
| 40 | _require_deep_gemm() |
| 41 | from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool |
| 42 | |
| 43 | assert _build(SparseAttnIndexerKpool, backend).topk_backend == backend |
| 44 | |
| 45 | |
| 46 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch") |
| 47 | @pytest.mark.parametrize("backend", ["auto", "aiter", "per_row", "torch"]) |
| 48 | def test_kpool_indexer_dispatches_through_shared_topk_backend_rocm( |
| 49 | backend: str, |
| 50 | ) -> None: |
| 51 | """The AMD kpool indexer must go through the same dispatcher, so the |
| 52 | AITER decode top-k is reachable from kernel_config.""" |
| 53 | from vllm.models.glm5next.amd.sparse_indexer import SparseAttnIndexerKpool |
| 54 | |
| 55 | assert _build(SparseAttnIndexerKpool, backend).topk_backend == backend |
| 56 | |
| 57 | |
| 58 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch") |
| 59 | @pytest.mark.parametrize( |
| 60 | "configured,full_cudagraph,expected", |
| 61 | [ |
| 62 | ("auto", False, "aiter"), |
| 63 | ("auto", True, "auto"), |
| 64 | ("aiter", True, "aiter"), |
| 65 | ("per_row", False, "per_row"), |
| 66 | ], |
| 67 | ) |
| 68 | def test_kpool_decode_topk_backend_never_auto_selects_aiter_under_full_cudagraph( |
| 69 | monkeypatch: pytest.MonkeyPatch, |
| 70 | configured: str, |
| 71 | full_cudagraph: bool, |
| 72 | expected: str, |
| 73 | ) -> None: |
| 74 | """A FULL cudagraph bakes in one kernel for every replay, so "auto" cannot |
| 75 | use the context length to pick AITER and must stay on the in-tree kernel. |
| 76 | Explicit backends keep working, and outside FULL the long-context heuristic |
| 77 | still narrows to AITER.""" |
| 78 | from vllm._aiter_ops import rocm_aiter_ops |
| 79 | from vllm.models.glm5next.amd import sparse_indexer |
| 80 | |
| 81 | monkeypatch.setattr( |
| 82 | rocm_aiter_ops, "is_indexer_top_k_enabled", lambda: True, raising=False |
| 83 | ) |
| 84 | monkeypatch.setattr( |
| 85 | rocm_aiter_ops, |
| 86 | "is_indexer_top_k_supported", |
| 87 | lambda **kwargs: True, |
| 88 | raising=False, |
| 89 | ) |
| 90 | |
| 91 | assert ( |
| 92 | sparse_indexer._kpool_decode_topk_backend( |
| 93 | configured, |
| 94 | num_rows=8, |
| 95 | max_valid_seq_len=64 * 1024, |
| 96 | select_k=512, |
| 97 | index_kpool=4, |
| 98 | full_cudagraph=full_cudagraph, |
| 99 | ) |
| 100 | == expected |
| 101 | ) |