[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>
Roberto L. Castro committed
Aug 31, 2026 at 22:44 UTC
f5c3cc240bc1fc0519694fe689edb87dee8dfd92
4 files changed
+33
-10
csrc/libtorch_stable/cooperative_topk.cu
+8
-6
@@ -16,7 +16,9 @@ template <uint32_t TopK, uint32_t CS>
16
void launch_cooperative_cluster(ct::CooperativeTopKParams<TopK>& params,
17
size_t smem, cudaStream_t stream) {
18
auto kernel = []() {
19
- if constexpr (CS == 16) {
19
+ if constexpr (CS == 2) {
20
+ return &ct::cooperative_topk_cs2<TopK>;
21
+ } else if constexpr (CS == 16) {
22
return &ct::cooperative_topk_cs16<TopK>;
23
} else if constexpr (CS == 8) {
24
return &ct::cooperative_topk_cs8<TopK>;
@@ -58,11 +60,9 @@ void launch_cooperative_topk_impl(const torch::stable::Tensor& logits,
60
const cudaStream_t stream = get_current_cuda_stream();
61
62
const uint32_t stride = static_cast<uint32_t>(logits.stride(0));
61
- // 32 = max clusters for CS=4 (32 x 4 = 128 CTAs = 66% of SMs, leaves
62
- // headroom)
63
STD_TORCH_CHECK(
64
- num_rows <= 32,
65
- "cooperative_topk supports <=32 rows; use persistent_topk for "
64
+ num_rows <= 64,
65
+ "cooperative_topk supports <=64 rows; use persistent_topk for "
66
"larger batches");
67
68
STD_TORCH_CHECK(stride % 4 == 0,
@@ -96,8 +96,10 @@ void launch_cooperative_topk_impl(const torch::stable::Tensor& logits,
96
launch_cooperative_cluster<TopK, 16>(params, ct::kSmemSize8, stream);
97
} else if (num_rows <= 8) {
98
launch_cooperative_cluster<TopK, 8>(params, ct::kSmemSize8, stream);
99
- } else {
99
+ } else if (num_rows <= 33) {
100
launch_cooperative_cluster<TopK, 4>(params, ct::kSmemSize4, stream);
101
+ } else {
102
+ launch_cooperative_cluster<TopK, 2>(params, ct::kSmemSize2, stream);
103
}
104
}
105
#endif // USE_ROCM
csrc/libtorch_stable/cooperative_topk.cuh
+10
-1
@@ -547,7 +547,8 @@ __device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
547
{}; // tracks the parity for mbarrier wait/arrive protocol
548
large_topk<TopK, CS, FusedSmem, true>(in, out, sl, phases, row_tie_ws);
549
} else {
550
- // Two-pass: only CS=4 in practice (CS=8 always fits in singlepass)
550
+ // Streaming path for cluster sizes whose per-block input exceeds the
551
+ // single-pass staging capacity.
552
auto* smem = reinterpret_cast<Smem4*>(sr);
553
if (tx < 2 * kStreamingStagesCS4) {
554
mbarrier_init(&smem->barrier[0][tx],
@@ -560,6 +561,13 @@ __device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
561
}
562
}
563
564
+template <uint32_t TopK>
565
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
566
+ __cluster_dims__(1, 2, 1)
567
+ cooperative_topk_cs2(CooperativeTopKParams<TopK> params) {
568
+ cooperative_topk_body<TopK, 2>(params);
569
+}
570
+
571
template <uint32_t TopK>
572
__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
573
__cluster_dims__(1, 4, 1)
@@ -586,6 +594,7 @@ constexpr size_t kSmemSize4_sp = sizeof(SmemSinglePass);
594
constexpr size_t kSmemSize4 =
595
(kSmemSize4_base > kSmemSize4_sp ? kSmemSize4_base : kSmemSize4_sp) +
596
sizeof(int32_t) * 2048 + 128;
597
+constexpr size_t kSmemSize2 = kSmemSize4;
598
constexpr size_t kSmemSize8 =
599
sizeof(SmemFused<kFusedStagesCS8>) + sizeof(int32_t) * 2048 + 128;
600
tests/kernels/test_top_k_per_row.py
+14
-2
@@ -60,9 +60,9 @@ def _run_topk_backend(
60
logits, lengths, indices, workspace, top_k, max_seq_len
61
)
62
elif backend == "cooperative_topk":
63
- if indices.shape[0] > 32:
63
+ if indices.shape[0] > 64:
64
pytest.skip(
65
- "cooperative_topk supports <=32 rows; "
65
+ "cooperative_topk supports <=64 rows; "
66
"persistent_topk covers larger batches"
67
)
68
if logits.stride(0) % 4 != 0:
@@ -881,6 +881,18 @@ def run_large_context_topk_test(
881
Torch: {torch_vals.sort(descending=True)[0][:10]}"""
882
883
884
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
885
+@pytest.mark.parametrize("top_k", [512, 1024, 2048])
886
+def test_cooperative_topk_cs2(top_k: int) -> None:
887
+ """The 64-row dispatch uses the two-CTA cooperative kernel."""
888
+ run_large_context_topk_test(
889
+ batch_size=64,
890
+ seq_lens=[65536] * 64,
891
+ top_k=top_k,
892
+ backend="cooperative_topk",
893
+ )
894
+
895
+
896
@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
897
@pytest.mark.parametrize(
898
"test_config",
vllm/model_executor/layers/sparse_attn_indexer.py
+1
-1
@@ -617,7 +617,7 @@ def sparse_attn_indexer(
617
use_cooperative_topk = (
618
current_platform.is_cuda()
619
and topk_tokens in (512, 1024, 2048)
620
- and num_rows <= 32
620
+ and num_rows <= 64
621
and logits.stride(0) % 4 == 0 # TMA 16-byte alignment
622
and current_platform.has_device_capability(90)
623
and not current_platform.is_device_capability_family(120)