mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-03 20:38:06 +00:00
[Perf][DSv4/DSv3.2] Add cluster-cooperative topK kernel for low-latency scenarios (#43008)
Signed-off-by: LopezCastroRoberto <[email protected]>
This commit is contained in:
Roberto L. Castro
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
// Cooperative cluster TopK for DeepSeek V3 sparse attention indexer.
|
||||
// See cooperative_topk.cuh for kernel implementation.
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#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 <uint32_t TopK, uint32_t CS>
|
||||
void launch_cooperative_cluster(ct::CooperativeTopKParams<TopK>& params,
|
||||
size_t smem, cudaStream_t stream) {
|
||||
auto kernel = []() {
|
||||
if constexpr (CS == 16) {
|
||||
return &ct::cooperative_topk_cs16<TopK>;
|
||||
} else if constexpr (CS == 8) {
|
||||
return &ct::cooperative_topk_cs8<TopK>;
|
||||
} else {
|
||||
static_assert(CS == 4, "unsupported cooperative_topk cluster size");
|
||||
return &ct::cooperative_topk_cs4<TopK>;
|
||||
}
|
||||
}();
|
||||
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 <uint32_t TopK>
|
||||
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<uint32_t>(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<TopK> params;
|
||||
params.input = logits.const_data_ptr<float>();
|
||||
params.output = output.mutable_data_ptr<int32_t>();
|
||||
params.lengths = lengths.const_data_ptr<int32_t>();
|
||||
params.num_rows = static_cast<uint32_t>(num_rows);
|
||||
params.stride = stride;
|
||||
params.tie_ws =
|
||||
reinterpret_cast<hist4096::Tie*>(workspace.mutable_data_ptr<uint8_t>());
|
||||
|
||||
constexpr uint32_t kTieWsPerRow =
|
||||
TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK;
|
||||
STD_TORCH_CHECK(
|
||||
workspace.size(0) >=
|
||||
static_cast<int64_t>(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<TopK, 16>(params, ct::kSmemSize8, stream);
|
||||
} else if (num_rows <= 8) {
|
||||
launch_cooperative_cluster<TopK, 8>(params, ct::kSmemSize8, stream);
|
||||
} else {
|
||||
launch_cooperative_cluster<TopK, 4>(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
|
||||
}
|
||||
@@ -0,0 +1,593 @@
|
||||
/*
|
||||
* Cooperative TopK kernel for DSA Indexer
|
||||
*/
|
||||
|
||||
#ifndef COOPERATIVE_TOPK_CUH_
|
||||
#define COOPERATIVE_TOPK_CUH_
|
||||
|
||||
#include <cooperative_groups.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda/ptx>
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
|
||||
#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 <uint32_t TopK = 1024>
|
||||
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 <uint32_t N>
|
||||
__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<kHistBits>(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 <uint32_t CS>
|
||||
__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<CS>(*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 <uint32_t TopK>
|
||||
__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 <typename SmemType, uint32_t kStages, uint32_t kBinBits,
|
||||
bool kIsScatter>
|
||||
__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<kBinBits>(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 <uint32_t kStages, uint32_t kPasses = 1>
|
||||
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<kFusedStagesCS8>;
|
||||
using Smem16 = SmemFused<kFusedStagesCS16>;
|
||||
using Smem4 = SmemFused<kStreamingStagesCS4, 2>;
|
||||
using SmemSinglePass = SmemFused<kMaxSinglePassStages>;
|
||||
|
||||
// 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 <uint32_t TopK, uint32_t CS, typename SmemType, bool kFused>
|
||||
__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<SmemType*>(smem_raw);
|
||||
int32_t* s_topk = reinterpret_cast<int32_t*>(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<SmemType, kStreamingStagesCS4, kHistBits, false>(
|
||||
row_input + my_start, my_len, 0, nullptr, phases, smem);
|
||||
}
|
||||
|
||||
// DSMEM all-reduce + find threshold
|
||||
dsmem_hist_reduce<CS>(
|
||||
smem->histogram); // each block histogram is summed across all CS blocks
|
||||
find_threshold<TopK>(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<SmemType, kStreamingStagesCS4, kHistBits, true>(
|
||||
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<TopK>(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<uint32_t>(TopK));
|
||||
hist4096::tie_handle_large<TopK>(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 <uint32_t TopK, uint32_t CS>
|
||||
__device__ void cooperative_topk_body(CooperativeTopKParams<TopK> 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<int32_t>(TopK)) {
|
||||
if (rank == 0) {
|
||||
for (uint32_t i = tx; i < TopK; i += hist4096::kBlockSize) {
|
||||
out[i] = (i < static_cast<uint32_t>(sl)) ? static_cast<int32_t>(i) : -1;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Short-Medium path: histogram_4096_topk on rank 0 only - all data fits in RF
|
||||
if (sl <= static_cast<int32_t>(hist4096::kHist4096MaxLen)) {
|
||||
if (rank == 0) {
|
||||
extern __shared__ uint8_t sr[];
|
||||
hist4096::histogram_4096_topk<TopK, 12>(
|
||||
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<kFusedStages>;
|
||||
|
||||
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<FusedSmem*>(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<TopK, CS, FusedSmem, true>(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<Smem4*>(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<TopK, CS, Smem4, false>(in, out, sl, hp, row_tie_ws);
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t TopK>
|
||||
__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
|
||||
__cluster_dims__(1, 4, 1)
|
||||
cooperative_topk_cs4(CooperativeTopKParams<TopK> params) {
|
||||
cooperative_topk_body<TopK, 4>(params);
|
||||
}
|
||||
|
||||
template <uint32_t TopK>
|
||||
__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
|
||||
__cluster_dims__(1, 8, 1)
|
||||
cooperative_topk_cs8(CooperativeTopKParams<TopK> params) {
|
||||
cooperative_topk_body<TopK, 8>(params);
|
||||
}
|
||||
|
||||
template <uint32_t TopK>
|
||||
__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
|
||||
__cluster_dims__(1, 16, 1)
|
||||
cooperative_topk_cs16(CooperativeTopKParams<TopK> params) {
|
||||
cooperative_topk_body<TopK, 16>(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<kFusedStagesCS8>) + sizeof(int32_t) * 2048 + 128;
|
||||
|
||||
} // namespace cooperative
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
#endif // COOPERATIVE_TOPK_CUH_
|
||||
@@ -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,
|
||||
|
||||
@@ -11,6 +11,8 @@
|
||||
#include <cub/cub.cuh>
|
||||
#include <cstdint>
|
||||
|
||||
#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 <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
|
||||
template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
|
||||
bool UsePredicatedShortLoads = false>
|
||||
__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<MAX_K, 12, 8>(score, dst, length,
|
||||
_smem_reg);
|
||||
} else {
|
||||
hist4096::histogram_4096_topk<MAX_K, 12, 8>(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<DType>(max_len);
|
||||
|
||||
#define DISPATCH_VEC_SIZE(VS) \
|
||||
if (vec_size == VS) { \
|
||||
auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>; \
|
||||
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<DType, IdType, VS, MAX_K, (VS != MAX_VEC)>; \
|
||||
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 <typename DType, typename IdType, uint32_t MAX_K = 2048>
|
||||
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<DType, IdType, MAX_K>(
|
||||
input, output_indices, lengths, num_rows, top_k_val, max_len, stream);
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
#endif // PERSISTENT_TOPK_CUH_
|
||||
|
||||
@@ -0,0 +1,554 @@
|
||||
/*
|
||||
* Shared 4096-bin single-CTA TopK helpers.
|
||||
*/
|
||||
|
||||
#ifndef TOPK_HISTOGRAM_4096_CUH_
|
||||
#define TOPK_HISTOGRAM_4096_CUH_
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
|
||||
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 <uint32_t kBits>
|
||||
__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<uint16_t>(~bits)
|
||||
: static_cast<uint16_t>(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 <uint32_t TopK>
|
||||
__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<TS*>(_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 <uint32_t TopK>
|
||||
__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<TS*>(_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 <uint32_t TopK, uint32_t HIST_BITS>
|
||||
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 <uint32_t TopK, uint32_t HIST_BITS,
|
||||
uint32_t VECS_PER_THREAD = kHist4096VecsPerThread,
|
||||
bool UsePredicatedLoads = false>
|
||||
__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<TopK, HIST_BITS>;
|
||||
auto* smem = static_cast<Smem*>(_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<uint4*>(
|
||||
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<uintptr_t>(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<const float4*>(scores + base);
|
||||
} else {
|
||||
load_float4_predicated(scores + base, static_cast<int>(base),
|
||||
static_cast<int>(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<const float4*>(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<const float*>(&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<HIST_BITS>(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<const float*>(&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<HIST_BITS>(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<uint32_t>(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<uint32_t>(
|
||||
__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<TopK>(smem->tie_buffer, num_ties, num_above, output, smem);
|
||||
} else {
|
||||
tie_handle_large<TopK>(smem->tie_buffer, num_ties, num_above, output,
|
||||
smem);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t TopK, uint32_t HIST_BITS,
|
||||
uint32_t VECS_PER_THREAD = kHist4096VecsPerThread>
|
||||
__device__ __noinline__ void histogram_4096_topk_predicated(
|
||||
const float* __restrict__ scores, int32_t* __restrict__ output,
|
||||
uint32_t length, void* _smem) {
|
||||
histogram_4096_topk<TopK, HIST_BITS, VECS_PER_THREAD, true>(scores, output,
|
||||
length, _smem);
|
||||
}
|
||||
|
||||
} // namespace topk_histogram_4096
|
||||
} // namespace vllm
|
||||
|
||||
#endif // TOPK_HISTOGRAM_4096_CUH_
|
||||
@@ -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",
|
||||
|
||||
@@ -14,6 +14,70 @@ TOP_K_VALUES = [2048, 3000]
|
||||
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}"
|
||||
)
|
||||
|
||||
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user