diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 4953b4d441ce..c5341a0f5188 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 kernels/test_top_k_per_row.py - label: Deepseek V4 Kernel Test (B200) key: deepseek-v4-kernel-test-b200 diff --git a/CMakeLists.txt b/CMakeLists.txt index e95fe38d3299..6d130f8bda2a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -382,6 +382,24 @@ 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_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1") + + endif() + endif() + if(VLLM_GPU_LANG STREQUAL "CUDA") SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") @@ -498,6 +516,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. @@ -1049,6 +1075,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/cooperative_topk.cu b/csrc/libtorch_stable/cooperative_topk.cu new file mode 100644 index 000000000000..f388a9e6c8e7 --- /dev/null +++ b/csrc/libtorch_stable/cooperative_topk.cu @@ -0,0 +1,146 @@ +// 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; +namespace hist4096 = vllm::topk_histogram_4096; +#endif + +#ifndef USE_ROCM +template +void launch_cooperative_cluster(ct::CooperativeTopKParams& params, + size_t smem, cudaStream_t stream) { + auto kernel = []() { + if constexpr (CS == 16) { + return &ct::cooperative_topk_cs16; + } 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; + } + }(); + if constexpr (CS > 8) { + cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, + 1); + } + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(params.num_rows, CS); + cfg.blockDim = dim3(hist4096::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) { + (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(); + + 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 <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK; + STD_TORCH_CHECK( + workspace.size(0) >= + static_cast(num_rows * kTieWsPerRow * sizeof(hist4096::Tie)), + "workspace too small"); + + 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) { + 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..b43b9b8447d6 --- /dev/null +++ b/csrc/libtorch_stable/cooperative_topk.cuh @@ -0,0 +1,593 @@ +/* + * Cooperative TopK kernel for DSA Indexer + */ + +#ifndef COOPERATIVE_TOPK_CUH_ +#define COOPERATIVE_TOPK_CUH_ + +#include +#include +#include +#include +#include +#include +#include + +#include "topk_histogram_4096.cuh" + +namespace vllm { +namespace cooperative { + +namespace hist4096 = topk_histogram_4096; + +constexpr uint32_t kHistBits = 10; +constexpr uint32_t kHistBins = 1 << kHistBits; +constexpr uint32_t kMaxTopK = 2048; + +constexpr uint32_t kElemPerStage = 16; +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; +// 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 + +template +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 + uint32_t num_rows, stride; +}; + +// ============================================================================ +// Cooperative helpers +// ============================================================================ + +// 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 hist4096::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); +} + +// ============================================================================ +// DSMEM histogram reduce +// ============================================================================ + +template +__device__ __forceinline__ void dsmem_hist_reduce(uint32_t* histogram) { + static_assert(kHistBins <= hist4096::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, + hist4096::MatchBin* match) { + const auto tx = threadIdx.x; + const auto li = tx % hist4096::kWarpSize, wi = tx / hist4096::kWarpSize; + const auto value = tx < kHistBins ? histogram[tx] : 0; + 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 = 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; + *match = {.bin = tx, .above_count = above, .equal_count = value}; + } + __syncthreads(); +} + +// 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 % hist4096::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 * hist4096::kBlockSize; + if (li >= sz) { + break; + } + const auto sc = smem->score_buffer[b][li]; + 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; + 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 < hist4096::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) hist4096::MatchBin match; + uint32_t warp_sum[hist4096::kNumWarps]; + union { + uint32_t histogram[kHistBins]; + hist4096::Tie tie_buffer[kMaxTopK]; + }; + 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, hist4096::Tie* tie_ws) { + const auto rank = blockIdx.y; // this block's position in cluster + const auto tx = threadIdx.x; + const auto lane = tx % hist4096::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 * hist4096::kBlockSize; + const auto li1 = tx + (i + 1) * hist4096::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 * hist4096::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 < hist4096::kMaxTies) { + smem->tie_buffer[p] = {gidx, score}; + } + } + } + } + __syncthreads(); + } else { + // Twopass scatter: re-stream data via TMA + uint32_t scatter_phases[kStreamingStagesCS4] = {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, hist4096::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 += 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 += 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 <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK)) { + tie_ws[tp] = hist4096::Tie{t.idx + my_start, t.score}; + } + } + + // 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 <= hist4096::kBlockSize) { + // copy ties from tie_ws back to smem, then refine + 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 += hist4096::kBlockSize) { + smem->tie_buffer[i] = hist4096::Tie{tie_ws[i].idx, tie_ws[i].score}; + } + __syncthreads(); + 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)); + hist4096::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 += hist4096::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(hist4096::kHist4096MaxLen)) { + if (rank == 0) { + extern __shared__ uint8_t sr[]; + hist4096::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) ? kFusedStagesCS16 + : (CS == 8) ? kFusedStagesCS8 + : 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) ? kFusedStagesCS16 + : (CS == 8) ? kFusedStagesCS8 + : kMaxSinglePassStages; + using FusedSmem = SmemFused; + + extern __shared__ uint8_t sr[]; + + 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); + 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 * kStreamingStagesCS4) { + mbarrier_init(&smem->barrier[0][tx], + 1); // init 2×2=4 barriers (2 passes × 2 stages) + } + __syncthreads(); + 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) { + cooperative_topk_body(params); +} + +template +__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) { + cooperative_topk_body(params); +} + +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 = + sizeof(SmemFused) + 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..85618feeb8a0 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 { @@ -935,8 +937,16 @@ __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 { + +namespace hist4096 = topk_histogram_4096; + +// ============================================================================ +// FilteredTopK — single CTA per row for bs > 32 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 // ============================================================================ @@ -963,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 @@ -1013,7 +1016,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 +1046,19 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) return; } + // Short path + if (length <= 32768) { + extern __shared__ uint8_t _smem_reg[]; + if constexpr (UsePredicatedShortLoads) { + hist4096::histogram_4096_topk_predicated(score, dst, length, + _smem_reg); + } else { + hist4096::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; @@ -1285,14 +1302,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 +1324,19 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input, return cudaSuccess; } +} // namespace filtered_topk + +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, + 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/topk_histogram_4096.cuh b/csrc/libtorch_stable/topk_histogram_4096.cuh new file mode 100644 index 000000000000..71c6c2cdf012 --- /dev/null +++ b/csrc/libtorch_stable/topk_histogram_4096.cuh @@ -0,0 +1,554 @@ +/* + * 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_ 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),