From 2848ae337b998f6f235d2e518cc95af799d158e6 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Fri, 19 Jun 2026 15:19:01 +0000 Subject: [PATCH 1/7] Rebase cooperative top-k onto libtorch_stable migration. Port cooperative cluster top-k kernels and launchers to csrc/libtorch_stable/, gate registration with VLLM_ENABLE_COOPERATIVE_TOPK, and route decode sparse indexer to cooperative_topk when eligible. Co-authored-by: Cursor Signed-off-by: LopezCastroRoberto --- CMakeLists.txt | 30 + csrc/libtorch_stable/cooperative_topk.cu | 156 +++ csrc/libtorch_stable/cooperative_topk.cuh | 1088 +++++++++++++++++ csrc/libtorch_stable/ops.h | 8 + csrc/libtorch_stable/persistent_topk.cuh | 631 +++++++++- csrc/libtorch_stable/torch_bindings.cpp | 9 + tests/kernels/test_top_k_per_row.py | 177 ++- .../layers/sparse_attn_indexer.py | 27 +- 8 files changed, 2079 insertions(+), 47 deletions(-) create mode 100644 csrc/libtorch_stable/cooperative_topk.cu create mode 100644 csrc/libtorch_stable/cooperative_topk.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index e95fe38d3299..c42cbccabaf4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -382,6 +382,36 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/custom_all_reduce.cu" "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") + if(VLLM_GPU_LANG STREQUAL "CUDA" AND + DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0) + + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS + "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS + "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + + if(COOPERATIVE_TOPK_ARCHS) + list(APPEND VLLM_STABLE_EXT_SRC + "csrc/libtorch_stable/cooperative_topk.cu") + list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1") + + cuda_archs_loose_intersection(COOPERATIVE_TOPK_SM90_ARCHS + "9.0a" "${COOPERATIVE_TOPK_ARCHS}") + if(COOPERATIVE_TOPK_SM90_ARCHS) + list(APPEND VLLM_GPU_FLAGS + "-DVLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY=1") + endif() + + set_source_files_properties( + "csrc/libtorch_stable/cooperative_topk.cu" + PROPERTIES CUDA_ARCHITECTURES "${COOPERATIVE_TOPK_ARCHS}") + endif() + endif() + if(VLLM_GPU_LANG STREQUAL "CUDA") SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") diff --git a/csrc/libtorch_stable/cooperative_topk.cu b/csrc/libtorch_stable/cooperative_topk.cu new file mode 100644 index 000000000000..97b7b38f42fe --- /dev/null +++ b/csrc/libtorch_stable/cooperative_topk.cu @@ -0,0 +1,156 @@ +// Cooperative cluster TopK for DeepSeek V3 sparse attention indexer. +// See cooperative_topk.cuh for kernel implementation. + +#include + +#include "torch_utils.h" + +#ifndef USE_ROCM + #include "cooperative_topk.cuh" +namespace ct = vllm::cooperative; +#endif + +#ifndef USE_ROCM +template +void launch_cooperative_cluster(ct::CooperativeTopKParams& params, + size_t smem, cudaStream_t stream) { + auto kernel = []() { + if constexpr (CS == 16) { + #ifndef VLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY + return &ct::cooperative_topk_cs16; + #else + static_assert(CS != 16, + "CS=16 cooperative_topk requires non-portable cluster " + "support"); + #endif + } else if constexpr (CS == 8) { + return &ct::cooperative_topk_cs8; + } else { + static_assert(CS == 4, "unsupported cooperative_topk cluster size"); + return &ct::cooperative_topk_cs4; + } + }(); + #ifndef VLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY + if constexpr (CS > 8) { + cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, + 1); + } + #endif + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(params.num_rows, CS); + cfg.blockDim = dim3(ct::kBlockSize); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim = {1, CS, 1}; + cfg.numAttrs = 1; + cfg.attrs = attrs; + cudaError_t err = cudaLaunchKernelEx(&cfg, kernel, params); + STD_TORCH_CHECK(err == cudaSuccess, "cooperative_topk launch failed: ", + cudaGetErrorString(err)); +} + +template +void launch_cooperative_topk_impl(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, + int64_t max_seq_len) { + const int64_t num_rows = logits.size(0); + const cudaStream_t stream = get_current_cuda_stream(); + + const uint32_t stride = static_cast(logits.stride(0)); + // 32 = max clusters for CS=4 (32 x 4 = 128 CTAs = 66% of SMs, leaves + // headroom) + STD_TORCH_CHECK( + num_rows <= 32, + "cooperative_topk supports <=32 rows; use persistent_topk for " + "larger batches"); + + STD_TORCH_CHECK( + stride % 4 == 0, + "cooperative_topk: stride must be multiple of 4 for TMA " + "alignment, got stride (max_model_len)=", + stride); + + STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); + STD_TORCH_CHECK( + workspace.scalar_type() == torch::headeronly::ScalarType::Byte, + "workspace must be uint8"); + + ct::CooperativeTopKParams params; + params.input = logits.const_data_ptr(); + params.output = output.mutable_data_ptr(); + params.lengths = lengths.const_data_ptr(); + params.num_rows = static_cast(num_rows); + params.stride = stride; + params.tie_ws = reinterpret_cast( + workspace.mutable_data_ptr()); + + constexpr uint32_t kTieWsPerRow = + TopK <= ct::kBlockSize ? ct::kMaxTies : TopK; + STD_TORCH_CHECK( + workspace.size(0) >= static_cast(num_rows * kTieWsPerRow * + sizeof(ct::Tie)), + "workspace too small"); + + #ifndef VLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY + if (num_rows <= 4) { + launch_cooperative_cluster(params, ct::kSmemSize8, stream); + } else if (num_rows <= 8) { + #else + if (num_rows <= 8) { + #endif + launch_cooperative_cluster(params, ct::kSmemSize8, stream); + } else { + launch_cooperative_cluster(params, ct::kSmemSize4, stream); + } +} +#endif // USE_ROCM + +void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len) { +#ifndef USE_ROCM + STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor"); + STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor"); + STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor"); + STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float, + "Only float32 supported"); + STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int, + "lengths must be int32"); + STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int, + "output must be int32"); + STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); + STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2, + "lengths must be 1D or 2D"); + STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous"); + STD_TORCH_CHECK(output.dim() == 2, "output must be 2D"); + const int64_t num_rows = logits.size(0); + STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch"); + STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k, + "output size mismatch"); + STD_TORCH_CHECK(k == 512 || k == 1024 || k == 2048, + "cooperative_topk supports k=512, k=1024, or k=2048, got k=", + k); + + if (k == 512) { + launch_cooperative_topk_impl<512>(logits, lengths, output, workspace, + max_seq_len); + } else if (k == 1024) { + launch_cooperative_topk_impl<1024>(logits, lengths, output, workspace, + max_seq_len); + } else { + launch_cooperative_topk_impl<2048>(logits, lengths, output, workspace, + max_seq_len); + } +#else + STD_TORCH_CHECK(false, "cooperative_topk is not supported on ROCm"); +#endif +} diff --git a/csrc/libtorch_stable/cooperative_topk.cuh b/csrc/libtorch_stable/cooperative_topk.cuh new file mode 100644 index 000000000000..cb5a608c1250 --- /dev/null +++ b/csrc/libtorch_stable/cooperative_topk.cuh @@ -0,0 +1,1088 @@ +/* + * Cooperative TopK kernel for DSA Indexer + */ + +#ifndef COOPERATIVE_TOPK_CUH_ +#define COOPERATIVE_TOPK_CUH_ + +#include +#include +#include +#include +#include +#include +#include + +namespace vllm { +namespace cooperative { + +// TopK is now a template parameter (512 or 1024) +constexpr uint32_t kBlockSize = 1024; +constexpr uint32_t kHistBits = 10; +constexpr uint32_t kHistBins = 1 << kHistBits; +constexpr uint32_t RADIX = 256; +constexpr uint32_t kMaxTies = 1024; +static_assert(kMaxTies <= kBlockSize, + "tie_handle requires kMaxTies <= kBlockSize"); +constexpr uint32_t kWarpSize = 32; +constexpr uint32_t kNumWarps = kBlockSize / kWarpSize; + +constexpr uint32_t kElemPerStage = 16; +constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize; // 16384 + +// CS=4: 2 TMA stages (double buffer), two-pass +constexpr uint32_t kNumStages4 = 2; +// CS=8: 4 TMA stages, single-pass (all data stays in smem) +constexpr uint32_t kNumStages8 = + 2; // 2 stages × 16K = 32K per block (enough for CS=8) +// Max data per block with CS=8: ceil(262144/8) = 32768 ≤ 4 × 8192 = 32768 +constexpr uint32_t kMaxSeqLen8 = kNumStages8 * kSizePerStage * 8; // 262144 +constexpr uint32_t kNumStages16 = 2; // double-buffer for CS=16 +constexpr uint32_t kMaxSeqLen16 = kNumStages16 * kSizePerStage * 16; // 524288 + +// Register path +constexpr uint32_t kHist4096Bits = 12; +constexpr uint32_t kHist4096Bins = 1 << kHist4096Bits; // 4096 +constexpr uint32_t kHist4096VecsPerThread = 4; +constexpr uint32_t kHist4096MaxLen = + kHist4096VecsPerThread * 4 * kBlockSize; // 16384 +constexpr uint32_t kHist4096Items = kHist4096Bins / kBlockSize; // 4 + +// CS=4 single-pass path +constexpr uint32_t kMaxSinglePassStages = 3; +constexpr uint32_t kMaxSinglePassPerBlock = + kMaxSinglePassStages * kSizePerStage; // 49152 + +constexpr uint32_t kStreamNumStages = 2; + +struct alignas(16) MatchBin { + uint32_t bin, above_count, equal_count; +}; +struct alignas(8) Tie { + uint32_t idx; + float score; +}; + +template +struct CooperativeTopKParams { + const float* __restrict__ input; + int32_t* __restrict__ output; + const int32_t* __restrict__ lengths; + Tie* __restrict__ tie_ws; // per-row tie workspace, see kTieWsPerRow + uint32_t num_rows, stride; +}; + +// ============================================================================ +// Common helpers +// ============================================================================ + +// converts the float32 score to a 32-bit ordered unsigned integer — the full +// precision key for radix sorting +__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits -> +// bin 0-4095) +template +__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) { + __half h = __float2half_rn(x); + uint16_t bits = __half_as_ushort(h); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) + : static_cast(bits | 0x8000); + return key >> (16 - kBits); +} + +// running sum within each warp — thread 0 gets its own value, thread 1 gets +// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc. +__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane, + uint32_t v) { +#pragma unroll + for (uint32_t o = 1; o < 32; o *= 2) { + uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o); + if (lane >= o) v += n; + } + return v; +} + +// Returns the sum of a value across all 32 threads in the warp, and every +// thread gets the same result SM90+ PTX instruction that does a hardware +// warp-wide reduction in a single instruction w.r.t. warp::reduce_sum(), which +// uses a __shfl_xor_sync butterfly tree (5 shuffles for 32 lanes) +__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) { + uint32_t r; + asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v)); + return r; +} + +// only CS adjacent lanes participate (sub-warp reduce), in opposite to +// warp_reduce_sum_full +template +__device__ __forceinline__ uint32_t warp_reduce_sum_subN(uint32_t v) { +#pragma unroll + for (uint32_t m = N >> 1; m > 0; m >>= 1) + v += __shfl_xor_sync(0xFFFFFFFF, v, m, 32); + return v; +} + +// ============================================================================ +// Helpers +// ============================================================================ + +__device__ __forceinline__ uint32_t extract_coarse_bin(float x) { + return extract_coarse_bin_N(x); +} + +__device__ __forceinline__ void mbarrier_init(uint64_t* a, uint32_t n) { + cuda::ptx::mbarrier_init(a, n); +} +__device__ __forceinline__ void mbarrier_wait(uint64_t* a, uint32_t p) { + while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed, + cuda::ptx::scope_cta, a, p)); +} +__device__ __forceinline__ void mbarrier_arrive_expect_tx(uint64_t* a, + uint32_t t) { + cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed, + cuda::ptx::scope_cta, + cuda::ptx::space_shared, a, t); +} +__device__ __forceinline__ void tma_load(void* d, const void* s, uint32_t n, + uint64_t* m) { + cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global, d, + s, n, m); +} + +// ============================================================================ +// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered +// key Each round narrows by 8 bits until ties are fully resolved +// ============================================================================ + +template +__device__ void tie_handle(const Tie* ties, uint32_t num_ties, + uint32_t num_above, int32_t* output, void* _smem) { + struct TS { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[RADIX]; + uint32_t warp_sum[kNumWarps]; + }; + auto* s = static_cast(_smem); + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize, wi = tx / kWarpSize; + + // Each thread loads one tie element. + const bool has = tx < num_ties; + const auto tie = has ? ties[tx] : Tie{0, 0.0f}; + const uint32_t key = convert_to_uint32_v2(tie.score); + + bool active = has; // tracks whether this thread's tie is still a candidate. + uint32_t remain = + TopK - num_above; // decreases each round as ties are resolved. + uint32_t wpos = TopK; // wpos will hold the final output position. + s->counter = 0; + __syncthreads(); + + // The 4-round radix loop - each round narrows by 8 bits until ties are fully + // resolved +#pragma unroll + for (int r = 0; r < 4; r++) { + uint32_t sh = 24 - r * 8; // round 0: bits 31-24, round 1: 23-16, etc. + uint32_t bin = (key >> sh) & 0xFF; // this tie's 8-bit bin for this round + + // Step 1: Build 256-bin histogram. + if (tx < RADIX) s->histogram[tx] = 0; + __syncthreads(); + if (active) atomicAdd(&s->histogram[bin], 1); + __syncthreads(); + + // Step 2: Prefix scan to find threshold + uint32_t hv = 0, wi2 = 0; + if (tx < RADIX) { + hv = s->histogram[tx]; + wi2 = warp_inclusive_sum(li, hv); + if (li == kWarpSize - 1) s->warp_sum[wi] = wi2; + } + __syncthreads(); + + if (tx < RADIX) { + auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; + auto tot = warp_reduce_sum_full(tmp); + auto inter = warp_reduce_sum_full(li < wi ? tmp : 0); + auto above = tot - (inter + wi2); + if (above < remain && above + hv >= remain) { + s->match = {tx, above, remain - above}; + } + } + __syncthreads(); + + // Step 3: Scatter + auto [thr, na, _] = s->match; // threshold bin, num above, unused + if (active) { + if (bin > thr) { + wpos = num_above + + atomicAdd(&s->counter, 1); // above -> place in output directly + active = false; + } else if (bin < thr) + active = false; // below -> discard + else if (r == 3) + wpos = TopK - atomicAdd(&s->match.equal_count, + -1u); // last round: place remaining + } + remain -= na; + if (!remain) break; // all ties resolved early + } + // Final write + if (wpos < TopK) output[wpos] = tie.idx; +} + +// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048). +// tie_handle assumes 1 tie per thread (max 1024). +// This version handles 2 ties per thread via kPerThread=2 +template +__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties, + uint32_t num_above, int32_t* output, + void* _smem) { + static_assert(TopK > kBlockSize); + struct TS { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[RADIX]; + uint32_t warp_sum[kNumWarps]; + }; + auto* s = static_cast(_smem); + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize; + const auto wi = tx / kWarpSize; + + constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize; + Tie my_ties[kPerThread]; + uint32_t keys[kPerThread]; + bool active[kPerThread]; + + for (uint32_t e = 0; e < kPerThread; e++) { + uint32_t idx = e * kBlockSize + tx; + if (idx < num_ties) { + my_ties[e] = ties[idx]; + keys[e] = convert_to_uint32_v2(ties[idx].score); + active[e] = true; + } else { + my_ties[e] = {0, 0.0f}; + keys[e] = 0; + active[e] = false; + } + } + + uint32_t remain = TopK - num_above; + s->counter = 0; + __syncthreads(); + + for (int r = 0; r < 4; r++) { + uint32_t sh = 24 - r * 8; + if (tx < RADIX) { + s->histogram[tx] = 0; + } + __syncthreads(); + + for (uint32_t e = 0; e < kPerThread; e++) { + if (active[e]) { + atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1); + } + } + __syncthreads(); + + uint32_t hv = 0; + if (tx < RADIX) { + hv = s->histogram[tx]; + auto wi2 = warp_inclusive_sum(li, hv); + if (li == kWarpSize - 1) { + s->warp_sum[wi] = wi2; + } + } + __syncthreads(); + if (tx < RADIX) { + auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; + auto total = warp_reduce_sum_full(tmp2); + auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0); + auto wi2 = warp_inclusive_sum(li, hv); + auto above = total - (inter + wi2); + if (above < remain && above + hv >= remain) { + s->match = { + .bin = tx, .above_count = above, .equal_count = remain - above}; + } + } + __syncthreads(); + + auto thr = s->match.bin; + auto na = s->match.above_count; + + for (uint32_t e = 0; e < kPerThread; e++) { + if (!active[e]) { + continue; + } + uint32_t bin = (keys[e] >> sh) & 0xFF; + if (bin > thr) { + uint32_t wpos = num_above + atomicAdd(&s->counter, 1); + if (wpos < TopK) { + output[wpos] = my_ties[e].idx; + } + active[e] = false; + } else if (bin < thr) { + active[e] = false; + } else if (r == 3) { + uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u); + if (wpos < TopK) { + output[wpos] = my_ties[e].idx; + } + } + } + + num_above += na; + remain -= na; + __syncthreads(); + s->counter = 0; + __syncthreads(); + } +} + +// ============================================================================ +// DSMEM histogram reduce +// ============================================================================ + +template +__device__ __forceinline__ void dsmem_hist_reduce(uint32_t* histogram) { + static_assert(kHistBins <= kBlockSize); + auto cluster = cooperative_groups::this_cluster(); + cluster.sync(); + const auto tx = threadIdx.x; + const auto rank = blockIdx.y; + constexpr auto kLocal = kHistBins / CS; + const auto off = kLocal * rank; + if (tx < kHistBins) { + const auto addr = &histogram[off + tx / CS]; + const auto src = cluster.map_shared_rank(addr, tx % CS); + *src = warp_reduce_sum_subN(*src); + } + cluster.sync(); +} + +// ============================================================================ +// Find threshold from reduced histogram +// ============================================================================ + +// NOTE: caller must ensure a cluster.sync() or __syncthreads() happened +// before calling this, so warp_sum writes are visible across warps. +// The first internal __syncthreads() is still needed for the warp_sum exchange. +template +__device__ __forceinline__ void find_threshold(uint32_t* histogram, + uint32_t* warp_sum, + uint32_t* counter_gt, + uint32_t* counter_eq, + MatchBin* match) { + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize, wi = tx / kWarpSize; + const auto value = tx < kHistBins ? histogram[tx] : 0; + const auto winc = warp_inclusive_sum(li, value); + if (li == kWarpSize - 1) warp_sum[wi] = winc; + __syncthreads(); + const auto tmp = warp_sum[li]; + const auto total = warp_reduce_sum_full(tmp); + auto pfx = warp_reduce_sum_full(li < wi ? tmp : 0) + winc; + const auto above = total - pfx; + if (tx < kHistBins && above < TopK && above + value >= TopK) { + *counter_gt = *counter_eq = 0; + *match = {.bin = tx, .above_count = above, .equal_count = value}; + } + __syncthreads(); +} + +// ============================================================================ +// Register-based single-CTA fast path for seq_len <= 16384 +// 4 float4 per thread × 1024 threads = 16384 elements max +// Uses 4096-bin (12-bit) histogram for better precision +// ============================================================================ + +template +struct Histogram4096Smem { + static constexpr uint32_t HIST_BINS = 1 << HIST_BITS; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + MatchBin match; + uint32_t warp_sum[kNumWarps]; + union { + uint32_t histogram[HIST_BINS]; + Tie tie_buffer[kMaxTies]; + }; + static_assert(sizeof(uint32_t) * HIST_BINS >= TopK * sizeof(Tie), + "histogram union must be large enough for TopK ties"); +}; + +template +__device__ void histogram_4096_topk(const float* __restrict__ scores, + int32_t* __restrict__ output, + uint32_t length, void* _smem) { + constexpr uint32_t HIST_BINS = 1 << HIST_BITS; + constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize; + static_assert(HIST_BINS >= kBlockSize, + "HIST_BITS must give >= kBlockSize bins"); + + using Smem = Histogram4096Smem; + auto* smem = static_cast(_smem); + const auto tx = threadIdx.x; + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + + // Phase 1: Load all data into RF + build histogram + float4 + vecs[VECS_PER_THREAD]; // 4 vectors x 4 floats = 16 elements per thread + if constexpr (ITEMS_PER_THREAD >= 4) { + // Zero the histogram (SMEM writes) + for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++) + reinterpret_cast( + smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] = + make_uint4(0, 0, 0, 0); + } else { + if (tx < HIST_BINS) smem->histogram[tx] = 0; + } + if (tx == 0) { + smem->counter_gt = 0; + smem->counter_eq = 0; + } +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { + const uint32_t base = (tx + v * kBlockSize) * 4; + if (base < length) { + // TODO (roberto): try if 256-bit loads are faster (GMEM->RF) + vecs[v] = *reinterpret_cast(scores + base); + } + } + __syncthreads(); + + // Build histogram from RF via atomic adds into the shared histogram + bool done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + atomicAdd(&smem->histogram[extract_coarse_bin_N(elems[e])], + 1); + } + } + } + __syncthreads(); + + // Phase 2: Prefix scan to find threshold bin + // Multi-element scan (4096 bins: 4 per thread) + uint32_t orig[ITEMS_PER_THREAD]; + uint32_t local_sum = 0; + + // Step 1: Each thread sums its 4 bins +#pragma unroll + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { + orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i]; + local_sum += orig[i]; + } + + // Step 2: Warp-level inclusive prefix sum on local_sum + const auto warp_inc = warp_inclusive_sum(lane_id, local_sum); + if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc; + __syncthreads(); + + // Step 3: Inter-warp prefix via redux.sync + const auto tmp = smem->warp_sum[lane_id]; + uint32_t prefix = warp_reduce_sum_full( + lane_id < warp_id ? tmp : 0); // sum of all prior warps + prefix += + warp_inc - local_sum; // exclusive prefix within this thread's position + + // Step 4: Find threshold - scan 4 bins, accumulate prefix +#pragma unroll + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { + prefix += orig[i]; + const auto above = length - prefix; // elements in bins ABOVE this one + if (above < TopK && above + orig[i] >= TopK) { + smem->match = {.bin = tx * ITEMS_PER_THREAD + i, + .above_count = above, + .equal_count = orig[i]}; + } + } + + __syncthreads(); + + // Phase 3: Scatter from registers + const auto [thr_bin, num_above, num_equal] = smem->match; + const bool need_tie = (num_equal + num_above > TopK); + + done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + const uint32_t bin = extract_coarse_bin_N(elems[e]); + if (bin > thr_bin) { + output[atomicAdd(&smem->counter_gt, 1)] = + idx; // above -> output directly + } else if (bin == thr_bin) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (!need_tie) { + if (pos + num_above < TopK) { + output[pos + num_above] = idx; // all fit + } + } else { + if (pos < TopK) { + smem->tie_buffer[pos] = {idx, elems[e]}; // store for refirement + } + } + } + // else: bin < thr_bin - discard (not in top-k) + } + } + } + + // Phase 4: Tie-breaking + if (!need_tie) return; + __syncthreads(); + + // Fast warp-ballot tie-breaking for small tie counts + const uint32_t num_ties = min(num_equal, static_cast(TopK)); + const uint32_t topk_remain = + TopK - num_above; // pick exactly remaining elements to fill topK + + auto is_greater = [](const Tie& a, const Tie& b) { + return (a.score > b.score) || (a.score == b.score && a.idx < b.idx); + }; + + if (num_ties <= kWarpSize) { + // <=32 ties - Use warp ballot + // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024 + // comparisons in one instruction per warp. O(1) work. + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + if (lane_id >= num_ties || warp_id >= num_ties) return; + const uint32_t mask = (1ull << num_ties) - 1u; + const auto tie = smem->tie_buffer[lane_id]; // each lane holds one tie + const auto target = + smem->tie_buffer[warp_id]; // each warp evaluates one candidate + const bool pred = + is_greater(tie, target); // compare all ties against target + const auto rank = static_cast( + __popc(__ballot_sync(mask, pred))); // count how many are greater + if (lane_id == 0 && rank < topk_remain) { + output[num_above + rank] = target.idx; // place at correct position + } + } else if (num_ties <= + kWarpSize * + 2) { // TODO (roberto): try to refactor this with <=32 case + // Same idea but each thread handles 2 tie elements + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + const auto lane1 = lane_id + kWarpSize; + const auto warp1 = warp_id + kWarpSize; + const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__}; + const auto tie0 = smem->tie_buffer[lane_id]; + const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid; + if (warp_id < num_ties) { + const auto target = smem->tie_buffer[warp_id]; + const auto r0 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); + const auto r1 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); + if (lane_id == 0 && r0 + r1 < topk_remain) + output[num_above + r0 + r1] = target.idx; + } + if (warp1 < num_ties) { + const auto target = smem->tie_buffer[warp1]; + const auto r0 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); + const auto r1 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); + if (lane_id == 0 && r0 + r1 < topk_remain) + output[num_above + r0 + r1] = target.idx; + } + } else { + // Large tie count: fall back to 4-round radix-256 sort + if constexpr (TopK <= kBlockSize) { + tie_handle(smem->tie_buffer, num_ties, num_above, output, smem); + } else { + tie_handle_large(smem->tie_buffer, num_ties, num_above, output, + smem); + } + } +} + +// ============================================================================ +// Streaming single-CTA path for seq_len 16K–32K +// TMA double-buffer, 4096-bin histogram, no cluster sync +struct StreamSmem { + uint64_t barrier[2][kStreamNumStages]; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + MatchBin match; + uint32_t warp_sum[kNumWarps]; + union { + uint32_t histogram[kHist4096Bins]; + Tie tie_buffer[kMaxTies]; + }; + alignas(128) float score_buffer[kStreamNumStages][kSizePerStage]; +}; + +// Streams data through shared memory in chunks, processing each chunk before +// loading the next overwrites each buffer after processing it (the epilogue +// prefetch loads the next chunk into the same slot) +template +__device__ void tma_stream_pass(const float* scores, uint32_t length, + uint32_t thr_bin, int32_t* indices, + uint32_t* phases, SmemType* smem) { + const auto tx = threadIdx.x; + const auto lane = tx % kWarpSize; + const auto ni = + (length + kSizePerStage - 1) / kSizePerStage; // total stages needed + const auto la = + (length + 3u) & ~3u; // length rounded up to float4 (TMA alignment) + const auto pass = + kIsScatter ? 1 : 0; // barrier dim: [0] for histogram, [1] for scatter + + // Prologue: issue initial TMA loads - prefill the pipeline + if (tx == 0) { +#pragma unroll + for (uint32_t i = 0; i < kStages; i++) { + if (i >= ni) { + break; + } + const auto o = i * kSizePerStage; + const auto sz = min(kSizePerStage, la - o) * sizeof(float); + tma_load(smem->score_buffer[i], scores + o, sz, + &smem->barrier[pass][i]); // cp.async.bulk is non-blocking + mbarrier_arrive_expect_tx(&smem->barrier[pass][i], sz); + } + } + + // Main loop: process stages + for (uint32_t it = 0; it < ni; it++) { + const auto b = it % kStages; // which buffer slot (0 or 1) + const auto o = it * kSizePerStage; + const auto sz = min(kSizePerStage, length - o); + + if (lane == 0) { + mbarrier_wait(&smem->barrier[pass][b], + phases[b] & 1); // wait for the data + } + phases[b]++; // advances the phase for next time this slot is reused + __syncwarp(); + +#pragma unroll + for (uint32_t i = 0; i < kElemPerStage; i++) { + const auto li = tx + i * kBlockSize; + if (li >= sz) { + break; + } + const auto sc = smem->score_buffer[b][li]; + const auto bn = extract_coarse_bin_N(sc); + if constexpr (kIsScatter) { // compile-time branch + // Scatter pass: place above-threshold and collect ties + const auto gi = o + li; + if (bn > thr_bin) { + indices[atomicAdd(&smem->counter_gt, 1)] = gi; + } else if (bn == thr_bin) { + const auto p = atomicAdd(&smem->counter_eq, 1); + if (p < kMaxTies) { + smem->tie_buffer[p] = {gi, sc}; + } + } + } else { + // Histogram pass: just count + atomicAdd(&smem->histogram[bn], 1); + } + } + __syncthreads(); // ensures all threads finished processing their buffer + // before next TMA load + + // Epilogue: issue next TMA load + if (tx == 0 && it + kStages < ni) { + const auto no = (it + kStages) * kSizePerStage; + const auto nsz = min(kSizePerStage, la - no) * sizeof(float); + tma_load(smem->score_buffer[b], scores + no, nsz, + &smem->barrier[pass][b]); + mbarrier_arrive_expect_tx(&smem->barrier[pass][b], nsz); + } + } +} + +// ============================================================================ +// Fused path: single TMA pass, rescan smem for scatter +// ============================================================================ + +// Fused shared memory layout for cluster cooperative paths. +// kPasses=1 for single-pass (CS=8, CS=4 singlepass), kPasses=2 for two-pass +// (CS=4). +template +struct SmemFused { + uint64_t barrier[kPasses][kStages]; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + alignas(128) MatchBin match; + uint32_t warp_sum[kNumWarps]; + union { + uint32_t histogram[kHistBins]; + Tie tie_buffer[kMaxTies]; + }; + alignas(128) float score_buffer[kStages][kSizePerStage]; +}; + +using Smem8 = SmemFused; +using Smem16 = SmemFused; +using Smem4 = SmemFused; +using SmemSinglePass = SmemFused; + +// Cluster-cooperative large path. +// kFused=true: all TMA stages resident, single-pass histogram + scatter (rescan +// from smem). kFused=false: TMA double-buffer streaming, two passes (histogram +// then scatter). +template +__device__ void large_topk(const float* __restrict__ row_input, + int32_t* __restrict__ row_output, uint32_t seq_len, + uint32_t* phases, Tie* tie_ws) { + const auto rank = blockIdx.y; // this block's position in cluster + const auto tx = threadIdx.x; + const auto lane = tx % kWarpSize; + + extern __shared__ uint8_t smem_raw[]; + auto* smem = reinterpret_cast(smem_raw); + int32_t* s_topk = reinterpret_cast(smem_raw + sizeof(SmemType)); + + // Partition row across cluster ranks + constexpr uint32_t kAlign = 4; + const auto units = + (seq_len + kAlign - 1) / kAlign; // float4-aligned element count + const auto base = units / CS, extra = units % CS; // elements per block + const auto lu = base + (rank < extra ? 1u : 0u); // remainder blocks + const auto ou = + rank * base + min(rank, extra); // this block's count (load-balanced) + const auto my_start = ou * kAlign; // global start offset + const auto my_len = min(my_start + lu * kAlign, seq_len) - + my_start; // actual length of this block + const auto num_iters = + (my_len + kSizePerStage - 1) / kSizePerStage; // TMA stages needed + const auto len_aligned = (my_len + 3u) & ~3u; + + if constexpr (kFused) { + // Fused init + TMA prologue + if (tx < kHistBins) { + smem->histogram[tx] = 0; // all threads zero histogram + } + if (tx == 0) { // thread 0 issues TMA - then all threads continue working + // until mbarrier sync + smem->counter_gt = 0; + smem->counter_eq = 0; + for (uint32_t i = 0; i < num_iters; i++) { + const auto off = i * kSizePerStage; + const auto sz = min(kSizePerStage, len_aligned - off) * sizeof(float); + tma_load(smem->score_buffer[i], row_input + my_start + off, sz, + &smem->barrier[0][i]); // cp.async.bulk of size kSizePerStage + // × sizeof(float) + mbarrier_arrive_expect_tx(&smem->barrier[0][i], sz); + } + } + __syncthreads(); + + // Histogram build. ILP unroll-by-2, no inter-stage sync + for (uint32_t iter = 0; iter < num_iters; iter++) { + const auto off = iter * kSizePerStage; + const auto sz = min(kSizePerStage, my_len - off); + if (lane == 0) { + mbarrier_wait(&smem->barrier[0][iter], + phases[iter] & 1); // wait for TMA + } + phases[iter]++; + __syncwarp(); +#pragma unroll + for (uint32_t i = 0; i < kElemPerStage; i += 2) { + const auto li0 = tx + i * kBlockSize; + const auto li1 = tx + (i + 1) * kBlockSize; + if (li0 >= sz) { + break; + } + const auto b0 = extract_coarse_bin(smem->score_buffer[iter][li0]); + if (li1 < sz) { + const auto b1 = extract_coarse_bin(smem->score_buffer[iter][li1]); + atomicAdd(&smem->histogram[b0], 1); + atomicAdd(&smem->histogram[b1], 1); + } else { + atomicAdd(&smem->histogram[b0], 1); + } + } + } + } else { + // Twopass: init then stream histogram pass + if (tx < kHistBins) { + smem->histogram[tx] = 0; + } + if (tx == 0) { + smem->counter_gt = 0; + smem->counter_eq = 0; + } + __syncthreads(); + tma_stream_pass( + row_input + my_start, my_len, 0, nullptr, phases, smem); + } + + // DSMEM all-reduce + find threshold + dsmem_hist_reduce( + smem->histogram); // each block histogram is summed across all CS blocks + find_threshold(smem->histogram, smem->warp_sum, &smem->counter_gt, + &smem->counter_eq, &smem->match); + + const auto thr = smem->match.bin; + + if constexpr (kFused) { + // Fused scatter: rescan score_buffer (still in smem) + for (uint32_t iter = 0; iter < num_iters; iter++) { + const auto off = iter * kSizePerStage; + const auto sz = min(kSizePerStage, my_len - off); +#pragma unroll + for (uint32_t i = 0; i < kElemPerStage; i++) { + const auto li = tx + i * kBlockSize; + if (li >= sz) { + break; + } + const auto score = smem->score_buffer[iter][li]; // still in smem + const auto bin = extract_coarse_bin(score); + const auto gidx = off + li; + if (bin > thr) { + s_topk[atomicAdd(&smem->counter_gt, 1)] = gidx; // above -> s_topk + } else if (bin == thr) { + const auto p = atomicAdd(&smem->counter_eq, + 1); // equal -> ties (later refinement) + if (p < kMaxTies) { + smem->tie_buffer[p] = {gidx, score}; + } + } + } + } + __syncthreads(); + } else { + // Twopass scatter: re-stream data via TMA + uint32_t scatter_phases[kNumStages4] = {0, 0}; + tma_stream_pass( + row_input + my_start, my_len, thr, s_topk, scatter_phases, smem); + } + + // Output collection via DSMEM prefix sum + constexpr uint32_t kAboveBits = 16; + constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1; + static_assert(kAboveMask >= TopK); + static_assert(kAboveMask >= kMaxSinglePassPerBlock, + "kAboveBits must cover max per-block element count"); + + const uint32_t la = smem->counter_gt; + const uint32_t le_full = smem->counter_eq; + const uint32_t le = + min(le_full, kMaxTies); // written smem tie_buffer entries + + __shared__ uint32_t s_local_counts[CS]; + __shared__ uint32_t s_prefix_packed; + __shared__ uint32_t s_total_above, s_total_equal; + + auto cluster = cooperative_groups::this_cluster(); + if (tx < CS) { + // Pack written tie counts into 32-bit: (equal << 16) | above. + // `le_full` may exceed the per-block tie buffer cap; using it here creates + // holes in tie_ws and can make TopK=2048 refine unwritten workspace slots. + const uint32_t packed = (le << kAboveBits) | la; + const auto dst = cluster.map_shared_rank(s_local_counts, tx); + dst[rank] = packed; // write my count to every block's s_local_counts[rank] + } + cluster.sync(); + + // Thread 0 computes serial prefix sum + if (tx == 0) { + uint32_t prefix = 0, ta = 0, te = 0; + for (uint32_t i = 0; i < CS; i++) { + if (i == rank) { + s_prefix_packed = prefix; // my prefix + } + ta += s_local_counts[i] & kAboveMask; // total above + te += s_local_counts[i] >> kAboveBits; // total equal + prefix += s_local_counts[i]; + } + s_total_above = ta; + s_total_equal = te; + } + __syncthreads(); + + const uint32_t prefix_above = s_prefix_packed & kAboveMask; + const uint32_t prefix_equal = s_prefix_packed >> kAboveBits; + + // Write to global output + for (uint32_t i = tx; i < la; i += kBlockSize) { + // indices are placed contiguously starting at prefix_above + row_output[prefix_above + i] = + s_topk[i] + my_start; // my_start: block-local -> row-global index + } + for (uint32_t i = tx; i < le; i += kBlockSize) { + const auto t = smem->tie_buffer[i]; + uint32_t p = s_total_above + prefix_equal + i; + if (p < TopK) { + row_output[p] = t.idx + my_start; + } + uint32_t tp = prefix_equal + i; + if (tp < (TopK <= kBlockSize ? kMaxTies : TopK)) { + tie_ws[tp] = Tie{t.idx + my_start, + t.score}; // Ties are copied for cross-block refinement + } + } + + // Tie refinement + cooperative_groups::this_cluster().sync(); + if (rank != 0) { // only rank 0 does tie refinement + return; + } + if (s_total_above + s_total_equal <= TopK) { // no ties to refine + return; + } + + // Tie-breaking uses FP32 (4-round radix sort) + if constexpr (TopK <= kBlockSize) { + // copy ties from tie_ws back to smem, then refine + const uint32_t num_ties = min(s_total_equal, kMaxTies); + // TODO (roberto): could vectorize with uint2 (8 bytes = exactly one Tie) + for (uint32_t i = tx; i < num_ties; i += kBlockSize) { + smem->tie_buffer[i] = Tie{tie_ws[i].idx, tie_ws[i].score}; + } + __syncthreads(); + tie_handle(smem->tie_buffer, num_ties, s_total_above, row_output, + smem); + } else { + // TopK=2048: process directly from tie_ws (GMEM) + const uint32_t num_ties = min(s_total_equal, static_cast(TopK)); + tie_handle_large(tie_ws, num_ties, s_total_above, row_output, smem); + } +} + +// ============================================================================ +// Adapted from https://github.com/sgl-project/sglang/pull/23600 +// sgl-project/sglang +// (python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/) +// ============================================================================ + +template +__device__ void cooperative_topk_body(CooperativeTopKParams params) { + const auto rank = blockIdx.y, row = blockIdx.x, tx = threadIdx.x; + const auto sl = params.lengths[row]; + int32_t* out = params.output + row * TopK; + const float* in = params.input + row * params.stride; + + // Trivial: seq_len <= TopK + if (sl <= static_cast(TopK)) { + if (rank == 0) { + for (uint32_t i = tx; i < TopK; i += kBlockSize) { + out[i] = (i < static_cast(sl)) ? static_cast(i) : -1; + } + } + return; + } + + // Short-Medium path: histogram_4096_topk on rank 0 only - all data fits in RF + if (sl <= static_cast(kHist4096MaxLen)) { + if (rank == 0) { + extern __shared__ uint8_t sr[]; + histogram_4096_topk(in, out, sl, + sr); // 4096-bin (12-bit) histogram + } + return; + } + + // Large path: init mbarriers + state, then dispatch fused or twopass + const uint32_t per_block = + (params.stride + CS - 1) / CS; // how many elements per block + constexpr uint32_t kFusedMax = ((CS == 16) ? kNumStages16 + : (CS == 8) ? kNumStages8 + : kMaxSinglePassStages) * + kSizePerStage; + const bool use_singlepass = + per_block <= + kFusedMax; // single pass or TMA streaming: histogram+scatter + + // Select smem type and stage count at compile time based on CS + constexpr uint32_t kFusedStages = (CS == 16) ? kNumStages16 + : (CS == 8) ? kNumStages8 + : kMaxSinglePassStages; + using FusedSmem = SmemFused; + + extern __shared__ uint8_t sr[]; + + constexpr uint32_t kTieWsPerRow = TopK <= kBlockSize ? kMaxTies : TopK; + Tie* row_tie_ws = params.tie_ws + row * kTieWsPerRow; + + if (use_singlepass) { + auto* smem = reinterpret_cast(sr); + const uint32_t sp_stages = (per_block + kSizePerStage - 1) / kSizePerStage; + if (tx < sp_stages) { + mbarrier_init(&smem->barrier[0][tx], + 1); // init 1 barrier per TMA stage - + // signal when async copies complete + } + __syncthreads(); + uint32_t phases[kFusedStages] = + {}; // tracks the parity for mbarrier wait/arrive protocol + large_topk(in, out, sl, phases, row_tie_ws); + } else { + // Two-pass: only CS=4 in practice (CS=8 always fits in singlepass) + auto* smem = reinterpret_cast(sr); + if (tx < 2 * kNumStages4) { + mbarrier_init(&smem->barrier[0][tx], + 1); // init 2×2=4 barriers (2 passes × 2 stages) + } + __syncthreads(); + uint32_t hp[kNumStages4] = {0, 0}; // histogram+scatter pass counters + large_topk(in, out, sl, hp, row_tie_ws); + } +} + +template +__global__ void __launch_bounds__(kBlockSize, 1) __cluster_dims__(1, 4, 1) + cooperative_topk_cs4(CooperativeTopKParams params) { + cooperative_topk_body(params); +} + +template +__global__ void __launch_bounds__(kBlockSize, 1) __cluster_dims__(1, 8, 1) + cooperative_topk_cs8(CooperativeTopKParams params) { + cooperative_topk_body(params); +} + +#ifndef VLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY +template +__global__ void __launch_bounds__(kBlockSize, 1) __cluster_dims__(1, 16, 1) + cooperative_topk_cs16(CooperativeTopKParams params) { + cooperative_topk_body(params); +} +#endif + +constexpr size_t kSmemSize4_base = + (sizeof(Smem4) > sizeof(StreamSmem) ? sizeof(Smem4) : sizeof(StreamSmem)); +constexpr size_t kSmemSize4_sp = sizeof(SmemSinglePass); +constexpr size_t kSmemSize4 = + (kSmemSize4_base > kSmemSize4_sp ? kSmemSize4_base : kSmemSize4_sp) + + sizeof(int32_t) * 2048 + 128; +constexpr size_t kSmemSize8 = + std::max(sizeof(SmemFused), sizeof(StreamSmem)) + + sizeof(int32_t) * 2048 + 128; + +} // namespace cooperative + +} // namespace vllm + +#endif // COOPERATIVE_TOPK_CUH_ diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 1cc8e8167a62..d60b68a5868d 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -343,6 +343,14 @@ void persistent_topk(const torch::stable::Tensor& logits, torch::stable::Tensor& workspace, int64_t k, int64_t max_seq_len); +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK +void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); +#endif + void selective_scan_fwd( const torch::stable::Tensor& u, const torch::stable::Tensor& delta, const torch::stable::Tensor& A, const torch::stable::Tensor& B, diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 6b25dc9940e8..1f1a4f19f214 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -126,10 +126,10 @@ struct RadixRowState { // ============================================================================ struct PersistentTopKParams { - const float* __restrict__ input; // [num_rows, stride] - int32_t* __restrict__ output; // [num_rows, top_k] - const int32_t* __restrict__ lengths; // [num_rows] - RadixRowState* row_states; // large path: per-group state + const float* __restrict__ input; // [num_rows, stride] + int32_t* __restrict__ output; // [num_rows, top_k] + int32_t* __restrict__ lengths; // [num_rows] + RadixRowState* row_states; // large path: per-group state uint32_t num_rows; uint32_t stride; uint32_t top_k; // actual k value for output stride @@ -935,8 +935,576 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) } // namespace persistent // ============================================================================ -// FlashInfer FilteredTopK (BS>32 dispatch) — float32 only. -// Extracted from flashinfer_topk.cuh. Lives in namespace vllm (not persistent). +// ============================================================================ +// Optimized FilteredTopK — single CTA per row for bs > 32. +// Kept with persistent_topk so the portable fallback owns the non-cluster path. +// ============================================================================ +namespace filtered_topk { + +constexpr uint32_t kBlockSize = 1024; +constexpr uint32_t kHistBits = 10; +constexpr uint32_t kHistBins = 1 << kHistBits; +constexpr uint32_t RADIX = 256; +constexpr uint32_t kMaxTies = 1024; +static_assert(kMaxTies <= kBlockSize, + "tie_handle requires kMaxTies <= kBlockSize"); +constexpr uint32_t kWarpSize = 32; +constexpr uint32_t kNumWarps = kBlockSize / kWarpSize; + +constexpr uint32_t kElemPerStage = 16; +constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize; // 16384 + +// CS=4: 2 TMA stages (double buffer), two-pass +constexpr uint32_t kNumStages4 = 2; +// CS=8: 4 TMA stages, single-pass (all data stays in smem) +constexpr uint32_t kNumStages8 = + 2; // 2 stages × 16K = 32K per block (enough for CS=8) +// Max data per block with CS=8: ceil(262144/8) = 32768 ≤ 4 × 8192 = 32768 +constexpr uint32_t kMaxSeqLen8 = kNumStages8 * kSizePerStage * 8; // 262144 +constexpr uint32_t kNumStages16 = 2; // double-buffer for CS=16 +constexpr uint32_t kMaxSeqLen16 = kNumStages16 * kSizePerStage * 16; // 524288 + +// Register path +constexpr uint32_t kHist4096Bits = 12; +constexpr uint32_t kHist4096Bins = 1 << kHist4096Bits; // 4096 +constexpr uint32_t kHist4096VecsPerThread = 4; +constexpr uint32_t kHist4096MaxLen = + kHist4096VecsPerThread * 4 * kBlockSize; // 16384 +constexpr uint32_t kHist4096Items = kHist4096Bins / kBlockSize; // 4 + +// CS=4 single-pass path +constexpr uint32_t kMaxSinglePassStages = 3; +constexpr uint32_t kMaxSinglePassPerBlock = + kMaxSinglePassStages * kSizePerStage; // 49152 + +constexpr uint32_t kStreamNumStages = 2; + +struct alignas(16) MatchBin { + uint32_t bin, above_count, equal_count; +}; +struct alignas(8) Tie { + uint32_t idx; + float score; +}; + +__device__ __forceinline__ void load_float4_predicated(const float* ptr, + int base, int seq_len, + float& v0, float& v1, + float& v2, float& v3) { + uint32_t r0, r1, r2, r3; + const int p0 = (base < seq_len); + const int p1 = (base + 1 < seq_len); + const int p2 = (base + 2 < seq_len); + const int p3 = (base + 3 < seq_len); + asm volatile( + "{\n" + " .reg .pred pr0, pr1, pr2, pr3;\n" + " setp.ne.u32 pr0, %4, 0;\n" + " setp.ne.u32 pr1, %5, 0;\n" + " setp.ne.u32 pr2, %6, 0;\n" + " setp.ne.u32 pr3, %7, 0;\n" + " mov.u32 %0, 0xFF800000;\n" + " mov.u32 %1, 0xFF800000;\n" + " mov.u32 %2, 0xFF800000;\n" + " mov.u32 %3, 0xFF800000;\n" + " @pr0 ld.global.cg.u32 %0, [%8];\n" + " @pr1 ld.global.cg.u32 %1, [%8+4];\n" + " @pr2 ld.global.cg.u32 %2, [%8+8];\n" + " @pr3 ld.global.cg.u32 %3, [%8+12];\n" + "}\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr)); + v0 = __uint_as_float(r0); + v1 = __uint_as_float(r1); + v2 = __uint_as_float(r2); + v3 = __uint_as_float(r3); +} + +// converts the float32 score to a 32-bit ordered unsigned integer — the full +// precision key for radix sorting +__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits -> +// bin 0-4095) +template +__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) { + __half h = __float2half_rn(x); + uint16_t bits = __half_as_ushort(h); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) + : static_cast(bits | 0x8000); + return key >> (16 - kBits); +} + +// running sum within each warp — thread 0 gets its own value, thread 1 gets +// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc. +__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane, + uint32_t v) { +#pragma unroll + for (uint32_t o = 1; o < 32; o *= 2) { + uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o); + if (lane >= o) v += n; + } + return v; +} + +// Returns the sum of a value across all 32 threads in the warp, and every +// thread gets the same result SM90+ PTX instruction that does a hardware +// warp-wide reduction in a single instruction w.r.t. warp::reduce_sum(), which +// uses a __shfl_xor_sync butterfly tree (5 shuffles for 32 lanes) +__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) { + uint32_t r; + asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v)); + return r; +} + +// ============================================================================ +// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered +// key Each round narrows by 8 bits until ties are fully resolved +// ============================================================================ + +template +__device__ void tie_handle(const Tie* ties, uint32_t num_ties, + uint32_t num_above, int32_t* output, void* _smem) { + struct TS { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[RADIX]; + uint32_t warp_sum[kNumWarps]; + }; + auto* s = static_cast(_smem); + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize, wi = tx / kWarpSize; + + // Each thread loads one tie element. + const bool has = tx < num_ties; + const auto tie = has ? ties[tx] : Tie{0, 0.0f}; + const uint32_t key = convert_to_uint32_v2(tie.score); + + bool active = has; // tracks whether this thread's tie is still a candidate. + uint32_t remain = + TopK - num_above; // decreases each round as ties are resolved. + uint32_t wpos = TopK; // wpos will hold the final output position. + s->counter = 0; + __syncthreads(); + + // The 4-round radix loop - each round narrows by 8 bits until ties are fully + // resolved +#pragma unroll + for (int r = 0; r < 4; r++) { + uint32_t sh = 24 - r * 8; // round 0: bits 31-24, round 1: 23-16, etc. + uint32_t bin = (key >> sh) & 0xFF; // this tie's 8-bit bin for this round + + // Step 1: Build 256-bin histogram. + if (tx < RADIX) s->histogram[tx] = 0; + __syncthreads(); + if (active) atomicAdd(&s->histogram[bin], 1); + __syncthreads(); + + // Step 2: Prefix scan to find threshold + uint32_t hv = 0, wi2 = 0; + if (tx < RADIX) { + hv = s->histogram[tx]; + wi2 = warp_inclusive_sum(li, hv); + if (li == kWarpSize - 1) s->warp_sum[wi] = wi2; + } + __syncthreads(); + + if (tx < RADIX) { + auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; + auto tot = warp_reduce_sum_full(tmp); + auto inter = warp_reduce_sum_full(li < wi ? tmp : 0); + auto above = tot - (inter + wi2); + if (above < remain && above + hv >= remain) { + s->match = {tx, above, remain - above}; + } + } + __syncthreads(); + + // Step 3: Scatter + auto [thr, na, _] = s->match; // threshold bin, num above, unused + if (active) { + if (bin > thr) { + wpos = num_above + + atomicAdd(&s->counter, 1); // above -> place in output directly + active = false; + } else if (bin < thr) + active = false; // below -> discard + else if (r == 3) + wpos = TopK - atomicAdd(&s->match.equal_count, + -1u); // last round: place remaining + } + remain -= na; + if (!remain) break; // all ties resolved early + } + // Final write + if (wpos < TopK) output[wpos] = tie.idx; +} + +// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048). +// tie_handle assumes 1 tie per thread (max 1024). +// This version handles 2 ties per thread via kPerThread=2 +template +__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties, + uint32_t num_above, int32_t* output, + void* _smem) { + static_assert(TopK > kBlockSize); + struct TS { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[RADIX]; + uint32_t warp_sum[kNumWarps]; + }; + auto* s = static_cast(_smem); + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize; + const auto wi = tx / kWarpSize; + + constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize; + Tie my_ties[kPerThread]; + uint32_t keys[kPerThread]; + bool active[kPerThread]; + + for (uint32_t e = 0; e < kPerThread; e++) { + uint32_t idx = e * kBlockSize + tx; + if (idx < num_ties) { + my_ties[e] = ties[idx]; + keys[e] = convert_to_uint32_v2(ties[idx].score); + active[e] = true; + } else { + my_ties[e] = {0, 0.0f}; + keys[e] = 0; + active[e] = false; + } + } + + uint32_t remain = TopK - num_above; + s->counter = 0; + __syncthreads(); + + for (int r = 0; r < 4; r++) { + uint32_t sh = 24 - r * 8; + if (tx < RADIX) { + s->histogram[tx] = 0; + } + __syncthreads(); + + for (uint32_t e = 0; e < kPerThread; e++) { + if (active[e]) { + atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1); + } + } + __syncthreads(); + + uint32_t hv = 0; + if (tx < RADIX) { + hv = s->histogram[tx]; + auto wi2 = warp_inclusive_sum(li, hv); + if (li == kWarpSize - 1) { + s->warp_sum[wi] = wi2; + } + } + __syncthreads(); + if (tx < RADIX) { + auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; + auto total = warp_reduce_sum_full(tmp2); + auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0); + auto wi2 = warp_inclusive_sum(li, hv); + auto above = total - (inter + wi2); + if (above < remain && above + hv >= remain) { + s->match = { + .bin = tx, .above_count = above, .equal_count = remain - above}; + } + } + __syncthreads(); + + auto thr = s->match.bin; + auto na = s->match.above_count; + + for (uint32_t e = 0; e < kPerThread; e++) { + if (!active[e]) { + continue; + } + uint32_t bin = (keys[e] >> sh) & 0xFF; + if (bin > thr) { + uint32_t wpos = num_above + atomicAdd(&s->counter, 1); + if (wpos < TopK) { + output[wpos] = my_ties[e].idx; + } + active[e] = false; + } else if (bin < thr) { + active[e] = false; + } else if (r == 3) { + uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u); + if (wpos < TopK) { + output[wpos] = my_ties[e].idx; + } + } + } + + num_above += na; + remain -= na; + __syncthreads(); + s->counter = 0; + __syncthreads(); + } +} + +// ============================================================================ +// Register-based single-CTA fast path for seq_len <= 16384 +// 4 float4 per thread × 1024 threads = 16384 elements max +// Uses 4096-bin (12-bit) histogram for better precision +// ============================================================================ + +template +struct Histogram4096Smem { + static constexpr uint32_t HIST_BINS = 1 << HIST_BITS; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + MatchBin match; + uint32_t warp_sum[kNumWarps]; + union { + uint32_t histogram[HIST_BINS]; + Tie tie_buffer[kMaxTies]; + }; + static_assert(sizeof(uint32_t) * HIST_BINS >= TopK * sizeof(Tie), + "histogram union must be large enough for TopK ties"); +}; + +template +__device__ void histogram_4096_topk(const float* __restrict__ scores, + int32_t* __restrict__ output, + uint32_t length, void* _smem) { + constexpr uint32_t HIST_BINS = 1 << HIST_BITS; + constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize; + static_assert(HIST_BINS >= kBlockSize, + "HIST_BITS must give >= kBlockSize bins"); + + using Smem = Histogram4096Smem; + auto* smem = static_cast(_smem); + const auto tx = threadIdx.x; + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + + // Phase 1: Load all data into RF + build histogram + float4 + vecs[VECS_PER_THREAD]; // 4 vectors x 4 floats = 16 elements per thread + if constexpr (ITEMS_PER_THREAD >= 4) { + // Zero the histogram (SMEM writes) + for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++) + reinterpret_cast( + smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] = + make_uint4(0, 0, 0, 0); + } else { + if (tx < HIST_BINS) smem->histogram[tx] = 0; + } + if (tx == 0) { + smem->counter_gt = 0; + smem->counter_eq = 0; + } + if constexpr (UsePredicatedLoads) { + const bool row_aligned = (reinterpret_cast(scores) & 0xFu) == 0; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { + const uint32_t base = (tx + v * kBlockSize) * 4; + if (base < length) { + if (row_aligned && base + 3 < length) { + vecs[v] = *reinterpret_cast(scores + base); + } else { + load_float4_predicated(scores + base, static_cast(base), + static_cast(length), vecs[v].x, vecs[v].y, + vecs[v].z, vecs[v].w); + } + } + } + } else { +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { + const uint32_t base = (tx + v * kBlockSize) * 4; + if (base < length) { + vecs[v] = *reinterpret_cast(scores + base); + } + } + } + __syncthreads(); + + // Build histogram from RF via atomic adds into the shared histogram + bool done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + atomicAdd(&smem->histogram[extract_coarse_bin_N(elems[e])], + 1); + } + } + } + __syncthreads(); + + // Phase 2: Prefix scan to find threshold bin + // Multi-element scan (4096 bins: 4 per thread) + uint32_t orig[ITEMS_PER_THREAD]; + uint32_t local_sum = 0; + + // Step 1: Each thread sums its 4 bins +#pragma unroll + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { + orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i]; + local_sum += orig[i]; + } + + // Step 2: Warp-level inclusive prefix sum on local_sum + const auto warp_inc = warp_inclusive_sum(lane_id, local_sum); + if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc; + __syncthreads(); + + // Step 3: Inter-warp prefix via redux.sync + const auto tmp = smem->warp_sum[lane_id]; + uint32_t prefix = warp_reduce_sum_full( + lane_id < warp_id ? tmp : 0); // sum of all prior warps + prefix += + warp_inc - local_sum; // exclusive prefix within this thread's position + + // Step 4: Find threshold - scan 4 bins, accumulate prefix +#pragma unroll + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { + prefix += orig[i]; + const auto above = length - prefix; // elements in bins ABOVE this one + if (above < TopK && above + orig[i] >= TopK) { + smem->match = {.bin = tx * ITEMS_PER_THREAD + i, + .above_count = above, + .equal_count = orig[i]}; + } + } + + __syncthreads(); + + // Phase 3: Scatter from registers + const auto [thr_bin, num_above, num_equal] = smem->match; + const bool need_tie = (num_equal + num_above > TopK); + + done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + const uint32_t bin = extract_coarse_bin_N(elems[e]); + if (bin > thr_bin) { + output[atomicAdd(&smem->counter_gt, 1)] = + idx; // above -> output directly + } else if (bin == thr_bin) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (!need_tie) { + if (pos + num_above < TopK) { + output[pos + num_above] = idx; // all fit + } + } else { + if (pos < TopK) { + smem->tie_buffer[pos] = {idx, elems[e]}; // store for refirement + } + } + } + // else: bin < thr_bin - discard (not in top-k) + } + } + } + + // Phase 4: Tie-breaking + if (!need_tie) return; + __syncthreads(); + + // Fast warp-ballot tie-breaking for small tie counts + const uint32_t num_ties = min(num_equal, static_cast(TopK)); + const uint32_t topk_remain = + TopK - num_above; // pick exactly remaining elements to fill topK + + auto is_greater = [](const Tie& a, const Tie& b) { + return (a.score > b.score) || (a.score == b.score && a.idx < b.idx); + }; + + if (num_ties <= kWarpSize) { + // <=32 ties - Use warp ballot + // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024 + // comparisons in one instruction per warp. O(1) work. + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + if (lane_id >= num_ties || warp_id >= num_ties) return; + const uint32_t mask = (1ull << num_ties) - 1u; + const auto tie = smem->tie_buffer[lane_id]; // each lane holds one tie + const auto target = + smem->tie_buffer[warp_id]; // each warp evaluates one candidate + const bool pred = + is_greater(tie, target); // compare all ties against target + const auto rank = static_cast( + __popc(__ballot_sync(mask, pred))); // count how many are greater + if (lane_id == 0 && rank < topk_remain) { + output[num_above + rank] = target.idx; // place at correct position + } + } else if (num_ties <= + kWarpSize * + 2) { // TODO (roberto): try to refactor this with <=32 case + // Same idea but each thread handles 2 tie elements + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + const auto lane1 = lane_id + kWarpSize; + const auto warp1 = warp_id + kWarpSize; + const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__}; + const auto tie0 = smem->tie_buffer[lane_id]; + const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid; + if (warp_id < num_ties) { + const auto target = smem->tie_buffer[warp_id]; + const auto r0 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); + const auto r1 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); + if (lane_id == 0 && r0 + r1 < topk_remain) + output[num_above + r0 + r1] = target.idx; + } + if (warp1 < num_ties) { + const auto target = smem->tie_buffer[warp1]; + const auto r0 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); + const auto r1 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); + if (lane_id == 0 && r0 + r1 < topk_remain) + output[num_above + r0 + r1] = target.idx; + } + } else { + // Large tie count: fall back to 4-round radix-256 sort + if constexpr (TopK <= kBlockSize) { + tie_handle(smem->tie_buffer, num_ties, num_above, output, smem); + } else { + tie_handle_large(smem->tie_buffer, num_ties, num_above, output, + smem); + } + } +} + +template +__device__ __noinline__ void histogram_4096_topk_predicated( + const float* __restrict__ scores, int32_t* __restrict__ output, + uint32_t length, void* _smem) { + histogram_4096_topk(scores, output, + length, _smem); +} + +// ============================================================================ +// FilteredTopK — single CTA per row for bs > 32 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 // ============================================================================ @@ -1013,7 +1581,8 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = * \tparam IdType Index type (int32_t) * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) */ -template +template __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) FilteredTopKUnifiedKernel(const DType* __restrict__ input, IdType* __restrict__ output, @@ -1042,6 +1611,18 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) return; } + // Short path + if (length <= 32768) { + extern __shared__ uint8_t _smem_reg[]; + if constexpr (UsePredicatedShortLoads) { + histogram_4096_topk_predicated(score, dst, length, + _smem_reg); + } else { + histogram_4096_topk(score, dst, length, _smem_reg); + } + return; + } + // Static shared memory alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128]; alignas(128) __shared__ int s_counter; @@ -1269,11 +1850,9 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) { } template -cudaError_t FilteredTopKRaggedTransform(const DType* input, - IdType* output_indices, - const IdType* lengths, - uint32_t num_rows, uint32_t top_k_val, - uint32_t max_len, +cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices, + IdType* lengths, uint32_t num_rows, + uint32_t top_k_val, uint32_t max_len, cudaStream_t stream = 0) { constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC; constexpr int MAX_VEC = 16 / sizeof(DType); @@ -1285,14 +1864,15 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input, const int vec_size = ComputeFilteredTopKVecSize(max_len); -#define DISPATCH_VEC_SIZE(VS) \ - if (vec_size == VS) { \ - auto kernel = FilteredTopKUnifiedKernel; \ - FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \ - FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \ - smem_size, stream)); \ - return cudaSuccess; \ +#define DISPATCH_VEC_SIZE(VS) \ + if (vec_size == VS) { \ + auto kernel = \ + FilteredTopKUnifiedKernel; \ + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \ + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \ + smem_size, stream)); \ + return cudaSuccess; \ } DISPATCH_VEC_SIZE(1) @@ -1306,6 +1886,17 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input, return cudaSuccess; } +} // namespace filtered_topk + +template +cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices, + IdType* lengths, uint32_t num_rows, + uint32_t top_k_val, uint32_t max_len, + cudaStream_t stream = 0) { + return filtered_topk::FilteredTopKRaggedTransform( + input, output_indices, lengths, num_rows, top_k_val, max_len, stream); +} + } // namespace vllm #endif // PERSISTENT_TOPK_CUH_ diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index d55c12d382a9..1be7217ce789 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -493,6 +493,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " "Tensor workspace, int k, int max_seq_len) -> ()"); +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.def( + "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); +#endif + // Activation ops ops.def( "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " @@ -711,6 +717,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); +#endif // Activation kernels (shared CUDA/ROCm) ops.impl("persistent_masked_m_silu_mul_quant", diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index 7b9c11495e8b..3a1ad0f0d23c 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -14,6 +14,70 @@ BATCH_SIZE = [1, 2, 2048] NEXT_N = [1, 8] DATA_GENERATION = ["random", "10LSBits"] +RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 + + +def _has_device_capability(major: int) -> bool: + return current_platform.is_cuda() and current_platform.has_device_capability(major) + + +COOPERATIVE_TOPK_BACKEND = pytest.param( + "cooperative_topk", + marks=pytest.mark.skipif( + not _has_device_capability(90), + reason="cooperative_topk requires SM90+", + ), +) +WORKSPACE_TOPK_BACKENDS = ["persistent_topk", COOPERATIVE_TOPK_BACKEND] +TOPK_BACKENDS = ["top_k_per_row_decode", *WORKSPACE_TOPK_BACKENDS] + + +def _run_topk_backend( + backend: str, + logits: torch.Tensor, + lengths: torch.Tensor, + indices: torch.Tensor, + top_k: int, + max_seq_len: int, + next_n: int = 1, +) -> None: + if backend == "top_k_per_row_decode": + torch.ops._C.top_k_per_row_decode( + logits, + next_n, + lengths, + indices, + indices.shape[0], + logits.stride(0), + logits.stride(1), + top_k, + ) + elif backend == "persistent_topk": + workspace = torch.empty( + RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda" + ) + torch.ops._C.persistent_topk( + logits, lengths, indices, workspace, top_k, max_seq_len + ) + elif backend == "cooperative_topk": + if indices.shape[0] > 32: + pytest.skip( + "cooperative_topk supports <=32 rows; " + "persistent_topk covers larger batches" + ) + if logits.stride(0) % 4 != 0: + pytest.skip( + "cooperative_topk requires row stride divisible by 4; " + "persistent_topk covers unaligned strides" + ) + workspace = torch.empty( + RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda" + ) + torch.ops._C.cooperative_topk( + logits, lengths, indices, workspace, top_k, max_seq_len + ) + else: + raise ValueError(f"Unknown top-k backend: {backend}") def create_random_logits( @@ -322,16 +386,19 @@ def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None: @pytest.mark.parametrize("clean_logits", [True, False]) @pytest.mark.parametrize("top_k", [2048]) @pytest.mark.parametrize("next_n", [1, 4]) +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() -def test_deepseek_persistent_topk( +def test_deepseek_workspace_topk( seq_len_range: tuple[int, int], test_id: str, clean_logits: bool, top_k: int, next_n: int, + backend: str, ) -> None: """ - Test persistent_topk with varying sequence lengths and speculative decoding. + Test workspace top-k backends with varying sequence lengths and speculative + decoding. Supports speculative decoding with next_n > 1. """ set_random_seed(42 if test_id == "short_sequences" else 43) @@ -347,6 +414,7 @@ def test_deepseek_persistent_topk( dtype=torch.int32, device="cuda", ) + seq_lens = (seq_lens + 3) & ~3 # align to 4 for TMA # Compute row boundaries for speculative decoding row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda") @@ -366,14 +434,11 @@ def test_deepseek_persistent_topk( offsets = torch.arange(next_n, device=logits.device, dtype=torch.int32) lengths = (seq_lens.unsqueeze(1) - next_n + 1 + offsets).flatten() - workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda") max_seq_len = int(seq_lens.max().item()) - torch.ops._C.persistent_topk( - logits, lengths, indices, workspace, top_k, max_seq_len - ) + _run_topk_backend(backend, logits, lengths, indices, top_k, max_seq_len, next_n) validate_topk_against_reference( - logits, indices, row_starts, row_ends, top_k, f"persistent_topk ({test_id})" + logits, indices, row_starts, row_ends, top_k, f"{backend} ({test_id})" ) @@ -383,9 +448,10 @@ def run_large_context_topk_test( top_k: int, data_type: str = "random", seed: int = 42, + backend: str = "cooperative_topk", ) -> None: """ - Helper to run persistent_topk kernel test with given parameters. + Helper to run a top-k backend test with given parameters. Args: batch_size: Number of rows/sequences @@ -393,6 +459,7 @@ def run_large_context_topk_test( top_k: Number of top elements to select data_type: Type of test data to generate seed: Random seed for reproducibility + backend: Top-k backend to test """ torch.set_default_device("cuda:0") set_random_seed(seed) @@ -449,11 +516,8 @@ def run_large_context_topk_test( # Create output tensor indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") - workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda") max_seq_len = max(seq_lens) - torch.ops._C.persistent_topk( - logits, lengths, indices, workspace, top_k, max_seq_len - ) + _run_topk_backend(backend, logits, lengths, indices, top_k, max_seq_len) torch.accelerator.synchronize() @@ -605,8 +669,9 @@ def run_large_context_topk_test( ), ], ) +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() -def test_persistent_topk_correctness(test_config: dict) -> None: +def test_workspace_topk_correctness(test_config: dict, backend: str) -> None: """ Comprehensive correctness tests covering: - Sequence length edge cases (trivial, boundary, varied) @@ -620,6 +685,7 @@ def test_persistent_topk_correctness(test_config: dict) -> None: seq_lens=test_config["seq_lens"], top_k=test_config["top_k"], data_type=test_config.get("data_type", "random"), + backend=backend, ) @@ -668,8 +734,9 @@ def test_persistent_topk_correctness(test_config: dict) -> None: ), ], ) +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() -def test_persistent_topk_algorithm_paths(test_config: dict) -> None: +def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None: """ Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2): - Batch size scalability (1, 4, 32, 256) @@ -680,12 +747,14 @@ def test_persistent_topk_algorithm_paths(test_config: dict) -> None: batch_size=test_config["batch_size"], seq_lens=[test_config["seq_len"]] * test_config["batch_size"], top_k=test_config["top_k"], + backend=backend, ) @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() -def test_persistent_topk_stress() -> None: +def test_workspace_topk_stress(backend: str) -> None: """ Stress test with random configurations to catch edge cases. Capped at 163840 (DeepSeek V3.2 max context) for realistic testing. @@ -700,16 +769,73 @@ def test_persistent_topk_stress() -> None: batch_size = torch.randint(1, 32, (1,)).item() # Random sequence lengths capped at DeepSeek V3.2 max context - seq_lens = torch.randint(100, 163840, (batch_size,)).tolist() + seq_lens_tensor = torch.randint(100, 163840, (batch_size,)) + if backend == "cooperative_topk": + seq_lens = ((seq_lens_tensor + 3) & ~3).tolist() + else: + seq_lens = seq_lens_tensor.tolist() run_large_context_topk_test( batch_size=batch_size, seq_lens=seq_lens, top_k=top_k, seed=seed, + backend=backend, ) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("backend", TOPK_BACKENDS) +@pytest.mark.parametrize("top_k", [512, 1024, 2048]) +@torch.inference_mode() +def test_deepseek_topk_backends_no_error_and_reference( + backend: str, + top_k: int, +) -> None: + """Exercise every production top-k backend on the same inputs.""" + run_large_context_topk_test( + batch_size=4, + seq_lens=[2049, 4097, 8191, 12000], + top_k=top_k, + data_type="random", + seed=123, + backend=backend, + ) + + +@pytest.mark.skipif(not _has_device_capability(90), reason="This test requires SM90+") +@torch.inference_mode() +def test_cooperative_topk_512_tie_workspace_is_per_row() -> None: + """Regression test for TopK=512 tie workspace row overlap.""" + torch.set_default_device("cuda:0") + + top_k = 512 + num_rows = 2 + stride = 65536 + lengths = torch.tensor([40960, 65536], dtype=torch.int32, device="cuda") + logits = torch.full( + (num_rows, stride), float("-inf"), dtype=torch.float32, device="cuda" + ) + + # Row 0 must never select these low indices: many better row-0 ties exist. + logits[0, :2048] = -10.0 + logits[0, 2048 : lengths[0]] = 1.0 + # Row 1 has higher exact tie scores. With the old row * TopK tie_ws stride, + # these row-1 ties could overwrite row 0's TopK=512 refinement workspace. + logits[1, : lengths[1]] = 2.0 + + indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") + workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda") + torch.ops._C.cooperative_topk(logits, lengths, indices, workspace, top_k, stride) + torch.accelerator.synchronize() + + row0 = indices[0].cpu() + assert torch.all(row0 >= 2048), ( + "cooperative_topk TopK=512 selected row-0 low-score indices, likely " + "from overlapping tie_ws rows" + ) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @pytest.mark.parametrize( "test_config", @@ -774,10 +900,11 @@ def test_persistent_topk_stress() -> None: ], ) @pytest.mark.parametrize("top_k", [512, 2048]) +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() -def test_persistent_topk(test_config: dict, top_k: int) -> None: +def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None: """ - Tests specific to the persistent_topk kernel: + Tests specific to workspace top-k backends: - Mixed medium/large rows in the same batch (dynamic per-row dispatch) - Boundary around LARGE_THRESHOLD (32K) - Trivial + medium + large rows in a single batch @@ -787,15 +914,17 @@ def test_persistent_topk(test_config: dict, top_k: int) -> None: seq_lens=test_config["seq_lens"], top_k=top_k, data_type=test_config.get("data_type", "random"), + backend=backend, ) @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @pytest.mark.parametrize("top_k", [512, 2048]) +@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS) @torch.inference_mode() -def test_persistent_topk_padded_stride(top_k: int) -> None: +def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None: """ - Test persistent_topk with padded logits (large stride, small seq_len) + Test workspace top-k backends with padded logits (large stride, small seq_len) to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits returns [B, max_model_len] with max_model_len=163840. """ @@ -818,11 +947,7 @@ def test_persistent_topk_padded_stride(top_k: int) -> None: lengths = torch.tensor(actual_seq_lens, dtype=torch.int32, device="cuda") indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda") - workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda") - - torch.ops._C.persistent_topk( - logits, lengths, indices, workspace, top_k, max(actual_seq_lens) - ) + _run_topk_backend(backend, logits, lengths, indices, top_k, max(actual_seq_lens)) torch.accelerator.synchronize() # Validate against torch.topk @@ -840,6 +965,6 @@ def test_persistent_topk_padded_stride(top_k: int) -> None: expected_vals = logits[i, expected].cpu().sort(descending=True)[0] actual_vals = logits[i, actual].cpu().sort(descending=True)[0] assert torch.allclose(expected_vals, actual_vals, rtol=1e-4, atol=1e-4), ( - f"Row {i}: persistent_topk with padded stride doesn't match. " + f"Row {i}: {backend} with padded stride doesn't match. " f"seq_len={sl}, stride={padded_stride}" ) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 45c5d5f78198..fe2b268cde66 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -334,7 +334,32 @@ def sparse_attn_indexer( num_rows = logits.shape[0] topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - if current_platform.is_cuda() and topk_tokens in (512, 1024, 2048): + use_cooperative_topk = ( + current_platform.is_cuda() + and topk_tokens in (512, 1024, 2048) + and num_rows <= 32 + and logits.stride(0) % 4 == 0 # TMA 16-byte alignment + and current_platform.has_device_capability(90) + ) + use_persistent_topk = current_platform.is_cuda() and topk_tokens in ( + 512, + 1024, + 2048, + ) + if use_cooperative_topk: + workspace_manager = current_workspace_manager() + (topk_workspace,) = workspace_manager.get_simultaneous( + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + ) + torch.ops._C.cooperative_topk( + logits, + seq_lens, + topk_indices, + topk_workspace, + topk_tokens, + attn_metadata_narrowed.max_seq_len, + ) + elif use_persistent_topk: workspace_manager = current_workspace_manager() (topk_workspace,) = workspace_manager.get_simultaneous( ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), From d867d1f88eee47d02c42053f0b4bb9e822566097 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Fri, 19 Jun 2026 16:36:08 +0000 Subject: [PATCH 2/7] fix rebase problems Signed-off-by: LopezCastroRoberto --- CMakeLists.txt | 4 ++++ csrc/libtorch_stable/persistent_topk.cuh | 14 +++++++++----- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index c42cbccabaf4..fc59214fbd5b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1079,6 +1079,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") target_compile_definitions(_C_stable_libtorch PRIVATE TORCH_TARGET_VERSION=0x020B000000000000ULL) target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA) + if(COOPERATIVE_TOPK_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + VLLM_ENABLE_COOPERATIVE_TOPK=1) + endif() # Needed by CUTLASS kernels target_compile_definitions(_C_stable_libtorch PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 1f1a4f19f214..8980c8a2cd57 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -128,7 +128,7 @@ struct RadixRowState { struct PersistentTopKParams { const float* __restrict__ input; // [num_rows, stride] int32_t* __restrict__ output; // [num_rows, top_k] - int32_t* __restrict__ lengths; // [num_rows] + const int32_t* __restrict__ lengths; // [num_rows] RadixRowState* row_states; // large path: per-group state uint32_t num_rows; uint32_t stride; @@ -1850,8 +1850,10 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) { } template -cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices, - IdType* lengths, uint32_t num_rows, +cudaError_t FilteredTopKRaggedTransform(const DType* input, + IdType* output_indices, + const IdType* lengths, + uint32_t num_rows, uint32_t top_k_val, uint32_t max_len, cudaStream_t stream = 0) { constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC; @@ -1889,8 +1891,10 @@ cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices, } // namespace filtered_topk template -cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices, - IdType* lengths, uint32_t num_rows, +cudaError_t FilteredTopKRaggedTransform(const DType* input, + IdType* output_indices, + const IdType* lengths, + uint32_t num_rows, uint32_t top_k_val, uint32_t max_len, cudaStream_t stream = 0) { return filtered_topk::FilteredTopKRaggedTransform( From b93ca2a9c610568331c7e3bbed6c5c10d8714708 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Fri, 19 Jun 2026 16:49:14 +0000 Subject: [PATCH 3/7] fix pre-commit Signed-off-by: LopezCastroRoberto --- csrc/libtorch_stable/cooperative_topk.cu | 35 ++++++++++++------------ csrc/libtorch_stable/persistent_topk.cuh | 14 +++++----- 2 files changed, 24 insertions(+), 25 deletions(-) diff --git a/csrc/libtorch_stable/cooperative_topk.cu b/csrc/libtorch_stable/cooperative_topk.cu index 97b7b38f42fe..02c3a2275026 100644 --- a/csrc/libtorch_stable/cooperative_topk.cu +++ b/csrc/libtorch_stable/cooperative_topk.cu @@ -50,8 +50,8 @@ void launch_cooperative_cluster(ct::CooperativeTopKParams& params, cfg.numAttrs = 1; cfg.attrs = attrs; cudaError_t err = cudaLaunchKernelEx(&cfg, kernel, params); - STD_TORCH_CHECK(err == cudaSuccess, "cooperative_topk launch failed: ", - cudaGetErrorString(err)); + STD_TORCH_CHECK(err == cudaSuccess, + "cooperative_topk launch failed: ", cudaGetErrorString(err)); } template @@ -71,11 +71,10 @@ void launch_cooperative_topk_impl(const torch::stable::Tensor& logits, "cooperative_topk supports <=32 rows; use persistent_topk for " "larger batches"); - STD_TORCH_CHECK( - stride % 4 == 0, - "cooperative_topk: stride must be multiple of 4 for TMA " - "alignment, got stride (max_model_len)=", - stride); + STD_TORCH_CHECK(stride % 4 == 0, + "cooperative_topk: stride must be multiple of 4 for TMA " + "alignment, got stride (max_model_len)=", + stride); STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); STD_TORCH_CHECK( @@ -88,14 +87,14 @@ void launch_cooperative_topk_impl(const torch::stable::Tensor& logits, params.lengths = lengths.const_data_ptr(); params.num_rows = static_cast(num_rows); params.stride = stride; - params.tie_ws = reinterpret_cast( - workspace.mutable_data_ptr()); + params.tie_ws = + reinterpret_cast(workspace.mutable_data_ptr()); constexpr uint32_t kTieWsPerRow = TopK <= ct::kBlockSize ? ct::kMaxTies : TopK; STD_TORCH_CHECK( - workspace.size(0) >= static_cast(num_rows * kTieWsPerRow * - sizeof(ct::Tie)), + workspace.size(0) >= + static_cast(num_rows * kTieWsPerRow * sizeof(ct::Tie)), "workspace too small"); #ifndef VLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY @@ -113,10 +112,10 @@ void launch_cooperative_topk_impl(const torch::stable::Tensor& logits, #endif // USE_ROCM void cooperative_topk(const torch::stable::Tensor& logits, - const torch::stable::Tensor& lengths, - torch::stable::Tensor& output, - torch::stable::Tensor& workspace, int64_t k, - int64_t max_seq_len) { + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len) { #ifndef USE_ROCM STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor"); STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor"); @@ -136,9 +135,9 @@ void cooperative_topk(const torch::stable::Tensor& logits, STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch"); STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k, "output size mismatch"); - STD_TORCH_CHECK(k == 512 || k == 1024 || k == 2048, - "cooperative_topk supports k=512, k=1024, or k=2048, got k=", - k); + STD_TORCH_CHECK( + k == 512 || k == 1024 || k == 2048, + "cooperative_topk supports k=512, k=1024, or k=2048, got k=", k); if (k == 512) { launch_cooperative_topk_impl<512>(logits, lengths, output, workspace, diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 8980c8a2cd57..70d157f1a221 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -126,10 +126,10 @@ struct RadixRowState { // ============================================================================ struct PersistentTopKParams { - const float* __restrict__ input; // [num_rows, stride] - int32_t* __restrict__ output; // [num_rows, top_k] + const float* __restrict__ input; // [num_rows, stride] + int32_t* __restrict__ output; // [num_rows, top_k] const int32_t* __restrict__ lengths; // [num_rows] - RadixRowState* row_states; // large path: per-group state + RadixRowState* row_states; // large path: per-group state uint32_t num_rows; uint32_t stride; uint32_t top_k; // actual k value for output stride @@ -1853,8 +1853,8 @@ template cudaError_t FilteredTopKRaggedTransform(const DType* input, IdType* output_indices, const IdType* lengths, - uint32_t num_rows, - uint32_t top_k_val, uint32_t max_len, + uint32_t num_rows, uint32_t top_k_val, + uint32_t max_len, cudaStream_t stream = 0) { constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC; constexpr int MAX_VEC = 16 / sizeof(DType); @@ -1894,8 +1894,8 @@ template cudaError_t FilteredTopKRaggedTransform(const DType* input, IdType* output_indices, const IdType* lengths, - uint32_t num_rows, - uint32_t top_k_val, uint32_t max_len, + uint32_t num_rows, uint32_t top_k_val, + uint32_t max_len, cudaStream_t stream = 0) { return filtered_topk::FilteredTopKRaggedTransform( input, output_indices, lengths, num_rows, top_k_val, max_len, stream); From 2ccaedacc465d05f302339da6193eeee5f5ed396 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Fri, 19 Jun 2026 17:23:14 +0000 Subject: [PATCH 4/7] fix CI Signed-off-by: LopezCastroRoberto --- CMakeLists.txt | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index fc59214fbd5b..037f529e497b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -395,8 +395,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(COOPERATIVE_TOPK_ARCHS) - list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/cooperative_topk.cu") list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1") cuda_archs_loose_intersection(COOPERATIVE_TOPK_SM90_ARCHS @@ -405,10 +403,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") list(APPEND VLLM_GPU_FLAGS "-DVLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY=1") endif() - - set_source_files_properties( - "csrc/libtorch_stable/cooperative_topk.cu" - PROPERTIES CUDA_ARCHITECTURES "${COOPERATIVE_TOPK_ARCHS}") endif() endif() @@ -528,6 +522,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") SRCS "${VLLM_STABLE_EXT_SRC}" CUDA_ARCHS "${CUDA_ARCHS}") + if(COOPERATIVE_TOPK_ARCHS) + list(APPEND VLLM_STABLE_EXT_SRC + "csrc/libtorch_stable/cooperative_topk.cu") + set_gencode_flags_for_srcs( + SRCS "csrc/libtorch_stable/cooperative_topk.cu" + CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}") + endif() + # Only build Marlin kernels if we are building for at least some compatible archs. # Keep building Marlin for 9.0 as there are some group sizes and shapes that # are not supported by Machete yet. From b317d150fad68e3c03c66122d5bfd9499fbe4aa8 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Tue, 23 Jun 2026 11:39:31 +0000 Subject: [PATCH 5/7] fix comments Co-authored-by: OpenAI Codex Signed-off-by: LopezCastroRoberto --- CMakeLists.txt | 6 - csrc/libtorch_stable/cooperative_topk.cu | 25 +- csrc/libtorch_stable/cooperative_topk.cuh | 643 ++----------------- csrc/libtorch_stable/persistent_topk.cuh | 575 +---------------- csrc/libtorch_stable/topk_histogram_4096.cuh | 555 ++++++++++++++++ 5 files changed, 639 insertions(+), 1165 deletions(-) create mode 100644 csrc/libtorch_stable/topk_histogram_4096.cuh diff --git a/CMakeLists.txt b/CMakeLists.txt index 037f529e497b..6d130f8bda2a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -397,12 +397,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(COOPERATIVE_TOPK_ARCHS) list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1") - cuda_archs_loose_intersection(COOPERATIVE_TOPK_SM90_ARCHS - "9.0a" "${COOPERATIVE_TOPK_ARCHS}") - if(COOPERATIVE_TOPK_SM90_ARCHS) - list(APPEND VLLM_GPU_FLAGS - "-DVLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY=1") - endif() endif() endif() diff --git a/csrc/libtorch_stable/cooperative_topk.cu b/csrc/libtorch_stable/cooperative_topk.cu index 02c3a2275026..f388a9e6c8e7 100644 --- a/csrc/libtorch_stable/cooperative_topk.cu +++ b/csrc/libtorch_stable/cooperative_topk.cu @@ -8,6 +8,7 @@ #ifndef USE_ROCM #include "cooperative_topk.cuh" namespace ct = vllm::cooperative; +namespace hist4096 = vllm::topk_histogram_4096; #endif #ifndef USE_ROCM @@ -16,13 +17,7 @@ void launch_cooperative_cluster(ct::CooperativeTopKParams& params, size_t smem, cudaStream_t stream) { auto kernel = []() { if constexpr (CS == 16) { - #ifndef VLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY return &ct::cooperative_topk_cs16; - #else - static_assert(CS != 16, - "CS=16 cooperative_topk requires non-portable cluster " - "support"); - #endif } else if constexpr (CS == 8) { return &ct::cooperative_topk_cs8; } else { @@ -30,18 +25,16 @@ void launch_cooperative_cluster(ct::CooperativeTopKParams& params, return &ct::cooperative_topk_cs4; } }(); - #ifndef VLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY if constexpr (CS > 8) { cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1); } - #endif cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); cudaLaunchConfig_t cfg = {}; cfg.gridDim = dim3(params.num_rows, CS); - cfg.blockDim = dim3(ct::kBlockSize); + cfg.blockDim = dim3(hist4096::kBlockSize); cfg.dynamicSmemBytes = smem; cfg.stream = stream; cudaLaunchAttribute attrs[1]; @@ -60,6 +53,7 @@ void launch_cooperative_topk_impl(const torch::stable::Tensor& logits, torch::stable::Tensor& output, torch::stable::Tensor& workspace, int64_t max_seq_len) { + (void)max_seq_len; // Kept for signature parity with persistent_topk. const int64_t num_rows = logits.size(0); const cudaStream_t stream = get_current_cuda_stream(); @@ -88,22 +82,19 @@ void launch_cooperative_topk_impl(const torch::stable::Tensor& logits, params.num_rows = static_cast(num_rows); params.stride = stride; params.tie_ws = - reinterpret_cast(workspace.mutable_data_ptr()); + reinterpret_cast(workspace.mutable_data_ptr()); constexpr uint32_t kTieWsPerRow = - TopK <= ct::kBlockSize ? ct::kMaxTies : TopK; + TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK; STD_TORCH_CHECK( workspace.size(0) >= - static_cast(num_rows * kTieWsPerRow * sizeof(ct::Tie)), + static_cast(num_rows * kTieWsPerRow * sizeof(hist4096::Tie)), "workspace too small"); - #ifndef VLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY - if (num_rows <= 4) { + const bool supports_cluster16 = get_device_prop()->major >= 10; + if (num_rows <= 4 && supports_cluster16) { launch_cooperative_cluster(params, ct::kSmemSize8, stream); } else if (num_rows <= 8) { - #else - if (num_rows <= 8) { - #endif launch_cooperative_cluster(params, ct::kSmemSize8, stream); } else { launch_cooperative_cluster(params, ct::kSmemSize4, stream); diff --git a/csrc/libtorch_stable/cooperative_topk.cuh b/csrc/libtorch_stable/cooperative_topk.cuh index cb5a608c1250..937a18fbf635 100644 --- a/csrc/libtorch_stable/cooperative_topk.cuh +++ b/csrc/libtorch_stable/cooperative_topk.cuh @@ -13,109 +13,44 @@ #include #include +#include "topk_histogram_4096.cuh" + namespace vllm { namespace cooperative { -// TopK is now a template parameter (512 or 1024) -constexpr uint32_t kBlockSize = 1024; +namespace hist4096 = topk_histogram_4096; + constexpr uint32_t kHistBits = 10; constexpr uint32_t kHistBins = 1 << kHistBits; -constexpr uint32_t RADIX = 256; -constexpr uint32_t kMaxTies = 1024; -static_assert(kMaxTies <= kBlockSize, - "tie_handle requires kMaxTies <= kBlockSize"); -constexpr uint32_t kWarpSize = 32; -constexpr uint32_t kNumWarps = kBlockSize / kWarpSize; +constexpr uint32_t kMaxTopK = 2048; constexpr uint32_t kElemPerStage = 16; -constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize; // 16384 - -// CS=4: 2 TMA stages (double buffer), two-pass -constexpr uint32_t kNumStages4 = 2; -// CS=8: 4 TMA stages, single-pass (all data stays in smem) -constexpr uint32_t kNumStages8 = - 2; // 2 stages × 16K = 32K per block (enough for CS=8) -// Max data per block with CS=8: ceil(262144/8) = 32768 ≤ 4 × 8192 = 32768 -constexpr uint32_t kMaxSeqLen8 = kNumStages8 * kSizePerStage * 8; // 262144 -constexpr uint32_t kNumStages16 = 2; // double-buffer for CS=16 -constexpr uint32_t kMaxSeqLen16 = kNumStages16 * kSizePerStage * 16; // 524288 +constexpr uint32_t kSizePerStage = kElemPerStage * hist4096::kBlockSize; // 16384 -// Register path -constexpr uint32_t kHist4096Bits = 12; -constexpr uint32_t kHist4096Bins = 1 << kHist4096Bits; // 4096 -constexpr uint32_t kHist4096VecsPerThread = 4; -constexpr uint32_t kHist4096MaxLen = - kHist4096VecsPerThread * 4 * kBlockSize; // 16384 -constexpr uint32_t kHist4096Items = kHist4096Bins / kBlockSize; // 4 +// CS=4 two-pass path uses two TMA stages as a double buffer. +constexpr uint32_t kStreamingStagesCS4 = 2; +// CS=8/16 fused paths keep all loaded TMA stages resident in smem. +constexpr uint32_t kFusedStagesCS8 = 2; +constexpr uint32_t kFusedStagesCS16 = 2; // CS=4 single-pass path constexpr uint32_t kMaxSinglePassStages = 3; constexpr uint32_t kMaxSinglePassPerBlock = kMaxSinglePassStages * kSizePerStage; // 49152 -constexpr uint32_t kStreamNumStages = 2; - -struct alignas(16) MatchBin { - uint32_t bin, above_count, equal_count; -}; -struct alignas(8) Tie { - uint32_t idx; - float score; -}; - template struct CooperativeTopKParams { const float* __restrict__ input; int32_t* __restrict__ output; const int32_t* __restrict__ lengths; - Tie* __restrict__ tie_ws; // per-row tie workspace, see kTieWsPerRow + hist4096::Tie* __restrict__ tie_ws; // per-row tie workspace, see kTieWsPerRow uint32_t num_rows, stride; }; // ============================================================================ -// Common helpers +// Cooperative helpers // ============================================================================ -// converts the float32 score to a 32-bit ordered unsigned integer — the full -// precision key for radix sorting -__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { - uint32_t bits = __float_as_uint(x); - return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); -} - -// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits -> -// bin 0-4095) -template -__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) { - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return key >> (16 - kBits); -} - -// running sum within each warp — thread 0 gets its own value, thread 1 gets -// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc. -__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane, - uint32_t v) { -#pragma unroll - for (uint32_t o = 1; o < 32; o *= 2) { - uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o); - if (lane >= o) v += n; - } - return v; -} - -// Returns the sum of a value across all 32 threads in the warp, and every -// thread gets the same result SM90+ PTX instruction that does a hardware -// warp-wide reduction in a single instruction w.r.t. warp::reduce_sum(), which -// uses a __shfl_xor_sync butterfly tree (5 shuffles for 32 lanes) -__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) { - uint32_t r; - asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v)); - return r; -} - // only CS adjacent lanes participate (sub-warp reduce), in opposite to // warp_reduce_sum_full template @@ -131,7 +66,7 @@ __device__ __forceinline__ uint32_t warp_reduce_sum_subN(uint32_t v) { // ============================================================================ __device__ __forceinline__ uint32_t extract_coarse_bin(float x) { - return extract_coarse_bin_N(x); + return hist4096::extract_coarse_bin_N(x); } __device__ __forceinline__ void mbarrier_init(uint64_t* a, uint32_t n) { @@ -153,205 +88,13 @@ __device__ __forceinline__ void tma_load(void* d, const void* s, uint32_t n, s, n, m); } -// ============================================================================ -// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered -// key Each round narrows by 8 bits until ties are fully resolved -// ============================================================================ - -template -__device__ void tie_handle(const Tie* ties, uint32_t num_ties, - uint32_t num_above, int32_t* output, void* _smem) { - struct TS { - alignas(128) uint32_t counter; - alignas(128) MatchBin match; - uint32_t histogram[RADIX]; - uint32_t warp_sum[kNumWarps]; - }; - auto* s = static_cast(_smem); - const auto tx = threadIdx.x; - const auto li = tx % kWarpSize, wi = tx / kWarpSize; - - // Each thread loads one tie element. - const bool has = tx < num_ties; - const auto tie = has ? ties[tx] : Tie{0, 0.0f}; - const uint32_t key = convert_to_uint32_v2(tie.score); - - bool active = has; // tracks whether this thread's tie is still a candidate. - uint32_t remain = - TopK - num_above; // decreases each round as ties are resolved. - uint32_t wpos = TopK; // wpos will hold the final output position. - s->counter = 0; - __syncthreads(); - - // The 4-round radix loop - each round narrows by 8 bits until ties are fully - // resolved -#pragma unroll - for (int r = 0; r < 4; r++) { - uint32_t sh = 24 - r * 8; // round 0: bits 31-24, round 1: 23-16, etc. - uint32_t bin = (key >> sh) & 0xFF; // this tie's 8-bit bin for this round - - // Step 1: Build 256-bin histogram. - if (tx < RADIX) s->histogram[tx] = 0; - __syncthreads(); - if (active) atomicAdd(&s->histogram[bin], 1); - __syncthreads(); - - // Step 2: Prefix scan to find threshold - uint32_t hv = 0, wi2 = 0; - if (tx < RADIX) { - hv = s->histogram[tx]; - wi2 = warp_inclusive_sum(li, hv); - if (li == kWarpSize - 1) s->warp_sum[wi] = wi2; - } - __syncthreads(); - - if (tx < RADIX) { - auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; - auto tot = warp_reduce_sum_full(tmp); - auto inter = warp_reduce_sum_full(li < wi ? tmp : 0); - auto above = tot - (inter + wi2); - if (above < remain && above + hv >= remain) { - s->match = {tx, above, remain - above}; - } - } - __syncthreads(); - - // Step 3: Scatter - auto [thr, na, _] = s->match; // threshold bin, num above, unused - if (active) { - if (bin > thr) { - wpos = num_above + - atomicAdd(&s->counter, 1); // above -> place in output directly - active = false; - } else if (bin < thr) - active = false; // below -> discard - else if (r == 3) - wpos = TopK - atomicAdd(&s->match.equal_count, - -1u); // last round: place remaining - } - remain -= na; - if (!remain) break; // all ties resolved early - } - // Final write - if (wpos < TopK) output[wpos] = tie.idx; -} - -// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048). -// tie_handle assumes 1 tie per thread (max 1024). -// This version handles 2 ties per thread via kPerThread=2 -template -__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties, - uint32_t num_above, int32_t* output, - void* _smem) { - static_assert(TopK > kBlockSize); - struct TS { - alignas(128) uint32_t counter; - alignas(128) MatchBin match; - uint32_t histogram[RADIX]; - uint32_t warp_sum[kNumWarps]; - }; - auto* s = static_cast(_smem); - const auto tx = threadIdx.x; - const auto li = tx % kWarpSize; - const auto wi = tx / kWarpSize; - - constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize; - Tie my_ties[kPerThread]; - uint32_t keys[kPerThread]; - bool active[kPerThread]; - - for (uint32_t e = 0; e < kPerThread; e++) { - uint32_t idx = e * kBlockSize + tx; - if (idx < num_ties) { - my_ties[e] = ties[idx]; - keys[e] = convert_to_uint32_v2(ties[idx].score); - active[e] = true; - } else { - my_ties[e] = {0, 0.0f}; - keys[e] = 0; - active[e] = false; - } - } - - uint32_t remain = TopK - num_above; - s->counter = 0; - __syncthreads(); - - for (int r = 0; r < 4; r++) { - uint32_t sh = 24 - r * 8; - if (tx < RADIX) { - s->histogram[tx] = 0; - } - __syncthreads(); - - for (uint32_t e = 0; e < kPerThread; e++) { - if (active[e]) { - atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1); - } - } - __syncthreads(); - - uint32_t hv = 0; - if (tx < RADIX) { - hv = s->histogram[tx]; - auto wi2 = warp_inclusive_sum(li, hv); - if (li == kWarpSize - 1) { - s->warp_sum[wi] = wi2; - } - } - __syncthreads(); - if (tx < RADIX) { - auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; - auto total = warp_reduce_sum_full(tmp2); - auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0); - auto wi2 = warp_inclusive_sum(li, hv); - auto above = total - (inter + wi2); - if (above < remain && above + hv >= remain) { - s->match = { - .bin = tx, .above_count = above, .equal_count = remain - above}; - } - } - __syncthreads(); - - auto thr = s->match.bin; - auto na = s->match.above_count; - - for (uint32_t e = 0; e < kPerThread; e++) { - if (!active[e]) { - continue; - } - uint32_t bin = (keys[e] >> sh) & 0xFF; - if (bin > thr) { - uint32_t wpos = num_above + atomicAdd(&s->counter, 1); - if (wpos < TopK) { - output[wpos] = my_ties[e].idx; - } - active[e] = false; - } else if (bin < thr) { - active[e] = false; - } else if (r == 3) { - uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u); - if (wpos < TopK) { - output[wpos] = my_ties[e].idx; - } - } - } - - num_above += na; - remain -= na; - __syncthreads(); - s->counter = 0; - __syncthreads(); - } -} - // ============================================================================ // DSMEM histogram reduce // ============================================================================ template __device__ __forceinline__ void dsmem_hist_reduce(uint32_t* histogram) { - static_assert(kHistBins <= kBlockSize); + static_assert(kHistBins <= hist4096::kBlockSize); auto cluster = cooperative_groups::this_cluster(); cluster.sync(); const auto tx = threadIdx.x; @@ -378,16 +121,16 @@ __device__ __forceinline__ void find_threshold(uint32_t* histogram, uint32_t* warp_sum, uint32_t* counter_gt, uint32_t* counter_eq, - MatchBin* match) { + hist4096::MatchBin* match) { const auto tx = threadIdx.x; - const auto li = tx % kWarpSize, wi = tx / kWarpSize; + const auto li = tx % hist4096::kWarpSize, wi = tx / hist4096::kWarpSize; const auto value = tx < kHistBins ? histogram[tx] : 0; - const auto winc = warp_inclusive_sum(li, value); - if (li == kWarpSize - 1) warp_sum[wi] = winc; + const auto winc = hist4096::warp_inclusive_sum(li, value); + if (li == hist4096::kWarpSize - 1) warp_sum[wi] = winc; __syncthreads(); const auto tmp = warp_sum[li]; - const auto total = warp_reduce_sum_full(tmp); - auto pfx = warp_reduce_sum_full(li < wi ? tmp : 0) + winc; + const auto total = hist4096::warp_reduce_sum_full(tmp); + auto pfx = hist4096::warp_reduce_sum_full(li < wi ? tmp : 0) + winc; const auto above = total - pfx; if (tx < kHistBins && above < TopK && above + value >= TopK) { *counter_gt = *counter_eq = 0; @@ -396,247 +139,6 @@ __device__ __forceinline__ void find_threshold(uint32_t* histogram, __syncthreads(); } -// ============================================================================ -// Register-based single-CTA fast path for seq_len <= 16384 -// 4 float4 per thread × 1024 threads = 16384 elements max -// Uses 4096-bin (12-bit) histogram for better precision -// ============================================================================ - -template -struct Histogram4096Smem { - static constexpr uint32_t HIST_BINS = 1 << HIST_BITS; - alignas(128) uint32_t counter_gt; - alignas(128) uint32_t counter_eq; - MatchBin match; - uint32_t warp_sum[kNumWarps]; - union { - uint32_t histogram[HIST_BINS]; - Tie tie_buffer[kMaxTies]; - }; - static_assert(sizeof(uint32_t) * HIST_BINS >= TopK * sizeof(Tie), - "histogram union must be large enough for TopK ties"); -}; - -template -__device__ void histogram_4096_topk(const float* __restrict__ scores, - int32_t* __restrict__ output, - uint32_t length, void* _smem) { - constexpr uint32_t HIST_BINS = 1 << HIST_BITS; - constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize; - static_assert(HIST_BINS >= kBlockSize, - "HIST_BITS must give >= kBlockSize bins"); - - using Smem = Histogram4096Smem; - auto* smem = static_cast(_smem); - const auto tx = threadIdx.x; - const auto lane_id = tx % kWarpSize; - const auto warp_id = tx / kWarpSize; - - // Phase 1: Load all data into RF + build histogram - float4 - vecs[VECS_PER_THREAD]; // 4 vectors x 4 floats = 16 elements per thread - if constexpr (ITEMS_PER_THREAD >= 4) { - // Zero the histogram (SMEM writes) - for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++) - reinterpret_cast( - smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] = - make_uint4(0, 0, 0, 0); - } else { - if (tx < HIST_BINS) smem->histogram[tx] = 0; - } - if (tx == 0) { - smem->counter_gt = 0; - smem->counter_eq = 0; - } -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { - const uint32_t base = (tx + v * kBlockSize) * 4; - if (base < length) { - // TODO (roberto): try if 256-bit loads are faster (GMEM->RF) - vecs[v] = *reinterpret_cast(scores + base); - } - } - __syncthreads(); - - // Build histogram from RF via atomic adds into the shared histogram - bool done = false; -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { - const float* elems = reinterpret_cast(&vecs[v]); -#pragma unroll - for (uint32_t e = 0; e < 4 && !done; e++) { - const uint32_t idx = (tx + v * kBlockSize) * 4 + e; - if (idx >= length) { - done = true; - } else { - atomicAdd(&smem->histogram[extract_coarse_bin_N(elems[e])], - 1); - } - } - } - __syncthreads(); - - // Phase 2: Prefix scan to find threshold bin - // Multi-element scan (4096 bins: 4 per thread) - uint32_t orig[ITEMS_PER_THREAD]; - uint32_t local_sum = 0; - - // Step 1: Each thread sums its 4 bins -#pragma unroll - for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { - orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i]; - local_sum += orig[i]; - } - - // Step 2: Warp-level inclusive prefix sum on local_sum - const auto warp_inc = warp_inclusive_sum(lane_id, local_sum); - if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc; - __syncthreads(); - - // Step 3: Inter-warp prefix via redux.sync - const auto tmp = smem->warp_sum[lane_id]; - uint32_t prefix = warp_reduce_sum_full( - lane_id < warp_id ? tmp : 0); // sum of all prior warps - prefix += - warp_inc - local_sum; // exclusive prefix within this thread's position - - // Step 4: Find threshold - scan 4 bins, accumulate prefix -#pragma unroll - for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { - prefix += orig[i]; - const auto above = length - prefix; // elements in bins ABOVE this one - if (above < TopK && above + orig[i] >= TopK) { - smem->match = {.bin = tx * ITEMS_PER_THREAD + i, - .above_count = above, - .equal_count = orig[i]}; - } - } - - __syncthreads(); - - // Phase 3: Scatter from registers - const auto [thr_bin, num_above, num_equal] = smem->match; - const bool need_tie = (num_equal + num_above > TopK); - - done = false; -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { - const float* elems = reinterpret_cast(&vecs[v]); -#pragma unroll - for (uint32_t e = 0; e < 4 && !done; e++) { - const uint32_t idx = (tx + v * kBlockSize) * 4 + e; - if (idx >= length) { - done = true; - } else { - const uint32_t bin = extract_coarse_bin_N(elems[e]); - if (bin > thr_bin) { - output[atomicAdd(&smem->counter_gt, 1)] = - idx; // above -> output directly - } else if (bin == thr_bin) { - const auto pos = atomicAdd(&smem->counter_eq, 1); - if (!need_tie) { - if (pos + num_above < TopK) { - output[pos + num_above] = idx; // all fit - } - } else { - if (pos < TopK) { - smem->tie_buffer[pos] = {idx, elems[e]}; // store for refirement - } - } - } - // else: bin < thr_bin - discard (not in top-k) - } - } - } - - // Phase 4: Tie-breaking - if (!need_tie) return; - __syncthreads(); - - // Fast warp-ballot tie-breaking for small tie counts - const uint32_t num_ties = min(num_equal, static_cast(TopK)); - const uint32_t topk_remain = - TopK - num_above; // pick exactly remaining elements to fill topK - - auto is_greater = [](const Tie& a, const Tie& b) { - return (a.score > b.score) || (a.score == b.score && a.idx < b.idx); - }; - - if (num_ties <= kWarpSize) { - // <=32 ties - Use warp ballot - // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024 - // comparisons in one instruction per warp. O(1) work. - const auto lane_id = tx % kWarpSize; - const auto warp_id = tx / kWarpSize; - if (lane_id >= num_ties || warp_id >= num_ties) return; - const uint32_t mask = (1ull << num_ties) - 1u; - const auto tie = smem->tie_buffer[lane_id]; // each lane holds one tie - const auto target = - smem->tie_buffer[warp_id]; // each warp evaluates one candidate - const bool pred = - is_greater(tie, target); // compare all ties against target - const auto rank = static_cast( - __popc(__ballot_sync(mask, pred))); // count how many are greater - if (lane_id == 0 && rank < topk_remain) { - output[num_above + rank] = target.idx; // place at correct position - } - } else if (num_ties <= - kWarpSize * - 2) { // TODO (roberto): try to refactor this with <=32 case - // Same idea but each thread handles 2 tie elements - const auto lane_id = tx % kWarpSize; - const auto warp_id = tx / kWarpSize; - const auto lane1 = lane_id + kWarpSize; - const auto warp1 = warp_id + kWarpSize; - const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__}; - const auto tie0 = smem->tie_buffer[lane_id]; - const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid; - if (warp_id < num_ties) { - const auto target = smem->tie_buffer[warp_id]; - const auto r0 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); - const auto r1 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); - if (lane_id == 0 && r0 + r1 < topk_remain) - output[num_above + r0 + r1] = target.idx; - } - if (warp1 < num_ties) { - const auto target = smem->tie_buffer[warp1]; - const auto r0 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); - const auto r1 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); - if (lane_id == 0 && r0 + r1 < topk_remain) - output[num_above + r0 + r1] = target.idx; - } - } else { - // Large tie count: fall back to 4-round radix-256 sort - if constexpr (TopK <= kBlockSize) { - tie_handle(smem->tie_buffer, num_ties, num_above, output, smem); - } else { - tie_handle_large(smem->tie_buffer, num_ties, num_above, output, - smem); - } - } -} - -// ============================================================================ -// Streaming single-CTA path for seq_len 16K–32K -// TMA double-buffer, 4096-bin histogram, no cluster sync -struct StreamSmem { - uint64_t barrier[2][kStreamNumStages]; - alignas(128) uint32_t counter_gt; - alignas(128) uint32_t counter_eq; - MatchBin match; - uint32_t warp_sum[kNumWarps]; - union { - uint32_t histogram[kHist4096Bins]; - Tie tie_buffer[kMaxTies]; - }; - alignas(128) float score_buffer[kStreamNumStages][kSizePerStage]; -}; - // Streams data through shared memory in chunks, processing each chunk before // loading the next overwrites each buffer after processing it (the epilogue // prefetch loads the next chunk into the same slot) @@ -646,7 +148,7 @@ __device__ void tma_stream_pass(const float* scores, uint32_t length, uint32_t thr_bin, int32_t* indices, uint32_t* phases, SmemType* smem) { const auto tx = threadIdx.x; - const auto lane = tx % kWarpSize; + const auto lane = tx % hist4096::kWarpSize; const auto ni = (length + kSizePerStage - 1) / kSizePerStage; // total stages needed const auto la = @@ -684,12 +186,12 @@ __device__ void tma_stream_pass(const float* scores, uint32_t length, #pragma unroll for (uint32_t i = 0; i < kElemPerStage; i++) { - const auto li = tx + i * kBlockSize; + const auto li = tx + i * hist4096::kBlockSize; if (li >= sz) { break; } const auto sc = smem->score_buffer[b][li]; - const auto bn = extract_coarse_bin_N(sc); + const auto bn = hist4096::extract_coarse_bin_N(sc); if constexpr (kIsScatter) { // compile-time branch // Scatter pass: place above-threshold and collect ties const auto gi = o + li; @@ -697,7 +199,7 @@ __device__ void tma_stream_pass(const float* scores, uint32_t length, indices[atomicAdd(&smem->counter_gt, 1)] = gi; } else if (bn == thr_bin) { const auto p = atomicAdd(&smem->counter_eq, 1); - if (p < kMaxTies) { + if (p < hist4096::kMaxTies) { smem->tie_buffer[p] = {gi, sc}; } } @@ -732,18 +234,18 @@ struct SmemFused { uint64_t barrier[kPasses][kStages]; alignas(128) uint32_t counter_gt; alignas(128) uint32_t counter_eq; - alignas(128) MatchBin match; - uint32_t warp_sum[kNumWarps]; + alignas(128) hist4096::MatchBin match; + uint32_t warp_sum[hist4096::kNumWarps]; union { uint32_t histogram[kHistBins]; - Tie tie_buffer[kMaxTies]; + hist4096::Tie tie_buffer[kMaxTopK]; }; alignas(128) float score_buffer[kStages][kSizePerStage]; }; -using Smem8 = SmemFused; -using Smem16 = SmemFused; -using Smem4 = SmemFused; +using Smem8 = SmemFused; +using Smem16 = SmemFused; +using Smem4 = SmemFused; using SmemSinglePass = SmemFused; // Cluster-cooperative large path. @@ -753,10 +255,10 @@ using SmemSinglePass = SmemFused; template __device__ void large_topk(const float* __restrict__ row_input, int32_t* __restrict__ row_output, uint32_t seq_len, - uint32_t* phases, Tie* tie_ws) { + uint32_t* phases, hist4096::Tie* tie_ws) { const auto rank = blockIdx.y; // this block's position in cluster const auto tx = threadIdx.x; - const auto lane = tx % kWarpSize; + const auto lane = tx % hist4096::kWarpSize; extern __shared__ uint8_t smem_raw[]; auto* smem = reinterpret_cast(smem_raw); @@ -809,8 +311,8 @@ __device__ void large_topk(const float* __restrict__ row_input, __syncwarp(); #pragma unroll for (uint32_t i = 0; i < kElemPerStage; i += 2) { - const auto li0 = tx + i * kBlockSize; - const auto li1 = tx + (i + 1) * kBlockSize; + const auto li0 = tx + i * hist4096::kBlockSize; + const auto li1 = tx + (i + 1) * hist4096::kBlockSize; if (li0 >= sz) { break; } @@ -834,7 +336,7 @@ __device__ void large_topk(const float* __restrict__ row_input, smem->counter_eq = 0; } __syncthreads(); - tma_stream_pass( + tma_stream_pass( row_input + my_start, my_len, 0, nullptr, phases, smem); } @@ -853,7 +355,7 @@ __device__ void large_topk(const float* __restrict__ row_input, const auto sz = min(kSizePerStage, my_len - off); #pragma unroll for (uint32_t i = 0; i < kElemPerStage; i++) { - const auto li = tx + i * kBlockSize; + const auto li = tx + i * hist4096::kBlockSize; if (li >= sz) { break; } @@ -865,7 +367,7 @@ __device__ void large_topk(const float* __restrict__ row_input, } else if (bin == thr) { const auto p = atomicAdd(&smem->counter_eq, 1); // equal -> ties (later refinement) - if (p < kMaxTies) { + if (p < hist4096::kMaxTies) { smem->tie_buffer[p] = {gidx, score}; } } @@ -874,8 +376,8 @@ __device__ void large_topk(const float* __restrict__ row_input, __syncthreads(); } else { // Twopass scatter: re-stream data via TMA - uint32_t scatter_phases[kNumStages4] = {0, 0}; - tma_stream_pass( + uint32_t scatter_phases[kStreamingStagesCS4] = {0, 0}; + tma_stream_pass( row_input + my_start, my_len, thr, s_topk, scatter_phases, smem); } @@ -889,7 +391,7 @@ __device__ void large_topk(const float* __restrict__ row_input, const uint32_t la = smem->counter_gt; const uint32_t le_full = smem->counter_eq; const uint32_t le = - min(le_full, kMaxTies); // written smem tie_buffer entries + min(le_full, hist4096::kMaxTies); // written smem tie_buffer entries __shared__ uint32_t s_local_counts[CS]; __shared__ uint32_t s_prefix_packed; @@ -926,21 +428,20 @@ __device__ void large_topk(const float* __restrict__ row_input, const uint32_t prefix_equal = s_prefix_packed >> kAboveBits; // Write to global output - for (uint32_t i = tx; i < la; i += kBlockSize) { + for (uint32_t i = tx; i < la; i += hist4096::kBlockSize) { // indices are placed contiguously starting at prefix_above row_output[prefix_above + i] = s_topk[i] + my_start; // my_start: block-local -> row-global index } - for (uint32_t i = tx; i < le; i += kBlockSize) { + for (uint32_t i = tx; i < le; i += hist4096::kBlockSize) { const auto t = smem->tie_buffer[i]; uint32_t p = s_total_above + prefix_equal + i; if (p < TopK) { row_output[p] = t.idx + my_start; } uint32_t tp = prefix_equal + i; - if (tp < (TopK <= kBlockSize ? kMaxTies : TopK)) { - tie_ws[tp] = Tie{t.idx + my_start, - t.score}; // Ties are copied for cross-block refinement + if (tp < (TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK)) { + tie_ws[tp] = hist4096::Tie{t.idx + my_start, t.score}; } } @@ -954,20 +455,21 @@ __device__ void large_topk(const float* __restrict__ row_input, } // Tie-breaking uses FP32 (4-round radix sort) - if constexpr (TopK <= kBlockSize) { + if constexpr (TopK <= hist4096::kBlockSize) { // copy ties from tie_ws back to smem, then refine - const uint32_t num_ties = min(s_total_equal, kMaxTies); + const uint32_t num_ties = min(s_total_equal, hist4096::kMaxTies); // TODO (roberto): could vectorize with uint2 (8 bytes = exactly one Tie) - for (uint32_t i = tx; i < num_ties; i += kBlockSize) { - smem->tie_buffer[i] = Tie{tie_ws[i].idx, tie_ws[i].score}; + for (uint32_t i = tx; i < num_ties; i += hist4096::kBlockSize) { + smem->tie_buffer[i] = hist4096::Tie{tie_ws[i].idx, tie_ws[i].score}; } __syncthreads(); - tie_handle(smem->tie_buffer, num_ties, s_total_above, row_output, - smem); + hist4096::tie_handle(smem->tie_buffer, num_ties, s_total_above, + row_output, smem); } else { // TopK=2048: process directly from tie_ws (GMEM) const uint32_t num_ties = min(s_total_equal, static_cast(TopK)); - tie_handle_large(tie_ws, num_ties, s_total_above, row_output, smem); + hist4096::tie_handle_large(tie_ws, num_ties, s_total_above, + row_output, smem); } } @@ -987,7 +489,7 @@ __device__ void cooperative_topk_body(CooperativeTopKParams params) { // Trivial: seq_len <= TopK if (sl <= static_cast(TopK)) { if (rank == 0) { - for (uint32_t i = tx; i < TopK; i += kBlockSize) { + for (uint32_t i = tx; i < TopK; i += hist4096::kBlockSize) { out[i] = (i < static_cast(sl)) ? static_cast(i) : -1; } } @@ -995,11 +497,11 @@ __device__ void cooperative_topk_body(CooperativeTopKParams params) { } // Short-Medium path: histogram_4096_topk on rank 0 only - all data fits in RF - if (sl <= static_cast(kHist4096MaxLen)) { + if (sl <= static_cast(hist4096::kHist4096MaxLen)) { if (rank == 0) { extern __shared__ uint8_t sr[]; - histogram_4096_topk(in, out, sl, - sr); // 4096-bin (12-bit) histogram + hist4096::histogram_4096_topk( + in, out, sl, sr); // 4096-bin (12-bit) histogram } return; } @@ -1007,8 +509,8 @@ __device__ void cooperative_topk_body(CooperativeTopKParams params) { // Large path: init mbarriers + state, then dispatch fused or twopass const uint32_t per_block = (params.stride + CS - 1) / CS; // how many elements per block - constexpr uint32_t kFusedMax = ((CS == 16) ? kNumStages16 - : (CS == 8) ? kNumStages8 + constexpr uint32_t kFusedMax = ((CS == 16) ? kFusedStagesCS16 + : (CS == 8) ? kFusedStagesCS8 : kMaxSinglePassStages) * kSizePerStage; const bool use_singlepass = @@ -1016,15 +518,16 @@ __device__ void cooperative_topk_body(CooperativeTopKParams params) { kFusedMax; // single pass or TMA streaming: histogram+scatter // Select smem type and stage count at compile time based on CS - constexpr uint32_t kFusedStages = (CS == 16) ? kNumStages16 - : (CS == 8) ? kNumStages8 + constexpr uint32_t kFusedStages = (CS == 16) ? kFusedStagesCS16 + : (CS == 8) ? kFusedStagesCS8 : kMaxSinglePassStages; using FusedSmem = SmemFused; extern __shared__ uint8_t sr[]; - constexpr uint32_t kTieWsPerRow = TopK <= kBlockSize ? kMaxTies : TopK; - Tie* row_tie_ws = params.tie_ws + row * kTieWsPerRow; + constexpr uint32_t kTieWsPerRow = + TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK; + hist4096::Tie* row_tie_ws = params.tie_ws + row * kTieWsPerRow; if (use_singlepass) { auto* smem = reinterpret_cast(sr); @@ -1041,45 +544,41 @@ __device__ void cooperative_topk_body(CooperativeTopKParams params) { } else { // Two-pass: only CS=4 in practice (CS=8 always fits in singlepass) auto* smem = reinterpret_cast(sr); - if (tx < 2 * kNumStages4) { + if (tx < 2 * kStreamingStagesCS4) { mbarrier_init(&smem->barrier[0][tx], 1); // init 2×2=4 barriers (2 passes × 2 stages) } __syncthreads(); - uint32_t hp[kNumStages4] = {0, 0}; // histogram+scatter pass counters + uint32_t hp[kStreamingStagesCS4] = {0, 0}; // histogram+scatter pass counters large_topk(in, out, sl, hp, row_tie_ws); } } template -__global__ void __launch_bounds__(kBlockSize, 1) __cluster_dims__(1, 4, 1) +__global__ void __launch_bounds__(hist4096::kBlockSize, 1) __cluster_dims__(1, 4, 1) cooperative_topk_cs4(CooperativeTopKParams params) { cooperative_topk_body(params); } template -__global__ void __launch_bounds__(kBlockSize, 1) __cluster_dims__(1, 8, 1) +__global__ void __launch_bounds__(hist4096::kBlockSize, 1) __cluster_dims__(1, 8, 1) cooperative_topk_cs8(CooperativeTopKParams params) { cooperative_topk_body(params); } -#ifndef VLLM_COOPERATIVE_TOPK_PORTABLE_CLUSTER_ONLY template -__global__ void __launch_bounds__(kBlockSize, 1) __cluster_dims__(1, 16, 1) +__global__ void __launch_bounds__(hist4096::kBlockSize, 1) __cluster_dims__(1, 16, 1) cooperative_topk_cs16(CooperativeTopKParams params) { cooperative_topk_body(params); } -#endif -constexpr size_t kSmemSize4_base = - (sizeof(Smem4) > sizeof(StreamSmem) ? sizeof(Smem4) : sizeof(StreamSmem)); +constexpr size_t kSmemSize4_base = sizeof(Smem4); constexpr size_t kSmemSize4_sp = sizeof(SmemSinglePass); constexpr size_t kSmemSize4 = (kSmemSize4_base > kSmemSize4_sp ? kSmemSize4_base : kSmemSize4_sp) + sizeof(int32_t) * 2048 + 128; constexpr size_t kSmemSize8 = - std::max(sizeof(SmemFused), sizeof(StreamSmem)) + - sizeof(int32_t) * 2048 + 128; + sizeof(SmemFused) + sizeof(int32_t) * 2048 + 128; } // namespace cooperative diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 70d157f1a221..38f290f40ae8 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -11,6 +11,8 @@ #include #include +#include "topk_histogram_4096.cuh" + namespace vllm { namespace persistent { @@ -941,567 +943,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) // ============================================================================ namespace filtered_topk { -constexpr uint32_t kBlockSize = 1024; -constexpr uint32_t kHistBits = 10; -constexpr uint32_t kHistBins = 1 << kHistBits; -constexpr uint32_t RADIX = 256; -constexpr uint32_t kMaxTies = 1024; -static_assert(kMaxTies <= kBlockSize, - "tie_handle requires kMaxTies <= kBlockSize"); -constexpr uint32_t kWarpSize = 32; -constexpr uint32_t kNumWarps = kBlockSize / kWarpSize; - -constexpr uint32_t kElemPerStage = 16; -constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize; // 16384 - -// CS=4: 2 TMA stages (double buffer), two-pass -constexpr uint32_t kNumStages4 = 2; -// CS=8: 4 TMA stages, single-pass (all data stays in smem) -constexpr uint32_t kNumStages8 = - 2; // 2 stages × 16K = 32K per block (enough for CS=8) -// Max data per block with CS=8: ceil(262144/8) = 32768 ≤ 4 × 8192 = 32768 -constexpr uint32_t kMaxSeqLen8 = kNumStages8 * kSizePerStage * 8; // 262144 -constexpr uint32_t kNumStages16 = 2; // double-buffer for CS=16 -constexpr uint32_t kMaxSeqLen16 = kNumStages16 * kSizePerStage * 16; // 524288 - -// Register path -constexpr uint32_t kHist4096Bits = 12; -constexpr uint32_t kHist4096Bins = 1 << kHist4096Bits; // 4096 -constexpr uint32_t kHist4096VecsPerThread = 4; -constexpr uint32_t kHist4096MaxLen = - kHist4096VecsPerThread * 4 * kBlockSize; // 16384 -constexpr uint32_t kHist4096Items = kHist4096Bins / kBlockSize; // 4 - -// CS=4 single-pass path -constexpr uint32_t kMaxSinglePassStages = 3; -constexpr uint32_t kMaxSinglePassPerBlock = - kMaxSinglePassStages * kSizePerStage; // 49152 - -constexpr uint32_t kStreamNumStages = 2; - -struct alignas(16) MatchBin { - uint32_t bin, above_count, equal_count; -}; -struct alignas(8) Tie { - uint32_t idx; - float score; -}; - -__device__ __forceinline__ void load_float4_predicated(const float* ptr, - int base, int seq_len, - float& v0, float& v1, - float& v2, float& v3) { - uint32_t r0, r1, r2, r3; - const int p0 = (base < seq_len); - const int p1 = (base + 1 < seq_len); - const int p2 = (base + 2 < seq_len); - const int p3 = (base + 3 < seq_len); - asm volatile( - "{\n" - " .reg .pred pr0, pr1, pr2, pr3;\n" - " setp.ne.u32 pr0, %4, 0;\n" - " setp.ne.u32 pr1, %5, 0;\n" - " setp.ne.u32 pr2, %6, 0;\n" - " setp.ne.u32 pr3, %7, 0;\n" - " mov.u32 %0, 0xFF800000;\n" - " mov.u32 %1, 0xFF800000;\n" - " mov.u32 %2, 0xFF800000;\n" - " mov.u32 %3, 0xFF800000;\n" - " @pr0 ld.global.cg.u32 %0, [%8];\n" - " @pr1 ld.global.cg.u32 %1, [%8+4];\n" - " @pr2 ld.global.cg.u32 %2, [%8+8];\n" - " @pr3 ld.global.cg.u32 %3, [%8+12];\n" - "}\n" - : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) - : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr)); - v0 = __uint_as_float(r0); - v1 = __uint_as_float(r1); - v2 = __uint_as_float(r2); - v3 = __uint_as_float(r3); -} - -// converts the float32 score to a 32-bit ordered unsigned integer — the full -// precision key for radix sorting -__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { - uint32_t bits = __float_as_uint(x); - return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); -} - -// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits -> -// bin 0-4095) -template -__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) { - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return key >> (16 - kBits); -} - -// running sum within each warp — thread 0 gets its own value, thread 1 gets -// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc. -__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane, - uint32_t v) { -#pragma unroll - for (uint32_t o = 1; o < 32; o *= 2) { - uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o); - if (lane >= o) v += n; - } - return v; -} - -// Returns the sum of a value across all 32 threads in the warp, and every -// thread gets the same result SM90+ PTX instruction that does a hardware -// warp-wide reduction in a single instruction w.r.t. warp::reduce_sum(), which -// uses a __shfl_xor_sync butterfly tree (5 shuffles for 32 lanes) -__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) { - uint32_t r; - asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v)); - return r; -} - -// ============================================================================ -// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered -// key Each round narrows by 8 bits until ties are fully resolved -// ============================================================================ - -template -__device__ void tie_handle(const Tie* ties, uint32_t num_ties, - uint32_t num_above, int32_t* output, void* _smem) { - struct TS { - alignas(128) uint32_t counter; - alignas(128) MatchBin match; - uint32_t histogram[RADIX]; - uint32_t warp_sum[kNumWarps]; - }; - auto* s = static_cast(_smem); - const auto tx = threadIdx.x; - const auto li = tx % kWarpSize, wi = tx / kWarpSize; - - // Each thread loads one tie element. - const bool has = tx < num_ties; - const auto tie = has ? ties[tx] : Tie{0, 0.0f}; - const uint32_t key = convert_to_uint32_v2(tie.score); - - bool active = has; // tracks whether this thread's tie is still a candidate. - uint32_t remain = - TopK - num_above; // decreases each round as ties are resolved. - uint32_t wpos = TopK; // wpos will hold the final output position. - s->counter = 0; - __syncthreads(); - - // The 4-round radix loop - each round narrows by 8 bits until ties are fully - // resolved -#pragma unroll - for (int r = 0; r < 4; r++) { - uint32_t sh = 24 - r * 8; // round 0: bits 31-24, round 1: 23-16, etc. - uint32_t bin = (key >> sh) & 0xFF; // this tie's 8-bit bin for this round - - // Step 1: Build 256-bin histogram. - if (tx < RADIX) s->histogram[tx] = 0; - __syncthreads(); - if (active) atomicAdd(&s->histogram[bin], 1); - __syncthreads(); - - // Step 2: Prefix scan to find threshold - uint32_t hv = 0, wi2 = 0; - if (tx < RADIX) { - hv = s->histogram[tx]; - wi2 = warp_inclusive_sum(li, hv); - if (li == kWarpSize - 1) s->warp_sum[wi] = wi2; - } - __syncthreads(); - - if (tx < RADIX) { - auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; - auto tot = warp_reduce_sum_full(tmp); - auto inter = warp_reduce_sum_full(li < wi ? tmp : 0); - auto above = tot - (inter + wi2); - if (above < remain && above + hv >= remain) { - s->match = {tx, above, remain - above}; - } - } - __syncthreads(); - - // Step 3: Scatter - auto [thr, na, _] = s->match; // threshold bin, num above, unused - if (active) { - if (bin > thr) { - wpos = num_above + - atomicAdd(&s->counter, 1); // above -> place in output directly - active = false; - } else if (bin < thr) - active = false; // below -> discard - else if (r == 3) - wpos = TopK - atomicAdd(&s->match.equal_count, - -1u); // last round: place remaining - } - remain -= na; - if (!remain) break; // all ties resolved early - } - // Final write - if (wpos < TopK) output[wpos] = tie.idx; -} - -// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048). -// tie_handle assumes 1 tie per thread (max 1024). -// This version handles 2 ties per thread via kPerThread=2 -template -__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties, - uint32_t num_above, int32_t* output, - void* _smem) { - static_assert(TopK > kBlockSize); - struct TS { - alignas(128) uint32_t counter; - alignas(128) MatchBin match; - uint32_t histogram[RADIX]; - uint32_t warp_sum[kNumWarps]; - }; - auto* s = static_cast(_smem); - const auto tx = threadIdx.x; - const auto li = tx % kWarpSize; - const auto wi = tx / kWarpSize; - - constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize; - Tie my_ties[kPerThread]; - uint32_t keys[kPerThread]; - bool active[kPerThread]; - - for (uint32_t e = 0; e < kPerThread; e++) { - uint32_t idx = e * kBlockSize + tx; - if (idx < num_ties) { - my_ties[e] = ties[idx]; - keys[e] = convert_to_uint32_v2(ties[idx].score); - active[e] = true; - } else { - my_ties[e] = {0, 0.0f}; - keys[e] = 0; - active[e] = false; - } - } - - uint32_t remain = TopK - num_above; - s->counter = 0; - __syncthreads(); - - for (int r = 0; r < 4; r++) { - uint32_t sh = 24 - r * 8; - if (tx < RADIX) { - s->histogram[tx] = 0; - } - __syncthreads(); - - for (uint32_t e = 0; e < kPerThread; e++) { - if (active[e]) { - atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1); - } - } - __syncthreads(); - - uint32_t hv = 0; - if (tx < RADIX) { - hv = s->histogram[tx]; - auto wi2 = warp_inclusive_sum(li, hv); - if (li == kWarpSize - 1) { - s->warp_sum[wi] = wi2; - } - } - __syncthreads(); - if (tx < RADIX) { - auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; - auto total = warp_reduce_sum_full(tmp2); - auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0); - auto wi2 = warp_inclusive_sum(li, hv); - auto above = total - (inter + wi2); - if (above < remain && above + hv >= remain) { - s->match = { - .bin = tx, .above_count = above, .equal_count = remain - above}; - } - } - __syncthreads(); - - auto thr = s->match.bin; - auto na = s->match.above_count; - - for (uint32_t e = 0; e < kPerThread; e++) { - if (!active[e]) { - continue; - } - uint32_t bin = (keys[e] >> sh) & 0xFF; - if (bin > thr) { - uint32_t wpos = num_above + atomicAdd(&s->counter, 1); - if (wpos < TopK) { - output[wpos] = my_ties[e].idx; - } - active[e] = false; - } else if (bin < thr) { - active[e] = false; - } else if (r == 3) { - uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u); - if (wpos < TopK) { - output[wpos] = my_ties[e].idx; - } - } - } - - num_above += na; - remain -= na; - __syncthreads(); - s->counter = 0; - __syncthreads(); - } -} - -// ============================================================================ -// Register-based single-CTA fast path for seq_len <= 16384 -// 4 float4 per thread × 1024 threads = 16384 elements max -// Uses 4096-bin (12-bit) histogram for better precision -// ============================================================================ - -template -struct Histogram4096Smem { - static constexpr uint32_t HIST_BINS = 1 << HIST_BITS; - alignas(128) uint32_t counter_gt; - alignas(128) uint32_t counter_eq; - MatchBin match; - uint32_t warp_sum[kNumWarps]; - union { - uint32_t histogram[HIST_BINS]; - Tie tie_buffer[kMaxTies]; - }; - static_assert(sizeof(uint32_t) * HIST_BINS >= TopK * sizeof(Tie), - "histogram union must be large enough for TopK ties"); -}; - -template -__device__ void histogram_4096_topk(const float* __restrict__ scores, - int32_t* __restrict__ output, - uint32_t length, void* _smem) { - constexpr uint32_t HIST_BINS = 1 << HIST_BITS; - constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize; - static_assert(HIST_BINS >= kBlockSize, - "HIST_BITS must give >= kBlockSize bins"); - - using Smem = Histogram4096Smem; - auto* smem = static_cast(_smem); - const auto tx = threadIdx.x; - const auto lane_id = tx % kWarpSize; - const auto warp_id = tx / kWarpSize; - - // Phase 1: Load all data into RF + build histogram - float4 - vecs[VECS_PER_THREAD]; // 4 vectors x 4 floats = 16 elements per thread - if constexpr (ITEMS_PER_THREAD >= 4) { - // Zero the histogram (SMEM writes) - for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++) - reinterpret_cast( - smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] = - make_uint4(0, 0, 0, 0); - } else { - if (tx < HIST_BINS) smem->histogram[tx] = 0; - } - if (tx == 0) { - smem->counter_gt = 0; - smem->counter_eq = 0; - } - if constexpr (UsePredicatedLoads) { - const bool row_aligned = (reinterpret_cast(scores) & 0xFu) == 0; -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { - const uint32_t base = (tx + v * kBlockSize) * 4; - if (base < length) { - if (row_aligned && base + 3 < length) { - vecs[v] = *reinterpret_cast(scores + base); - } else { - load_float4_predicated(scores + base, static_cast(base), - static_cast(length), vecs[v].x, vecs[v].y, - vecs[v].z, vecs[v].w); - } - } - } - } else { -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { - const uint32_t base = (tx + v * kBlockSize) * 4; - if (base < length) { - vecs[v] = *reinterpret_cast(scores + base); - } - } - } - __syncthreads(); - - // Build histogram from RF via atomic adds into the shared histogram - bool done = false; -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { - const float* elems = reinterpret_cast(&vecs[v]); -#pragma unroll - for (uint32_t e = 0; e < 4 && !done; e++) { - const uint32_t idx = (tx + v * kBlockSize) * 4 + e; - if (idx >= length) { - done = true; - } else { - atomicAdd(&smem->histogram[extract_coarse_bin_N(elems[e])], - 1); - } - } - } - __syncthreads(); - - // Phase 2: Prefix scan to find threshold bin - // Multi-element scan (4096 bins: 4 per thread) - uint32_t orig[ITEMS_PER_THREAD]; - uint32_t local_sum = 0; - - // Step 1: Each thread sums its 4 bins -#pragma unroll - for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { - orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i]; - local_sum += orig[i]; - } - - // Step 2: Warp-level inclusive prefix sum on local_sum - const auto warp_inc = warp_inclusive_sum(lane_id, local_sum); - if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc; - __syncthreads(); - - // Step 3: Inter-warp prefix via redux.sync - const auto tmp = smem->warp_sum[lane_id]; - uint32_t prefix = warp_reduce_sum_full( - lane_id < warp_id ? tmp : 0); // sum of all prior warps - prefix += - warp_inc - local_sum; // exclusive prefix within this thread's position - - // Step 4: Find threshold - scan 4 bins, accumulate prefix -#pragma unroll - for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { - prefix += orig[i]; - const auto above = length - prefix; // elements in bins ABOVE this one - if (above < TopK && above + orig[i] >= TopK) { - smem->match = {.bin = tx * ITEMS_PER_THREAD + i, - .above_count = above, - .equal_count = orig[i]}; - } - } - - __syncthreads(); - - // Phase 3: Scatter from registers - const auto [thr_bin, num_above, num_equal] = smem->match; - const bool need_tie = (num_equal + num_above > TopK); - - done = false; -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { - const float* elems = reinterpret_cast(&vecs[v]); -#pragma unroll - for (uint32_t e = 0; e < 4 && !done; e++) { - const uint32_t idx = (tx + v * kBlockSize) * 4 + e; - if (idx >= length) { - done = true; - } else { - const uint32_t bin = extract_coarse_bin_N(elems[e]); - if (bin > thr_bin) { - output[atomicAdd(&smem->counter_gt, 1)] = - idx; // above -> output directly - } else if (bin == thr_bin) { - const auto pos = atomicAdd(&smem->counter_eq, 1); - if (!need_tie) { - if (pos + num_above < TopK) { - output[pos + num_above] = idx; // all fit - } - } else { - if (pos < TopK) { - smem->tie_buffer[pos] = {idx, elems[e]}; // store for refirement - } - } - } - // else: bin < thr_bin - discard (not in top-k) - } - } - } - - // Phase 4: Tie-breaking - if (!need_tie) return; - __syncthreads(); - - // Fast warp-ballot tie-breaking for small tie counts - const uint32_t num_ties = min(num_equal, static_cast(TopK)); - const uint32_t topk_remain = - TopK - num_above; // pick exactly remaining elements to fill topK - - auto is_greater = [](const Tie& a, const Tie& b) { - return (a.score > b.score) || (a.score == b.score && a.idx < b.idx); - }; - - if (num_ties <= kWarpSize) { - // <=32 ties - Use warp ballot - // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024 - // comparisons in one instruction per warp. O(1) work. - const auto lane_id = tx % kWarpSize; - const auto warp_id = tx / kWarpSize; - if (lane_id >= num_ties || warp_id >= num_ties) return; - const uint32_t mask = (1ull << num_ties) - 1u; - const auto tie = smem->tie_buffer[lane_id]; // each lane holds one tie - const auto target = - smem->tie_buffer[warp_id]; // each warp evaluates one candidate - const bool pred = - is_greater(tie, target); // compare all ties against target - const auto rank = static_cast( - __popc(__ballot_sync(mask, pred))); // count how many are greater - if (lane_id == 0 && rank < topk_remain) { - output[num_above + rank] = target.idx; // place at correct position - } - } else if (num_ties <= - kWarpSize * - 2) { // TODO (roberto): try to refactor this with <=32 case - // Same idea but each thread handles 2 tie elements - const auto lane_id = tx % kWarpSize; - const auto warp_id = tx / kWarpSize; - const auto lane1 = lane_id + kWarpSize; - const auto warp1 = warp_id + kWarpSize; - const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__}; - const auto tie0 = smem->tie_buffer[lane_id]; - const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid; - if (warp_id < num_ties) { - const auto target = smem->tie_buffer[warp_id]; - const auto r0 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); - const auto r1 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); - if (lane_id == 0 && r0 + r1 < topk_remain) - output[num_above + r0 + r1] = target.idx; - } - if (warp1 < num_ties) { - const auto target = smem->tie_buffer[warp1]; - const auto r0 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); - const auto r1 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); - if (lane_id == 0 && r0 + r1 < topk_remain) - output[num_above + r0 + r1] = target.idx; - } - } else { - // Large tie count: fall back to 4-round radix-256 sort - if constexpr (TopK <= kBlockSize) { - tie_handle(smem->tie_buffer, num_ties, num_above, output, smem); - } else { - tie_handle_large(smem->tie_buffer, num_ties, num_above, output, - smem); - } - } -} - -template -__device__ __noinline__ void histogram_4096_topk_predicated( - const float* __restrict__ scores, int32_t* __restrict__ output, - uint32_t length, void* _smem) { - histogram_4096_topk(scores, output, - length, _smem); -} +namespace hist4096 = topk_histogram_4096; // ============================================================================ // FilteredTopK — single CTA per row for bs > 32 @@ -1531,13 +973,6 @@ struct vec_t { data[i] = ptr[i]; } } - - FLASHINFER_INLINE void cast_store(T* ptr) const { -#pragma unroll - for (size_t i = 0; i < N; ++i) { - ptr[i] = data[i]; - } - } }; #undef FLASHINFER_INLINE @@ -1615,10 +1050,10 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) if (length <= 32768) { extern __shared__ uint8_t _smem_reg[]; if constexpr (UsePredicatedShortLoads) { - histogram_4096_topk_predicated(score, dst, length, + hist4096::histogram_4096_topk_predicated(score, dst, length, _smem_reg); } else { - histogram_4096_topk(score, dst, length, _smem_reg); + hist4096::histogram_4096_topk(score, dst, length, _smem_reg); } return; } diff --git a/csrc/libtorch_stable/topk_histogram_4096.cuh b/csrc/libtorch_stable/topk_histogram_4096.cuh new file mode 100644 index 000000000000..92b02538c686 --- /dev/null +++ b/csrc/libtorch_stable/topk_histogram_4096.cuh @@ -0,0 +1,555 @@ +/* + * Shared 4096-bin single-CTA TopK helpers. + */ + +#ifndef TOPK_HISTOGRAM_4096_CUH_ +#define TOPK_HISTOGRAM_4096_CUH_ + +#include +#include +#include + +namespace vllm { +namespace topk_histogram_4096 { + +constexpr uint32_t kBlockSize = 1024; +constexpr uint32_t RADIX = 256; +constexpr uint32_t kMaxTies = 1024; +static_assert(kMaxTies <= kBlockSize, + "tie_handle requires kMaxTies <= kBlockSize"); +constexpr uint32_t kWarpSize = 32; +constexpr uint32_t kNumWarps = kBlockSize / kWarpSize; + +// Register path +constexpr uint32_t kHist4096VecsPerThread = 4; +constexpr uint32_t kHist4096MaxLen = + kHist4096VecsPerThread * 4 * kBlockSize; // 16384 + +struct alignas(16) MatchBin { + uint32_t bin, above_count, equal_count; +}; +struct alignas(8) Tie { + uint32_t idx; + float score; +}; + +__device__ __forceinline__ void load_float4_predicated(const float* ptr, + int base, int seq_len, + float& v0, float& v1, + float& v2, float& v3) { + uint32_t r0, r1, r2, r3; + const int p0 = (base < seq_len); + const int p1 = (base + 1 < seq_len); + const int p2 = (base + 2 < seq_len); + const int p3 = (base + 3 < seq_len); + asm volatile( + "{\n" + " .reg .pred pr0, pr1, pr2, pr3;\n" + " setp.ne.u32 pr0, %4, 0;\n" + " setp.ne.u32 pr1, %5, 0;\n" + " setp.ne.u32 pr2, %6, 0;\n" + " setp.ne.u32 pr3, %7, 0;\n" + " mov.u32 %0, 0xFF800000;\n" + " mov.u32 %1, 0xFF800000;\n" + " mov.u32 %2, 0xFF800000;\n" + " mov.u32 %3, 0xFF800000;\n" + " @pr0 ld.global.cg.u32 %0, [%8];\n" + " @pr1 ld.global.cg.u32 %1, [%8+4];\n" + " @pr2 ld.global.cg.u32 %2, [%8+8];\n" + " @pr3 ld.global.cg.u32 %3, [%8+12];\n" + "}\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr)); + v0 = __uint_as_float(r0); + v1 = __uint_as_float(r1); + v2 = __uint_as_float(r2); + v3 = __uint_as_float(r3); +} + +// converts the float32 score to a 32-bit ordered unsigned integer — the full +// precision key for radix sorting +__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits -> +// bin 0-4095) +template +__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) { + __half h = __float2half_rn(x); + uint16_t bits = __half_as_ushort(h); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) + : static_cast(bits | 0x8000); + return key >> (16 - kBits); +} + +// running sum within each warp — thread 0 gets its own value, thread 1 gets +// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc. +__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane, + uint32_t v) { +#pragma unroll + for (uint32_t o = 1; o < 32; o *= 2) { + uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o); + if (lane >= o) v += n; + } + return v; +} + +// Returns the sum of a value across all 32 threads in the warp, and every +// thread gets the same result SM90+ PTX instruction that does a hardware +// warp-wide reduction in a single instruction w.r.t. warp::reduce_sum(), which +// uses a __shfl_xor_sync butterfly tree (5 shuffles for 32 lanes) +__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) { + uint32_t r; + asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v)); + return r; +} + +// ============================================================================ +// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered +// key Each round narrows by 8 bits until ties are fully resolved +// ============================================================================ + +template +__device__ void tie_handle(const Tie* ties, uint32_t num_ties, + uint32_t num_above, int32_t* output, void* _smem) { + struct TS { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[RADIX]; + uint32_t warp_sum[kNumWarps]; + }; + auto* s = static_cast(_smem); + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize, wi = tx / kWarpSize; + + // Each thread loads one tie element. + const bool has = tx < num_ties; + const auto tie = has ? ties[tx] : Tie{0, 0.0f}; + const uint32_t key = convert_to_uint32_v2(tie.score); + + bool active = has; // tracks whether this thread's tie is still a candidate. + uint32_t remain = + TopK - num_above; // decreases each round as ties are resolved. + uint32_t wpos = TopK; // wpos will hold the final output position. + s->counter = 0; + __syncthreads(); + + // The 4-round radix loop - each round narrows by 8 bits until ties are fully + // resolved +#pragma unroll + for (int r = 0; r < 4; r++) { + uint32_t sh = 24 - r * 8; // round 0: bits 31-24, round 1: 23-16, etc. + uint32_t bin = (key >> sh) & 0xFF; // this tie's 8-bit bin for this round + + // Step 1: Build 256-bin histogram. + if (tx < RADIX) s->histogram[tx] = 0; + __syncthreads(); + if (active) atomicAdd(&s->histogram[bin], 1); + __syncthreads(); + + // Step 2: Prefix scan to find threshold + uint32_t hv = 0, wi2 = 0; + if (tx < RADIX) { + hv = s->histogram[tx]; + wi2 = warp_inclusive_sum(li, hv); + if (li == kWarpSize - 1) s->warp_sum[wi] = wi2; + } + __syncthreads(); + + if (tx < RADIX) { + auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; + auto tot = warp_reduce_sum_full(tmp); + auto inter = warp_reduce_sum_full(li < wi ? tmp : 0); + auto above = tot - (inter + wi2); + if (above < remain && above + hv >= remain) { + s->match = {tx, above, remain - above}; + } + } + __syncthreads(); + + // Step 3: Scatter + auto [thr, na, _] = s->match; // threshold bin, num above, unused + if (active) { + if (bin > thr) { + wpos = num_above + + atomicAdd(&s->counter, 1); // above -> place in output directly + active = false; + } else if (bin < thr) + active = false; // below -> discard + else if (r == 3) + wpos = TopK - atomicAdd(&s->match.equal_count, + -1u); // last round: place remaining + } + remain -= na; + if (!remain) break; // all ties resolved early + } + // Final write + if (wpos < TopK) output[wpos] = tie.idx; +} + +// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048). +// tie_handle assumes 1 tie per thread (max 1024). +// This version handles 2 ties per thread via kPerThread=2 +template +__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties, + uint32_t num_above, int32_t* output, + void* _smem) { + static_assert(TopK > kBlockSize); + struct TS { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[RADIX]; + uint32_t warp_sum[kNumWarps]; + }; + auto* s = static_cast(_smem); + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize; + const auto wi = tx / kWarpSize; + + constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize; + Tie my_ties[kPerThread]; + uint32_t keys[kPerThread]; + bool active[kPerThread]; + + for (uint32_t e = 0; e < kPerThread; e++) { + uint32_t idx = e * kBlockSize + tx; + if (idx < num_ties) { + my_ties[e] = ties[idx]; + keys[e] = convert_to_uint32_v2(ties[idx].score); + active[e] = true; + } else { + my_ties[e] = {0, 0.0f}; + keys[e] = 0; + active[e] = false; + } + } + + uint32_t remain = TopK - num_above; + s->counter = 0; + __syncthreads(); + + for (int r = 0; r < 4; r++) { + uint32_t sh = 24 - r * 8; + if (tx < RADIX) { + s->histogram[tx] = 0; + } + __syncthreads(); + + for (uint32_t e = 0; e < kPerThread; e++) { + if (active[e]) { + atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1); + } + } + __syncthreads(); + + uint32_t hv = 0; + if (tx < RADIX) { + hv = s->histogram[tx]; + auto wi2 = warp_inclusive_sum(li, hv); + if (li == kWarpSize - 1) { + s->warp_sum[wi] = wi2; + } + } + __syncthreads(); + if (tx < RADIX) { + auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; + auto total = warp_reduce_sum_full(tmp2); + auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0); + auto wi2 = warp_inclusive_sum(li, hv); + auto above = total - (inter + wi2); + if (above < remain && above + hv >= remain) { + s->match = { + .bin = tx, .above_count = above, .equal_count = remain - above}; + } + } + __syncthreads(); + + auto thr = s->match.bin; + auto na = s->match.above_count; + + for (uint32_t e = 0; e < kPerThread; e++) { + if (!active[e]) { + continue; + } + uint32_t bin = (keys[e] >> sh) & 0xFF; + if (bin > thr) { + uint32_t wpos = num_above + atomicAdd(&s->counter, 1); + if (wpos < TopK) { + output[wpos] = my_ties[e].idx; + } + active[e] = false; + } else if (bin < thr) { + active[e] = false; + } else if (r == 3) { + uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u); + if (wpos < TopK) { + output[wpos] = my_ties[e].idx; + } + } + } + + num_above += na; + remain -= na; + __syncthreads(); + s->counter = 0; + __syncthreads(); + } +} + +// ============================================================================ +// Register-based single-CTA fast path for seq_len <= 16384 +// 4 float4 per thread × 1024 threads = 16384 elements max +// Uses 4096-bin (12-bit) histogram for better precision +// ============================================================================ + +template +struct Histogram4096Smem { + static constexpr uint32_t HIST_BINS = 1 << HIST_BITS; + static constexpr uint32_t TIE_CAPACITY = + TopK > kMaxTies ? TopK : kMaxTies; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + MatchBin match; + uint32_t warp_sum[kNumWarps]; + union { + uint32_t histogram[HIST_BINS]; + Tie tie_buffer[TIE_CAPACITY]; + }; +}; + +template +__device__ void histogram_4096_topk(const float* __restrict__ scores, + int32_t* __restrict__ output, + uint32_t length, void* _smem) { + constexpr uint32_t HIST_BINS = 1 << HIST_BITS; + constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize; + static_assert(HIST_BINS >= kBlockSize, + "HIST_BITS must give >= kBlockSize bins"); + + using Smem = Histogram4096Smem; + auto* smem = static_cast(_smem); + const auto tx = threadIdx.x; + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + + // Phase 1: Load all data into RF + build histogram + float4 + vecs[VECS_PER_THREAD]; // 4 vectors x 4 floats = 16 elements per thread + if constexpr (ITEMS_PER_THREAD >= 4) { + // Zero the histogram (SMEM writes) + for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++) + reinterpret_cast( + smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] = + make_uint4(0, 0, 0, 0); + } else { + if (tx < HIST_BINS) smem->histogram[tx] = 0; + } + if (tx == 0) { + smem->counter_gt = 0; + smem->counter_eq = 0; + } + if constexpr (UsePredicatedLoads) { + const bool row_aligned = (reinterpret_cast(scores) & 0xFu) == 0; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { + const uint32_t base = (tx + v * kBlockSize) * 4; + if (base < length) { + if (row_aligned && base + 3 < length) { + vecs[v] = *reinterpret_cast(scores + base); + } else { + load_float4_predicated(scores + base, static_cast(base), + static_cast(length), vecs[v].x, vecs[v].y, + vecs[v].z, vecs[v].w); + } + } + } + } else { +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { + const uint32_t base = (tx + v * kBlockSize) * 4; + if (base < length) { + vecs[v] = *reinterpret_cast(scores + base); + } + } + } + __syncthreads(); + + // Build histogram from RF via atomic adds into the shared histogram + bool done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + atomicAdd(&smem->histogram[extract_coarse_bin_N(elems[e])], + 1); + } + } + } + __syncthreads(); + + // Phase 2: Prefix scan to find threshold bin + // Multi-element scan (4096 bins: 4 per thread) + uint32_t orig[ITEMS_PER_THREAD]; + uint32_t local_sum = 0; + + // Step 1: Each thread sums its 4 bins +#pragma unroll + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { + orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i]; + local_sum += orig[i]; + } + + // Step 2: Warp-level inclusive prefix sum on local_sum + const auto warp_inc = warp_inclusive_sum(lane_id, local_sum); + if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc; + __syncthreads(); + + // Step 3: Inter-warp prefix via redux.sync + const auto tmp = smem->warp_sum[lane_id]; + uint32_t prefix = warp_reduce_sum_full( + lane_id < warp_id ? tmp : 0); // sum of all prior warps + prefix += + warp_inc - local_sum; // exclusive prefix within this thread's position + + // Step 4: Find threshold - scan 4 bins, accumulate prefix +#pragma unroll + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { + prefix += orig[i]; + const auto above = length - prefix; // elements in bins ABOVE this one + if (above < TopK && above + orig[i] >= TopK) { + smem->match = {.bin = tx * ITEMS_PER_THREAD + i, + .above_count = above, + .equal_count = orig[i]}; + } + } + + __syncthreads(); + + // Phase 3: Scatter from registers + const auto [thr_bin, num_above, num_equal] = smem->match; + const bool need_tie = (num_equal + num_above > TopK); + + done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + const uint32_t bin = extract_coarse_bin_N(elems[e]); + if (bin > thr_bin) { + output[atomicAdd(&smem->counter_gt, 1)] = + idx; // above -> output directly + } else if (bin == thr_bin) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (!need_tie) { + if (pos + num_above < TopK) { + output[pos + num_above] = idx; // all fit + } + } else { + if (pos < TopK) { + smem->tie_buffer[pos] = {idx, elems[e]}; // store for refirement + } + } + } + // else: bin < thr_bin - discard (not in top-k) + } + } + } + + // Phase 4: Tie-breaking + if (!need_tie) return; + __syncthreads(); + + // Fast warp-ballot tie-breaking for small tie counts + const uint32_t num_ties = min(num_equal, static_cast(TopK)); + const uint32_t topk_remain = + TopK - num_above; // pick exactly remaining elements to fill topK + + auto is_greater = [](const Tie& a, const Tie& b) { + return (a.score > b.score) || (a.score == b.score && a.idx < b.idx); + }; + + if (num_ties <= kWarpSize) { + // <=32 ties - Use warp ballot + // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024 + // comparisons in one instruction per warp. O(1) work. + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + if (lane_id >= num_ties || warp_id >= num_ties) return; + const uint32_t mask = (1ull << num_ties) - 1u; + const auto tie = smem->tie_buffer[lane_id]; // each lane holds one tie + const auto target = + smem->tie_buffer[warp_id]; // each warp evaluates one candidate + const bool pred = + is_greater(tie, target); // compare all ties against target + const auto rank = static_cast( + __popc(__ballot_sync(mask, pred))); // count how many are greater + if (lane_id == 0 && rank < topk_remain) { + output[num_above + rank] = target.idx; // place at correct position + } + } else if (num_ties <= + kWarpSize * + 2) { // TODO (roberto): try to refactor this with <=32 case + // Same idea but each thread handles 2 tie elements + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + const auto lane1 = lane_id + kWarpSize; + const auto warp1 = warp_id + kWarpSize; + const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__}; + const auto tie0 = smem->tie_buffer[lane_id]; + const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid; + if (warp_id < num_ties) { + const auto target = smem->tie_buffer[warp_id]; + const auto r0 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); + const auto r1 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); + if (lane_id == 0 && r0 + r1 < topk_remain) + output[num_above + r0 + r1] = target.idx; + } + if (warp1 < num_ties) { + const auto target = smem->tie_buffer[warp1]; + const auto r0 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); + const auto r1 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); + if (lane_id == 0 && r0 + r1 < topk_remain) + output[num_above + r0 + r1] = target.idx; + } + } else { + // Large tie count: fall back to 4-round radix-256 sort + if constexpr (TopK <= kBlockSize) { + tie_handle(smem->tie_buffer, num_ties, num_above, output, smem); + } else { + tie_handle_large(smem->tie_buffer, num_ties, num_above, output, + smem); + } + } +} + +template +__device__ __noinline__ void histogram_4096_topk_predicated( + const float* __restrict__ scores, int32_t* __restrict__ output, + uint32_t length, void* _smem) { + histogram_4096_topk(scores, output, + length, _smem); +} + +} // namespace topk_histogram_4096 +} // namespace vllm + +#endif // TOPK_HISTOGRAM_4096_CUH_ From 147a3b2b8eade18aea39dd382b0b2d27b8512d40 Mon Sep 17 00:00:00 2001 From: LopezCastroRoberto Date: Tue, 23 Jun 2026 12:12:55 +0000 Subject: [PATCH 6/7] fix pre-commit Signed-off-by: LopezCastroRoberto --- .buildkite/test_areas/kernels.yaml | 2 ++ csrc/libtorch_stable/cooperative_topk.cuh | 24 ++++++++++++-------- csrc/libtorch_stable/persistent_topk.cuh | 5 ++-- csrc/libtorch_stable/topk_histogram_4096.cuh | 3 +-- 4 files changed, 21 insertions(+), 13 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 4953b4d441ce..387daacf700f 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -47,8 +47,10 @@ steps: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations commands: - pytest -v -s kernels/test_fused_deepseek_v4_*.py + - pytest -v -s tests/kernels/test_top_k_per_row.py - label: Deepseek V4 Kernel Test (B200) key: deepseek-v4-kernel-test-b200 diff --git a/csrc/libtorch_stable/cooperative_topk.cuh b/csrc/libtorch_stable/cooperative_topk.cuh index 937a18fbf635..b43b9b8447d6 100644 --- a/csrc/libtorch_stable/cooperative_topk.cuh +++ b/csrc/libtorch_stable/cooperative_topk.cuh @@ -25,7 +25,8 @@ constexpr uint32_t kHistBins = 1 << kHistBits; constexpr uint32_t kMaxTopK = 2048; constexpr uint32_t kElemPerStage = 16; -constexpr uint32_t kSizePerStage = kElemPerStage * hist4096::kBlockSize; // 16384 +constexpr uint32_t kSizePerStage = + kElemPerStage * hist4096::kBlockSize; // 16384 // CS=4 two-pass path uses two TMA stages as a double buffer. constexpr uint32_t kStreamingStagesCS4 = 2; @@ -43,7 +44,8 @@ struct CooperativeTopKParams { const float* __restrict__ input; int32_t* __restrict__ output; const int32_t* __restrict__ lengths; - hist4096::Tie* __restrict__ tie_ws; // per-row tie workspace, see kTieWsPerRow + hist4096::Tie* __restrict__ tie_ws; // per-row tie workspace, see + // kTieWsPerRow uint32_t num_rows, stride; }; @@ -549,26 +551,30 @@ __device__ void cooperative_topk_body(CooperativeTopKParams params) { 1); // init 2×2=4 barriers (2 passes × 2 stages) } __syncthreads(); - uint32_t hp[kStreamingStagesCS4] = {0, 0}; // histogram+scatter pass counters + uint32_t hp[kStreamingStagesCS4] = {0, + 0}; // histogram+scatter pass counters large_topk(in, out, sl, hp, row_tie_ws); } } template -__global__ void __launch_bounds__(hist4096::kBlockSize, 1) __cluster_dims__(1, 4, 1) - cooperative_topk_cs4(CooperativeTopKParams params) { +__global__ void __launch_bounds__(hist4096::kBlockSize, 1) + __cluster_dims__(1, 4, 1) + cooperative_topk_cs4(CooperativeTopKParams params) { cooperative_topk_body(params); } template -__global__ void __launch_bounds__(hist4096::kBlockSize, 1) __cluster_dims__(1, 8, 1) - cooperative_topk_cs8(CooperativeTopKParams params) { +__global__ void __launch_bounds__(hist4096::kBlockSize, 1) + __cluster_dims__(1, 8, 1) + cooperative_topk_cs8(CooperativeTopKParams params) { cooperative_topk_body(params); } template -__global__ void __launch_bounds__(hist4096::kBlockSize, 1) __cluster_dims__(1, 16, 1) - cooperative_topk_cs16(CooperativeTopKParams params) { +__global__ void __launch_bounds__(hist4096::kBlockSize, 1) + __cluster_dims__(1, 16, 1) + cooperative_topk_cs16(CooperativeTopKParams params) { cooperative_topk_body(params); } diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 38f290f40ae8..85618feeb8a0 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -1051,9 +1051,10 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) extern __shared__ uint8_t _smem_reg[]; if constexpr (UsePredicatedShortLoads) { hist4096::histogram_4096_topk_predicated(score, dst, length, - _smem_reg); + _smem_reg); } else { - hist4096::histogram_4096_topk(score, dst, length, _smem_reg); + hist4096::histogram_4096_topk(score, dst, length, + _smem_reg); } return; } diff --git a/csrc/libtorch_stable/topk_histogram_4096.cuh b/csrc/libtorch_stable/topk_histogram_4096.cuh index 92b02538c686..71c6c2cdf012 100644 --- a/csrc/libtorch_stable/topk_histogram_4096.cuh +++ b/csrc/libtorch_stable/topk_histogram_4096.cuh @@ -307,8 +307,7 @@ __device__ void tie_handle_large(const Tie* ties, uint32_t num_ties, template struct Histogram4096Smem { static constexpr uint32_t HIST_BINS = 1 << HIST_BITS; - static constexpr uint32_t TIE_CAPACITY = - TopK > kMaxTies ? TopK : kMaxTies; + static constexpr uint32_t TIE_CAPACITY = TopK > kMaxTies ? TopK : kMaxTies; alignas(128) uint32_t counter_gt; alignas(128) uint32_t counter_eq; MatchBin match; From d29c385b2999a5dde10c8a506f4cc0db29777237 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Tue, 23 Jun 2026 21:07:38 +0000 Subject: [PATCH 7/7] minor Signed-off-by: Woosuk Kwon --- .buildkite/test_areas/kernels.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 387daacf700f..c5341a0f5188 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -50,7 +50,7 @@ steps: - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations commands: - pytest -v -s kernels/test_fused_deepseek_v4_*.py - - pytest -v -s tests/kernels/test_top_k_per_row.py + - pytest -v -s kernels/test_top_k_per_row.py - label: Deepseek V4 Kernel Test (B200) key: deepseek-v4-kernel-test-b200