diff --git a/csrc/libtorch_stable/cooperative_topk.cu b/csrc/libtorch_stable/cooperative_topk.cu index f388a9e6c8e7..6c1ea6c0ee33 100644 --- a/csrc/libtorch_stable/cooperative_topk.cu +++ b/csrc/libtorch_stable/cooperative_topk.cu @@ -16,7 +16,9 @@ template void launch_cooperative_cluster(ct::CooperativeTopKParams& params, size_t smem, cudaStream_t stream) { auto kernel = []() { - if constexpr (CS == 16) { + if constexpr (CS == 2) { + return &ct::cooperative_topk_cs2; + } else if constexpr (CS == 16) { return &ct::cooperative_topk_cs16; } else if constexpr (CS == 8) { return &ct::cooperative_topk_cs8; @@ -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(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(params, ct::kSmemSize8, stream); } else if (num_rows <= 8) { launch_cooperative_cluster(params, ct::kSmemSize8, stream); - } else { + } else if (num_rows <= 33) { launch_cooperative_cluster(params, ct::kSmemSize4, stream); + } else { + launch_cooperative_cluster(params, ct::kSmemSize2, stream); } } #endif // USE_ROCM diff --git a/csrc/libtorch_stable/cooperative_topk.cuh b/csrc/libtorch_stable/cooperative_topk.cuh index 3312c127a762..f34664a66f7d 100644 --- a/csrc/libtorch_stable/cooperative_topk.cuh +++ b/csrc/libtorch_stable/cooperative_topk.cuh @@ -547,7 +547,8 @@ __device__ void cooperative_topk_body(CooperativeTopKParams params) { {}; // tracks the parity for mbarrier wait/arrive protocol large_topk(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(sr); if (tx < 2 * kStreamingStagesCS4) { mbarrier_init(&smem->barrier[0][tx], @@ -560,6 +561,13 @@ __device__ void cooperative_topk_body(CooperativeTopKParams params) { } } +template +__global__ void __launch_bounds__(hist4096::kBlockSize, 1) + __cluster_dims__(1, 2, 1) + cooperative_topk_cs2(CooperativeTopKParams params) { + cooperative_topk_body(params); +} + template __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) + sizeof(int32_t) * 2048 + 128; diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index bb283e26a25a..a1e27d2e50a3 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/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", diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index bad32f520ab5..bb3572b17876 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/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)