[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>

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)

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,