[Perf][Kernel] Tune cooperative topk for medium batch-sizes (#53382)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Co-authored-by: Codex <codex@openai.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

csrc/libtorch_stable/cooperative_topk.cu

@@ -16,7 +16,9 @@ template <uint32_t TopK, uint32_t CS>
 void launch_cooperative_cluster(ct::CooperativeTopKParams<TopK>& params,
                                 size_t smem, cudaStream_t stream) {
   auto kernel = []() {
-    if constexpr (CS == 16) {
+    if constexpr (CS == 2) {
+      return &ct::cooperative_topk_cs2<TopK>;
+    } else if constexpr (CS == 16) {
       return &ct::cooperative_topk_cs16<TopK>;
     } else if constexpr (CS == 8) {
       return &ct::cooperative_topk_cs8<TopK>;
@@ -58,11 +60,9 @@ void launch_cooperative_topk_impl(const torch::stable::Tensor& logits,
   const cudaStream_t stream = get_current_cuda_stream();
 
   const uint32_t stride = static_cast<uint32_t>(logits.stride(0));
-  // 32 = max clusters for CS=4 (32 x 4 = 128 CTAs = 66% of SMs, leaves
-  // headroom)
   STD_TORCH_CHECK(
-      num_rows <= 32,
-      "cooperative_topk supports <=32 rows; use persistent_topk for "
+      num_rows <= 64,
+      "cooperative_topk supports <=64 rows; use persistent_topk for "
       "larger batches");
 
   STD_TORCH_CHECK(stride % 4 == 0,
@@ -96,8 +96,10 @@ void launch_cooperative_topk_impl(const torch::stable::Tensor& logits,
     launch_cooperative_cluster<TopK, 16>(params, ct::kSmemSize8, stream);
   } else if (num_rows <= 8) {
     launch_cooperative_cluster<TopK, 8>(params, ct::kSmemSize8, stream);
-  } else {
+  } else if (num_rows <= 33) {
     launch_cooperative_cluster<TopK, 4>(params, ct::kSmemSize4, stream);
+  } else {
+    launch_cooperative_cluster<TopK, 2>(params, ct::kSmemSize2, stream);
   }
 }
 #endif  // USE_ROCM

csrc/libtorch_stable/cooperative_topk.cuh

@@ -547,7 +547,8 @@ __device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
         {};  // tracks the parity for mbarrier wait/arrive protocol
     large_topk<TopK, CS, FusedSmem, true>(in, out, sl, phases, row_tie_ws);
   } else {
-    // Two-pass: only CS=4 in practice (CS=8 always fits in singlepass)
+    // Streaming path for cluster sizes whose per-block input exceeds the
+    // single-pass staging capacity.
     auto* smem = reinterpret_cast<Smem4*>(sr);
     if (tx < 2 * kStreamingStagesCS4) {
       mbarrier_init(&smem->barrier[0][tx],
@@ -560,6 +561,13 @@ __device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
   }
 }
 
+template <uint32_t TopK>
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
+    __cluster_dims__(1, 2, 1)
+        cooperative_topk_cs2(CooperativeTopKParams<TopK> params) {
+  cooperative_topk_body<TopK, 2>(params);
+}
+
 template <uint32_t TopK>
 __global__ void __launch_bounds__(hist4096::kBlockSize, 1)
     __cluster_dims__(1, 4, 1)
@@ -586,6 +594,7 @@ constexpr size_t kSmemSize4_sp = sizeof(SmemSinglePass);
 constexpr size_t kSmemSize4 =
     (kSmemSize4_base > kSmemSize4_sp ? kSmemSize4_base : kSmemSize4_sp) +
     sizeof(int32_t) * 2048 + 128;
+constexpr size_t kSmemSize2 = kSmemSize4;
 constexpr size_t kSmemSize8 =
     sizeof(SmemFused<kFusedStagesCS8>) + sizeof(int32_t) * 2048 + 128;
 

tests/kernels/test_top_k_per_row.py

@@ -60,9 +60,9 @@ def _run_topk_backend(
             logits, lengths, indices, workspace, top_k, max_seq_len
         )
     elif backend == "cooperative_topk":
-        if indices.shape[0] > 32:
+        if indices.shape[0] > 64:
             pytest.skip(
-                "cooperative_topk supports <=32 rows; "
+                "cooperative_topk supports <=64 rows; "
                 "persistent_topk covers larger batches"
             )
         if logits.stride(0) % 4 != 0:
@@ -881,6 +881,18 @@ def run_large_context_topk_test(
             Torch: {torch_vals.sort(descending=True)[0][:10]}"""
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("top_k", [512, 1024, 2048])
+def test_cooperative_topk_cs2(top_k: int) -> None:
+    """The 64-row dispatch uses the two-CTA cooperative kernel."""
+    run_large_context_topk_test(
+        batch_size=64,
+        seq_lens=[65536] * 64,
+        top_k=top_k,
+        backend="cooperative_topk",
+    )
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
 @pytest.mark.parametrize(
     "test_config",

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -617,7 +617,7 @@ def sparse_attn_indexer(
         use_cooperative_topk = (
             current_platform.is_cuda()
             and topk_tokens in (512, 1024, 2048)
-            and num_rows <= 32
+            and num_rows <= 64
             and logits.stride(0) % 4 == 0  # TMA 16-byte alignment
             and current_platform.has_device_capability(90)
             and not current_platform.is_device_capability_family(120)