@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 855cd4d787

[Perf][DSv4/DSv3.2] Add cluster-cooperative topK kernel for low-latency scenarios (#43008)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

Roberto L. Castro committed Jun 24, 2026 at 01:11 UTC 855cd4d787608b9bdecdc491d51fefadd3fc67dd
10 files changed +1568 -45
.buildkite/test_areas/kernels.yaml
+2
@@ -47,8 +47,10 @@ steps:
47 - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
48 - vllm/models/deepseek_v4/common/ops/
49 - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
50 + - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations
51 commands:
52 - pytest -v -s kernels/test_fused_deepseek_v4_*.py
53 + - pytest -v -s kernels/test_top_k_per_row.py
54
55 - label: Deepseek V4 Kernel Test (B200)
56 key: deepseek-v4-kernel-test-b200
CMakeLists.txt
+30
@@ -382,6 +382,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
382 "csrc/libtorch_stable/custom_all_reduce.cu"
383 "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
384
385 + if(VLLM_GPU_LANG STREQUAL "CUDA" AND
386 + DEFINED CMAKE_CUDA_COMPILER_VERSION AND
387 + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0)
388 +
389 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
390 + cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
391 + "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
392 + else()
393 + cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
394 + "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
395 + endif()
396 +
397 + if(COOPERATIVE_TOPK_ARCHS)
398 + list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1")
399 +
400 + endif()
401 + endif()
402 +
403 if(VLLM_GPU_LANG STREQUAL "CUDA")
404 SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
405
@@ -498,6 +516,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
516 SRCS "${VLLM_STABLE_EXT_SRC}"
517 CUDA_ARCHS "${CUDA_ARCHS}")
518
519 + if(COOPERATIVE_TOPK_ARCHS)
520 + list(APPEND VLLM_STABLE_EXT_SRC
521 + "csrc/libtorch_stable/cooperative_topk.cu")
522 + set_gencode_flags_for_srcs(
523 + SRCS "csrc/libtorch_stable/cooperative_topk.cu"
524 + CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}")
525 + endif()
526 +
527 # Only build Marlin kernels if we are building for at least some compatible archs.
528 # Keep building Marlin for 9.0 as there are some group sizes and shapes that
529 # are not supported by Machete yet.
@@ -1049,6 +1075,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1075 target_compile_definitions(_C_stable_libtorch PRIVATE
1076 TORCH_TARGET_VERSION=0x020B000000000000ULL)
1077 target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
1078 + if(COOPERATIVE_TOPK_ARCHS)
1079 + target_compile_definitions(_C_stable_libtorch PRIVATE
1080 + VLLM_ENABLE_COOPERATIVE_TOPK=1)
1081 + endif()
1082 # Needed by CUTLASS kernels
1083 target_compile_definitions(_C_stable_libtorch PRIVATE
1084 CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
csrc/libtorch_stable/cooperative_topk.cu new
+146
@@ -0,0 +1,146 @@
1 +// Cooperative cluster TopK for DeepSeek V3 sparse attention indexer.
2 +// See cooperative_topk.cuh for kernel implementation.
3 +
4 +#include <cuda_runtime.h>
5 +
6 +#include "torch_utils.h"
7 +
8 +#ifndef USE_ROCM
9 + #include "cooperative_topk.cuh"
10 +namespace ct = vllm::cooperative;
11 +namespace hist4096 = vllm::topk_histogram_4096;
12 +#endif
13 +
14 +#ifndef USE_ROCM
15 +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) {
20 + return &ct::cooperative_topk_cs16<TopK>;
21 + } else if constexpr (CS == 8) {
22 + return &ct::cooperative_topk_cs8<TopK>;
23 + } else {
24 + static_assert(CS == 4, "unsupported cooperative_topk cluster size");
25 + return &ct::cooperative_topk_cs4<TopK>;
26 + }
27 + }();
28 + if constexpr (CS > 8) {
29 + cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed,
30 + 1);
31 + }
32 + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
33 + smem);
34 +
35 + cudaLaunchConfig_t cfg = {};
36 + cfg.gridDim = dim3(params.num_rows, CS);
37 + cfg.blockDim = dim3(hist4096::kBlockSize);
38 + cfg.dynamicSmemBytes = smem;
39 + cfg.stream = stream;
40 + cudaLaunchAttribute attrs[1];
41 + attrs[0].id = cudaLaunchAttributeClusterDimension;
42 + attrs[0].val.clusterDim = {1, CS, 1};
43 + cfg.numAttrs = 1;
44 + cfg.attrs = attrs;
45 + cudaError_t err = cudaLaunchKernelEx(&cfg, kernel, params);
46 + STD_TORCH_CHECK(err == cudaSuccess,
47 + "cooperative_topk launch failed: ", cudaGetErrorString(err));
48 +}
49 +
50 +template <uint32_t TopK>
51 +void launch_cooperative_topk_impl(const torch::stable::Tensor& logits,
52 + const torch::stable::Tensor& lengths,
53 + torch::stable::Tensor& output,
54 + torch::stable::Tensor& workspace,
55 + int64_t max_seq_len) {
56 + (void)max_seq_len; // Kept for signature parity with persistent_topk.
57 + const int64_t num_rows = logits.size(0);
58 + const cudaStream_t stream = get_current_cuda_stream();
59 +
60 + 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 "
66 + "larger batches");
67 +
68 + STD_TORCH_CHECK(stride % 4 == 0,
69 + "cooperative_topk: stride must be multiple of 4 for TMA "
70 + "alignment, got stride (max_model_len)=",
71 + stride);
72 +
73 + STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
74 + STD_TORCH_CHECK(
75 + workspace.scalar_type() == torch::headeronly::ScalarType::Byte,
76 + "workspace must be uint8");
77 +
78 + ct::CooperativeTopKParams<TopK> params;
79 + params.input = logits.const_data_ptr<float>();
80 + params.output = output.mutable_data_ptr<int32_t>();
81 + params.lengths = lengths.const_data_ptr<int32_t>();
82 + params.num_rows = static_cast<uint32_t>(num_rows);
83 + params.stride = stride;
84 + params.tie_ws =
85 + reinterpret_cast<hist4096::Tie*>(workspace.mutable_data_ptr<uint8_t>());
86 +
87 + constexpr uint32_t kTieWsPerRow =
88 + TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK;
89 + STD_TORCH_CHECK(
90 + workspace.size(0) >=
91 + static_cast<int64_t>(num_rows * kTieWsPerRow * sizeof(hist4096::Tie)),
92 + "workspace too small");
93 +
94 + const bool supports_cluster16 = get_device_prop()->major >= 10;
95 + if (num_rows <= 4 && supports_cluster16) {
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 {
100 + launch_cooperative_cluster<TopK, 4>(params, ct::kSmemSize4, stream);
101 + }
102 +}
103 +#endif // USE_ROCM
104 +
105 +void cooperative_topk(const torch::stable::Tensor& logits,
106 + const torch::stable::Tensor& lengths,
107 + torch::stable::Tensor& output,
108 + torch::stable::Tensor& workspace, int64_t k,
109 + int64_t max_seq_len) {
110 +#ifndef USE_ROCM
111 + STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
112 + STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
113 + STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
114 + STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float,
115 + "Only float32 supported");
116 + STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int,
117 + "lengths must be int32");
118 + STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int,
119 + "output must be int32");
120 + STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
121 + STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
122 + "lengths must be 1D or 2D");
123 + STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
124 + STD_TORCH_CHECK(output.dim() == 2, "output must be 2D");
125 + const int64_t num_rows = logits.size(0);
126 + STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
127 + STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
128 + "output size mismatch");
129 + STD_TORCH_CHECK(
130 + k == 512 || k == 1024 || k == 2048,
131 + "cooperative_topk supports k=512, k=1024, or k=2048, got k=", k);
132 +
133 + if (k == 512) {
134 + launch_cooperative_topk_impl<512>(logits, lengths, output, workspace,
135 + max_seq_len);
136 + } else if (k == 1024) {
137 + launch_cooperative_topk_impl<1024>(logits, lengths, output, workspace,
138 + max_seq_len);
139 + } else {
140 + launch_cooperative_topk_impl<2048>(logits, lengths, output, workspace,
141 + max_seq_len);
142 + }
143 +#else
144 + STD_TORCH_CHECK(false, "cooperative_topk is not supported on ROCm");
145 +#endif
146 +}
csrc/libtorch_stable/cooperative_topk.cuh new
+593
@@ -0,0 +1,593 @@
1 +/*
2 + * Cooperative TopK kernel for DSA Indexer
3 + */
4 +
5 +#ifndef COOPERATIVE_TOPK_CUH_
6 +#define COOPERATIVE_TOPK_CUH_
7 +
8 +#include <cooperative_groups.h>
9 +#include <cuda.h>
10 +#include <cuda_fp16.h>
11 +#include <cuda_runtime.h>
12 +#include <cuda/ptx>
13 +#include <algorithm>
14 +#include <cstdint>
15 +
16 +#include "topk_histogram_4096.cuh"
17 +
18 +namespace vllm {
19 +namespace cooperative {
20 +
21 +namespace hist4096 = topk_histogram_4096;
22 +
23 +constexpr uint32_t kHistBits = 10;
24 +constexpr uint32_t kHistBins = 1 << kHistBits;
25 +constexpr uint32_t kMaxTopK = 2048;
26 +
27 +constexpr uint32_t kElemPerStage = 16;
28 +constexpr uint32_t kSizePerStage =
29 + kElemPerStage * hist4096::kBlockSize; // 16384
30 +
31 +// CS=4 two-pass path uses two TMA stages as a double buffer.
32 +constexpr uint32_t kStreamingStagesCS4 = 2;
33 +// CS=8/16 fused paths keep all loaded TMA stages resident in smem.
34 +constexpr uint32_t kFusedStagesCS8 = 2;
35 +constexpr uint32_t kFusedStagesCS16 = 2;
36 +
37 +// CS=4 single-pass path
38 +constexpr uint32_t kMaxSinglePassStages = 3;
39 +constexpr uint32_t kMaxSinglePassPerBlock =
40 + kMaxSinglePassStages * kSizePerStage; // 49152
41 +
42 +template <uint32_t TopK = 1024>
43 +struct CooperativeTopKParams {
44 + const float* __restrict__ input;
45 + int32_t* __restrict__ output;
46 + const int32_t* __restrict__ lengths;
47 + hist4096::Tie* __restrict__ tie_ws; // per-row tie workspace, see
48 + // kTieWsPerRow
49 + uint32_t num_rows, stride;
50 +};
51 +
52 +// ============================================================================
53 +// Cooperative helpers
54 +// ============================================================================
55 +
56 +// only CS adjacent lanes participate (sub-warp reduce), in opposite to
57 +// warp_reduce_sum_full
58 +template <uint32_t N>
59 +__device__ __forceinline__ uint32_t warp_reduce_sum_subN(uint32_t v) {
60 +#pragma unroll
61 + for (uint32_t m = N >> 1; m > 0; m >>= 1)
62 + v += __shfl_xor_sync(0xFFFFFFFF, v, m, 32);
63 + return v;
64 +}
65 +
66 +// ============================================================================
67 +// Helpers
68 +// ============================================================================
69 +
70 +__device__ __forceinline__ uint32_t extract_coarse_bin(float x) {
71 + return hist4096::extract_coarse_bin_N<kHistBits>(x);
72 +}
73 +
74 +__device__ __forceinline__ void mbarrier_init(uint64_t* a, uint32_t n) {
75 + cuda::ptx::mbarrier_init(a, n);
76 +}
77 +__device__ __forceinline__ void mbarrier_wait(uint64_t* a, uint32_t p) {
78 + while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed,
79 + cuda::ptx::scope_cta, a, p));
80 +}
81 +__device__ __forceinline__ void mbarrier_arrive_expect_tx(uint64_t* a,
82 + uint32_t t) {
83 + cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed,
84 + cuda::ptx::scope_cta,
85 + cuda::ptx::space_shared, a, t);
86 +}
87 +__device__ __forceinline__ void tma_load(void* d, const void* s, uint32_t n,
88 + uint64_t* m) {
89 + cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global, d,
90 + s, n, m);
91 +}
92 +
93 +// ============================================================================
94 +// DSMEM histogram reduce
95 +// ============================================================================
96 +
97 +template <uint32_t CS>
98 +__device__ __forceinline__ void dsmem_hist_reduce(uint32_t* histogram) {
99 + static_assert(kHistBins <= hist4096::kBlockSize);
100 + auto cluster = cooperative_groups::this_cluster();
101 + cluster.sync();
102 + const auto tx = threadIdx.x;
103 + const auto rank = blockIdx.y;
104 + constexpr auto kLocal = kHistBins / CS;
105 + const auto off = kLocal * rank;
106 + if (tx < kHistBins) {
107 + const auto addr = &histogram[off + tx / CS];
108 + const auto src = cluster.map_shared_rank(addr, tx % CS);
109 + *src = warp_reduce_sum_subN<CS>(*src);
110 + }
111 + cluster.sync();
112 +}
113 +
114 +// ============================================================================
115 +// Find threshold from reduced histogram
116 +// ============================================================================
117 +
118 +// NOTE: caller must ensure a cluster.sync() or __syncthreads() happened
119 +// before calling this, so warp_sum writes are visible across warps.
120 +// The first internal __syncthreads() is still needed for the warp_sum exchange.
121 +template <uint32_t TopK>
122 +__device__ __forceinline__ void find_threshold(uint32_t* histogram,
123 + uint32_t* warp_sum,
124 + uint32_t* counter_gt,
125 + uint32_t* counter_eq,
126 + hist4096::MatchBin* match) {
127 + const auto tx = threadIdx.x;
128 + const auto li = tx % hist4096::kWarpSize, wi = tx / hist4096::kWarpSize;
129 + const auto value = tx < kHistBins ? histogram[tx] : 0;
130 + const auto winc = hist4096::warp_inclusive_sum(li, value);
131 + if (li == hist4096::kWarpSize - 1) warp_sum[wi] = winc;
132 + __syncthreads();
133 + const auto tmp = warp_sum[li];
134 + const auto total = hist4096::warp_reduce_sum_full(tmp);
135 + auto pfx = hist4096::warp_reduce_sum_full(li < wi ? tmp : 0) + winc;
136 + const auto above = total - pfx;
137 + if (tx < kHistBins && above < TopK && above + value >= TopK) {
138 + *counter_gt = *counter_eq = 0;
139 + *match = {.bin = tx, .above_count = above, .equal_count = value};
140 + }
141 + __syncthreads();
142 +}
143 +
144 +// Streams data through shared memory in chunks, processing each chunk before
145 +// loading the next overwrites each buffer after processing it (the epilogue
146 +// prefetch loads the next chunk into the same slot)
147 +template <typename SmemType, uint32_t kStages, uint32_t kBinBits,
148 + bool kIsScatter>
149 +__device__ void tma_stream_pass(const float* scores, uint32_t length,
150 + uint32_t thr_bin, int32_t* indices,
151 + uint32_t* phases, SmemType* smem) {
152 + const auto tx = threadIdx.x;
153 + const auto lane = tx % hist4096::kWarpSize;
154 + const auto ni =
155 + (length + kSizePerStage - 1) / kSizePerStage; // total stages needed
156 + const auto la =
157 + (length + 3u) & ~3u; // length rounded up to float4 (TMA alignment)
158 + const auto pass =
159 + kIsScatter ? 1 : 0; // barrier dim: [0] for histogram, [1] for scatter
160 +
161 + // Prologue: issue initial TMA loads - prefill the pipeline
162 + if (tx == 0) {
163 +#pragma unroll
164 + for (uint32_t i = 0; i < kStages; i++) {
165 + if (i >= ni) {
166 + break;
167 + }
168 + const auto o = i * kSizePerStage;
169 + const auto sz = min(kSizePerStage, la - o) * sizeof(float);
170 + tma_load(smem->score_buffer[i], scores + o, sz,
171 + &smem->barrier[pass][i]); // cp.async.bulk is non-blocking
172 + mbarrier_arrive_expect_tx(&smem->barrier[pass][i], sz);
173 + }
174 + }
175 +
176 + // Main loop: process stages
177 + for (uint32_t it = 0; it < ni; it++) {
178 + const auto b = it % kStages; // which buffer slot (0 or 1)
179 + const auto o = it * kSizePerStage;
180 + const auto sz = min(kSizePerStage, length - o);
181 +
182 + if (lane == 0) {
183 + mbarrier_wait(&smem->barrier[pass][b],
184 + phases[b] & 1); // wait for the data
185 + }
186 + phases[b]++; // advances the phase for next time this slot is reused
187 + __syncwarp();
188 +
189 +#pragma unroll
190 + for (uint32_t i = 0; i < kElemPerStage; i++) {
191 + const auto li = tx + i * hist4096::kBlockSize;
192 + if (li >= sz) {
193 + break;
194 + }
195 + const auto sc = smem->score_buffer[b][li];
196 + const auto bn = hist4096::extract_coarse_bin_N<kBinBits>(sc);
197 + if constexpr (kIsScatter) { // compile-time branch
198 + // Scatter pass: place above-threshold and collect ties
199 + const auto gi = o + li;
200 + if (bn > thr_bin) {
201 + indices[atomicAdd(&smem->counter_gt, 1)] = gi;
202 + } else if (bn == thr_bin) {
203 + const auto p = atomicAdd(&smem->counter_eq, 1);
204 + if (p < hist4096::kMaxTies) {
205 + smem->tie_buffer[p] = {gi, sc};
206 + }
207 + }
208 + } else {
209 + // Histogram pass: just count
210 + atomicAdd(&smem->histogram[bn], 1);
211 + }
212 + }
213 + __syncthreads(); // ensures all threads finished processing their buffer
214 + // before next TMA load
215 +
216 + // Epilogue: issue next TMA load
217 + if (tx == 0 && it + kStages < ni) {
218 + const auto no = (it + kStages) * kSizePerStage;
219 + const auto nsz = min(kSizePerStage, la - no) * sizeof(float);
220 + tma_load(smem->score_buffer[b], scores + no, nsz,
221 + &smem->barrier[pass][b]);
222 + mbarrier_arrive_expect_tx(&smem->barrier[pass][b], nsz);
223 + }
224 + }
225 +}
226 +
227 +// ============================================================================
228 +// Fused path: single TMA pass, rescan smem for scatter
229 +// ============================================================================
230 +
231 +// Fused shared memory layout for cluster cooperative paths.
232 +// kPasses=1 for single-pass (CS=8, CS=4 singlepass), kPasses=2 for two-pass
233 +// (CS=4).
234 +template <uint32_t kStages, uint32_t kPasses = 1>
235 +struct SmemFused {
236 + uint64_t barrier[kPasses][kStages];
237 + alignas(128) uint32_t counter_gt;
238 + alignas(128) uint32_t counter_eq;
239 + alignas(128) hist4096::MatchBin match;
240 + uint32_t warp_sum[hist4096::kNumWarps];
241 + union {
242 + uint32_t histogram[kHistBins];
243 + hist4096::Tie tie_buffer[kMaxTopK];
244 + };
245 + alignas(128) float score_buffer[kStages][kSizePerStage];
246 +};
247 +
248 +using Smem8 = SmemFused<kFusedStagesCS8>;
249 +using Smem16 = SmemFused<kFusedStagesCS16>;
250 +using Smem4 = SmemFused<kStreamingStagesCS4, 2>;
251 +using SmemSinglePass = SmemFused<kMaxSinglePassStages>;
252 +
253 +// Cluster-cooperative large path.
254 +// kFused=true: all TMA stages resident, single-pass histogram + scatter (rescan
255 +// from smem). kFused=false: TMA double-buffer streaming, two passes (histogram
256 +// then scatter).
257 +template <uint32_t TopK, uint32_t CS, typename SmemType, bool kFused>
258 +__device__ void large_topk(const float* __restrict__ row_input,
259 + int32_t* __restrict__ row_output, uint32_t seq_len,
260 + uint32_t* phases, hist4096::Tie* tie_ws) {
261 + const auto rank = blockIdx.y; // this block's position in cluster
262 + const auto tx = threadIdx.x;
263 + const auto lane = tx % hist4096::kWarpSize;
264 +
265 + extern __shared__ uint8_t smem_raw[];
266 + auto* smem = reinterpret_cast<SmemType*>(smem_raw);
267 + int32_t* s_topk = reinterpret_cast<int32_t*>(smem_raw + sizeof(SmemType));
268 +
269 + // Partition row across cluster ranks
270 + constexpr uint32_t kAlign = 4;
271 + const auto units =
272 + (seq_len + kAlign - 1) / kAlign; // float4-aligned element count
273 + const auto base = units / CS, extra = units % CS; // elements per block
274 + const auto lu = base + (rank < extra ? 1u : 0u); // remainder blocks
275 + const auto ou =
276 + rank * base + min(rank, extra); // this block's count (load-balanced)
277 + const auto my_start = ou * kAlign; // global start offset
278 + const auto my_len = min(my_start + lu * kAlign, seq_len) -
279 + my_start; // actual length of this block
280 + const auto num_iters =
281 + (my_len + kSizePerStage - 1) / kSizePerStage; // TMA stages needed
282 + const auto len_aligned = (my_len + 3u) & ~3u;
283 +
284 + if constexpr (kFused) {
285 + // Fused init + TMA prologue
286 + if (tx < kHistBins) {
287 + smem->histogram[tx] = 0; // all threads zero histogram
288 + }
289 + if (tx == 0) { // thread 0 issues TMA - then all threads continue working
290 + // until mbarrier sync
291 + smem->counter_gt = 0;
292 + smem->counter_eq = 0;
293 + for (uint32_t i = 0; i < num_iters; i++) {
294 + const auto off = i * kSizePerStage;
295 + const auto sz = min(kSizePerStage, len_aligned - off) * sizeof(float);
296 + tma_load(smem->score_buffer[i], row_input + my_start + off, sz,
297 + &smem->barrier[0][i]); // cp.async.bulk of size kSizePerStage
298 + // × sizeof(float)
299 + mbarrier_arrive_expect_tx(&smem->barrier[0][i], sz);
300 + }
301 + }
302 + __syncthreads();
303 +
304 + // Histogram build. ILP unroll-by-2, no inter-stage sync
305 + for (uint32_t iter = 0; iter < num_iters; iter++) {
306 + const auto off = iter * kSizePerStage;
307 + const auto sz = min(kSizePerStage, my_len - off);
308 + if (lane == 0) {
309 + mbarrier_wait(&smem->barrier[0][iter],
310 + phases[iter] & 1); // wait for TMA
311 + }
312 + phases[iter]++;
313 + __syncwarp();
314 +#pragma unroll
315 + for (uint32_t i = 0; i < kElemPerStage; i += 2) {
316 + const auto li0 = tx + i * hist4096::kBlockSize;
317 + const auto li1 = tx + (i + 1) * hist4096::kBlockSize;
318 + if (li0 >= sz) {
319 + break;
320 + }
321 + const auto b0 = extract_coarse_bin(smem->score_buffer[iter][li0]);
322 + if (li1 < sz) {
323 + const auto b1 = extract_coarse_bin(smem->score_buffer[iter][li1]);
324 + atomicAdd(&smem->histogram[b0], 1);
325 + atomicAdd(&smem->histogram[b1], 1);
326 + } else {
327 + atomicAdd(&smem->histogram[b0], 1);
328 + }
329 + }
330 + }
331 + } else {
332 + // Twopass: init then stream histogram pass
333 + if (tx < kHistBins) {
334 + smem->histogram[tx] = 0;
335 + }
336 + if (tx == 0) {
337 + smem->counter_gt = 0;
338 + smem->counter_eq = 0;
339 + }
340 + __syncthreads();
341 + tma_stream_pass<SmemType, kStreamingStagesCS4, kHistBits, false>(
342 + row_input + my_start, my_len, 0, nullptr, phases, smem);
343 + }
344 +
345 + // DSMEM all-reduce + find threshold
346 + dsmem_hist_reduce<CS>(
347 + smem->histogram); // each block histogram is summed across all CS blocks
348 + find_threshold<TopK>(smem->histogram, smem->warp_sum, &smem->counter_gt,
349 + &smem->counter_eq, &smem->match);
350 +
351 + const auto thr = smem->match.bin;
352 +
353 + if constexpr (kFused) {
354 + // Fused scatter: rescan score_buffer (still in smem)
355 + for (uint32_t iter = 0; iter < num_iters; iter++) {
356 + const auto off = iter * kSizePerStage;
357 + const auto sz = min(kSizePerStage, my_len - off);
358 +#pragma unroll
359 + for (uint32_t i = 0; i < kElemPerStage; i++) {
360 + const auto li = tx + i * hist4096::kBlockSize;
361 + if (li >= sz) {
362 + break;
363 + }
364 + const auto score = smem->score_buffer[iter][li]; // still in smem
365 + const auto bin = extract_coarse_bin(score);
366 + const auto gidx = off + li;
367 + if (bin > thr) {
368 + s_topk[atomicAdd(&smem->counter_gt, 1)] = gidx; // above -> s_topk
369 + } else if (bin == thr) {
370 + const auto p = atomicAdd(&smem->counter_eq,
371 + 1); // equal -> ties (later refinement)
372 + if (p < hist4096::kMaxTies) {
373 + smem->tie_buffer[p] = {gidx, score};
374 + }
375 + }
376 + }
377 + }
378 + __syncthreads();
379 + } else {
380 + // Twopass scatter: re-stream data via TMA
381 + uint32_t scatter_phases[kStreamingStagesCS4] = {0, 0};
382 + tma_stream_pass<SmemType, kStreamingStagesCS4, kHistBits, true>(
383 + row_input + my_start, my_len, thr, s_topk, scatter_phases, smem);
384 + }
385 +
386 + // Output collection via DSMEM prefix sum
387 + constexpr uint32_t kAboveBits = 16;
388 + constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1;
389 + static_assert(kAboveMask >= TopK);
390 + static_assert(kAboveMask >= kMaxSinglePassPerBlock,
391 + "kAboveBits must cover max per-block element count");
392 +
393 + const uint32_t la = smem->counter_gt;
394 + const uint32_t le_full = smem->counter_eq;
395 + const uint32_t le =
396 + min(le_full, hist4096::kMaxTies); // written smem tie_buffer entries
397 +
398 + __shared__ uint32_t s_local_counts[CS];
399 + __shared__ uint32_t s_prefix_packed;
400 + __shared__ uint32_t s_total_above, s_total_equal;
401 +
402 + auto cluster = cooperative_groups::this_cluster();
403 + if (tx < CS) {
404 + // Pack written tie counts into 32-bit: (equal << 16) | above.
405 + // `le_full` may exceed the per-block tie buffer cap; using it here creates
406 + // holes in tie_ws and can make TopK=2048 refine unwritten workspace slots.
407 + const uint32_t packed = (le << kAboveBits) | la;
408 + const auto dst = cluster.map_shared_rank(s_local_counts, tx);
409 + dst[rank] = packed; // write my count to every block's s_local_counts[rank]
410 + }
411 + cluster.sync();
412 +
413 + // Thread 0 computes serial prefix sum
414 + if (tx == 0) {
415 + uint32_t prefix = 0, ta = 0, te = 0;
416 + for (uint32_t i = 0; i < CS; i++) {
417 + if (i == rank) {
418 + s_prefix_packed = prefix; // my prefix
419 + }
420 + ta += s_local_counts[i] & kAboveMask; // total above
421 + te += s_local_counts[i] >> kAboveBits; // total equal
422 + prefix += s_local_counts[i];
423 + }
424 + s_total_above = ta;
425 + s_total_equal = te;
426 + }
427 + __syncthreads();
428 +
429 + const uint32_t prefix_above = s_prefix_packed & kAboveMask;
430 + const uint32_t prefix_equal = s_prefix_packed >> kAboveBits;
431 +
432 + // Write to global output
433 + for (uint32_t i = tx; i < la; i += hist4096::kBlockSize) {
434 + // indices are placed contiguously starting at prefix_above
435 + row_output[prefix_above + i] =
436 + s_topk[i] + my_start; // my_start: block-local -> row-global index
437 + }
438 + for (uint32_t i = tx; i < le; i += hist4096::kBlockSize) {
439 + const auto t = smem->tie_buffer[i];
440 + uint32_t p = s_total_above + prefix_equal + i;
441 + if (p < TopK) {
442 + row_output[p] = t.idx + my_start;
443 + }
444 + uint32_t tp = prefix_equal + i;
445 + if (tp < (TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK)) {
446 + tie_ws[tp] = hist4096::Tie{t.idx + my_start, t.score};
447 + }
448 + }
449 +
450 + // Tie refinement
451 + cooperative_groups::this_cluster().sync();
452 + if (rank != 0) { // only rank 0 does tie refinement
453 + return;
454 + }
455 + if (s_total_above + s_total_equal <= TopK) { // no ties to refine
456 + return;
457 + }
458 +
459 + // Tie-breaking uses FP32 (4-round radix sort)
460 + if constexpr (TopK <= hist4096::kBlockSize) {
461 + // copy ties from tie_ws back to smem, then refine
462 + const uint32_t num_ties = min(s_total_equal, hist4096::kMaxTies);
463 + // TODO (roberto): could vectorize with uint2 (8 bytes = exactly one Tie)
464 + for (uint32_t i = tx; i < num_ties; i += hist4096::kBlockSize) {
465 + smem->tie_buffer[i] = hist4096::Tie{tie_ws[i].idx, tie_ws[i].score};
466 + }
467 + __syncthreads();
468 + hist4096::tie_handle<TopK>(smem->tie_buffer, num_ties, s_total_above,
469 + row_output, smem);
470 + } else {
471 + // TopK=2048: process directly from tie_ws (GMEM)
472 + const uint32_t num_ties = min(s_total_equal, static_cast<uint32_t>(TopK));
473 + hist4096::tie_handle_large<TopK>(tie_ws, num_ties, s_total_above,
474 + row_output, smem);
475 + }
476 +}
477 +
478 +// ============================================================================
479 +// Adapted from https://github.com/sgl-project/sglang/pull/23600
480 +// sgl-project/sglang
481 +// (python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/)
482 +// ============================================================================
483 +
484 +template <uint32_t TopK, uint32_t CS>
485 +__device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
486 + const auto rank = blockIdx.y, row = blockIdx.x, tx = threadIdx.x;
487 + const auto sl = params.lengths[row];
488 + int32_t* out = params.output + row * TopK;
489 + const float* in = params.input + row * params.stride;
490 +
491 + // Trivial: seq_len <= TopK
492 + if (sl <= static_cast<int32_t>(TopK)) {
493 + if (rank == 0) {
494 + for (uint32_t i = tx; i < TopK; i += hist4096::kBlockSize) {
495 + out[i] = (i < static_cast<uint32_t>(sl)) ? static_cast<int32_t>(i) : -1;
496 + }
497 + }
498 + return;
499 + }
500 +
501 + // Short-Medium path: histogram_4096_topk on rank 0 only - all data fits in RF
502 + if (sl <= static_cast<int32_t>(hist4096::kHist4096MaxLen)) {
503 + if (rank == 0) {
504 + extern __shared__ uint8_t sr[];
505 + hist4096::histogram_4096_topk<TopK, 12>(
506 + in, out, sl, sr); // 4096-bin (12-bit) histogram
507 + }
508 + return;
509 + }
510 +
511 + // Large path: init mbarriers + state, then dispatch fused or twopass
512 + const uint32_t per_block =
513 + (params.stride + CS - 1) / CS; // how many elements per block
514 + constexpr uint32_t kFusedMax = ((CS == 16) ? kFusedStagesCS16
515 + : (CS == 8) ? kFusedStagesCS8
516 + : kMaxSinglePassStages) *
517 + kSizePerStage;
518 + const bool use_singlepass =
519 + per_block <=
520 + kFusedMax; // single pass or TMA streaming: histogram+scatter
521 +
522 + // Select smem type and stage count at compile time based on CS
523 + constexpr uint32_t kFusedStages = (CS == 16) ? kFusedStagesCS16
524 + : (CS == 8) ? kFusedStagesCS8
525 + : kMaxSinglePassStages;
526 + using FusedSmem = SmemFused<kFusedStages>;
527 +
528 + extern __shared__ uint8_t sr[];
529 +
530 + constexpr uint32_t kTieWsPerRow =
531 + TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK;
532 + hist4096::Tie* row_tie_ws = params.tie_ws + row * kTieWsPerRow;
533 +
534 + if (use_singlepass) {
535 + auto* smem = reinterpret_cast<FusedSmem*>(sr);
536 + const uint32_t sp_stages = (per_block + kSizePerStage - 1) / kSizePerStage;
537 + if (tx < sp_stages) {
538 + mbarrier_init(&smem->barrier[0][tx],
539 + 1); // init 1 barrier per TMA stage -
540 + // signal when async copies complete
541 + }
542 + __syncthreads();
543 + uint32_t phases[kFusedStages] =
544 + {}; // tracks the parity for mbarrier wait/arrive protocol
545 + large_topk<TopK, CS, FusedSmem, true>(in, out, sl, phases, row_tie_ws);
546 + } else {
547 + // Two-pass: only CS=4 in practice (CS=8 always fits in singlepass)
548 + auto* smem = reinterpret_cast<Smem4*>(sr);
549 + if (tx < 2 * kStreamingStagesCS4) {
550 + mbarrier_init(&smem->barrier[0][tx],
551 + 1); // init 2×2=4 barriers (2 passes × 2 stages)
552 + }
553 + __syncthreads();
554 + uint32_t hp[kStreamingStagesCS4] = {0,
555 + 0}; // histogram+scatter pass counters
556 + large_topk<TopK, CS, Smem4, false>(in, out, sl, hp, row_tie_ws);
557 + }
558 +}
559 +
560 +template <uint32_t TopK>
561 +__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
562 + __cluster_dims__(1, 4, 1)
563 + cooperative_topk_cs4(CooperativeTopKParams<TopK> params) {
564 + cooperative_topk_body<TopK, 4>(params);
565 +}
566 +
567 +template <uint32_t TopK>
568 +__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
569 + __cluster_dims__(1, 8, 1)
570 + cooperative_topk_cs8(CooperativeTopKParams<TopK> params) {
571 + cooperative_topk_body<TopK, 8>(params);
572 +}
573 +
574 +template <uint32_t TopK>
575 +__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
576 + __cluster_dims__(1, 16, 1)
577 + cooperative_topk_cs16(CooperativeTopKParams<TopK> params) {
578 + cooperative_topk_body<TopK, 16>(params);
579 +}
580 +
581 +constexpr size_t kSmemSize4_base = sizeof(Smem4);
582 +constexpr size_t kSmemSize4_sp = sizeof(SmemSinglePass);
583 +constexpr size_t kSmemSize4 =
584 + (kSmemSize4_base > kSmemSize4_sp ? kSmemSize4_base : kSmemSize4_sp) +
585 + sizeof(int32_t) * 2048 + 128;
586 +constexpr size_t kSmemSize8 =
587 + sizeof(SmemFused<kFusedStagesCS8>) + sizeof(int32_t) * 2048 + 128;
588 +
589 +} // namespace cooperative
590 +
591 +} // namespace vllm
592 +
593 +#endif // COOPERATIVE_TOPK_CUH_
csrc/libtorch_stable/ops.h
+8
@@ -343,6 +343,14 @@ void persistent_topk(const torch::stable::Tensor& logits,
343 torch::stable::Tensor& workspace, int64_t k,
344 int64_t max_seq_len);
345
346 +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
347 +void cooperative_topk(const torch::stable::Tensor& logits,
348 + const torch::stable::Tensor& lengths,
349 + torch::stable::Tensor& output,
350 + torch::stable::Tensor& workspace, int64_t k,
351 + int64_t max_seq_len);
352 +#endif
353 +
354 void selective_scan_fwd(
355 const torch::stable::Tensor& u, const torch::stable::Tensor& delta,
356 const torch::stable::Tensor& A, const torch::stable::Tensor& B,
csrc/libtorch_stable/persistent_topk.cuh
+49 -18
@@ -11,6 +11,8 @@
11 #include <cub/cub.cuh>
12 #include <cstdint>
13
14 +#include "topk_histogram_4096.cuh"
15 +
16 namespace vllm {
17 namespace persistent {
18
@@ -935,8 +937,16 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
937 } // namespace persistent
938
939 // ============================================================================
938 -// FlashInfer FilteredTopK (BS>32 dispatch) — float32 only.
939 -// Extracted from flashinfer_topk.cuh. Lives in namespace vllm (not persistent).
940 +// ============================================================================
941 +// Optimized FilteredTopK — single CTA per row for bs > 32.
942 +// Kept with persistent_topk so the portable fallback owns the non-cluster path.
943 +// ============================================================================
944 +namespace filtered_topk {
945 +
946 +namespace hist4096 = topk_histogram_4096;
947 +
948 +// ============================================================================
949 +// FilteredTopK — single CTA per row for bs > 32
950 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
951 // ============================================================================
952
@@ -963,13 +973,6 @@ struct vec_t {
973 data[i] = ptr[i];
974 }
975 }
966 -
967 - FLASHINFER_INLINE void cast_store(T* ptr) const {
968 -#pragma unroll
969 - for (size_t i = 0; i < N; ++i) {
970 - ptr[i] = data[i];
971 - }
972 - }
976 };
977 #undef FLASHINFER_INLINE
978
@@ -1013,7 +1016,8 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
1016 * \tparam IdType Index type (int32_t)
1017 * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
1018 */
1016 -template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
1019 +template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
1020 + bool UsePredicatedShortLoads = false>
1021 __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
1022 FilteredTopKUnifiedKernel(const DType* __restrict__ input,
1023 IdType* __restrict__ output,
@@ -1042,6 +1046,19 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
1046 return;
1047 }
1048
1049 + // Short path
1050 + if (length <= 32768) {
1051 + extern __shared__ uint8_t _smem_reg[];
1052 + if constexpr (UsePredicatedShortLoads) {
1053 + hist4096::histogram_4096_topk_predicated<MAX_K, 12, 8>(score, dst, length,
1054 + _smem_reg);
1055 + } else {
1056 + hist4096::histogram_4096_topk<MAX_K, 12, 8>(score, dst, length,
1057 + _smem_reg);
1058 + }
1059 + return;
1060 + }
1061 +
1062 // Static shared memory
1063 alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
1064 alignas(128) __shared__ int s_counter;
@@ -1285,14 +1302,15 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input,
1302
1303 const int vec_size = ComputeFilteredTopKVecSize<DType>(max_len);
1304
1288 -#define DISPATCH_VEC_SIZE(VS) \
1289 - if (vec_size == VS) { \
1290 - auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>; \
1291 - FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \
1292 - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \
1293 - FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
1294 - smem_size, stream)); \
1295 - return cudaSuccess; \
1305 +#define DISPATCH_VEC_SIZE(VS) \
1306 + if (vec_size == VS) { \
1307 + auto kernel = \
1308 + FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K, (VS != MAX_VEC)>; \
1309 + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \
1310 + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \
1311 + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
1312 + smem_size, stream)); \
1313 + return cudaSuccess; \
1314 }
1315
1316 DISPATCH_VEC_SIZE(1)
@@ -1306,6 +1324,19 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input,
1324 return cudaSuccess;
1325 }
1326
1327 +} // namespace filtered_topk
1328 +
1329 +template <typename DType, typename IdType, uint32_t MAX_K = 2048>
1330 +cudaError_t FilteredTopKRaggedTransform(const DType* input,
1331 + IdType* output_indices,
1332 + const IdType* lengths,
1333 + uint32_t num_rows, uint32_t top_k_val,
1334 + uint32_t max_len,
1335 + cudaStream_t stream = 0) {
1336 + return filtered_topk::FilteredTopKRaggedTransform<DType, IdType, MAX_K>(
1337 + input, output_indices, lengths, num_rows, top_k_val, max_len, stream);
1338 +}
1339 +
1340 } // namespace vllm
1341
1342 #endif // PERSISTENT_TOPK_CUH_
csrc/libtorch_stable/topk_histogram_4096.cuh new
+554
@@ -0,0 +1,554 @@
1 +/*
2 + * Shared 4096-bin single-CTA TopK helpers.
3 + */
4 +
5 +#ifndef TOPK_HISTOGRAM_4096_CUH_
6 +#define TOPK_HISTOGRAM_4096_CUH_
7 +
8 +#include <cuda_fp16.h>
9 +#include <cuda_runtime.h>
10 +#include <cstdint>
11 +
12 +namespace vllm {
13 +namespace topk_histogram_4096 {
14 +
15 +constexpr uint32_t kBlockSize = 1024;
16 +constexpr uint32_t RADIX = 256;
17 +constexpr uint32_t kMaxTies = 1024;
18 +static_assert(kMaxTies <= kBlockSize,
19 + "tie_handle requires kMaxTies <= kBlockSize");
20 +constexpr uint32_t kWarpSize = 32;
21 +constexpr uint32_t kNumWarps = kBlockSize / kWarpSize;
22 +
23 +// Register path
24 +constexpr uint32_t kHist4096VecsPerThread = 4;
25 +constexpr uint32_t kHist4096MaxLen =
26 + kHist4096VecsPerThread * 4 * kBlockSize; // 16384
27 +
28 +struct alignas(16) MatchBin {
29 + uint32_t bin, above_count, equal_count;
30 +};
31 +struct alignas(8) Tie {
32 + uint32_t idx;
33 + float score;
34 +};
35 +
36 +__device__ __forceinline__ void load_float4_predicated(const float* ptr,
37 + int base, int seq_len,
38 + float& v0, float& v1,
39 + float& v2, float& v3) {
40 + uint32_t r0, r1, r2, r3;
41 + const int p0 = (base < seq_len);
42 + const int p1 = (base + 1 < seq_len);
43 + const int p2 = (base + 2 < seq_len);
44 + const int p3 = (base + 3 < seq_len);
45 + asm volatile(
46 + "{\n"
47 + " .reg .pred pr0, pr1, pr2, pr3;\n"
48 + " setp.ne.u32 pr0, %4, 0;\n"
49 + " setp.ne.u32 pr1, %5, 0;\n"
50 + " setp.ne.u32 pr2, %6, 0;\n"
51 + " setp.ne.u32 pr3, %7, 0;\n"
52 + " mov.u32 %0, 0xFF800000;\n"
53 + " mov.u32 %1, 0xFF800000;\n"
54 + " mov.u32 %2, 0xFF800000;\n"
55 + " mov.u32 %3, 0xFF800000;\n"
56 + " @pr0 ld.global.cg.u32 %0, [%8];\n"
57 + " @pr1 ld.global.cg.u32 %1, [%8+4];\n"
58 + " @pr2 ld.global.cg.u32 %2, [%8+8];\n"
59 + " @pr3 ld.global.cg.u32 %3, [%8+12];\n"
60 + "}\n"
61 + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
62 + : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr));
63 + v0 = __uint_as_float(r0);
64 + v1 = __uint_as_float(r1);
65 + v2 = __uint_as_float(r2);
66 + v3 = __uint_as_float(r3);
67 +}
68 +
69 +// converts the float32 score to a 32-bit ordered unsigned integer — the full
70 +// precision key for radix sorting
71 +__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
72 + uint32_t bits = __float_as_uint(x);
73 + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
74 +}
75 +
76 +// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits ->
77 +// bin 0-4095)
78 +template <uint32_t kBits>
79 +__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) {
80 + __half h = __float2half_rn(x);
81 + uint16_t bits = __half_as_ushort(h);
82 + uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
83 + : static_cast<uint16_t>(bits | 0x8000);
84 + return key >> (16 - kBits);
85 +}
86 +
87 +// running sum within each warp — thread 0 gets its own value, thread 1 gets
88 +// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc.
89 +__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane,
90 + uint32_t v) {
91 +#pragma unroll
92 + for (uint32_t o = 1; o < 32; o *= 2) {
93 + uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o);
94 + if (lane >= o) v += n;
95 + }
96 + return v;
97 +}
98 +
99 +// Returns the sum of a value across all 32 threads in the warp, and every
100 +// thread gets the same result SM90+ PTX instruction that does a hardware
101 +// warp-wide reduction in a single instruction w.r.t. warp::reduce_sum(), which
102 +// uses a __shfl_xor_sync butterfly tree (5 shuffles for 32 lanes)
103 +__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) {
104 + uint32_t r;
105 + asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v));
106 + return r;
107 +}
108 +
109 +// ============================================================================
110 +// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered
111 +// key Each round narrows by 8 bits until ties are fully resolved
112 +// ============================================================================
113 +
114 +template <uint32_t TopK>
115 +__device__ void tie_handle(const Tie* ties, uint32_t num_ties,
116 + uint32_t num_above, int32_t* output, void* _smem) {
117 + struct TS {
118 + alignas(128) uint32_t counter;
119 + alignas(128) MatchBin match;
120 + uint32_t histogram[RADIX];
121 + uint32_t warp_sum[kNumWarps];
122 + };
123 + auto* s = static_cast<TS*>(_smem);
124 + const auto tx = threadIdx.x;
125 + const auto li = tx % kWarpSize, wi = tx / kWarpSize;
126 +
127 + // Each thread loads one tie element.
128 + const bool has = tx < num_ties;
129 + const auto tie = has ? ties[tx] : Tie{0, 0.0f};
130 + const uint32_t key = convert_to_uint32_v2(tie.score);
131 +
132 + bool active = has; // tracks whether this thread's tie is still a candidate.
133 + uint32_t remain =
134 + TopK - num_above; // decreases each round as ties are resolved.
135 + uint32_t wpos = TopK; // wpos will hold the final output position.
136 + s->counter = 0;
137 + __syncthreads();
138 +
139 + // The 4-round radix loop - each round narrows by 8 bits until ties are fully
140 + // resolved
141 +#pragma unroll
142 + for (int r = 0; r < 4; r++) {
143 + uint32_t sh = 24 - r * 8; // round 0: bits 31-24, round 1: 23-16, etc.
144 + uint32_t bin = (key >> sh) & 0xFF; // this tie's 8-bit bin for this round
145 +
146 + // Step 1: Build 256-bin histogram.
147 + if (tx < RADIX) s->histogram[tx] = 0;
148 + __syncthreads();
149 + if (active) atomicAdd(&s->histogram[bin], 1);
150 + __syncthreads();
151 +
152 + // Step 2: Prefix scan to find threshold
153 + uint32_t hv = 0, wi2 = 0;
154 + if (tx < RADIX) {
155 + hv = s->histogram[tx];
156 + wi2 = warp_inclusive_sum(li, hv);
157 + if (li == kWarpSize - 1) s->warp_sum[wi] = wi2;
158 + }
159 + __syncthreads();
160 +
161 + if (tx < RADIX) {
162 + auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0;
163 + auto tot = warp_reduce_sum_full(tmp);
164 + auto inter = warp_reduce_sum_full(li < wi ? tmp : 0);
165 + auto above = tot - (inter + wi2);
166 + if (above < remain && above + hv >= remain) {
167 + s->match = {tx, above, remain - above};
168 + }
169 + }
170 + __syncthreads();
171 +
172 + // Step 3: Scatter
173 + auto [thr, na, _] = s->match; // threshold bin, num above, unused
174 + if (active) {
175 + if (bin > thr) {
176 + wpos = num_above +
177 + atomicAdd(&s->counter, 1); // above -> place in output directly
178 + active = false;
179 + } else if (bin < thr)
180 + active = false; // below -> discard
181 + else if (r == 3)
182 + wpos = TopK - atomicAdd(&s->match.equal_count,
183 + -1u); // last round: place remaining
184 + }
185 + remain -= na;
186 + if (!remain) break; // all ties resolved early
187 + }
188 + // Final write
189 + if (wpos < TopK) output[wpos] = tie.idx;
190 +}
191 +
192 +// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048).
193 +// tie_handle assumes 1 tie per thread (max 1024).
194 +// This version handles 2 ties per thread via kPerThread=2
195 +template <uint32_t TopK>
196 +__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties,
197 + uint32_t num_above, int32_t* output,
198 + void* _smem) {
199 + static_assert(TopK > kBlockSize);
200 + struct TS {
201 + alignas(128) uint32_t counter;
202 + alignas(128) MatchBin match;
203 + uint32_t histogram[RADIX];
204 + uint32_t warp_sum[kNumWarps];
205 + };
206 + auto* s = static_cast<TS*>(_smem);
207 + const auto tx = threadIdx.x;
208 + const auto li = tx % kWarpSize;
209 + const auto wi = tx / kWarpSize;
210 +
211 + constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize;
212 + Tie my_ties[kPerThread];
213 + uint32_t keys[kPerThread];
214 + bool active[kPerThread];
215 +
216 + for (uint32_t e = 0; e < kPerThread; e++) {
217 + uint32_t idx = e * kBlockSize + tx;
218 + if (idx < num_ties) {
219 + my_ties[e] = ties[idx];
220 + keys[e] = convert_to_uint32_v2(ties[idx].score);
221 + active[e] = true;
222 + } else {
223 + my_ties[e] = {0, 0.0f};
224 + keys[e] = 0;
225 + active[e] = false;
226 + }
227 + }
228 +
229 + uint32_t remain = TopK - num_above;
230 + s->counter = 0;
231 + __syncthreads();
232 +
233 + for (int r = 0; r < 4; r++) {
234 + uint32_t sh = 24 - r * 8;
235 + if (tx < RADIX) {
236 + s->histogram[tx] = 0;
237 + }
238 + __syncthreads();
239 +
240 + for (uint32_t e = 0; e < kPerThread; e++) {
241 + if (active[e]) {
242 + atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1);
243 + }
244 + }
245 + __syncthreads();
246 +
247 + uint32_t hv = 0;
248 + if (tx < RADIX) {
249 + hv = s->histogram[tx];
250 + auto wi2 = warp_inclusive_sum(li, hv);
251 + if (li == kWarpSize - 1) {
252 + s->warp_sum[wi] = wi2;
253 + }
254 + }
255 + __syncthreads();
256 + if (tx < RADIX) {
257 + auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0;
258 + auto total = warp_reduce_sum_full(tmp2);
259 + auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0);
260 + auto wi2 = warp_inclusive_sum(li, hv);
261 + auto above = total - (inter + wi2);
262 + if (above < remain && above + hv >= remain) {
263 + s->match = {
264 + .bin = tx, .above_count = above, .equal_count = remain - above};
265 + }
266 + }
267 + __syncthreads();
268 +
269 + auto thr = s->match.bin;
270 + auto na = s->match.above_count;
271 +
272 + for (uint32_t e = 0; e < kPerThread; e++) {
273 + if (!active[e]) {
274 + continue;
275 + }
276 + uint32_t bin = (keys[e] >> sh) & 0xFF;
277 + if (bin > thr) {
278 + uint32_t wpos = num_above + atomicAdd(&s->counter, 1);
279 + if (wpos < TopK) {
280 + output[wpos] = my_ties[e].idx;
281 + }
282 + active[e] = false;
283 + } else if (bin < thr) {
284 + active[e] = false;
285 + } else if (r == 3) {
286 + uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u);
287 + if (wpos < TopK) {
288 + output[wpos] = my_ties[e].idx;
289 + }
290 + }
291 + }
292 +
293 + num_above += na;
294 + remain -= na;
295 + __syncthreads();
296 + s->counter = 0;
297 + __syncthreads();
298 + }
299 +}
300 +
301 +// ============================================================================
302 +// Register-based single-CTA fast path for seq_len <= 16384
303 +// 4 float4 per thread × 1024 threads = 16384 elements max
304 +// Uses 4096-bin (12-bit) histogram for better precision
305 +// ============================================================================
306 +
307 +template <uint32_t TopK, uint32_t HIST_BITS>
308 +struct Histogram4096Smem {
309 + static constexpr uint32_t HIST_BINS = 1 << HIST_BITS;
310 + static constexpr uint32_t TIE_CAPACITY = TopK > kMaxTies ? TopK : kMaxTies;
311 + alignas(128) uint32_t counter_gt;
312 + alignas(128) uint32_t counter_eq;
313 + MatchBin match;
314 + uint32_t warp_sum[kNumWarps];
315 + union {
316 + uint32_t histogram[HIST_BINS];
317 + Tie tie_buffer[TIE_CAPACITY];
318 + };
319 +};
320 +
321 +template <uint32_t TopK, uint32_t HIST_BITS,
322 + uint32_t VECS_PER_THREAD = kHist4096VecsPerThread,
323 + bool UsePredicatedLoads = false>
324 +__device__ void histogram_4096_topk(const float* __restrict__ scores,
325 + int32_t* __restrict__ output,
326 + uint32_t length, void* _smem) {
327 + constexpr uint32_t HIST_BINS = 1 << HIST_BITS;
328 + constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize;
329 + static_assert(HIST_BINS >= kBlockSize,
330 + "HIST_BITS must give >= kBlockSize bins");
331 +
332 + using Smem = Histogram4096Smem<TopK, HIST_BITS>;
333 + auto* smem = static_cast<Smem*>(_smem);
334 + const auto tx = threadIdx.x;
335 + const auto lane_id = tx % kWarpSize;
336 + const auto warp_id = tx / kWarpSize;
337 +
338 + // Phase 1: Load all data into RF + build histogram
339 + float4
340 + vecs[VECS_PER_THREAD]; // 4 vectors x 4 floats = 16 elements per thread
341 + if constexpr (ITEMS_PER_THREAD >= 4) {
342 + // Zero the histogram (SMEM writes)
343 + for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++)
344 + reinterpret_cast<uint4*>(
345 + smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] =
346 + make_uint4(0, 0, 0, 0);
347 + } else {
348 + if (tx < HIST_BINS) smem->histogram[tx] = 0;
349 + }
350 + if (tx == 0) {
351 + smem->counter_gt = 0;
352 + smem->counter_eq = 0;
353 + }
354 + if constexpr (UsePredicatedLoads) {
355 + const bool row_aligned = (reinterpret_cast<uintptr_t>(scores) & 0xFu) == 0;
356 +#pragma unroll
357 + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) {
358 + const uint32_t base = (tx + v * kBlockSize) * 4;
359 + if (base < length) {
360 + if (row_aligned && base + 3 < length) {
361 + vecs[v] = *reinterpret_cast<const float4*>(scores + base);
362 + } else {
363 + load_float4_predicated(scores + base, static_cast<int>(base),
364 + static_cast<int>(length), vecs[v].x, vecs[v].y,
365 + vecs[v].z, vecs[v].w);
366 + }
367 + }
368 + }
369 + } else {
370 +#pragma unroll
371 + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) {
372 + const uint32_t base = (tx + v * kBlockSize) * 4;
373 + if (base < length) {
374 + vecs[v] = *reinterpret_cast<const float4*>(scores + base);
375 + }
376 + }
377 + }
378 + __syncthreads();
379 +
380 + // Build histogram from RF via atomic adds into the shared histogram
381 + bool done = false;
382 +#pragma unroll
383 + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) {
384 + const float* elems = reinterpret_cast<const float*>(&vecs[v]);
385 +#pragma unroll
386 + for (uint32_t e = 0; e < 4 && !done; e++) {
387 + const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
388 + if (idx >= length) {
389 + done = true;
390 + } else {
391 + atomicAdd(&smem->histogram[extract_coarse_bin_N<HIST_BITS>(elems[e])],
392 + 1);
393 + }
394 + }
395 + }
396 + __syncthreads();
397 +
398 + // Phase 2: Prefix scan to find threshold bin
399 + // Multi-element scan (4096 bins: 4 per thread)
400 + uint32_t orig[ITEMS_PER_THREAD];
401 + uint32_t local_sum = 0;
402 +
403 + // Step 1: Each thread sums its 4 bins
404 +#pragma unroll
405 + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) {
406 + orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i];
407 + local_sum += orig[i];
408 + }
409 +
410 + // Step 2: Warp-level inclusive prefix sum on local_sum
411 + const auto warp_inc = warp_inclusive_sum(lane_id, local_sum);
412 + if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc;
413 + __syncthreads();
414 +
415 + // Step 3: Inter-warp prefix via redux.sync
416 + const auto tmp = smem->warp_sum[lane_id];
417 + uint32_t prefix = warp_reduce_sum_full(
418 + lane_id < warp_id ? tmp : 0); // sum of all prior warps
419 + prefix +=
420 + warp_inc - local_sum; // exclusive prefix within this thread's position
421 +
422 + // Step 4: Find threshold - scan 4 bins, accumulate prefix
423 +#pragma unroll
424 + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) {
425 + prefix += orig[i];
426 + const auto above = length - prefix; // elements in bins ABOVE this one
427 + if (above < TopK && above + orig[i] >= TopK) {
428 + smem->match = {.bin = tx * ITEMS_PER_THREAD + i,
429 + .above_count = above,
430 + .equal_count = orig[i]};
431 + }
432 + }
433 +
434 + __syncthreads();
435 +
436 + // Phase 3: Scatter from registers
437 + const auto [thr_bin, num_above, num_equal] = smem->match;
438 + const bool need_tie = (num_equal + num_above > TopK);
439 +
440 + done = false;
441 +#pragma unroll
442 + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) {
443 + const float* elems = reinterpret_cast<const float*>(&vecs[v]);
444 +#pragma unroll
445 + for (uint32_t e = 0; e < 4 && !done; e++) {
446 + const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
447 + if (idx >= length) {
448 + done = true;
449 + } else {
450 + const uint32_t bin = extract_coarse_bin_N<HIST_BITS>(elems[e]);
451 + if (bin > thr_bin) {
452 + output[atomicAdd(&smem->counter_gt, 1)] =
453 + idx; // above -> output directly
454 + } else if (bin == thr_bin) {
455 + const auto pos = atomicAdd(&smem->counter_eq, 1);
456 + if (!need_tie) {
457 + if (pos + num_above < TopK) {
458 + output[pos + num_above] = idx; // all fit
459 + }
460 + } else {
461 + if (pos < TopK) {
462 + smem->tie_buffer[pos] = {idx, elems[e]}; // store for refirement
463 + }
464 + }
465 + }
466 + // else: bin < thr_bin - discard (not in top-k)
467 + }
468 + }
469 + }
470 +
471 + // Phase 4: Tie-breaking
472 + if (!need_tie) return;
473 + __syncthreads();
474 +
475 + // Fast warp-ballot tie-breaking for small tie counts
476 + const uint32_t num_ties = min(num_equal, static_cast<uint32_t>(TopK));
477 + const uint32_t topk_remain =
478 + TopK - num_above; // pick exactly remaining elements to fill topK
479 +
480 + auto is_greater = [](const Tie& a, const Tie& b) {
481 + return (a.score > b.score) || (a.score == b.score && a.idx < b.idx);
482 + };
483 +
484 + if (num_ties <= kWarpSize) {
485 + // <=32 ties - Use warp ballot
486 + // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024
487 + // comparisons in one instruction per warp. O(1) work.
488 + const auto lane_id = tx % kWarpSize;
489 + const auto warp_id = tx / kWarpSize;
490 + if (lane_id >= num_ties || warp_id >= num_ties) return;
491 + const uint32_t mask = (1ull << num_ties) - 1u;
492 + const auto tie = smem->tie_buffer[lane_id]; // each lane holds one tie
493 + const auto target =
494 + smem->tie_buffer[warp_id]; // each warp evaluates one candidate
495 + const bool pred =
496 + is_greater(tie, target); // compare all ties against target
497 + const auto rank = static_cast<uint32_t>(
498 + __popc(__ballot_sync(mask, pred))); // count how many are greater
499 + if (lane_id == 0 && rank < topk_remain) {
500 + output[num_above + rank] = target.idx; // place at correct position
501 + }
502 + } else if (num_ties <=
503 + kWarpSize *
504 + 2) { // TODO (roberto): try to refactor this with <=32 case
505 + // Same idea but each thread handles 2 tie elements
506 + const auto lane_id = tx % kWarpSize;
507 + const auto warp_id = tx / kWarpSize;
508 + const auto lane1 = lane_id + kWarpSize;
509 + const auto warp1 = warp_id + kWarpSize;
510 + const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__};
511 + const auto tie0 = smem->tie_buffer[lane_id];
512 + const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid;
513 + if (warp_id < num_ties) {
514 + const auto target = smem->tie_buffer[warp_id];
515 + const auto r0 =
516 + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target)));
517 + const auto r1 =
518 + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target)));
519 + if (lane_id == 0 && r0 + r1 < topk_remain)
520 + output[num_above + r0 + r1] = target.idx;
521 + }
522 + if (warp1 < num_ties) {
523 + const auto target = smem->tie_buffer[warp1];
524 + const auto r0 =
525 + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target)));
526 + const auto r1 =
527 + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target)));
528 + if (lane_id == 0 && r0 + r1 < topk_remain)
529 + output[num_above + r0 + r1] = target.idx;
530 + }
531 + } else {
532 + // Large tie count: fall back to 4-round radix-256 sort
533 + if constexpr (TopK <= kBlockSize) {
534 + tie_handle<TopK>(smem->tie_buffer, num_ties, num_above, output, smem);
535 + } else {
536 + tie_handle_large<TopK>(smem->tie_buffer, num_ties, num_above, output,
537 + smem);
538 + }
539 + }
540 +}
541 +
542 +template <uint32_t TopK, uint32_t HIST_BITS,
543 + uint32_t VECS_PER_THREAD = kHist4096VecsPerThread>
544 +__device__ __noinline__ void histogram_4096_topk_predicated(
545 + const float* __restrict__ scores, int32_t* __restrict__ output,
546 + uint32_t length, void* _smem) {
547 + histogram_4096_topk<TopK, HIST_BITS, VECS_PER_THREAD, true>(scores, output,
548 + length, _smem);
549 +}
550 +
551 +} // namespace topk_histogram_4096
552 +} // namespace vllm
553 +
554 +#endif // TOPK_HISTOGRAM_4096_CUH_
csrc/libtorch_stable/torch_bindings.cpp
+9
@@ -493,6 +493,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
493 "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
494 "Tensor workspace, int k, int max_seq_len) -> ()");
495
496 +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
497 + ops.def(
498 + "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, "
499 + "Tensor workspace, int k, int max_seq_len) -> ()");
500 +#endif
501 +
502 // Activation ops
503 ops.def(
504 "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
@@ -711,6 +717,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
717 ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill));
718 ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode));
719 ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
720 +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
721 + ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk));
722 +#endif
723
724 // Activation kernels (shared CUDA/ROCm)
725 ops.impl("persistent_masked_m_silu_mul_quant",
tests/kernels/test_top_k_per_row.py
+151 -26
@@ -14,6 +14,70 @@ TOP_K_VALUES = [2048, 3000]
14 BATCH_SIZE = [1, 2, 2048]
15 NEXT_N = [1, 8]
16 DATA_GENERATION = ["random", "10LSBits"]
17 +RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
18 +
19 +
20 +def _has_device_capability(major: int) -> bool:
21 + return current_platform.is_cuda() and current_platform.has_device_capability(major)
22 +
23 +
24 +COOPERATIVE_TOPK_BACKEND = pytest.param(
25 + "cooperative_topk",
26 + marks=pytest.mark.skipif(
27 + not _has_device_capability(90),
28 + reason="cooperative_topk requires SM90+",
29 + ),
30 +)
31 +WORKSPACE_TOPK_BACKENDS = ["persistent_topk", COOPERATIVE_TOPK_BACKEND]
32 +TOPK_BACKENDS = ["top_k_per_row_decode", *WORKSPACE_TOPK_BACKENDS]
33 +
34 +
35 +def _run_topk_backend(
36 + backend: str,
37 + logits: torch.Tensor,
38 + lengths: torch.Tensor,
39 + indices: torch.Tensor,
40 + top_k: int,
41 + max_seq_len: int,
42 + next_n: int = 1,
43 +) -> None:
44 + if backend == "top_k_per_row_decode":
45 + torch.ops._C.top_k_per_row_decode(
46 + logits,
47 + next_n,
48 + lengths,
49 + indices,
50 + indices.shape[0],
51 + logits.stride(0),
52 + logits.stride(1),
53 + top_k,
54 + )
55 + elif backend == "persistent_topk":
56 + workspace = torch.empty(
57 + RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda"
58 + )
59 + torch.ops._C.persistent_topk(
60 + logits, lengths, indices, workspace, top_k, max_seq_len
61 + )
62 + elif backend == "cooperative_topk":
63 + if indices.shape[0] > 32:
64 + pytest.skip(
65 + "cooperative_topk supports <=32 rows; "
66 + "persistent_topk covers larger batches"
67 + )
68 + if logits.stride(0) % 4 != 0:
69 + pytest.skip(
70 + "cooperative_topk requires row stride divisible by 4; "
71 + "persistent_topk covers unaligned strides"
72 + )
73 + workspace = torch.empty(
74 + RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda"
75 + )
76 + torch.ops._C.cooperative_topk(
77 + logits, lengths, indices, workspace, top_k, max_seq_len
78 + )
79 + else:
80 + raise ValueError(f"Unknown top-k backend: {backend}")
81
82
83 def create_random_logits(
@@ -322,16 +386,19 @@ def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None:
386 @pytest.mark.parametrize("clean_logits", [True, False])
387 @pytest.mark.parametrize("top_k", [2048])
388 @pytest.mark.parametrize("next_n", [1, 4])
389 +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
390 @torch.inference_mode()
326 -def test_deepseek_persistent_topk(
391 +def test_deepseek_workspace_topk(
392 seq_len_range: tuple[int, int],
393 test_id: str,
394 clean_logits: bool,
395 top_k: int,
396 next_n: int,
397 + backend: str,
398 ) -> None:
399 """
334 - Test persistent_topk with varying sequence lengths and speculative decoding.
400 + Test workspace top-k backends with varying sequence lengths and speculative
401 + decoding.
402 Supports speculative decoding with next_n > 1.
403 """
404 set_random_seed(42 if test_id == "short_sequences" else 43)
@@ -347,6 +414,7 @@ def test_deepseek_persistent_topk(
414 dtype=torch.int32,
415 device="cuda",
416 )
417 + seq_lens = (seq_lens + 3) & ~3 # align to 4 for TMA
418
419 # Compute row boundaries for speculative decoding
420 row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")
@@ -366,14 +434,11 @@ def test_deepseek_persistent_topk(
434 offsets = torch.arange(next_n, device=logits.device, dtype=torch.int32)
435 lengths = (seq_lens.unsqueeze(1) - next_n + 1 + offsets).flatten()
436
369 - workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
437 max_seq_len = int(seq_lens.max().item())
371 - torch.ops._C.persistent_topk(
372 - logits, lengths, indices, workspace, top_k, max_seq_len
373 - )
438 + _run_topk_backend(backend, logits, lengths, indices, top_k, max_seq_len, next_n)
439
440 validate_topk_against_reference(
376 - logits, indices, row_starts, row_ends, top_k, f"persistent_topk ({test_id})"
441 + logits, indices, row_starts, row_ends, top_k, f"{backend} ({test_id})"
442 )
443
444
@@ -383,9 +448,10 @@ def run_large_context_topk_test(
448 top_k: int,
449 data_type: str = "random",
450 seed: int = 42,
451 + backend: str = "cooperative_topk",
452 ) -> None:
453 """
388 - Helper to run persistent_topk kernel test with given parameters.
454 + Helper to run a top-k backend test with given parameters.
455
456 Args:
457 batch_size: Number of rows/sequences
@@ -393,6 +459,7 @@ def run_large_context_topk_test(
459 top_k: Number of top elements to select
460 data_type: Type of test data to generate
461 seed: Random seed for reproducibility
462 + backend: Top-k backend to test
463 """
464 torch.set_default_device("cuda:0")
465 set_random_seed(seed)
@@ -449,11 +516,8 @@ def run_large_context_topk_test(
516 # Create output tensor
517 indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
518
452 - workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
519 max_seq_len = max(seq_lens)
454 - torch.ops._C.persistent_topk(
455 - logits, lengths, indices, workspace, top_k, max_seq_len
456 - )
520 + _run_topk_backend(backend, logits, lengths, indices, top_k, max_seq_len)
521
522 torch.accelerator.synchronize()
523
@@ -605,8 +669,9 @@ def run_large_context_topk_test(
669 ),
670 ],
671 )
672 +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
673 @torch.inference_mode()
609 -def test_persistent_topk_correctness(test_config: dict) -> None:
674 +def test_workspace_topk_correctness(test_config: dict, backend: str) -> None:
675 """
676 Comprehensive correctness tests covering:
677 - Sequence length edge cases (trivial, boundary, varied)
@@ -620,6 +685,7 @@ def test_persistent_topk_correctness(test_config: dict) -> None:
685 seq_lens=test_config["seq_lens"],
686 top_k=test_config["top_k"],
687 data_type=test_config.get("data_type", "random"),
688 + backend=backend,
689 )
690
691
@@ -668,8 +734,9 @@ def test_persistent_topk_correctness(test_config: dict) -> None:
734 ),
735 ],
736 )
737 +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
738 @torch.inference_mode()
672 -def test_persistent_topk_algorithm_paths(test_config: dict) -> None:
739 +def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None:
740 """
741 Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2):
742 - Batch size scalability (1, 4, 32, 256)
@@ -680,12 +747,14 @@ def test_persistent_topk_algorithm_paths(test_config: dict) -> None:
747 batch_size=test_config["batch_size"],
748 seq_lens=[test_config["seq_len"]] * test_config["batch_size"],
749 top_k=test_config["top_k"],
750 + backend=backend,
751 )
752
753
754 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
755 +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
756 @torch.inference_mode()
688 -def test_persistent_topk_stress() -> None:
757 +def test_workspace_topk_stress(backend: str) -> None:
758 """
759 Stress test with random configurations to catch edge cases.
760 Capped at 163840 (DeepSeek V3.2 max context) for realistic testing.
@@ -700,16 +769,73 @@ def test_persistent_topk_stress() -> None:
769 batch_size = torch.randint(1, 32, (1,)).item()
770
771 # Random sequence lengths capped at DeepSeek V3.2 max context
703 - seq_lens = torch.randint(100, 163840, (batch_size,)).tolist()
772 + seq_lens_tensor = torch.randint(100, 163840, (batch_size,))
773 + if backend == "cooperative_topk":
774 + seq_lens = ((seq_lens_tensor + 3) & ~3).tolist()
775 + else:
776 + seq_lens = seq_lens_tensor.tolist()
777
778 run_large_context_topk_test(
779 batch_size=batch_size,
780 seq_lens=seq_lens,
781 top_k=top_k,
782 seed=seed,
783 + backend=backend,
784 )
785
786
787 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
788 +@pytest.mark.parametrize("backend", TOPK_BACKENDS)
789 +@pytest.mark.parametrize("top_k", [512, 1024, 2048])
790 +@torch.inference_mode()
791 +def test_deepseek_topk_backends_no_error_and_reference(
792 + backend: str,
793 + top_k: int,
794 +) -> None:
795 + """Exercise every production top-k backend on the same inputs."""
796 + run_large_context_topk_test(
797 + batch_size=4,
798 + seq_lens=[2049, 4097, 8191, 12000],
799 + top_k=top_k,
800 + data_type="random",
801 + seed=123,
802 + backend=backend,
803 + )
804 +
805 +
806 +@pytest.mark.skipif(not _has_device_capability(90), reason="This test requires SM90+")
807 +@torch.inference_mode()
808 +def test_cooperative_topk_512_tie_workspace_is_per_row() -> None:
809 + """Regression test for TopK=512 tie workspace row overlap."""
810 + torch.set_default_device("cuda:0")
811 +
812 + top_k = 512
813 + num_rows = 2
814 + stride = 65536
815 + lengths = torch.tensor([40960, 65536], dtype=torch.int32, device="cuda")
816 + logits = torch.full(
817 + (num_rows, stride), float("-inf"), dtype=torch.float32, device="cuda"
818 + )
819 +
820 + # Row 0 must never select these low indices: many better row-0 ties exist.
821 + logits[0, :2048] = -10.0
822 + logits[0, 2048 : lengths[0]] = 1.0
823 + # Row 1 has higher exact tie scores. With the old row * TopK tie_ws stride,
824 + # these row-1 ties could overwrite row 0's TopK=512 refinement workspace.
825 + logits[1, : lengths[1]] = 2.0
826 +
827 + indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
828 + workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda")
829 + torch.ops._C.cooperative_topk(logits, lengths, indices, workspace, top_k, stride)
830 + torch.accelerator.synchronize()
831 +
832 + row0 = indices[0].cpu()
833 + assert torch.all(row0 >= 2048), (
834 + "cooperative_topk TopK=512 selected row-0 low-score indices, likely "
835 + "from overlapping tie_ws rows"
836 + )
837 +
838 +
839 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
840 @pytest.mark.parametrize(
841 "test_config",
@@ -774,10 +900,11 @@ def test_persistent_topk_stress() -> None:
900 ],
901 )
902 @pytest.mark.parametrize("top_k", [512, 2048])
903 +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
904 @torch.inference_mode()
778 -def test_persistent_topk(test_config: dict, top_k: int) -> None:
905 +def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None:
906 """
780 - Tests specific to the persistent_topk kernel:
907 + Tests specific to workspace top-k backends:
908 - Mixed medium/large rows in the same batch (dynamic per-row dispatch)
909 - Boundary around LARGE_THRESHOLD (32K)
910 - Trivial + medium + large rows in a single batch
@@ -787,15 +914,17 @@ def test_persistent_topk(test_config: dict, top_k: int) -> None:
914 seq_lens=test_config["seq_lens"],
915 top_k=top_k,
916 data_type=test_config.get("data_type", "random"),
917 + backend=backend,
918 )
919
920
921 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
922 @pytest.mark.parametrize("top_k", [512, 2048])
923 +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
924 @torch.inference_mode()
796 -def test_persistent_topk_padded_stride(top_k: int) -> None:
925 +def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None:
926 """
798 - Test persistent_topk with padded logits (large stride, small seq_len)
927 + Test workspace top-k backends with padded logits (large stride, small seq_len)
928 to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits
929 returns [B, max_model_len] with max_model_len=163840.
930 """
@@ -818,11 +947,7 @@ def test_persistent_topk_padded_stride(top_k: int) -> None:
947
948 lengths = torch.tensor(actual_seq_lens, dtype=torch.int32, device="cuda")
949 indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda")
821 - workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
822 -
823 - torch.ops._C.persistent_topk(
824 - logits, lengths, indices, workspace, top_k, max(actual_seq_lens)
825 - )
950 + _run_topk_backend(backend, logits, lengths, indices, top_k, max(actual_seq_lens))
951 torch.accelerator.synchronize()
952
953 # Validate against torch.topk
@@ -840,6 +965,6 @@ def test_persistent_topk_padded_stride(top_k: int) -> None:
965 expected_vals = logits[i, expected].cpu().sort(descending=True)[0]
966 actual_vals = logits[i, actual].cpu().sort(descending=True)[0]
967 assert torch.allclose(expected_vals, actual_vals, rtol=1e-4, atol=1e-4), (
843 - f"Row {i}: persistent_topk with padded stride doesn't match. "
968 + f"Row {i}: {backend} with padded stride doesn't match. "
969 f"seq_len={sl}, stride={padded_stride}"
970 )
vllm/model_executor/layers/sparse_attn_indexer.py
+26 -1
@@ -334,7 +334,32 @@ def sparse_attn_indexer(
334 num_rows = logits.shape[0]
335 topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
336
337 - if current_platform.is_cuda() and topk_tokens in (512, 1024, 2048):
337 + use_cooperative_topk = (
338 + current_platform.is_cuda()
339 + and topk_tokens in (512, 1024, 2048)
340 + and num_rows <= 32
341 + and logits.stride(0) % 4 == 0 # TMA 16-byte alignment
342 + and current_platform.has_device_capability(90)
343 + )
344 + use_persistent_topk = current_platform.is_cuda() and topk_tokens in (
345 + 512,
346 + 1024,
347 + 2048,
348 + )
349 + if use_cooperative_topk:
350 + workspace_manager = current_workspace_manager()
351 + (topk_workspace,) = workspace_manager.get_simultaneous(
352 + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
353 + )
354 + torch.ops._C.cooperative_topk(
355 + logits,
356 + seq_lens,
357 + topk_indices,
358 + topk_workspace,
359 + topk_tokens,
360 + attn_metadata_narrowed.max_seq_len,
361 + )
362 + elif use_persistent_topk:
363 workspace_manager = current_workspace_manager()
364 (topk_workspace,) = workspace_manager.get_simultaneous(
365 ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),