[Perf][GLM5.3-Flash] Use cooperative top-k for small GLM decode batches (#57327)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Chauncey committed
Sep 18, 2026 at 10:31 UTC
2bbdfcfcfea79db5a3a4a32a573e0998fd6a58eb
2 files changed
+63
-1
tests/models/glm5next/test_sparse_indexer_topk_dispatch.py
new
+41
@@ -0,0 +1,41 @@
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."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.models.glm5next.nvidia.sparse_indexer import _use_cooperative_topk
9
+from vllm.platforms import current_platform
10
+
11
+pytestmark = pytest.mark.skipif(
12
+ not current_platform.is_cuda(), reason="CUDA-only dispatch"
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)
21
+
22
+
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)
vllm/models/glm5next/nvidia/sparse_indexer.py
+22
-1
@@ -38,6 +38,22 @@ if current_platform.is_cuda_alike():
38
39
logger = init_logger(__name__)
40
41
+
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
+
57
# kpool write helper: form pools from the current token batch and compress them
58
# into the index K cache via the fused Triton kernel.
59
@@ -573,7 +589,12 @@ def sparse_attn_indexer_kpool(
589
(topk_workspace,) = workspace_manager.get_simultaneous(
590
((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
591
)
576
- torch.ops._C.persistent_topk(
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,