diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index db49b8091b..b6c469181f 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -1871,17 +1871,32 @@ def test_sparse_indexer_topk_backend_resolution() -> None: has_coop = _has_cooperative_topk() has_fi = _has_flashinfer_topk() - # "auto" is exactly the pre-existing chain: cooperative -> persistent -> - # per_row. It must never select the opt-in backends by itself. + # "auto" is the chain cooperative -> persistent -> per_row, with the + # cooperative/persistent switch at cooperative_topk's cluster-wave row + # limit. It must never select the opt-in backends by itself. if has_coop: - assert resolve("auto") == "cooperative" + from vllm.model_executor.layers.indexer_topk import ( + AUTO_COOPERATIVE_MAX_ROWS, + ) + + # cooperative for every batch within its cluster-wave row limit... + assert ( + resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS) == "cooperative" + ) assert resolve("auto", k=512, num_rows=8) == "cooperative" - # Past cooperative's 64-row limit -> persistent, even where - # DeepSelect would be applicable. + # ...persistent past the row limit, even where DeepSelect would be + # applicable. + assert ( + resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS + 1) + == "persistent" + ) assert resolve("auto", k=512, num_rows=128) == "persistent" - # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: - # still cooperative at <= 64 rows... - assert resolve("auto", unaligned_logits, num_rows=64) == "cooperative" + # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: the + # same row-limit logic applies. + assert ( + resolve("auto", unaligned_logits, num_rows=AUTO_COOPERATIVE_MAX_ROWS) + == "cooperative" + ) # ...and persistent past the row limit. assert resolve("auto", unaligned_logits, num_rows=128) == "persistent" # Unsupported topk (outside {512, 1024, 2048} for the workspace kernels) diff --git a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py index 83dfdc9f6f..9f94da459c 100644 --- a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py +++ b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Focused unit tests for the GLM decode top-k dispatch.""" +"""Tests for the GLM kpool indexer's top-k backend wiring.""" import pytest import torch -from vllm.models.glm5next.nvidia.sparse_indexer import _use_cooperative_topk +from vllm.config import VllmConfig, set_current_vllm_config from vllm.platforms import current_platform pytestmark = pytest.mark.skipif( @@ -13,29 +13,31 @@ pytestmark = pytest.mark.skipif( ) -@pytest.mark.parametrize("num_rows", [1, 64], ids=["1row", "64rows"]) -@torch.inference_mode() -def test_cooperative_topk_is_selected_for_supported_batch(num_rows: int): - logits = torch.empty(num_rows, 8192, dtype=torch.float32, device="cuda") - assert _use_cooperative_topk(logits, 512, num_rows) +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( - ("num_rows", "select_k", "stride_padding"), - [ - (65, 512, 0), - (1, 256, 0), - (1, 512, 1), - ], - ids=["too_many_rows", "unsupported_k", "unaligned_stride"], -) -@torch.inference_mode() -def test_cooperative_topk_falls_back_on_unsupported_batch( - num_rows: int, select_k: int, stride_padding: int -): - logits = torch.empty( - num_rows, 8192 + stride_padding, dtype=torch.float32, device="cuda" - ) - if stride_padding: - logits = logits[:, :8192] - assert not _use_cooperative_topk(logits, select_k, num_rows) +@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 diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py index 2ddf710703..c3422445b2 100644 --- a/vllm/model_executor/layers/indexer_topk.py +++ b/vllm/model_executor/layers/indexer_topk.py @@ -14,6 +14,12 @@ from vllm.v1.worker.workspace import current_workspace_manager RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 +# cooperative_topk's hard row limit: one cluster wave covers at most this many +# rows. "auto" uses cooperative_topk for every batch within the limit and falls +# back to persistent_topk past it; explicit ``cooperative`` requests are +# validated against the same bound. +AUTO_COOPERATIVE_MAX_ROWS = 64 + # --------------------------------------------------------------------------- # DeepSelect (vllm._deepselect_C) # --------------------------------------------------------------------------- @@ -209,8 +215,13 @@ class SparseIndexerTopk(torch.nn.Module): def _resolve_auto( self, logits: torch.Tensor, topk_tokens: int, num_rows: int ) -> str: - """The pre-existing priority chain: cooperative -> persistent -> - per_row. deep_select/flashinfer/torch are opt-in only.""" + """The priority chain: cooperative -> persistent -> per_row. + deep_select/flashinfer/torch are opt-in only. + + cooperative_topk is preferred whenever it is applicable, i.e. within + its AUTO_COOPERATIVE_MAX_ROWS row limit; larger batches go to + persistent_topk. + """ if not self._cooperative_constraints(logits, topk_tokens, num_rows): return "cooperative" if self._is_cuda and topk_tokens in (512, 1024, 2048): @@ -228,8 +239,10 @@ class SparseIndexerTopk(torch.nn.Module): failures.append( f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}" ) - if num_rows > 64: - failures.append(f"num_rows must be <= 64, got {num_rows}") + if num_rows > AUTO_COOPERATIVE_MAX_ROWS: + failures.append( + f"num_rows must be <= {AUTO_COOPERATIVE_MAX_ROWS}, got {num_rows}" + ) if logits.stride(0) % 4 != 0: failures.append( f"logits.stride(0) must be divisible by 4, got {logits.stride(0)}" @@ -293,13 +306,18 @@ class SparseIndexerTopk(torch.nn.Module): (topk_workspace,) = current_workspace_manager().get_simultaneous( ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), ) + # persistent_topk's last argument is the host-side "max seq_len + # across rows" bound (it gates the sampled_topk path and clamps + # per-row lengths), not the padded row width: logits is sized to + # max_model_len, so passing its width would incorrectly enable the + # long-row kernels on short-context batches. torch.ops._C.persistent_topk( logits, seq_lens, topk_indices, topk_workspace, topk_tokens, - logits.shape[1], + max_seq_len, ) elif backend == "flashinfer": # Deferred: importing flashinfer initializes CUDA at import time. diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index f7204949fc..c80e442015 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -10,6 +10,7 @@ from vllm.config import get_current_vllm_config_or_none from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.layers.indexer_topk import get_indexer_topk from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, _build_decode_scatter_indices, @@ -39,21 +40,6 @@ if current_platform.is_cuda_alike(): logger = init_logger(__name__) -def _use_cooperative_topk( - logits: torch.Tensor, - select_k: int, - num_rows: int, -) -> bool: - return ( - current_platform.is_cuda() - and select_k in (512, 1024, 2048) - and num_rows <= 64 - and logits.stride(0) % 4 == 0 - and current_platform.has_device_capability(90) - and not current_platform.is_device_capability_family(120) - ) - - # kpool write helper: form pools from the current token batch and compress them # into the index K cache via the fused Triton kernel. @@ -134,6 +120,7 @@ def sparse_attn_indexer_kpool( # path and when the tail cache is disabled. tail_kv_cache: torch.Tensor | None = None, tail_prefix: str | None = None, + topk_backend: str = "auto", ) -> torch.Tensor: # careful! this will be None in dummy run attn_metadata = get_forward_context().attn_metadata @@ -584,35 +571,18 @@ def sparse_attn_indexer_kpool( else: topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - if current_platform.is_cuda() and select_k in (512, 1024, 2048): - workspace_manager = current_workspace_manager() - (topk_workspace,) = workspace_manager.get_simultaneous( - ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), - ) - topk_op = ( - torch.ops._C.cooperative_topk - if _use_cooperative_topk(logits, select_k, num_rows) - else torch.ops._C.persistent_topk - ) - topk_op( - logits, - seq_lens, - topk_dst, - topk_workspace, - select_k, - attn_metadata_narrowed.max_seq_len, - ) - else: - torch.ops._C.top_k_per_row_decode( - logits, - next_n, - seq_lens, - topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), - select_k, - ) + # Shared dispatcher with the DSA sparse indexer, so the kpool select + # honors kernel_config.sparse_indexer_topk_backend and picks between + # cooperative/persistent by batch size instead of always taking the + # same kernel. + get_indexer_topk(topk_backend)( + logits, + seq_lens, + next_n, + topk_dst, + select_k, + attn_metadata_narrowed.max_seq_len, + ) # Resolve to token-level indices in the output buffer. if index_kpool > 1: @@ -692,6 +662,11 @@ class SparseAttnIndexerKpool(CustomOp): "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed." ) _cfg = get_current_vllm_config_or_none() + self.topk_backend = ( + _cfg.kernel_config.sparse_indexer_topk_backend + if _cfg is not None + else "auto" + ) _parallel = _cfg.parallel_config if _cfg is not None else None if ( _parallel is not None @@ -766,4 +741,5 @@ class SparseAttnIndexerKpool(CustomOp): positions, self.tail_cache.kv_cache if self.tail_cache is not None else None, self.tail_cache.prefix if self.tail_cache is not None else None, + self.topk_backend, )