iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[GLM-5.3-Flash] Route kpool indexer top-k through the shared SparseIndexerTopk dispatcher (#57546)

Browse Source 
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
This commit is contained in:
Chauncey
2026-09-20 11:17:59 +00:00
committed by GitHub
parent f648eed23d
commit bf01fc4a31
4 changed files with 94 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/test_top_k_per_row.py
+23 -8
View File
@@ -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)
tests/models/glm5next/test_sparse_indexer_topk_dispatch.py
+28 -26
View File
@@ -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
vllm/model_executor/layers/indexer_topk.py
+23 -5
View File
@@ -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.
vllm/models/glm5next/nvidia/sparse_indexer.py
+20 -44
View File
@@ -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,
)