From 2bbdfcfcfea79db5a3a4a32a573e0998fd6a58eb Mon Sep 17 00:00:00 2001
From: Chauncey <chaunceyjiang@gmail.com>
Date: Fri, 18 Sep 2026 10:31:35 +0800
Subject: [PATCH] [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>
---
 .../test_sparse_indexer_topk_dispatch.py      | 41 +++++++++++++++++++
 vllm/models/glm5next/nvidia/sparse_indexer.py | 23 ++++++++++-
 2 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 tests/models/glm5next/test_sparse_indexer_topk_dispatch.py

diff --git a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py
new file mode 100644
index 0000000000..83dfdc9f6f
--- /dev/null
+++ b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py
@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Focused unit tests for the GLM decode top-k dispatch."""
+
+import pytest
+import torch
+
+from vllm.models.glm5next.nvidia.sparse_indexer import _use_cooperative_topk
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="CUDA-only dispatch"
+)
+
+
+@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)
+
+
+@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)
diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py
index 1fc42c9e41..9dab659d4c 100644
--- a/vllm/models/glm5next/nvidia/sparse_indexer.py
+++ b/vllm/models/glm5next/nvidia/sparse_indexer.py
@@ -38,6 +38,22 @@ 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.
 
@@ -573,7 +589,12 @@ def sparse_attn_indexer_kpool(
             (topk_workspace,) = workspace_manager.get_simultaneous(
                 ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
             )
-            torch.ops._C.persistent_topk(
+            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,