| 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 | // Clamp at 0: `sl` is compared signed here but cast to uint32 below, so a |
| 488 | // negative length would otherwise emit indices 0..TopK-1 as valid instead |
| 489 | // of the -1 padding. |
| 490 | const int32_t sl = params.lengths[row] > 0 ? params.lengths[row] : 0; |
| 491 | int32_t* out = params.output + row * TopK; |
| 492 | const float* in = params.input + row * params.stride; |
| 493 | |
| 494 | // Trivial: seq_len <= TopK |
| 495 | if (sl <= static_cast<int32_t>(TopK)) { |
| 496 | if (rank == 0) { |
| 497 | for (uint32_t i = tx; i < TopK; i += hist4096::kBlockSize) { |
| 498 | out[i] = (i < static_cast<uint32_t>(sl)) ? static_cast<int32_t>(i) : -1; |
| 499 | } |
| 500 | } |
| 501 | return; |
| 502 | } |
| 503 | |
| 504 | // Short-Medium path: histogram_4096_topk on rank 0 only - all data fits in RF |
| 505 | if (sl <= static_cast<int32_t>(hist4096::kHist4096MaxLen)) { |
| 506 | if (rank == 0) { |
| 507 | extern __shared__ uint8_t sr[]; |
| 508 | hist4096::histogram_4096_topk<TopK, 12>( |
| 509 | in, out, sl, sr); // 4096-bin (12-bit) histogram |
| 510 | } |
| 511 | return; |
| 512 | } |
| 513 | |
| 514 | // Large path: init mbarriers + state, then dispatch fused or twopass |
| 515 | const uint32_t per_block = |
| 516 | (params.stride + CS - 1) / CS; // how many elements per block |
| 517 | constexpr uint32_t kFusedMax = ((CS == 16) ? kFusedStagesCS16 |
| 518 | : (CS == 8) ? kFusedStagesCS8 |
| 519 | : kMaxSinglePassStages) * |
| 520 | kSizePerStage; |
| 521 | const bool use_singlepass = |
| 522 | per_block <= |
| 523 | kFusedMax; // single pass or TMA streaming: histogram+scatter |
| 524 | |
| 525 | // Select smem type and stage count at compile time based on CS |
| 526 | constexpr uint32_t kFusedStages = (CS == 16) ? kFusedStagesCS16 |
| 527 | : (CS == 8) ? kFusedStagesCS8 |
| 528 | : kMaxSinglePassStages; |
| 529 | using FusedSmem = SmemFused<kFusedStages>; |
| 530 | |
| 531 | extern __shared__ uint8_t sr[]; |
| 532 | |
| 533 | constexpr uint32_t kTieWsPerRow = |
| 534 | TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK; |
| 535 | hist4096::Tie* row_tie_ws = params.tie_ws + row * kTieWsPerRow; |
| 536 | |
| 537 | if (use_singlepass) { |
| 538 | auto* smem = reinterpret_cast<FusedSmem*>(sr); |
| 539 | const uint32_t sp_stages = (per_block + kSizePerStage - 1) / kSizePerStage; |
| 540 | if (tx < sp_stages) { |
| 541 | mbarrier_init(&smem->barrier[0][tx], |
| 542 | 1); // init 1 barrier per TMA stage - |
| 543 | // signal when async copies complete |
| 544 | } |
| 545 | __syncthreads(); |
| 546 | uint32_t phases[kFusedStages] = |
| 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 | // 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], |
| 555 | 1); // init 2×2=4 barriers (2 passes × 2 stages) |
| 556 | } |
| 557 | __syncthreads(); |
| 558 | uint32_t hp[kStreamingStagesCS4] = {0, |
| 559 | 0}; // histogram+scatter pass counters |
| 560 | large_topk<TopK, CS, Smem4, false>(in, out, sl, hp, row_tie_ws); |
| 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) |
| 574 | cooperative_topk_cs4(CooperativeTopKParams<TopK> params) { |
| 575 | cooperative_topk_body<TopK, 4>(params); |
| 576 | } |
| 577 | |
| 578 | template <uint32_t TopK> |
| 579 | __global__ void __launch_bounds__(hist4096::kBlockSize, 1) |
| 580 | __cluster_dims__(1, 8, 1) |
| 581 | cooperative_topk_cs8(CooperativeTopKParams<TopK> params) { |
| 582 | cooperative_topk_body<TopK, 8>(params); |
| 583 | } |
| 584 | |
| 585 | template <uint32_t TopK> |
| 586 | __global__ void __launch_bounds__(hist4096::kBlockSize, 1) |
| 587 | __cluster_dims__(1, 16, 1) |
| 588 | cooperative_topk_cs16(CooperativeTopKParams<TopK> params) { |
| 589 | cooperative_topk_body<TopK, 16>(params); |
| 590 | } |
| 591 | |
| 592 | constexpr size_t kSmemSize4_base = sizeof(Smem4); |
| 593 | 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 | |
| 601 | } // namespace cooperative |
| 602 | |
| 603 | } // namespace vllm |
| 604 | |
| 605 | #endif // COOPERATIVE_TOPK_CUH_ |