[GLM-5.3-Flash] Route kpool indexer top-k through the shared SparseIndexerTopk dispatcher (#57546)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Sep 20, 2026 at 19:17 UTC
bf01fc4a313cc90863c470f37a26e41fd7b16fd1
4 files changed
+94
-83
tests/kernels/test_top_k_per_row.py
+23
-8
@@ -1871,17 +1871,32 @@ def test_sparse_indexer_topk_backend_resolution() -> None:
1871
has_coop = _has_cooperative_topk()
1872
has_fi = _has_flashinfer_topk()
1873
1874
- # "auto" is exactly the pre-existing chain: cooperative -> persistent ->
1875
- # per_row. It must never select the opt-in backends by itself.
1874
+ # "auto" is the chain cooperative -> persistent -> per_row, with the
1875
+ # cooperative/persistent switch at cooperative_topk's cluster-wave row
1876
+ # limit. It must never select the opt-in backends by itself.
1877
if has_coop:
1877
- assert resolve("auto") == "cooperative"
1878
+ from vllm.model_executor.layers.indexer_topk import (
1879
+ AUTO_COOPERATIVE_MAX_ROWS,
1880
+ )
1881
+
1882
+ # cooperative for every batch within its cluster-wave row limit...
1883
+ assert (
1884
+ resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS) == "cooperative"
1885
+ )
1886
assert resolve("auto", k=512, num_rows=8) == "cooperative"
1879
- # Past cooperative's 64-row limit -> persistent, even where
1880
- # DeepSelect would be applicable.
1887
+ # ...persistent past the row limit, even where DeepSelect would be
1888
+ # applicable.
1889
+ assert (
1890
+ resolve("auto", k=512, num_rows=AUTO_COOPERATIVE_MAX_ROWS + 1)
1891
+ == "persistent"
1892
+ )
1893
assert resolve("auto", k=512, num_rows=128) == "persistent"
1882
- # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned:
1883
- # still cooperative at <= 64 rows...
1884
- assert resolve("auto", unaligned_logits, num_rows=64) == "cooperative"
1894
+ # 16B-aligned (cooperative TMA ok) but not DeepSelect-aligned: the
1895
+ # same row-limit logic applies.
1896
+ assert (
1897
+ resolve("auto", unaligned_logits, num_rows=AUTO_COOPERATIVE_MAX_ROWS)
1898
+ == "cooperative"
1899
+ )
1900
# ...and persistent past the row limit.
1901
assert resolve("auto", unaligned_logits, num_rows=128) == "persistent"
1902
# Unsupported topk (outside {512, 1024, 2048} for the workspace kernels)
tests/models/glm5next/test_sparse_indexer_topk_dispatch.py
+28
-26
@@ -1,11 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Focused unit tests for the GLM decode top-k dispatch."""
3
+"""Tests for the GLM kpool indexer's top-k backend wiring."""
4
5
import pytest
6
import torch
7
8
-from vllm.models.glm5next.nvidia.sparse_indexer import _use_cooperative_topk
8
+from vllm.config import VllmConfig, set_current_vllm_config
9
from vllm.platforms import current_platform
10
11
pytestmark = pytest.mark.skipif(
@@ -13,29 +13,31 @@ pytestmark = pytest.mark.skipif(
13
)
14
15
16
-@pytest.mark.parametrize("num_rows", [1, 64], ids=["1row", "64rows"])
17
-@torch.inference_mode()
18
-def test_cooperative_topk_is_selected_for_supported_batch(num_rows: int):
19
- logits = torch.empty(num_rows, 8192, dtype=torch.float32, device="cuda")
20
- assert _use_cooperative_topk(logits, 512, num_rows)
16
+def _require_deep_gemm() -> None:
17
+ from vllm.utils.deep_gemm import has_deep_gemm
18
19
+ if not has_deep_gemm():
20
+ pytest.skip("kpool indexer requires DeepGEMM")
21
23
-@pytest.mark.parametrize(
24
- ("num_rows", "select_k", "stride_padding"),
25
- [
26
- (65, 512, 0),
27
- (1, 256, 0),
28
- (1, 512, 1),
29
- ],
30
- ids=["too_many_rows", "unsupported_k", "unaligned_stride"],
31
-)
32
-@torch.inference_mode()
33
-def test_cooperative_topk_falls_back_on_unsupported_batch(
34
- num_rows: int, select_k: int, stride_padding: int
35
-):
36
- logits = torch.empty(
37
- num_rows, 8192 + stride_padding, dtype=torch.float32, device="cuda"
38
- )
39
- if stride_padding:
40
- logits = logits[:, :8192]
41
- assert not _use_cooperative_topk(logits, select_k, num_rows)
22
+
23
+@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"])
24
+def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None:
25
+ """The kpool indexer must read kernel_config.sparse_indexer_topk_backend
26
+ and hand it to the shared SparseIndexerTopk dispatcher, rather than
27
+ hard-coding its own cooperative/persistent/per_row choice."""
28
+ _require_deep_gemm()
29
+ from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool
30
+
31
+ cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend})
32
+ with set_current_vllm_config(cfg):
33
+ op = SparseAttnIndexerKpool(
34
+ k_cache=None,
35
+ quant_block_size=128,
36
+ scale_fmt="ue8m0",
37
+ topk_tokens=2048,
38
+ head_dim=128,
39
+ max_pool_len=4096,
40
+ max_total_seq_len=8192,
41
+ topk_indices_buffer=torch.empty(8, 2176, dtype=torch.int32, device="cuda"),
42
+ )
43
+ assert op.topk_backend == backend
vllm/model_executor/layers/indexer_topk.py
+23
-5
@@ -14,6 +14,12 @@ from vllm.v1.worker.workspace import current_workspace_manager
14
15
RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
16
17
+# cooperative_topk's hard row limit: one cluster wave covers at most this many
18
+# rows. "auto" uses cooperative_topk for every batch within the limit and falls
19
+# back to persistent_topk past it; explicit ``cooperative`` requests are
20
+# validated against the same bound.
21
+AUTO_COOPERATIVE_MAX_ROWS = 64
22
+
23
# ---------------------------------------------------------------------------
24
# DeepSelect (vllm._deepselect_C)
25
# ---------------------------------------------------------------------------
@@ -209,8 +215,13 @@ class SparseIndexerTopk(torch.nn.Module):
215
def _resolve_auto(
216
self, logits: torch.Tensor, topk_tokens: int, num_rows: int
217
) -> str:
212
- """The pre-existing priority chain: cooperative -> persistent ->
213
- per_row. deep_select/flashinfer/torch are opt-in only."""
218
+ """The priority chain: cooperative -> persistent -> per_row.
219
+ deep_select/flashinfer/torch are opt-in only.
220
+
221
+ cooperative_topk is preferred whenever it is applicable, i.e. within
222
+ its AUTO_COOPERATIVE_MAX_ROWS row limit; larger batches go to
223
+ persistent_topk.
224
+ """
225
if not self._cooperative_constraints(logits, topk_tokens, num_rows):
226
return "cooperative"
227
if self._is_cuda and topk_tokens in (512, 1024, 2048):
@@ -228,8 +239,10 @@ class SparseIndexerTopk(torch.nn.Module):
239
failures.append(
240
f"topk_tokens must be in (512, 1024, 2048), got {topk_tokens}"
241
)
231
- if num_rows > 64:
232
- failures.append(f"num_rows must be <= 64, got {num_rows}")
242
+ if num_rows > AUTO_COOPERATIVE_MAX_ROWS:
243
+ failures.append(
244
+ f"num_rows must be <= {AUTO_COOPERATIVE_MAX_ROWS}, got {num_rows}"
245
+ )
246
if logits.stride(0) % 4 != 0:
247
failures.append(
248
f"logits.stride(0) must be divisible by 4, got {logits.stride(0)}"
@@ -293,13 +306,18 @@ class SparseIndexerTopk(torch.nn.Module):
306
(topk_workspace,) = current_workspace_manager().get_simultaneous(
307
((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
308
)
309
+ # persistent_topk's last argument is the host-side "max seq_len
310
+ # across rows" bound (it gates the sampled_topk path and clamps
311
+ # per-row lengths), not the padded row width: logits is sized to
312
+ # max_model_len, so passing its width would incorrectly enable the
313
+ # long-row kernels on short-context batches.
314
torch.ops._C.persistent_topk(
315
logits,
316
seq_lens,
317
topk_indices,
318
topk_workspace,
319
topk_tokens,
302
- logits.shape[1],
320
+ max_seq_len,
321
)
322
elif backend == "flashinfer":
323
# Deferred: importing flashinfer initializes CUDA at import time.
vllm/models/glm5next/nvidia/sparse_indexer.py
+20
-44
@@ -10,6 +10,7 @@ from vllm.config import get_current_vllm_config_or_none
10
from vllm.forward_context import get_forward_context
11
from vllm.logger import init_logger
12
from vllm.model_executor.custom_op import CustomOp
13
+from vllm.model_executor.layers.indexer_topk import get_indexer_topk
14
from vllm.models.glm5next.common.sparse_indexer import (
15
RADIX_TOPK_WORKSPACE_SIZE,
16
_build_decode_scatter_indices,
@@ -39,21 +40,6 @@ if current_platform.is_cuda_alike():
40
logger = init_logger(__name__)
41
42
42
-def _use_cooperative_topk(
43
- logits: torch.Tensor,
44
- select_k: int,
45
- num_rows: int,
46
-) -> bool:
47
- return (
48
- current_platform.is_cuda()
49
- and select_k in (512, 1024, 2048)
50
- and num_rows <= 64
51
- and logits.stride(0) % 4 == 0
52
- and current_platform.has_device_capability(90)
53
- and not current_platform.is_device_capability_family(120)
54
- )
55
-
56
-
43
# kpool write helper: form pools from the current token batch and compress them
44
# into the index K cache via the fused Triton kernel.
45
@@ -134,6 +120,7 @@ def sparse_attn_indexer_kpool(
120
# path and when the tail cache is disabled.
121
tail_kv_cache: torch.Tensor | None = None,
122
tail_prefix: str | None = None,
123
+ topk_backend: str = "auto",
124
) -> torch.Tensor:
125
# careful! this will be None in dummy run
126
attn_metadata = get_forward_context().attn_metadata
@@ -584,35 +571,18 @@ def sparse_attn_indexer_kpool(
571
else:
572
topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
573
587
- if current_platform.is_cuda() and select_k in (512, 1024, 2048):
588
- workspace_manager = current_workspace_manager()
589
- (topk_workspace,) = workspace_manager.get_simultaneous(
590
- ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
591
- )
592
- topk_op = (
593
- torch.ops._C.cooperative_topk
594
- if _use_cooperative_topk(logits, select_k, num_rows)
595
- else torch.ops._C.persistent_topk
596
- )
597
- topk_op(
598
- logits,
599
- seq_lens,
600
- topk_dst,
601
- topk_workspace,
602
- select_k,
603
- attn_metadata_narrowed.max_seq_len,
604
- )
605
- else:
606
- torch.ops._C.top_k_per_row_decode(
607
- logits,
608
- next_n,
609
- seq_lens,
610
- topk_dst,
611
- num_rows,
612
- logits.stride(0),
613
- logits.stride(1),
614
- select_k,
615
- )
574
+ # Shared dispatcher with the DSA sparse indexer, so the kpool select
575
+ # honors kernel_config.sparse_indexer_topk_backend and picks between
576
+ # cooperative/persistent by batch size instead of always taking the
577
+ # same kernel.
578
+ get_indexer_topk(topk_backend)(
579
+ logits,
580
+ seq_lens,
581
+ next_n,
582
+ topk_dst,
583
+ select_k,
584
+ attn_metadata_narrowed.max_seq_len,
585
+ )
586
587
# Resolve to token-level indices in the output buffer.
588
if index_kpool > 1:
@@ -692,6 +662,11 @@ class SparseAttnIndexerKpool(CustomOp):
662
"Sparse Attention Indexer CUDA op requires DeepGEMM to be installed."
663
)
664
_cfg = get_current_vllm_config_or_none()
665
+ self.topk_backend = (
666
+ _cfg.kernel_config.sparse_indexer_topk_backend
667
+ if _cfg is not None
668
+ else "auto"
669
+ )
670
_parallel = _cfg.parallel_config if _cfg is not None else None
671
if (
672
_parallel is not None
@@ -766,4 +741,5 @@ class SparseAttnIndexerKpool(CustomOp):
741
positions,
742
self.tail_cache.kv_cache if self.tail_cache is not None else None,
743
self.tail_cache.prefix if self.tail_cache is not None else None,
744
+ self.topk_backend,
745
)