minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048) (#26788)

Browse Source 
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
2026-07-05 20:23:39 -07:00
committed by GitHub
co-authored by Claude
parent c9ceab34cf
commit c016c6f355
11 changed files with 1483 additions and 1344 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh
+290 -325
View File
@@ -1,193 +1,293 @@
/**
* \file topk_v2.cuh
* \brief TopK kernel for DeepSeek v4.
* Adapted from
* 1:
* https://github.com/vllm-project/vllm/blob/a8c6ee9b787d273916206a29b77feebadb80c368/csrc/persistent_topk.cuh
* 2:
* https://github.com/flashinfer-ai/flashinfer/blob/c2b4db2b1a84448d802f0e6ac445243312bd6a4c/include/flashinfer/topk.cuh
* DarkSharpness never took a detailed look at these 2 implementation, but his claude code did.
* So we add credit to the reference implementations.
*/
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/topk/cluster.cuh>
#include <sgl_kernel/deepseek_v4/topk/register.cuh>
#include <sgl_kernel/deepseek_v4/topk/streaming.cuh>
#include <sgl_kernel/deepseek_v4/topk_impl.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/object.h>
#include <cfloat>
#include <bit>
#include <cstdint>
#include <iterator>
namespace {
#ifndef SGL_TOPK
#define SGL_TOPK 512
#endif
namespace impl = device::topk;
using impl::TopKProblem;
inline constexpr uint32_t K = SGL_TOPK;
using Register2 = impl::TopKRegister<2>; // <= 8192, register-resident, 1 read
using Register4 = impl::TopKRegister<4>; // <= 16384, register-resident, 1 read
using Streaming = impl::TopKStreaming;
using Cluster = impl::TopKCluster<8>;
template <auto* f, size_t kMaxDynamicSMEM>
void setup_kernel_smem_once(host::DebugInfo where = {}) {
[[maybe_unused]]
static const auto result = [] {
const auto fptr = std::bit_cast<const void*>(f);
return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
}();
host::RuntimeDeviceCheck(result, where);
}
constexpr uint32_t kBlockSize = impl::TopKConfig::kBlockSize;
constexpr uint32_t kOccupancy = impl::TopKConfig::kOccupancy;
constexpr uint32_t kMaxTopK = impl::TopKConfig::kMaxTopK;
constexpr uint32_t kClusterSize = Cluster::kClusterSize;
constexpr uint32_t kReg2MaxSeqLen = Register2::kMaxSeqLen; // 8192
constexpr uint32_t kReg4MaxSeqLen = Register4::kMaxSeqLen; // 16384
namespace impl = device::top512;
using Large = impl::ClusterTopK<K>;
using Medium = impl::StreamingTopK<K>;
using Small = impl::RegisterTopK<K>;
#define TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, kOccupancy)
#define CLUSTER_TOPK_KERNEL TOPK_KERNEL __cluster_dims__(1, kClusterSize, 1)
using Metadata = Large::Metadata;
constexpr uint32_t kBlockSize = impl::kBlockSize;
constexpr uint32_t kNumClusters = 15; // based on hardware limits
constexpr uint32_t kClusterSize = Large::kClusterSize;
constexpr uint32_t kMax2PassLength = Small::kMax2PassLength;
constexpr uint32_t kMaxSupportedLength = Large::kMaxLength;
constexpr uint32_t kClusterFloor = 65536;
constexpr uint32_t kClusterMaxBatch = 512;
constexpr uint32_t kNumPersistentClusters = 15 * kOccupancy;
/// Common metadata lives at metadata[0] (first row of the [batch_size+1, 4] tensor).
/// Per-item metadata starts at metadata[1..batch_size]. The plan kernel writes both.
struct alignas(16) GlobalMetadata {
uint32_t cluster_threshold; // decided per-batch in plan kernel
uint32_t num_cluster_items; // N = number of items routed to the cluster path
uint32_t reserved[2];
/// Metadata tensor rows (each 8 B / 2 int32). Row 0 is the global plan result;
/// rows 1..N are the (batch_id, seq_len) of items routed to the cluster pool.
struct alignas(8) GlobalMetadata {
uint32_t cluster_threshold;
uint32_t num_cluster_items; // N = number of items routed to the cluster pool
};
static_assert(sizeof(GlobalMetadata) == sizeof(Metadata), "layout: row 0 must occupy one Metadata-sized slot");
struct alignas(8) PlanItem {
uint32_t batch_id;
uint32_t seq_len;
};
static_assert(sizeof(GlobalMetadata) == 2 * sizeof(int32_t) && sizeof(PlanItem) == sizeof(GlobalMetadata));
// optimize occupancy for prefill
#define SMALL_TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, 2)
// cluster at y dim
#define LARGE_CLUSTER __cluster_dims__(1, kClusterSize, 1)
// stage-1 is persistent cluster, and shared memory usage is huge (can not 2)
#define LARGE_TOPK_STAGE_1 __global__ __launch_bounds__(kBlockSize, 1) LARGE_CLUSTER
// stage-2 is non-persistent non-cluster, with less shared memory and higher occupancy
#define LARGE_TOPK_STAGE_2 __global__ __launch_bounds__(kBlockSize, 2)
// fused into 1 stage when batch-size <= kNumPersistentClusters
#define FUSED_COMBINE_KERNEL __global__ __launch_bounds__(kBlockSize, 1) LARGE_CLUSTER
// plan runs once as a single block before the combine kernels
#define PLAN_KERNEL __global__ __launch_bounds__(kBlockSize, 1)
struct TopKParams {
const uint32_t* __restrict__ seq_lens;
struct TopKLaunchParams {
const float* __restrict__ scores;
const int32_t* __restrict__ seq_lens;
const int32_t* __restrict__ page_table;
int32_t* __restrict__ page_indices;
int32_t* __restrict__ raw_indices; // optional raw (pre-transform) indices output; nullptr if unused
const PlanItem* __restrict__ metadata; // [0]=GlobalMetadata, [1+i]=PlanItem
int64_t score_stride;
int64_t page_table_stride;
uint8_t* __restrict__ workspace; // [batch, kWorkspaceBytes] -- internally allocated
/// Pointer to the full metadata tensor: metadata[0] is GlobalMetadata, metadata[1..]
/// are per-item entries (at most kNumClusters * rounds of them).
const Metadata* __restrict__ metadata = nullptr;
int64_t workspace_stride; // bytes per batch
uint32_t batch_size;
uint32_t topk;
uint32_t page_bits;
uint32_t cluster_floor; // seq_len > this routes to the cluster path (batch-aware, host-set)
SGL_DEVICE const float* get_scores(const uint32_t batch_id) const {
return scores + batch_id * score_stride;
SGL_DEVICE const GlobalMetadata& global() const {
return *reinterpret_cast<const GlobalMetadata*>(metadata);
}
SGL_DEVICE impl::TransformParams get_transform(const uint32_t batch_id, int32_t* indices) const {
return {
SGL_DEVICE uint32_t cluster_threshold() const {
return global().cluster_threshold;
}
SGL_DEVICE const PlanItem& item(uint32_t i) const {
return metadata[1 + i];
}
SGL_DEVICE int32_t* get_output_ptr(uint32_t batch_id) const {
return page_indices + batch_id * static_cast<int64_t>(topk);
}
SGL_DEVICE TopKProblem problem(uint32_t batch_id, uint32_t seq_len) const {
const auto k = static_cast<int64_t>(topk);
return TopKProblem{
.in = scores + batch_id * score_stride,
.out = page_indices + batch_id * k,
.raw_out = raw_indices != nullptr ? raw_indices + batch_id * k : nullptr,
.page_table = page_table + batch_id * page_table_stride,
.indices_in = indices,
.indices_out = page_indices + batch_id * K,
.topk = topk,
.seq_len = seq_len,
.page_bits = page_bits,
};
}
SGL_DEVICE const GlobalMetadata& get_global_metadata() const {
return *reinterpret_cast<const GlobalMetadata*>(metadata);
}
SGL_DEVICE const Metadata& get_item_metadata(uint32_t work_id) const {
return metadata[1 + work_id]; // +1 to skip the GlobalMetadata row
SGL_DEVICE TopKProblem problem(uint32_t batch_id) const {
return this->problem(batch_id, static_cast<uint32_t>(seq_lens[batch_id]));
}
};
SGL_DEVICE uint2 partition_work(uint32_t length, uint32_t rank) {
constexpr uint32_t kTMAAlign = 4;
const auto total_units = (length + kTMAAlign - 1) / kTMAAlign;
const auto base = total_units / kClusterSize;
const auto extra = total_units % kClusterSize;
const auto local_units = base + (rank < extra ? 1u : 0u);
const auto offset_units = rank * base + min(rank, extra);
const auto offset = offset_units * kTMAAlign;
const auto finish = min(offset + local_units * kTMAAlign, length);
return {offset, finish - offset};
/**
* \brief Persistent cluster kernel for the long items. It will handle long inputs.
* The short items are handled by the separate topk_kernel.
*/
template <bool kPDL>
CLUSTER_TOPK_KERNEL void topk_persistent_cluster_kernel(const __grid_constant__ TopKLaunchParams params) {
device::enable_smem_spilling();
__shared__ impl::MaxSmem<Cluster::Smem> smem;
const uint32_t num_cluster_items = params.global().num_cluster_items;
device::PDLWaitPrimary<kPDL>();
device::PDLTriggerSecondary<kPDL>();
#pragma unroll 1
for (uint32_t w = blockIdx.x; w < num_cluster_items; w += kNumPersistentClusters) {
const auto it = params.item(w);
const auto problem = params.problem(it.batch_id, it.seq_len);
Cluster::forward<false>(problem, &smem);
__syncthreads();
}
}
/// Persistent scheduler. A single block:
/// 1. Decides a cluster_threshold from the real seq_lens distribution (or
/// uses the caller-supplied `static_cluster_threshold` when non-zero).
/// 2. Writes that threshold + N into metadata[0] (the GlobalMetadata row).
/// 3. Compacts items with seq_len > threshold into metadata[1..N+1), laid out
/// to match the persistent consumer's round-robin stride (kNumClusters).
/// Entries for clusters that get no work are zero-filled.
PLAN_KERNEL void topk_plan(
template <typename F>
SGL_DEVICE void for_each_item(uint32_t topk, const F& f) {
constexpr uint32_t kNumElems = kMaxTopK / kBlockSize;
#pragma unroll
for (uint32_t i = 0; i < kNumElems; ++i) {
if (const auto tx = i * kBlockSize + threadIdx.x; tx < topk) {
__builtin_assume(tx < kMaxTopK);
f(tx, i);
}
}
}
template <bool kPDL>
SGL_DEVICE void trivial_transform(const TopKProblem& problem) {
device::PDLWaitPrimary<kPDL>();
device::PDLTriggerSecondary<kPDL>();
for_each_item(problem.topk, [&](uint32_t tx, uint32_t) {
problem.transform_output(tx, tx < problem.seq_len ? static_cast<int32_t>(tx) : -1);
});
}
SGL_DEVICE void problem_transform(TopKProblem& problem, int32_t* output_ptr) {
static_assert(kMaxTopK % kBlockSize == 0);
constexpr uint32_t kNumElems = kMaxTopK / kBlockSize;
int32_t source_index[kNumElems];
for_each_item(problem.topk, [&](uint32_t tx, uint32_t i) { source_index[i] = problem.out[tx]; });
problem.out = output_ptr;
for_each_item(problem.topk, [&](uint32_t tx, uint32_t i) { problem.transform_output(tx, source_index[i]); });
}
/**
* \brief Main kernel for the short items and epilogue of long items.
* \tparam kPDL whether to use PDL to synchronize with the cluster kernel (if any)
* \tparam kLevel:
* - Level 0: max_seq_len <= 8192 -> trivial + register<2>
* - Level 1: max_seq_len <= 16384 -> trivial + register<4>
* - Level 2: max_seq_len <= cluster_floor -> trivial + register<4> + streaming
* - Level 3: max_seq_len > cluster_floor -> + epilogue process of cluster path
*/
template <bool kPDL, int kLevel>
TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams params) {
device::enable_smem_spilling();
auto problem = params.problem(blockIdx.x);
constexpr uint32_t kU32Max = std::numeric_limits<uint32_t>::max();
__shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
__shared__ int32_t topk_indices[kMaxTopK];
problem.out = topk_indices;
constexpr bool kHandleCluster = (kLevel == 3);
// non-trivial path: dispatch based on level and seq_len
const auto cluster_threshold = kHandleCluster ? params.cluster_threshold() : kU32Max;
if constexpr (kLevel == 0) {
__builtin_assume(problem.seq_len <= kReg2MaxSeqLen);
Register2::forward<kPDL>(problem, &smem);
} else if constexpr (kLevel == 1) {
__builtin_assume(problem.seq_len <= kReg4MaxSeqLen);
Register4::forward<kPDL>(problem, &smem); // max_seq_len <= 16384 guarantees seq <= 16384
} else {
static_assert(kLevel == 2 || kLevel == 3, "we only support level = 0,1,2,3 now");
// if using cluster, we can delay the PDL wait
constexpr bool kPDLEarly = kPDL && !kHandleCluster;
constexpr bool kPDLFinal = kPDL && kHandleCluster;
if (problem.seq_len <= kReg4MaxSeqLen) {
Register4::forward<kPDLEarly>(problem, &smem);
} else if (problem.seq_len <= cluster_threshold) {
Streaming::forward<kPDLEarly>(problem, &smem);
} else { // cluster path do nothing here
problem.out = params.get_output_ptr(blockIdx.x);
}
device::PDLWaitPrimary<kPDLFinal>();
}
// page-table transform pass (gathers kept out of the hot scatter loop),
// then trigger the dependent kernel only after the full output is written.
device::PDLTriggerSecondary<kPDL>();
__syncthreads();
problem_transform(problem, params.get_output_ptr(blockIdx.x));
}
template <bool kPDL>
CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLaunchParams params) {
device::enable_smem_spilling();
auto problem = params.problem(blockIdx.x);
__shared__ impl::MaxSmem<Streaming::Smem, Cluster::Smem> smem;
if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
__shared__ int32_t topk_indices[kMaxTopK];
problem.out = topk_indices;
// randomly elect one worker rank to avoid workload imbalance
const auto worker_rank = blockIdx.x % kClusterSize;
// for small batch, we will fuse in the cluster case
if (problem.seq_len <= kReg4MaxSeqLen) {
if (blockIdx.y == worker_rank) Register4::forward<kPDL>(problem, &smem);
} else if (problem.seq_len <= params.cluster_floor) {
if (blockIdx.y == worker_rank) Streaming::forward<kPDL>(problem, &smem);
} else {
auto cluster = cooperative_groups::this_cluster();
problem.out = cluster.map_shared_rank(topk_indices, worker_rank);
Cluster::forward<kPDL>(problem, &smem); // write to peer's output shared memory
cluster.sync();
}
device::PDLWaitPrimary<kPDL>();
__syncthreads();
if (blockIdx.y == worker_rank) problem_transform(problem, params.get_output_ptr(blockIdx.x));
}
// --- Plan: choose cluster_threshold from the seq_len distribution -----------
__global__ __launch_bounds__(kBlockSize, 1) void topk_plan(
const uint32_t* __restrict__ seq_lens,
Metadata* __restrict__ metadata,
PlanItem* __restrict__ metadata, // [0]=GlobalMetadata, [1+i]=PlanItem
const uint32_t batch_size,
const uint32_t static_cluster_threshold) {
// Candidate thresholds, strictly increasing. Picked to give the auto-heuristic
// reasonable granularity without needing a full sort. Must all be >= kMax2PassLength.
// Candidate (threshold T_j, cap_j) pairs, T strictly increasing. The plan lowers
// cluster_threshold to T_j while #(items with seq_len > T_j) <= cap_j, so cap_j
// bounds how many long items go to the persistent pool. The pool runs N items in
// ceil(N / kNumPersistentClusters) waves; the longer the seq the more waves pay
// off (streaming a single block over a long item is very slow), so cap_j is the
// measured cluster-vs-streaming crossover (B200, occ2) and GROWS with T -- a flat
// cap = pool size only fits the shortest (~98K, one-wave) bucket. (Plan is tunable.)
struct Pair {
uint32_t threshold;
uint32_t max_batch_size;
};
/// NOTE: only tuned on B200
constexpr Pair kCandidates[] = {
{32768, 30},
{40960, 45},
{49152, 45},
{65536, 60},
{98304, 60},
{131072, 75},
{196608, 90},
{262144, 105},
{65536, 30}, // (65536,98304]: ~1 pool wave, streams beyond 30
{98304, 48}, // (98304,131072]
{131072, 60}, // (131072,196608]
{196608, 80}, // (196608,262144]
{262144, 112}, // (262144,393216]
{393216, 128}, // (393216,inf): longest -- worth many pool waves; a top
// threshold here lets overloaded ~280-393K batches still stream
};
constexpr uint32_t kNumCandidates = std::size(kCandidates);
constexpr uint32_t kMinBatchSize = kCandidates[0].max_batch_size;
static_assert(kCandidates[0].threshold == kMax2PassLength);
static_assert(kCandidates[kNumCandidates - 1].threshold == kMaxSupportedLength);
static_assert(kCandidates[0].threshold == kClusterFloor);
__shared__ uint32_t s_count; // final N after compaction
__shared__ uint32_t s_counts[kNumCandidates];
__shared__ uint32_t s_threshold;
__shared__ uint32_t s_count;
const auto tx = threadIdx.x;
if (tx == 0) s_count = 0;
if (tx < kNumCandidates) s_counts[tx] = 0;
if (tx == 0) s_count = 0;
__syncthreads();
// --- Phase 1: decide threshold ------------------------------------------
if (static_cluster_threshold > 0) {
if (tx == 0) s_threshold = static_cluster_threshold;
} else if (batch_size <= kMinBatchSize) {
if (tx == 0) s_threshold = kMax2PassLength; // always prefer cluster
} else {
// Count items above each candidate threshold. Monotonically non-increasing in T.
for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
const uint32_t sl = seq_lens[i];
assert(sl <= kMaxSupportedLength);
uint32_t count = 0;
#pragma unroll
for (uint32_t j = 0; j < kNumCandidates; ++j) {
count += (sl > kCandidates[j].threshold ? 1 : 0);
}
if (count > 0) {
atomicAdd(&s_counts[count - 1], 1);
}
if (count > 0) atomicAdd(&s_counts[count - 1], 1);
}
__syncthreads();
if (tx == 0) {
uint32_t accum = 0;
uint32_t chosen = kMaxSupportedLength;
uint32_t chosen = kCandidates[kNumCandidates - 1].threshold;
#pragma unroll
for (uint32_t i = 0; i < kNumCandidates; ++i) {
const auto j = kNumCandidates - 1 - i;
accum += s_counts[j];
/// NOTE: `accum` increasing, while `max_batch_size` decreasing
accum += s_counts[j]; // # items with seq_len > kCandidates[j].threshold
if (accum > kCandidates[j].max_batch_size) break;
chosen = kCandidates[j].threshold;
}
@@ -195,163 +295,25 @@ PLAN_KERNEL void topk_plan(
}
}
__syncthreads();
// sanity check: below 2 pass threshold, must fits in small path
const auto cluster_threshold = max(s_threshold, kMax2PassLength);
const auto cluster_threshold = max(s_threshold, kClusterFloor);
// --- Phase 2: compact items with seq_len > threshold into metadata[1..] -
// Per-item rows live at metadata[1 + pos]; metadata[0] is the GlobalMetadata row.
// Compact items with seq_len > threshold into metadata[1..N]: their batch ids
// are the work list the persistent cluster pool fetches.
for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
const uint32_t sl = seq_lens[i];
if (sl > cluster_threshold) {
const auto pos = atomicAdd(&s_count, 1);
metadata[1 + pos] = {i, sl, false};
metadata[1 + pos] = {i, sl};
}
}
__syncthreads();
const auto N = s_count;
// --- Phase 3: has_next + sentinels + GlobalMetadata ---------------------
for (uint32_t i = tx; i < N; i += kBlockSize) {
if (i + kNumClusters < N) metadata[1 + i].has_next = true;
}
// Zero-fill the first kNumClusters sentinel slots that got no valid entry.
if (tx < kNumClusters && tx >= N) metadata[1 + tx] = {0, 0, false};
// Write global metadata (row 0).
if (tx == 0) {
auto* g = reinterpret_cast<GlobalMetadata*>(metadata);
*g = {
.cluster_threshold = cluster_threshold,
.num_cluster_items = N,
.reserved = {0, 0},
};
*g = {.cluster_threshold = cluster_threshold, .num_cluster_items = s_count};
}
}
SMALL_TOPK_KERNEL void // short context
topk_short_transform(const __grid_constant__ TopKParams params) {
alignas(128) extern __shared__ uint8_t smem[];
__shared__ int32_t s_topk_indices[K];
const auto batch_id = blockIdx.x;
const auto seq_len = params.seq_lens[batch_id];
const auto transform = params.get_transform(batch_id, s_topk_indices);
// trivial case
if (seq_len <= K) {
impl::trivial_transform(transform, seq_len, K);
} else {
Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem, /*use_pdl=*/true);
device::PDLTriggerSecondary<true>();
Small::transform(transform);
}
}
LARGE_TOPK_STAGE_1 void // long context, middle to large batch size
topk_combine_preprocess(const __grid_constant__ TopKParams params) {
alignas(128) extern __shared__ uint8_t smem[];
__shared__ int32_t s_topk_indices[K];
uint32_t work_id = blockIdx.x;
uint32_t batch_id;
uint32_t seq_len;
bool has_next;
uint32_t length;
uint32_t offset;
const auto cluster_rank = blockIdx.y;
const auto prefetch_metadata = [&] {
const auto metadata = params.get_item_metadata(work_id);
batch_id = metadata.batch_id;
seq_len = metadata.seq_len;
has_next = metadata.has_next;
work_id += kNumClusters; // advance to the next item for this cluster
};
const auto launch_prologue = [&] {
const auto partition = partition_work(seq_len, cluster_rank);
offset = partition.x;
length = partition.y;
Large::stage1_prologue(params.get_scores(batch_id) + offset, length, smem);
};
device::PDLWaitPrimary<true>();
device::PDLTriggerSecondary<true>();
prefetch_metadata();
if (seq_len == 0) return;
Large::stage1_init(smem);
launch_prologue();
while (true) {
const auto this_length = length;
const auto this_offset = offset;
const auto need_prefetch = has_next;
const auto transform = params.get_transform(batch_id, s_topk_indices);
const auto ws = params.workspace + batch_id * params.workspace_stride;
if (need_prefetch) prefetch_metadata();
Large::stage1(s_topk_indices, this_length, smem, /*reuse=*/true);
if (need_prefetch) launch_prologue();
Large::stage1_epilogue(transform, this_offset, ws, smem);
if (!need_prefetch) break;
}
}
LARGE_TOPK_STAGE_2 void // long context, middle to large batch size
topk_combine_transform(const __grid_constant__ TopKParams params) {
alignas(128) extern __shared__ uint8_t smem[];
__shared__ int32_t s_topk_indices[K];
const auto batch_id = blockIdx.x;
const auto seq_len = params.seq_lens[batch_id];
const auto cluster_threshold = params.get_global_metadata().cluster_threshold;
const auto transform = params.get_transform(batch_id, s_topk_indices);
if (seq_len <= K) {
impl::trivial_transform(transform, seq_len, K);
} else if (seq_len <= kMax2PassLength) {
if (seq_len <= Small::kMax1PassLength) {
Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem);
} else {
__syncwarp();
Small::run<true>(params.get_scores(batch_id), s_topk_indices, seq_len, smem);
}
Small::transform(transform);
} else if (seq_len <= cluster_threshold) {
Medium::run(params.get_scores(batch_id), seq_len, s_topk_indices, smem);
Medium::transform(transform, smem);
} else {
const auto ws = params.workspace + batch_id * params.workspace_stride;
device::PDLWaitPrimary<true>();
Large::transform(transform, ws, smem);
}
}
FUSED_COMBINE_KERNEL void // long context, small batch size
topk_fused_transform(const __grid_constant__ TopKParams params) {
alignas(128) extern __shared__ uint8_t smem[];
__shared__ int32_t s_topk_indices[K];
const auto batch_id = blockIdx.x;
const auto cluster_rank = blockIdx.y;
const auto seq_len = params.seq_lens[batch_id];
const auto transform = params.get_transform(batch_id, s_topk_indices);
if (seq_len <= K) {
if (cluster_rank != 0) return; // only first rank work
impl::trivial_transform(transform, seq_len, K);
} else if (seq_len <= Small::kMax1PassLength) {
if (cluster_rank != 0) return; // only first rank work
Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem, /*use_pdl=*/true);
Small::transform(transform);
} else {
const auto [offset, length] = partition_work(seq_len, cluster_rank);
const auto ws = params.workspace + batch_id * params.workspace_stride;
Large::stage1_init(smem);
device::PDLWaitPrimary<true>();
Large::stage1_prologue(params.get_scores(batch_id) + offset, length, smem);
Large::stage1(s_topk_indices, length, smem);
Large::stage1_epilogue(transform, offset, ws, smem);
cooperative_groups::this_cluster().sync();
if (cluster_rank != 0) return; // only first rank do the stage-2
Large::transform(transform, ws, smem);
}
}
struct CombinedTopKKernel {
static constexpr auto kStage1SMEM = sizeof(Large::Smem) + 128;
static constexpr auto kStage2SMEM = std::max(sizeof(Small::Smem), sizeof(Medium::Smem)) + 128;
struct TopKKernel {
static void plan( //
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::TensorView metadata,
@@ -362,25 +324,22 @@ struct CombinedTopKKernel {
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B}) //
TensorMatcher({B}) // seq_lens
.with_dtype<int32_t>()
.with_device(device_)
.verify(seq_lens);
TensorMatcher({Bp1, 4}) //
TensorMatcher({Bp1, 2}) // metadata: [0]=GlobalMetadata, [1..N]=PlanItem(batch_id, seq_len)
.with_dtype<int32_t>()
.with_device(device_)
.verify(metadata);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1);
if (batch_size <= kNumClusters) return; // metadata unused in fused path
RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
const auto device = device_.unwrap();
constexpr auto kernel = topk_plan;
LaunchKernel(1, kBlockSize, device)( //
kernel,
static_cast<uint32_t*>(seq_lens.data_ptr()),
static_cast<Metadata*>(metadata.data_ptr()),
topk_plan,
static_cast<const uint32_t*>(seq_lens.data_ptr()),
static_cast<PlanItem*>(metadata.data_ptr()),
batch_size,
static_cluster_threshold);
}
@@ -391,101 +350,107 @@ struct CombinedTopKKernel {
const tvm::ffi::TensorView page_table,
const tvm::ffi::TensorView page_indices,
const uint32_t page_size,
const tvm::ffi::TensorView workspace,
const tvm::ffi::TensorView metadata) {
const tvm::ffi::TensorView metadata,
const tvm::ffi::Optional<tvm::ffi::TensorView> raw_indices) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto Bp1 = SymbolicSize{"batch_size_plus_1"};
auto L = SymbolicSize{"max_seq_len"};
auto S = SymbolicSize{"score_stride"};
auto P = SymbolicSize{"page_table_stride"};
auto W = SymbolicSize{"workspace_stride"};
constexpr auto D = Large::kWorkspaceInts;
auto K = SymbolicSize{"topk"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, L}) //
TensorMatcher({B, L}) // score
.with_strides({S, 1})
.with_dtype<float>()
.with_device(device_)
.verify(scores);
TensorMatcher({B}) //
TensorMatcher({B}) // seq_lens
.with_dtype<int32_t>()
.with_device(device_)
.verify(seq_lens);
TensorMatcher({B, -1}) //
TensorMatcher({B, -1}) // page_table
.with_strides({P, 1})
.with_dtype<int32_t>()
.with_device(device_)
.verify(page_table);
TensorMatcher({B, K}) //
TensorMatcher({B, K}) // page_indices
.with_dtype<int32_t>()
.with_device(device_)
.verify(page_indices);
TensorMatcher({B, D}) //
.with_strides({W, 1})
.with_dtype<int32_t>()
.with_device(device_)
.verify(workspace);
TensorMatcher({Bp1, 4}) //
TensorMatcher({Bp1, 2}) // metadata: [0]=GlobalMetadata, [1..N]=PlanItem(batch_id, seq_len)
.with_dtype<int32_t>()
.with_device(device_)
.verify(metadata);
int32_t* raw_indices_ptr = nullptr;
if (raw_indices.has_value()) {
TensorMatcher({B, K}).with_dtype<int32_t>().with_device(device_).verify(raw_indices.value());
raw_indices_ptr = static_cast<int32_t*>(raw_indices.value().data_ptr());
}
RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2");
RuntimeCheck(S.unwrap() % 4 == 0, "score_stride must be a multiple of 4 (16-byte vectorized load)");
RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
const auto topk = static_cast<uint32_t>(K.unwrap());
RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 2048]");
const auto page_bits = static_cast<uint32_t>(std::countr_zero(page_size));
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto max_seq_len = static_cast<uint32_t>(L.unwrap());
const auto device = device_.unwrap();
RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2");
RuntimeCheck(S.unwrap() % 4 == 0, "score_stride must be a multiple of 4 (TMA 16-byte alignment)");
RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
// NOTE: this should be fixed later
// RuntimeCheck(max_seq_len <= kMaxSupportedLength, max_seq_len, " exceeds the maximum supported length");
const auto params = TopKParams{
.seq_lens = static_cast<uint32_t*>(seq_lens.data_ptr()),
.scores = static_cast<float*>(scores.data_ptr()),
.page_table = static_cast<int32_t*>(page_table.data_ptr()),
// The fused kernel runs one 8-block cluster per batch element, and B200 fits one
// wave of exactly 15 such clusters (occ2). For batch <= 15 it stays latency-bound,
// so the 8-way split beats streaming from a much lower seq (measured crossover
// ~36-40K); batch 16 spills into a 2nd wave (+25%) and keeps the 64K floor.
// The floor is chosen on the host per launch.
constexpr uint32_t kClusterFloorSmall = 32768;
constexpr uint32_t kSmallBatchLowFloor = 15;
const auto params = TopKLaunchParams{
.scores = static_cast<const float*>(scores.data_ptr()),
.seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
.page_table = static_cast<const int32_t*>(page_table.data_ptr()),
.page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
.raw_indices = raw_indices_ptr,
.metadata = static_cast<const PlanItem*>(metadata.data_ptr()),
.score_stride = S.unwrap(),
.page_table_stride = P.unwrap(),
.workspace = static_cast<uint8_t*>(workspace.data_ptr()),
.metadata = static_cast<const Metadata*>(metadata.data_ptr()),
.workspace_stride = W.unwrap() * static_cast<int64_t>(sizeof(int32_t)),
.batch_size = batch_size,
.topk = topk,
.page_bits = page_bits,
.cluster_floor = (batch_size <= kSmallBatchLowFloor) ? kClusterFloorSmall : kClusterFloor,
};
if (max_seq_len <= Small::kMax1PassLength) {
// All items fit in the short path -- no stage-1 needed
constexpr auto kernel = topk_short_transform;
setup_kernel_smem_once<kernel, kStage2SMEM>();
LaunchKernel(batch_size, kBlockSize, device, kStage2SMEM) //
.enable_pdl(true)(kernel, params);
} else {
// Some items may be large -- launch stage-1 + main
if (batch_size <= kNumClusters) {
// can fuse into 1 stage
constexpr auto kernel = topk_fused_transform;
constexpr auto kSMEM = std::max(kStage1SMEM, kStage2SMEM);
setup_kernel_smem_once<kernel, kSMEM>();
LaunchKernel({batch_size, kClusterSize}, kBlockSize, device, kSMEM)
.enable_cluster({1, kClusterSize})
.enable_pdl(true)(kernel, params);
const bool use_cluster = (max_seq_len > params.cluster_floor) && (batch_size <= kClusterMaxBatch);
constexpr bool kUsePDL = true;
if (use_cluster) {
if (batch_size <= kNumPersistentClusters) {
LaunchKernel({batch_size, kClusterSize}, kBlockSize, device)
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
.launch(topk_small_batch_kernel<kUsePDL>, params);
} else {
// stage 1 + stage 2
constexpr auto kernel_stage_1 = topk_combine_preprocess;
setup_kernel_smem_once<kernel_stage_1, kStage1SMEM>();
const auto num_clusters = std::min(batch_size, kNumClusters);
LaunchKernel({num_clusters, kClusterSize}, kBlockSize, device, kStage1SMEM)
.enable_cluster({1, kClusterSize})
.enable_pdl(true)(kernel_stage_1, params);
constexpr auto kernel_stage_2 = topk_combine_transform;
setup_kernel_smem_once<kernel_stage_2, kStage2SMEM>();
LaunchKernel(batch_size, kBlockSize, device, kStage2SMEM) //
.enable_pdl(true)(kernel_stage_2, params);
const uint32_t num_clusters = std::min(batch_size, kNumPersistentClusters);
LaunchKernel({num_clusters, kClusterSize}, kBlockSize, device)
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
.launch(topk_persistent_cluster_kernel<kUsePDL>, params);
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/3>, params);
}
} else if (max_seq_len <= kReg2MaxSeqLen) {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/0>, params);
} else if (max_seq_len <= kReg4MaxSeqLen) {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/1>, params);
} else {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/2>, params);
}
}
};
python/sglang/jit_kernel/dsv4/topk.py
+13 -12
View File
@@ -29,15 +29,16 @@ def _jit_topk_v1_module(topk: int):
@cache_once
def _jit_topk_v2_module(topk: int):
def _jit_topk_v2_module():
# v2 is universal: topk (<= 2048) is a runtime argument, not a compile-time
# constant, so a single module serves every k.
return load_jit(
make_name(f"topk_v2_{topk}"),
make_name("topk_v2"),
cuda_files=["deepseek_v4/topk_v2.cuh"],
cuda_wrappers=[
("topk_transform", "CombinedTopKKernel::transform"),
("topk_plan", "CombinedTopKKernel::plan"),
("topk_transform", "TopKKernel::transform"),
("topk_plan", "TopKKernel::plan"),
],
extra_cuda_cflags=[f"-DSGL_TOPK={topk}"],
)
@@ -60,12 +61,13 @@ def topk_transform_512(
)
_WORKSPACE_INTS_PER_BATCH = 2 + 1024 * 2
_PLAN_METADATA_INTS_PER_BATCH = 4
# metadata is (batch+1, 2) int32: row 0 = {cluster_threshold, num_cluster_items};
# rows 1..N = {batch_id, seq_len} of items routed to the persistent cluster pool.
_PLAN_METADATA_INTS_PER_BATCH = 2
def plan_topk_v2(seq_lens: torch.Tensor, static_threshold: int = 0) -> torch.Tensor:
module = _jit_topk_v2_module(512) # does not matter
module = _jit_topk_v2_module()
bs = seq_lens.shape[0]
metadata = seq_lens.new_empty(bs + 1, _PLAN_METADATA_INTS_PER_BATCH)
module.topk_plan(seq_lens, metadata, static_threshold)
@@ -79,16 +81,15 @@ def topk_transform_512_v2(
out_page_indices: torch.Tensor,
page_size: int,
metadata: torch.Tensor,
out_raw_indices: Optional[torch.Tensor] = None,
) -> None:
module = _jit_topk_v2_module(out_page_indices.shape[1])
bs = scores.shape[0]
workspace = seq_lens.new_empty(bs, _WORKSPACE_INTS_PER_BATCH)
module = _jit_topk_v2_module()
module.topk_transform(
scores,
seq_lens,
page_tables,
out_page_indices,
page_size,
workspace,
metadata,
out_raw_indices,
)
python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/cluster.cuh
-257
View File
@@ -1,257 +0,0 @@
#pragma once
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include "common.cuh"
#include "ptx.cuh"
#include <cooperative_groups.h>
#include <cstdint>
namespace device::top512 {
template <uint32_t K>
struct ClusterTopK {
static constexpr uint32_t kClusterSize = 8;
static constexpr uint32_t kHistBits = 10;
static constexpr uint32_t kHistBins = 1 << kHistBits;
static constexpr uint32_t kRadixBins = 256;
static constexpr uint32_t kElemPerStage = 8;
static constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize;
static constexpr uint32_t kNumStages = 4;
static constexpr uint32_t kMaxLength = kClusterSize * kNumStages * kSizePerStage;
static constexpr uint32_t kStoreLane = kBlockSize - 1;
static constexpr uint32_t kAboveBits = 11;
// ---------------------------------------------------------------------------
// Shared memory layouts
// ---------------------------------------------------------------------------
struct Smem {
uint64_t barrier[kNumStages];
uint32_t local_above_equal[kClusterSize];
uint32_t prefix_above_equal;
alignas(128) uint32_t counter_gt;
alignas(128) uint32_t counter_eq;
alignas(128) MatchBin match;
alignas(128) uint32_t warp_sum[kNumWarps];
uint32_t histogram[kHistBins];
alignas(128) float score_buffer[kNumStages][kSizePerStage];
Tie tie_buffer[kMaxTies];
};
struct alignas(16) Metadata {
uint32_t batch_id;
uint32_t seq_len;
bool has_next;
};
struct WorkSpace {
uint2 metadata; // {num_above, num_ties}
Tie ties[kMaxTies];
};
static constexpr uint32_t kWorkspaceInts = sizeof(WorkSpace) / sizeof(uint32_t);
// ---------------------------------------------------------------------------
// Stage 1: histogram + cluster reduce + find threshold + scatter
// ---------------------------------------------------------------------------
SGL_DEVICE static void stage1_init(void* _smem) {
const auto tx = threadIdx.x;
__builtin_assume(tx < kBlockSize);
const auto smem = static_cast<Smem*>(_smem);
if (tx < kHistBins) smem->histogram[tx] = 0;
if (tx < kNumStages) ptx::mbarrier_init(&smem->barrier[tx], 1);
__syncthreads();
}
SGL_DEVICE static void stage1_prologue(const float* scores, uint32_t length, void* _smem) {
if (threadIdx.x == 0) {
const auto smem = static_cast<Smem*>(_smem);
const auto num_stages = (length + kSizePerStage - 1) / kSizePerStage;
const auto length_aligned = (length + 3u) & ~3u; // align to 4 for TMA
#pragma unroll
for (uint32_t stage = 0; stage < kNumStages; stage++) {
if (stage >= num_stages) break;
const auto offset = stage * kSizePerStage;
const auto size = min(kSizePerStage, length_aligned - offset);
const auto size_bytes = size * sizeof(float);
const auto bar = &smem->barrier[stage];
ptx::tma_load(smem->score_buffer[stage], scores + offset, size_bytes, bar);
ptx::mbarrier_arrive_expect_tx(bar, size_bytes);
}
}
}
SGL_DEVICE static void stage1(int32_t* indices, uint32_t length, void* _smem, bool reuse = false) {
const auto smem = static_cast<Smem*>(_smem);
const auto tx = threadIdx.x;
__builtin_assume(tx < kBlockSize);
const auto lane_id = tx % kWarpThreads;
const auto warp_id = tx / kWarpThreads;
// Initialize shared memory histogram, counters, and barriers
#pragma unroll
for (uint32_t stage = 0; stage < kNumStages; stage++) {
const auto offset = stage * kSizePerStage;
if (offset >= length) break;
const auto size = min(kSizePerStage, length - offset);
if (lane_id == 0) ptx::mbarrier_wait(&smem->barrier[stage], 0);
__syncwarp();
#pragma unroll
for (uint32_t i = 0; i < kElemPerStage; ++i) {
const auto idx = tx + i * kBlockSize;
if (idx >= size) break;
const auto score = smem->score_buffer[stage][idx];
const auto bin = extract_coarse_bin<kHistBits>(score);
atomicAdd(&smem->histogram[bin], 1);
}
}
static_assert(kHistBins <= kBlockSize);
// 2-shot all-reduce
{
auto cluster = cooperative_groups::this_cluster();
cluster.sync();
const auto cluster_rank = blockIdx.y;
const auto kLocalSize = kHistBins / kClusterSize;
const auto offset = kLocalSize * cluster_rank;
const auto src_tx = tx / kClusterSize;
const auto src_rank = tx % kClusterSize;
if (tx < kHistBins) {
const auto addr = &smem->histogram[offset + src_tx];
const auto src_addr = cluster.map_shared_rank(addr, src_rank);
*src_addr = warp::reduce_sum<kClusterSize>(*src_addr);
}
cluster.sync();
}
// now each block holds the whole histogram, find the threshold bin
{
const auto value = tx < kHistBins ? smem->histogram[tx] : 0;
const auto warp_inc = warp_inclusive_sum(lane_id, value);
if (lane_id == kWarpThreads - 1) {
smem->warp_sum[warp_id] = warp_inc;
}
__syncthreads();
const auto tmp = smem->warp_sum[lane_id];
// total_length = sum of all bins in the globally-reduced histogram
// (problem.length is block-local; after cluster reduction we need the global total)
const auto total_length = warp::reduce_sum(tmp);
uint32_t prefix_sum = warp::reduce_sum(lane_id < warp_id ? tmp : 0);
prefix_sum += warp_inc;
const auto above = total_length - prefix_sum;
if (tx < kHistBins && above < K && above + value >= K) {
smem->counter_gt = smem->counter_eq = 0;
smem->match = {
.bin = tx,
.above_count = above,
.equal_count = value,
};
}
__syncthreads();
}
const auto [thr_bin, num_above, num_equal] = smem->match;
// write above and equal results to global memory
#pragma unroll
for (uint32_t stage = 0; stage < kNumStages; stage++) {
const auto offset = stage * kSizePerStage;
if (offset >= length) break;
#pragma unroll
for (uint32_t i = 0; i < kElemPerStage; ++i) {
const auto buf_idx = tx + i * kBlockSize;
const auto global_idx = offset + buf_idx;
if (global_idx >= length) break;
const auto score = smem->score_buffer[stage][buf_idx];
const auto bin = extract_coarse_bin<kHistBits>(score);
if (bin > thr_bin) {
indices[atomicAdd(&smem->counter_gt, 1)] = global_idx;
} else if (bin == thr_bin) {
const auto pos = atomicAdd(&smem->counter_eq, 1);
if (pos < kMaxTies) smem->tie_buffer[pos] = {global_idx, score};
}
}
}
if (reuse) {
const auto num_stages = (length + kSizePerStage - 1) / kSizePerStage;
if (tx < kHistBins) smem->histogram[tx] = 0;
if (tx < num_stages) ptx::mbarrier_arrive(&smem->barrier[tx]);
}
__syncthreads();
}
// ---------------------------------------------------------------------------
// Stage 1 epilogue: cross-block prefix sum + page translate + tie store
// ---------------------------------------------------------------------------
SGL_DEVICE static void stage1_epilogue(const TransformParams params, const uint32_t offset, void* _ws, void* _smem) {
auto cluster = cooperative_groups::this_cluster();
const auto smem = static_cast<Smem*>(_smem);
const auto tx = threadIdx.x;
const auto local_above = smem->counter_gt;
const auto local_equal = smem->counter_eq;
const auto cluster_rank = blockIdx.y;
constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1;
static_assert(kAboveMask >= K);
// Pack local counts -- NO alignment rounding (contiguous layout)
static_assert(kMaxTies <= kBlockSize);
const auto idx_above = tx < local_above ? params.indices_in[tx] : 0;
const auto tie_value = tx < local_equal ? smem->tie_buffer[tx] : Tie{0, 0.0f};
// push to remote shared memory, can reduce latency of reading remote
if (tx < kClusterSize) {
const auto value = (local_equal << kAboveBits) | local_above;
const auto dst_addr = cluster.map_shared_rank(smem->local_above_equal, tx);
dst_addr[cluster_rank] = value;
}
// after this last sync, only read local shared memory
// so that it is safe when peer rank has already exited the kernel
cluster.sync();
if (tx < kClusterSize) {
const auto value = tx < cluster_rank ? smem->local_above_equal[tx] : 0;
const auto kActiveMask = (1u << kClusterSize) - 1;
smem->prefix_above_equal = warp::reduce_sum<kClusterSize>(value, kActiveMask);
}
__syncthreads();
const auto prefix_packed = smem->prefix_above_equal;
const auto prefix_above = prefix_packed & kAboveMask;
const auto prefix_equal = prefix_packed >> kAboveBits;
// Page-translate above elements
if (tx < local_above) {
params.write(tx + prefix_above, idx_above + offset);
}
// Contiguous tie store via regular global writes (no TMA, no gaps)
const auto ws = static_cast<WorkSpace*>(_ws);
if (tx < local_equal && tx + prefix_equal < kMaxTies) {
ws->ties[tx + prefix_equal] = {tie_value.idx + offset, tie_value.score};
}
// Block 0 writes global metadata {num_above, num_ties}
if (cluster_rank == kClusterSize - 1 && tx == 0) {
const auto sum_above = prefix_above + local_above;
const auto sum_equal = prefix_equal + local_equal;
ws->metadata = make_uint2(sum_above, sum_equal);
}
}
SGL_DEVICE static void transform(const TransformParams params, const void* _ws, void* _smem) {
const auto ws = static_cast<const WorkSpace*>(_ws);
const auto meta = &ws->metadata;
const auto [num_above, num_equal] = *meta;
if (num_above >= K || num_equal == 0) return;
const auto clamped_ties = min(num_equal, kMaxTies);
tie_handle_transform(ws->ties, clamped_ties, num_above, K, params, _smem);
}
};
} // namespace device::top512
python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/common.cuh
-176
View File
@@ -1,176 +0,0 @@
#pragma once
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <cstdint>
namespace device::top512 {
inline constexpr uint32_t kMaxTopK = 1024;
inline constexpr uint32_t kBlockSize = 1024;
inline constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads;
inline constexpr uint32_t kMaxTies = 1024; // == kBlockSize: 1 element per thread in stage2
static constexpr uint32_t kRadixBins = 256;
static_assert(kMaxTopK <= kBlockSize && kMaxTies <= kBlockSize);
// always use float4 to load from global memory
using Vec4 = AlignedVector<float, 4>;
SGL_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, uint32_t i, uint32_t page_bits) {
const uint32_t mask = (1u << page_bits) - 1u;
return (page_table[i >> page_bits] << page_bits) | (i & mask);
}
struct TransformParams {
const int32_t* __restrict__ page_table;
const int32_t* __restrict__ indices_in;
int32_t* __restrict__ indices_out;
uint32_t page_bits;
SGL_DEVICE void transform(const uint32_t idx) const {
indices_out[idx] = page_to_indices(page_table, indices_in[idx], page_bits);
}
SGL_DEVICE void write(const uint32_t dst, const uint32_t src) const {
indices_out[dst] = page_to_indices(page_table, src, page_bits);
}
};
struct alignas(16) MatchBin {
uint32_t bin;
uint32_t above_count;
uint32_t equal_count;
};
struct alignas(8) Tie {
uint32_t idx;
float score;
};
struct TieHandleSmem {
alignas(128) uint32_t counter; // output position counter
alignas(128) MatchBin match;
uint32_t histogram[kRadixBins]; // 256-bin radix histogram
uint32_t warp_sum[kNumWarps]; // for 2-pass prefix sum
};
template <uint32_t kBits>
SGL_DEVICE uint32_t extract_coarse_bin(float x) {
static_assert(0 < kBits && kBits < 15);
const auto hx = cast<fp16_t>(x);
const uint16_t bits = *reinterpret_cast<const uint16_t*>(&hx);
const uint16_t key = (bits & 0x8000) ? ~bits : bits | 0x8000;
return key >> (16 - kBits);
}
SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
static_assert(kWarpThreads == 32);
#pragma unroll
for (uint32_t offset = 1; offset < 32; offset *= 2) {
uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
if (lane_id >= offset) val += n;
}
return val;
}
/// Order-preserving float32 -> uint32 for radix select
SGL_DEVICE uint32_t extract_exact_bin(float x) {
uint32_t bits = __float_as_uint(x);
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
SGL_DEVICE void trivial_transform(const TransformParams& params, uint32_t length, uint32_t K) {
if (const auto tx = threadIdx.x; tx < length) {
params.write(tx, tx);
} else if (tx < K) {
params.indices_out[tx] = -1;
}
}
SGL_DEVICE void tie_handle_transform(
const Tie* __restrict__ ties, //
const uint32_t num_ties,
const uint32_t num_above,
const uint32_t K,
const TransformParams params,
void* _smem) {
auto* smem = static_cast<TieHandleSmem*>(_smem);
const auto tx = threadIdx.x;
const auto lane_id = tx % kWarpThreads;
const auto warp_id = tx / kWarpThreads;
// Each thread loads one element (or becomes inactive)
const bool has_elem = tx < num_ties;
const auto tie = has_elem ? ties[tx] : Tie{0, 0.0f};
const uint32_t key = extract_exact_bin(tie.score);
const uint32_t idx = tie.idx;
bool active = has_elem;
uint32_t topk_remain = K - num_above;
uint32_t write_pos = K;
smem->counter = 0;
__syncthreads();
// Number of warps covering the 256-bin histogram (256/32 = 8)
constexpr uint32_t kRadixWarps = kRadixBins / kWarpThreads;
#pragma unroll
for (int round = 0; round < 4; round++) {
const uint32_t shift = 24 - round * 8;
const uint32_t bin = (key >> shift) & 0xFFu;
// 1. Build histogram
if (tx < kRadixBins) smem->histogram[tx] = 0;
__syncthreads();
if (active) atomicAdd(&smem->histogram[bin], 1);
__syncthreads();
// 2. v2-style 2-pass prefix sum on 256 bins
// Only first 256 threads (8 warps) carry histogram bins.
// Other threads get hist_val=0 and harmless prefix results.
uint32_t hist_val = 0;
uint32_t warp_inc = 0;
if (tx < kRadixBins) {
hist_val = smem->histogram[tx];
warp_inc = warp_inclusive_sum(lane_id, hist_val);
if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc;
}
__syncthreads();
if (tx < kRadixBins) {
// Inter-warp prefix (only first kHistWarps warp totals matter)
const auto tmp = (lane_id < kRadixWarps) ? smem->warp_sum[lane_id] : 0;
const auto total = warp::reduce_sum(tmp);
const auto inter = warp::reduce_sum(lane_id < warp_id ? tmp : 0);
const auto prefix = inter + warp_inc; // inclusive prefix through this bin
const auto above = total - prefix; // elements in bins ABOVE this one
// 3. Find threshold bin
if (above < topk_remain && above + hist_val >= topk_remain) {
smem->match = {tx, above, topk_remain - above};
}
}
__syncthreads();
const auto [thr, n_above, _] = smem->match;
// 4. Scatter
if (active) {
if (bin > thr) {
write_pos = num_above + atomicAdd(&smem->counter, 1);
active = false;
} else if (bin < thr) {
active = false;
} else if (round == 3) {
write_pos = K - atomicAdd(&smem->match.equal_count, -1u);
}
// my_bin == thr && round < 3: stay active for next round
}
topk_remain -= n_above;
if (topk_remain == 0) break;
}
if (write_pos < K) params.write(write_pos, idx);
}
} // namespace device::top512
python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/ptx.cuh
-54
View File
@@ -1,54 +0,0 @@
#pragma once
#include <sgl_kernel/utils.cuh>
#include <cuda/ptx>
#include <cstdint>
namespace device::top512 {
namespace ptx {
SGL_DEVICE void mbarrier_wait(uint64_t* addr, uint32_t phase) {
while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed, cuda::ptx::scope_cta, addr, phase))
;
}
SGL_DEVICE void mbarrier_init(uint64_t* addr, uint32_t arrives) {
cuda::ptx::mbarrier_init(addr, arrives);
}
SGL_DEVICE void mbarrier_arrive_expect_tx(uint64_t* addr, uint32_t tx) {
cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed, cuda::ptx::scope_cta, cuda::ptx::space_shared, addr, tx);
}
SGL_DEVICE void mbarrier_arrive(uint64_t* addr) {
cuda::ptx::mbarrier_arrive(cuda::ptx::sem_relaxed, cuda::ptx::scope_cta, cuda::ptx::space_shared, addr);
}
SGL_DEVICE void tma_load(void* dst, const void* src, uint32_t num_bytes, uint64_t* mbar) {
cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global, dst, src, num_bytes, mbar);
}
SGL_DEVICE uint32_t elect_sync() {
uint32_t pred = 0;
asm volatile(
"{\n\t"
".reg .pred %%px;\n\t"
"elect.sync _|%%px, %1;\n\t"
"@%%px mov.s32 %0, 1;\n\t"
"}"
: "+r"(pred)
: "r"(0xFFFFFFFF));
return pred;
}
SGL_DEVICE bool elect_sync_cta(uint32_t tx) {
const auto warp_id = tx / 32;
const auto uniform_warp_id = __shfl_sync(0xFFFFFFFF, warp_id, 0);
return (uniform_warp_id == 0 && elect_sync());
}
} // namespace ptx
} // namespace device::top512
python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/register.cuh
-302
View File
@@ -1,302 +0,0 @@
#pragma once
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include "common.cuh"
#include "ptx.cuh"
#include <cfloat>
#include <cstdint>
namespace device::top512 {
template <uint32_t K>
struct RegisterTopK {
static constexpr uint32_t kHistBits = 12;
static constexpr uint32_t kHistBins = 1 << kHistBits;
static constexpr uint32_t kVecsPerThread = 4;
static constexpr uint32_t kMaxTolerance = 0;
static constexpr uint32_t kMax1PassLength = kVecsPerThread * 4 * kBlockSize;
static constexpr uint32_t kMaxExtraLength = kMax1PassLength;
static constexpr uint32_t kMax2PassLength = kMax1PassLength + kMaxExtraLength;
struct Smem {
using HistVec = AlignedVector<uint32_t, kHistBins / kBlockSize>;
alignas(128) uint32_t counter_gt;
alignas(128) uint32_t counter_eq;
uint64_t mbarrier; // for cp.async
MatchBin match;
uint32_t warp_sum[kNumWarps];
union {
uint32_t histogram[kHistBins];
HistVec histogram_vec[kBlockSize];
Tie tie_buffer[kMaxTies];
};
alignas(16) float score_buffer[kMaxExtraLength];
};
template <bool kIs2Pass = false>
SGL_DEVICE static void
run(const float* scores, //
int32_t* indices,
const uint32_t length,
void* _smem,
const bool use_pdl = false) {
const auto smem = static_cast<Smem*>(_smem);
const auto tx = threadIdx.x;
const auto lane_id = tx % kWarpThreads;
const auto warp_id = tx / kWarpThreads;
// Initialize shared memory histogram
{
typename Smem::HistVec hist_vec;
hist_vec.fill(0);
smem->histogram_vec[tx] = hist_vec;
if (tx == 0) {
smem->counter_gt = smem->counter_eq = 0;
if constexpr (kIs2Pass) {
ptx::mbarrier_init(&smem->mbarrier, 1);
}
}
__syncthreads();
}
if (use_pdl) device::PDLWaitPrimary<true>();
// Load scores into registers
Vec4 local[kVecsPerThread];
#pragma unroll
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
const uint32_t base = (tx + v * kBlockSize) * 4;
if (base >= length) break;
local[v].load(scores, tx + v * kBlockSize);
}
// Fetch the next chunk of scores
if constexpr (kIs2Pass) {
if (ptx::elect_sync_cta(tx)) {
const auto length_aligned = (length + 3u - kMax1PassLength) & ~3u;
const auto size_bytes = length_aligned * sizeof(float);
ptx::tma_load(smem->score_buffer, scores + kMax1PassLength, size_bytes, &smem->mbarrier);
ptx::mbarrier_arrive_expect_tx(&smem->mbarrier, size_bytes);
}
__syncwarp(); // avoid warp divergence on
}
// Accumulate histogram via shared-memory atomics
#pragma unroll
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
#pragma unroll
for (uint32_t e = 0; e < 4; ++e) {
if constexpr (!kIs2Pass) {
const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
if (idx >= length) goto LABEL_ACC_FINISH;
}
atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(local[v][e])], 1);
}
}
if constexpr (kIs2Pass) {
// 16K ~ 32K. `i` is a float4 index
if (lane_id == 0) ptx::mbarrier_wait(&smem->mbarrier, 0);
__syncwarp();
for (uint32_t i = tx; i + kMax1PassLength < length; i += kBlockSize) {
const auto val = smem->score_buffer[i];
atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(val)], 1);
}
}
[[maybe_unused]] LABEL_ACC_FINISH:
__syncthreads();
// Phase 2: Exclusive prefix scan -> find threshold bin
{
constexpr uint32_t kItems = kHistBins / kBlockSize;
uint32_t orig[kItems];
const auto hist_vec = smem->histogram_vec[tx];
uint32_t tmp_local_sum = 0;
#pragma unroll
for (uint32_t i = 0; i < kItems; ++i) {
orig[i] = hist_vec[i];
tmp_local_sum += orig[i];
}
const auto warp_inc = warp_inclusive_sum(lane_id, tmp_local_sum);
const auto warp_exc = warp_inc - tmp_local_sum;
if (lane_id == kWarpThreads - 1) {
smem->warp_sum[warp_id] = warp_inc;
}
__syncthreads();
const auto tmp = smem->warp_sum[lane_id];
// Exactly one bin satisfies: above < K && above + count >= K
uint32_t prefix_sum = warp::reduce_sum(lane_id < warp_id ? tmp : 0);
prefix_sum += warp_exc;
#pragma unroll
for (uint32_t i = 0; i < kItems; ++i) {
prefix_sum += orig[i];
const auto above = length - prefix_sum;
if (above < K && above + orig[i] >= K) {
smem->match = {
.bin = tx * kItems + i,
.above_count = above,
.equal_count = orig[i],
};
}
}
__syncthreads();
}
const auto [thr_bin, num_above, num_equal] = smem->match;
// Phase 3: Scatter
// Elements strictly above threshold go directly to output.
// Tied elements: simple path admits first-come; tiebreak path collects into tie_buffer.
const bool need_tiebreak = (num_equal + num_above > K + kMaxTolerance);
const auto topk_indices = indices;
const auto tie_buffer = smem->tie_buffer;
#pragma unroll
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
#pragma unroll
for (uint32_t e = 0; e < 4; ++e) {
const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
if constexpr (!kIs2Pass) {
if (idx >= length) goto LABEL_SCATTER_DONE;
}
const uint32_t bin = extract_coarse_bin<kHistBits>(local[v][e]);
if (bin > thr_bin) {
topk_indices[atomicAdd(&smem->counter_gt, 1)] = idx;
} else if (bin == thr_bin) {
const auto pos = atomicAdd(&smem->counter_eq, 1);
if (need_tiebreak) {
if (pos < kMaxTies) {
tie_buffer[pos] = {.idx = idx, .score = local[v][e]};
}
} else {
if (const auto which = pos + num_above; which < K) {
topk_indices[which] = idx;
}
}
}
}
// prefetch the next scores
if constexpr (kIs2Pass) {
local[v].load(smem->score_buffer, tx + v * kBlockSize);
}
}
// 16K ~ 32K, already in registers: similar loop as above but read from smem->score_buffer
if constexpr (kIs2Pass) {
#pragma unroll
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
#pragma unroll
for (uint32_t e = 0; e < 4; ++e) {
const uint32_t idx = (tx + v * kBlockSize) * 4 + e + kMax1PassLength;
if (idx >= length) goto LABEL_SCATTER_DONE;
const uint32_t bin = extract_coarse_bin<kHistBits>(local[v][e]);
if (bin > thr_bin) {
topk_indices[atomicAdd(&smem->counter_gt, 1)] = idx;
} else if (bin == thr_bin) {
const auto pos = atomicAdd(&smem->counter_eq, 1);
if (need_tiebreak) {
if (pos < kMaxTies) {
tie_buffer[pos] = {.idx = idx, .score = local[v][e]};
}
} else {
if (const auto which = pos + num_above; which < K) {
topk_indices[which] = idx;
}
}
}
}
}
}
[[maybe_unused]] LABEL_SCATTER_DONE:
if (!need_tiebreak) return;
// Phase 4: Tie-breaking within the threshold bin.
// Assume num_ties <= kBlockSize (at most 1 block of ties).
// Each thread takes one tied element, computes its rank (number of
// elements with strictly higher score, breaking exact float ties by
// original index), and writes to output if rank < topk_remain.
__syncthreads();
static_assert(kMaxTies <= kBlockSize);
const uint32_t num_ties = min(num_equal, kMaxTies);
const uint32_t topk_remain = K - num_above;
const 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 <= kWarpThreads) {
static_assert(kWarpThreads <= kNumWarps);
if (lane_id >= num_ties || warp_id >= num_ties) return; // some threads are idle
/// NOTE: use long long to avoid mask overflow when num_ties == 32
const uint32_t mask = (1ull << num_ties) - 1u;
const auto tie = tie_buffer[lane_id];
const auto target_tie = tie_buffer[warp_id];
const bool pred = is_greater(tie, target_tie);
const auto rank = static_cast<uint32_t>(__popc(__ballot_sync(mask, pred)));
if (lane_id == 0 && rank < topk_remain) {
topk_indices[num_above + rank] = target_tie.idx;
}
} else if (num_ties <= kWarpThreads * 2) {
// 64 x 64 topk implementation: each thread takes 2 elements
const auto lane_id_1 = lane_id + kWarpThreads;
const auto warp_id_1 = warp_id + kWarpThreads;
const auto invalid = Tie{.idx = 0xFFFFFFFF, .score = -FLT_MAX};
const auto tie_0 = tie_buffer[lane_id];
const auto tie_1 = lane_id_1 < num_ties ? tie_buffer[lane_id_1] : invalid;
if (true) {
const auto target = tie_buffer[warp_id];
const bool pred_0 = is_greater(tie_0, target);
const bool pred_1 = is_greater(tie_1, target);
const auto rank_0 = static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_0)));
const auto rank_1 = static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_1)));
const auto rank = rank_0 + rank_1;
if (lane_id == 0 && rank < topk_remain) {
topk_indices[num_above + rank] = target.idx;
}
}
if (warp_id_1 < num_ties) {
const auto target = tie_buffer[warp_id_1];
const bool pred_0 = is_greater(tie_0, target);
const bool pred_1 = is_greater(tie_1, target);
const auto rank_0 = static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_0)));
const auto rank_1 = static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_1)));
const auto rank = rank_0 + rank_1;
if (lane_id == 0 && rank < topk_remain) {
topk_indices[num_above + rank] = target.idx;
}
}
} else {
/// NOTE: Based on my observation, this path is very rarely reached
[[unlikely]];
// Block-level: each thread reads from tie_buffer in shared memory
for (auto i = warp_id; i < num_ties; i += kNumWarps) {
const auto target_tie = tie_buffer[i];
uint32_t local_rank = 0;
for (auto j = lane_id; j < num_ties; j += kWarpThreads) {
const auto tie = tie_buffer[j];
if (is_greater(tie, target_tie)) local_rank++;
}
// sum the rank across the warp
const auto rank = warp::reduce_sum(local_rank);
if (lane_id == 0 && rank < topk_remain) {
topk_indices[num_above + rank] = target_tie.idx;
}
}
}
}
SGL_DEVICE static void transform(const TransformParams params) {
__syncthreads();
if (const auto tx = threadIdx.x; tx < K) params.transform(tx);
}
};
} // namespace device::top512
python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/streaming.cuh
-213
View File
@@ -1,213 +0,0 @@
#pragma once
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include "common.cuh"
#include "ptx.cuh"
#include <cfloat>
#include <cstdint>
namespace device::top512 {
template <uint32_t K>
struct StreamingTopK {
static constexpr uint32_t kHistBits = 12;
static constexpr uint32_t kHistBins = 1 << kHistBits;
static constexpr uint32_t kRadixBins = 256;
static constexpr uint32_t kElemPerStage = 8;
static constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize;
static constexpr uint32_t kNumStages = 2; // double buffer
static constexpr uint32_t kHistItems = kHistBins / kBlockSize; // 4
static_assert(kHistItems * kBlockSize == kHistBins);
using HistVec = AlignedVector<uint32_t, kHistItems>;
struct Smem {
uint64_t barrier[2][kNumStages];
alignas(128) uint32_t counter_gt;
alignas(128) uint32_t counter_eq;
alignas(128) MatchBin match;
alignas(128) uint32_t warp_sum[kNumWarps];
union {
uint32_t histogram[kHistBins];
HistVec histogram_vec[kBlockSize];
Tie tie_buffer[kMaxTies];
};
union {
float score_buffer[kNumStages][kSizePerStage];
TieHandleSmem stage2; // reuse smem for tie handling in phase D
};
};
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
/// NOTE: length must be 4-aligned since we load 4 floats/thread. Caller should round up.
template <bool kIsScatter>
SGL_DEVICE static void issue_tma(const float* scores, uint32_t stage, uint32_t length, Smem* smem) {
const auto buf_idx = stage % kNumStages;
const auto offset = stage * kSizePerStage;
const auto size = min(kSizePerStage, length - offset);
const auto size_bytes = size * sizeof(float);
const auto bar = &smem->barrier[kIsScatter][buf_idx];
ptx::tma_load(smem->score_buffer[buf_idx], scores + offset, size_bytes, bar);
ptx::mbarrier_arrive_expect_tx(bar, size_bytes);
}
// ---------------------------------------------------------------------------
// Unified streaming pass. Used for both phase A (kIsScatter=false) and
// phase C (kIsScatter=true). Each buffer is reused across iterations via the
// reuse-arrive trick (same pattern as ClusterTopKImpl::stage1).
// ---------------------------------------------------------------------------
template <bool kIsScatter>
SGL_DEVICE static void stream_pass(
const float* scores,
const uint32_t length,
const uint32_t thr_bin, // ignored when !kIsScatter
int32_t* s_topk_indices, // ignored when !kIsScatter
Smem* smem) {
const auto tx = threadIdx.x;
const auto num_iters = (length + kSizePerStage - 1) / kSizePerStage;
const auto lane_id = tx % kWarpThreads;
// Initial double-buffer TMA prologue.
const auto length_aligned = (length + 3u) & ~3u;
if (tx == 0) {
#pragma unroll
for (uint32_t i = 0; i < kNumStages; i++) {
if (i >= num_iters) break;
issue_tma<kIsScatter>(scores, i, length_aligned, smem);
}
}
for (uint32_t iter = 0; iter < num_iters; iter++) {
const auto buf_idx = iter % kNumStages;
const auto offset = iter * kSizePerStage;
const auto this_size = min(kSizePerStage, length - offset);
if (lane_id == 1) {
const auto phase_bit = (iter / kNumStages) & 1;
ptx::mbarrier_wait(&smem->barrier[kIsScatter][buf_idx], phase_bit);
}
__syncwarp();
#pragma unroll
for (uint32_t i = 0; i < kElemPerStage; i++) {
const auto local_idx = tx + i * kBlockSize;
if (local_idx >= this_size) break;
const auto score = smem->score_buffer[buf_idx][local_idx];
const auto bin = extract_coarse_bin<kHistBits>(score);
if constexpr (kIsScatter) {
const auto global_idx = offset + local_idx;
if (bin > thr_bin) {
const auto pos = atomicAdd(&smem->counter_gt, 1);
if (pos < K) s_topk_indices[pos] = global_idx;
} else if (bin == thr_bin) {
const auto pos = atomicAdd(&smem->counter_eq, 1);
if (pos < kMaxTies) smem->tie_buffer[pos] = {global_idx, score};
}
} else {
atomicAdd(&smem->histogram[bin], 1);
}
}
__syncthreads();
if (tx == 0) {
if (const auto next_iter = iter + kNumStages; next_iter < num_iters) {
issue_tma<kIsScatter>(scores, next_iter, length_aligned, smem);
}
}
}
}
// ---------------------------------------------------------------------------
// Phase B: find the threshold bin via a warp-level prefix scan.
// Same structure as SmallTopKImpl's phase 2 (4 bins/thread, warp_sum relay).
// ---------------------------------------------------------------------------
SGL_DEVICE static void find_threshold(uint32_t length, Smem* smem) {
const auto tx = threadIdx.x;
const auto lane_id = tx % kWarpThreads;
const auto warp_id = tx / kWarpThreads;
uint32_t orig[kHistItems];
const auto hist_vec = smem->histogram_vec[tx];
uint32_t local_sum = 0;
#pragma unroll
for (uint32_t i = 0; i < kHistItems; ++i) {
orig[i] = hist_vec[i];
local_sum += orig[i];
}
const auto warp_inc = warp_inclusive_sum(lane_id, local_sum);
const auto warp_exc = warp_inc - local_sum;
if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc;
__syncthreads();
const auto tmp = smem->warp_sum[lane_id];
uint32_t prefix_sum = warp::reduce_sum(lane_id < warp_id ? tmp : 0);
prefix_sum += warp_exc;
#pragma unroll
for (uint32_t i = 0; i < kHistItems; ++i) {
prefix_sum += orig[i];
const auto above = length - prefix_sum;
if (above < K && above + orig[i] >= K) {
smem->match = {
.bin = tx * kHistItems + i,
.above_count = above,
.equal_count = orig[i],
};
}
}
__syncthreads();
}
SGL_DEVICE static void run(const float* scores, const uint32_t length, int32_t* topk_indices, void* _smem) {
const auto smem = static_cast<Smem*>(_smem);
const auto tx = threadIdx.x;
__builtin_assume(tx < kBlockSize);
// Init histogram, barriers, counters.
{
HistVec zero;
zero.fill(0);
smem->histogram_vec[tx] = zero;
if (tx < 2 * kNumStages) {
const auto base_barrier = &smem->barrier[0][0];
ptx::mbarrier_init(&base_barrier[tx], 1);
}
if (tx == 0) {
smem->counter_gt = 0;
smem->counter_eq = 0;
}
__syncthreads();
}
// Phase A: histogram pass (pipelined TMA stream).
stream_pass<false>(scores, length, 0, nullptr, smem);
// Phase B: locate threshold bin & re-init barriers
find_threshold(length, smem);
// Phase C: scatter pass.
stream_pass<true>(scores, length, smem->match.bin, topk_indices, smem);
}
SGL_DEVICE static void transform(const TransformParams params, void* _smem) {
// Phase D: page-translate above entries, then refine ties.
const auto smem = static_cast<Smem*>(_smem);
const auto tx = threadIdx.x;
const auto num_above = smem->match.above_count;
if (tx < num_above) params.transform(tx);
const auto num_equal = smem->counter_eq;
if (num_above >= K || num_equal == 0) return;
const auto clamped_ties = min(num_equal, kMaxTies);
tie_handle_transform(smem->tie_buffer, clamped_ties, num_above, K, params, &smem->stage2);
}
};
} // namespace device::top512
python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh
+752
View File
@@ -0,0 +1,752 @@
/// \file topk_impl.cuh
/// \brief DeepSeek-V4 (DSA indexer) top-k implementation classes.
///
/// This header holds ONLY the device-side implementation classes + helpers; the
/// `__global__` kernels and the host dispatcher live in csrc/deepseek_v4/topk_v2.cuh.
///
/// Design notes:
/// - top-k (`topk`) is a *runtime* value (<= kMaxTopK = 2048), never a
/// compile-time constant.
/// - the output is the page-table transform of the selected raw indices
/// (`TopKProblem::emit` then `transform_output`).
/// - each block reads its own `seq_len` (per-batch ragged lengths) -- the host
/// launches one universal kernel and dispatches per block.
/// - the cluster size is fixed at 8 (dynamic persistent clusters are hard).
///
/// Algorithm: fp16 coarse histogram -> threshold bin -> fp32-boundary collect ->
/// exact radix tie-break.
#pragma once
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <cfloat>
#include <cooperative_groups.h>
#include <cstdint>
#include <limits>
namespace device::topk {
namespace cg = cooperative_groups;
/// sgl_kernel names the warp size `kWarpThreads`; alias it locally as `kWarpSize`.
inline constexpr uint32_t kWarpSize = kWarpThreads;
// ---------------------------------------------------------------------------
// Shared-memory storage sized/aligned for several impl `Smem` types
// ---------------------------------------------------------------------------
/// Compile-time max over a non-empty pack (avoids an <algorithm> dependency).
template <typename T>
constexpr T ct_max(T a) {
return a;
}
template <typename T, typename... Ts>
constexpr T ct_max(T a, Ts... rest) {
const T m = ct_max(rest...);
return a > m ? a : m;
}
/// Static shared-memory buffer sized + aligned to hold any one of the given
/// impl `Smem` types. A kernel that dispatches across several paths (e.g. the
/// fused small-batch kernel runs either Streaming or Cluster; the main kernel
/// runs any of Register2/Register4/Streaming) declares one
/// `__shared__ MaxSmem<...> smem` and hands `&smem` to whichever forward() it
/// calls -- instead of hand-picking "the largest" type and relying on it
/// staying the largest. `&smem` converts to the `void*` the forwards expect;
/// the buffer is aligned to the strictest member, so the cast is well-aligned.
template <typename... Smems>
struct MaxSmem {
static constexpr size_t kSize = ct_max(sizeof(Smems)...);
static constexpr size_t kAlign = ct_max(alignof(Smems)...);
alignas(kAlign) uint8_t storage[kSize];
};
// ---------------------------------------------------------------------------
// Order-preserving float -> integer key extraction
// ---------------------------------------------------------------------------
SGL_DEVICE uint32_t extract_exact_bin(float x) {
uint32_t bits = __float_as_uint(x);
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
template <uint32_t kBits>
SGL_DEVICE uint32_t extract_coarse_bin(float x) {
static_assert(0 < kBits && kBits < 15);
const auto hx = cast<fp16_t>(x);
const uint16_t bits = *reinterpret_cast<const uint16_t*>(&hx);
const uint16_t key = (bits & 0x8000) ? ~bits : bits | 0x8000;
return key >> (16 - kBits);
}
// Smallest fp32 value `v` for which `extract_coarse_bin<kBits>(v) >= bin`, i.e. the
// lower fp32 boundary of coarse bin `bin`. Because `extract_coarse_bin` is monotonic
// non-decreasing in its argument, the collect pass can classify an element with two
// fp32 comparisons against these boundaries instead of recomputing the fp16 bin --
// removing the F2F conversion and bit-twiddle from the (compute-bound) second pass.
// Returns -inf for bin 0 (everything qualifies) and +inf for bins past the top.
template <uint32_t kBits>
SGL_DEVICE float coarse_bin_lower_bound(uint32_t bin) {
if (bin == 0) return -FLT_MAX;
if (bin >= (1u << kBits)) return FLT_MAX;
constexpr uint32_t kShift = 16 - kBits;
const uint32_t key = bin << kShift; // ordered16 key at the low edge of `bin`
// ordered16 -> fp16 value (inverse of the transform in extract_coarse_bin)
const auto to_val = [](uint32_t okey) -> float {
const uint16_t ob = static_cast<uint16_t>(okey);
const uint16_t hb = (ob & 0x8000) ? static_cast<uint16_t>(ob ^ 0x8000) : static_cast<uint16_t>(~ob);
return cast<float>(*reinterpret_cast<const fp16_t*>(&hb));
};
// fp16 rounds to nearest, so the fp32 boundary is the midpoint between the fp16
// value at this key and the next-lower fp16 value (ordered key - 1).
return 0.5f * (to_val(key) + to_val(key - 1));
}
SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
#pragma unroll
for (uint32_t offset = 1; offset < 32; offset *= 2) {
uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
if (lane_id >= offset) val += n;
}
return val;
}
SGL_DEVICE uint32_t warp_sum_bool(bool pred, uint32_t mask = 0xFFFFFFFF) {
return __popc(__ballot_sync(mask, pred));
}
struct alignas(8) TieValue {
float value;
uint32_t idx;
inline static constexpr TieValue invalid() {
return TieValue{-FLT_MAX, 0xFFFFFFFFu};
}
};
// ---------------------------------------------------------------------------
// Per-batch problem description + page-table transform sink
// ---------------------------------------------------------------------------
SGL_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, uint32_t i, uint32_t page_bits) {
const uint32_t mask = (1u << page_bits) - 1u;
return (page_table[i >> page_bits] << page_bits) | (i & mask);
}
/// One batch element's worth of work. `emit(pos, raw_idx)` writes the selected raw
/// index to output slot `pos`; `transform_output` then applies the page-table
/// transform in a separate pass (and records the raw index in `raw_out` if set).
struct TopKProblem {
const float* __restrict__ in;
int32_t* __restrict__ out; // page_indices [topk]
int32_t* __restrict__ raw_out; // optional raw (pre-transform) indices [topk]; nullptr if unused
const int32_t* __restrict__ page_table;
uint32_t topk;
uint32_t seq_len;
uint32_t page_bits;
// Write the raw selected index; the page-table transform is applied afterwards
// by transform_output() in a separate, pipelined pass. Keeping the per-element
// page_table gather off the atomic-serialized scatter loop is measurably faster
// for both short and long context.
SGL_DEVICE void emit(uint32_t pos, uint32_t raw_idx) const {
out[pos] = static_cast<int32_t>(raw_idx);
}
SGL_DEVICE void transform_output(uint32_t t, int32_t raw) const {
if (raw_out != nullptr) raw_out[t] = raw;
out[t] = raw < 0 ? -1 : page_to_indices(page_table, raw, page_bits);
}
};
// ---------------------------------------------------------------------------
// Shared configuration + tie handling (exact radix select on the threshold bin)
// ---------------------------------------------------------------------------
struct TopKConfig {
static constexpr uint32_t kMaxTopK = 2048;
static constexpr uint32_t kBlockSize = 1024;
static constexpr uint32_t kOccupancy = 2;
static constexpr uint32_t kNumWarps = kBlockSize / kWarpSize;
static constexpr uint32_t kMaxNumTie = 1024;
static constexpr uint32_t kRadixSize = 1 << 8;
static constexpr uint32_t kTopKItems = (kMaxTopK + kBlockSize - 1) / kBlockSize;
static_assert(kMaxNumTie <= kBlockSize && kBlockSize % kNumWarps == 0);
struct TieHandleSmem {
struct alignas(16) MatchBin {
uint32_t bin;
uint32_t above_count;
uint32_t equal_count;
uint32_t _pad = 0;
};
alignas(128) uint32_t counter;
alignas(128) uint32_t counter_final;
MatchBin match;
uint32_t warp_sum[kNumWarps];
uint32_t histogram[2][kRadixSize];
};
/// Resolve the threshold bin's ties exactly. `base` is the number of strictly
/// "above" elements already emitted (final output starts at slot `base`);
/// `topk` here is the number of remaining slots to fill (== global_topk - base).
SGL_DEVICE static void handle_tie( //
const TieValue* tie_buffer,
const TopKProblem& problem,
const uint32_t base,
const uint32_t num_ties,
const uint32_t topk,
TieHandleSmem* smem) {
constexpr auto is_greater = [](const TieValue& a, const TieValue& b) {
return (a.value > b.value) || (a.value == b.value && a.idx < b.idx);
};
const auto tx = threadIdx.x;
const auto lane_id = tx % kWarpSize;
const auto warp_id = tx / kWarpSize;
static_assert(kNumWarps == kWarpSize);
if (num_ties <= topk) {
if (tx < num_ties) problem.emit(base + tx, tie_buffer[tx].idx);
} else if (num_ties <= kWarpSize) {
if (lane_id >= num_ties || warp_id >= num_ties) return; // some threads are idle
/// NOTE: use long long to avoid mask overflow when num_tie == 32
const uint32_t mask = (1ull << num_ties) - 1u;
const auto tie = tie_buffer[lane_id];
const auto target = tie_buffer[warp_id];
const auto rank = warp_sum_bool(is_greater(tie, target), mask);
if (lane_id == 0 && rank < topk) problem.emit(base + rank, target.idx);
} else if (num_ties <= kWarpSize * 2) {
// 64 x 64 topk implementation: each thread takes 2 elements
const auto warp_id_0 = warp_id;
const auto warp_id_1 = warp_id + kWarpSize;
const auto lane_id_1 = lane_id + kWarpSize;
const auto invalid = TieValue::invalid();
const auto tie_0 = tie_buffer[lane_id];
const auto tie_1 = lane_id_1 < num_ties ? tie_buffer[lane_id_1] : invalid;
const auto target_0 = tie_buffer[warp_id_0];
const auto target_1 = tie_buffer[warp_id_1];
if (true) { // NOTE: warp_id_0 <= kNumWarps < num_ties
const auto rank_0 = warp_sum_bool(is_greater(tie_0, target_0));
const auto rank_1 = warp_sum_bool(is_greater(tie_1, target_0));
const auto rank = rank_0 + rank_1;
if (lane_id == 0 && rank < topk) problem.emit(base + rank, target_0.idx);
}
if (warp_id_1 < num_ties) {
const auto rank_0 = warp_sum_bool(is_greater(tie_0, target_1));
const auto rank_1 = warp_sum_bool(is_greater(tie_1, target_1));
const auto rank = rank_0 + rank_1;
if (lane_id == 0 && rank < topk) problem.emit(base + rank, target_1.idx);
}
} else if (num_ties <= kWarpSize * 4) {
// 128 x 128 topk implementation: each thread takes 4 elements and does local sort + merge
const auto invalid = TieValue::invalid();
const TieValue tie[] = {
tie_buffer[lane_id + 0 * kWarpSize],
tie_buffer[lane_id + 1 * kWarpSize],
lane_id + 2 * kWarpSize < num_ties ? tie_buffer[lane_id + 2 * kWarpSize] : invalid,
lane_id + 3 * kWarpSize < num_ties ? tie_buffer[lane_id + 3 * kWarpSize] : invalid,
};
const TieValue target[] = {
tie_buffer[warp_id + 0 * kWarpSize],
tie_buffer[warp_id + 1 * kWarpSize],
tie_buffer[warp_id + 2 * kWarpSize],
tie_buffer[warp_id + 3 * kWarpSize],
};
#pragma unroll
for (int i = 0; i < 4; ++i) {
if (i >= 2 && warp_id + i * kWarpSize >= num_ties) break;
uint32_t rank = 0;
#pragma unroll
for (int j = 0; j < 4; ++j) {
rank += warp_sum_bool(is_greater(tie[j], target[i]));
}
if (lane_id == 0 && rank < topk) problem.emit(base + rank, target[i].idx);
}
} else {
// Each thread loads one element (or becomes inactive)
bool active = tx < num_ties;
const auto tie = active ? tie_buffer[tx] : TieValue::invalid();
const uint32_t key = extract_exact_bin(tie.value);
const uint32_t idx = tie.idx;
uint32_t topk_remain = topk;
uint32_t write_pos = topk;
if (tx < kRadixSize) smem->histogram[0][tx] = 0;
if (tx == kRadixSize) smem->counter = smem->counter_final = 0;
__syncthreads();
uint32_t total_active = num_ties;
#pragma unroll
for (int round = 0; round < 4; round++) {
const uint32_t shift = 24 - round * 8;
const uint32_t bin = (key >> shift) & 0xFFu;
const auto hist_idx = round % 2;
const auto histogram = smem->histogram[hist_idx];
if (active) {
atomicAdd(&histogram[bin], 1);
}
if (round < 3 && tx < kRadixSize) {
smem->histogram[hist_idx ^ 1][tx] = 0;
}
__syncthreads();
uint32_t hist_val = 0;
uint32_t warp_inc = 0;
if (tx < kRadixSize) {
hist_val = histogram[tx];
warp_inc = warp_inclusive_sum(lane_id, hist_val);
if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc;
}
__syncthreads();
if (tx < kRadixSize) {
const auto inter = warp::reduce_sum(lane_id < warp_id ? smem->warp_sum[lane_id] : 0);
const auto prefix = inter + warp_inc; // inclusive prefix through this bin
const auto above = total_active - prefix; // elements in bins ABOVE this one
// 3. Find threshold bin
if (above < topk_remain && above + hist_val >= topk_remain) {
smem->match = {tx, above, hist_val};
}
}
__syncthreads();
const auto [threshold_bin, above_count, equal_count, __] = smem->match;
if (round < 3) total_active = equal_count;
topk_remain -= above_count;
// 4. Scatter
if (active) {
if (bin > threshold_bin) {
write_pos = atomicAdd(&smem->counter, 1);
active = false;
} else if (bin < threshold_bin) {
active = false;
} else if (round == 3) {
write_pos = topk - topk_remain + atomicAdd(&smem->counter_final, 1);
}
// my_bin == thr && round < 3: stay active for next round
}
if (round == 3 || topk_remain == 0) break;
}
if (write_pos < topk) problem.emit(base + write_pos, idx);
}
}
};
// ---------------------------------------------------------------------------
// Radix base: histogram storage + input iteration + threshold-bin search
// ---------------------------------------------------------------------------
template <uint32_t kHistBits_>
struct TopKRadixBase : TopKConfig {
static constexpr uint32_t kVecSize = 4;
static constexpr uint32_t kHistBits = kHistBits_;
static constexpr uint32_t kHistSize = 1 << kHistBits;
using vec_t = AlignedVector<float, kVecSize>;
struct Smem {
using kHistVec = AlignedVector<uint32_t, kHistSize / kBlockSize>;
alignas(128) uint32_t count_eq;
alignas(128) uint32_t count_gt;
uint32_t threshold_bin;
uint32_t warp_sum[kNumWarps];
union {
TieHandleSmem tie_handle_smem;
uint32_t histogram[kHistSize];
kHistVec hist_vecs[kBlockSize];
};
TieValue tie_values[kMaxNumTie];
};
protected:
template <typename F>
SGL_DEVICE static void for_each_input(const float* __restrict__ in, uint32_t seq_len, F&& fn) {
const auto tx = threadIdx.x;
const uint32_t num_full = seq_len / kVecSize; // fully-in-bounds vectors
vec_t next_vec;
uint32_t vi = tx;
if (vi < num_full) next_vec.load(in, vi);
while (vi < num_full) {
const auto cur = next_vec;
const auto base = vi * kVecSize;
vi += kBlockSize;
if (vi < num_full) next_vec.load(in, vi);
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j) {
fn(cur[j], base + j);
}
}
// Tail: at most one partial vector, `rem` in [0, kVecSize).
static_assert(kVecSize <= kBlockSize); // ensure tail correctness
const uint32_t tail_start = num_full * kVecSize;
if (tx < seq_len - tail_start) {
const auto idx = tail_start + tx;
fn(in[idx], idx);
}
}
SGL_DEVICE static void find_threshold(const uint32_t topk, const uint32_t seq_len, Smem* smem) {
const auto tx = threadIdx.x;
constexpr uint32_t kItems = kHistSize / kBlockSize;
uint32_t orig[kItems];
const auto hist_vec = smem->hist_vecs[tx];
uint32_t tmp_local_sum = 0;
#pragma unroll
for (uint32_t i = 0; i < kItems; ++i) {
orig[i] = hist_vec[i];
tmp_local_sum += orig[i];
}
const auto lane_id = tx % kWarpSize;
const auto warp_id = tx / kWarpSize;
const auto warp_inc = warp_inclusive_sum(lane_id, tmp_local_sum);
const auto warp_exc = warp_inc - tmp_local_sum;
if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc;
__syncthreads();
const auto tmp = smem->warp_sum[lane_id];
// Exactly one bin satisfies: above < K && above + count >= K
uint32_t prefix_sum = warp::reduce_sum(lane_id < warp_id ? tmp : 0);
prefix_sum += warp_exc;
#pragma unroll
for (uint32_t i = 0; i < kItems; ++i) {
prefix_sum += orig[i];
const auto above = seq_len - prefix_sum;
if (above < topk && above + orig[i] >= topk) {
smem->threshold_bin = tx * kItems + i;
}
}
__syncthreads();
}
};
// ---------------------------------------------------------------------------
// Register path: scores stay resident in registers across both passes (read
// once). Templated on kLocalVecs so the caller picks the smallest covering
// kernel -- a larger kLocalVecs raises kMaxSeqLen but its fixed-unrolled loop
// wastes work on shorter sequences.
// ---------------------------------------------------------------------------
template <uint32_t kLocalVecs_>
struct TopKRegister : TopKRadixBase<12> {
static constexpr uint32_t kLocalVecs = kLocalVecs_;
static constexpr uint32_t kMaxSeqLen = kBlockSize * kVecSize * kLocalVecs;
using Smem = typename TopKRadixBase<12>::Smem;
template <bool kUsePDL>
SGL_DEVICE static void forward(const TopKProblem problem, void* _smem) {
const auto tx = threadIdx.x;
const auto smem = static_cast<Smem*>(_smem);
{
Smem::kHistVec hist_vec;
hist_vec.fill(0);
smem->hist_vecs[tx] = hist_vec;
}
if (tx == 0) {
smem->count_eq = 0;
smem->count_gt = 0;
}
__syncthreads();
PDLWaitPrimary<kUsePDL>();
// A vector `vi` is fully in bounds iff vi < num_full; only full vectors are
// vector-loaded (16B aligned, never straddling seq_len). The <kVecSize tail is
// a scalar remainder on the LAST lanes (which own the fewest full vectors, so
// it overlaps the busy lanes' extra vector). The full path has no per-element
// bounds check, keeping register pressure low enough to hold all vectors.
const uint32_t num_full = problem.seq_len / kVecSize;
const uint32_t tail_start = num_full * kVecSize;
const uint32_t tail = problem.seq_len - tail_start;
// Phase 1: load full vectors + build histogram
vec_t local_vecs[kLocalVecs];
#pragma unroll
for (uint32_t i = 0; i < kLocalVecs; ++i) {
const auto vi = tx + kBlockSize * i;
if (vi >= num_full) break;
local_vecs[i].load(problem.in, vi);
}
#pragma unroll
for (uint32_t i = 0; i < kLocalVecs; ++i) {
const auto vi = tx + kBlockSize * i;
if (vi >= num_full) break;
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j)
atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(local_vecs[i][j])], 1);
}
if (tx >= kBlockSize - tail) {
const uint32_t idx = tail_start + tx - (kBlockSize - tail);
atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(problem.in[idx])], 1);
}
__syncthreads();
// Phase 2: Find the threshold bin
find_threshold(problem.topk, problem.seq_len, smem);
// Phase 3: collect by two fp32 boundaries (raw indices; transform applied later)
const auto topk = problem.topk;
const auto threshold_bin = smem->threshold_bin;
const auto v_hi = coarse_bin_lower_bound<kHistBits>(threshold_bin + 1);
const auto v_lo = coarse_bin_lower_bound<kHistBits>(threshold_bin);
const auto collect = [&](float val, uint32_t idx) {
if (val >= v_hi) {
const auto pos = atomicAdd(&smem->count_gt, 1);
if (pos < topk) [[likely]]
problem.emit(pos, idx);
} else if (val >= v_lo) {
const auto count_eq = atomicAdd(&smem->count_eq, 1);
if (count_eq < kMaxNumTie) [[likely]]
smem->tie_values[count_eq] = {val, idx};
}
};
#pragma unroll
for (uint32_t i = 0; i < kLocalVecs; ++i) {
const auto vi = tx + kBlockSize * i;
const auto base = vi * kVecSize;
if (vi >= num_full) break;
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j)
collect(local_vecs[i][j], base + j);
}
if (tx >= kBlockSize - tail) {
const uint32_t idx = tail_start + tx - (kBlockSize - tail);
collect(problem.in[idx], idx);
}
// Phase 4: Handle ties.
__syncthreads();
const auto above_count = smem->count_gt;
const auto equal_count = smem->count_eq;
const auto remain_topk = above_count < topk ? topk - above_count : 0;
const auto tie_count = min(equal_count, kMaxNumTie);
handle_tie(smem->tie_values, problem, above_count, tie_count, remain_topk, &smem->tie_handle_smem);
}
};
// ---------------------------------------------------------------------------
// Streaming path: seq_len > 8192 -- two vectorized passes over global memory
// ---------------------------------------------------------------------------
struct TopKStreaming : TopKRegister<2> {
public:
static constexpr uint32_t kMaxSeqLen = std::numeric_limits<uint32_t>::max();
template <bool kUsePDL>
SGL_DEVICE static void forward(const TopKProblem problem, void* _smem) {
const auto tx = threadIdx.x;
const auto smem = static_cast<Smem*>(_smem);
{
Smem::kHistVec hist_vec;
hist_vec.fill(0);
smem->hist_vecs[tx] = hist_vec;
}
if (tx == 0) {
smem->count_eq = 0;
smem->count_gt = 0;
}
__syncthreads();
PDLWaitPrimary<kUsePDL>();
// Phase 1: Load and build histogram
for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t) {
const auto bin = extract_coarse_bin<kHistBits>(val);
atomicAdd(&smem->histogram[bin], 1);
});
__syncthreads();
// Phase 2: Find the threshold bin
find_threshold(problem.topk, problem.seq_len, smem);
// Phase 3: Collect candidates and sort. Classify by two fp32 boundaries derived
// from the threshold bin instead of recomputing the fp16 bin per element: an
// element is "above" iff val >= v_hi (bin > threshold) and a "tie" iff
// v_lo <= val < v_hi (bin == threshold). This drops the F2F + bit-twiddle from
// the second full pass over the input.
const auto threshold_bin = smem->threshold_bin;
const float v_hi = coarse_bin_lower_bound<kHistBits>(threshold_bin + 1);
const float v_lo = coarse_bin_lower_bound<kHistBits>(threshold_bin);
const auto topk = problem.topk;
for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) {
if (val >= v_hi) {
const auto pos = atomicAdd(&smem->count_gt, 1);
if (pos < topk) [[likely]] {
problem.emit(pos, idx);
}
} else if (val >= v_lo) {
const auto count_eq = atomicAdd(&smem->count_eq, 1);
if (count_eq < kMaxNumTie) [[likely]] {
smem->tie_values[count_eq] = {val, idx};
}
}
});
// Phase 4: Handle ties. Drive the output layout from the *collect* counts so it
// is self-consistent with the fp32 classification above (rather than the fp16
// histogram counts), even if rounding moves a boundary element between the
// "above" and "tie" sets. above_count is < topk by the threshold-bin invariant,
// so the count_gt guard above effectively never triggers.
__syncthreads();
const auto above_count = smem->count_gt;
const auto equal_count = smem->count_eq;
const auto remain_topk = above_count < topk ? topk - above_count : 0;
const auto tie_count = min(equal_count, kMaxNumTie);
handle_tie(smem->tie_values, problem, above_count, tie_count, remain_topk, &smem->tie_handle_smem);
}
};
// ---------------------------------------------------------------------------
// Cluster path: very long seq_len, small batch. `kClusterSize` blocks cooperate
// on one batch element via distributed shared memory (one cluster per element).
// ---------------------------------------------------------------------------
template <uint32_t kClusterSize_>
struct TopKCluster : TopKRadixBase<10> {
public:
static constexpr uint32_t kClusterSize = kClusterSize_;
static constexpr uint32_t kMaxSeqLen = std::numeric_limits<uint32_t>::max();
using Base = TopKRadixBase<10>;
struct Smem : Base::Smem {
using kHistVec = Base::Smem::kHistVec;
uint32_t start_eq_local, start_gt_local;
int32_t tmp_out[kMaxTopK];
};
// Process ONE batch element (one cluster). NO PDL and NO trailing barrier --
// the persistent kernel does PDLWaitPrimary once before its item loop and a
// cluster.sync() after each forward(). Writes raw indices to out; the kernel's
// transform pass applies the page-table transform.
template <bool kUsePDL>
SGL_DEVICE static void forward(TopKProblem problem, void* _smem) {
const auto tx = threadIdx.x;
const auto smem = static_cast<Smem*>(_smem);
const auto cluster = cg::this_cluster();
const auto this_rank = blockIdx.y;
const bool is_primary = (this_rank == 0);
constexpr uint32_t kAlignElems = kWarpSize * kVecSize;
const uint32_t chunk_size = div_ceil(problem.seq_len, kClusterSize * kAlignElems) * kAlignElems;
const uint32_t chunk_start = min(this_rank * chunk_size, problem.seq_len);
const uint32_t chunk_finish = min(chunk_start + chunk_size, problem.seq_len);
const uint32_t local_seq_len = chunk_finish - chunk_start;
problem.in += chunk_start;
{
typename Smem::kHistVec hist_vec;
hist_vec.fill(0);
smem->hist_vecs[tx] = hist_vec;
}
if (tx == 0) {
smem->count_eq = 0;
smem->count_gt = 0;
}
__syncthreads();
PDLWaitPrimary<kUsePDL>();
// Phase 1: Load and build histogram over this rank's contiguous chunk.
for_each_input(problem.in, local_seq_len, [&](float val, uint32_t) {
const auto bin = extract_coarse_bin<kHistBits>(val);
atomicAdd(&smem->histogram[bin], 1);
});
__syncthreads();
// Phase 1.5: reduce the histogram across the cluster
{
// 1-shot all-reduce: each rank owns kPartition consecutive bins;
// for each owned bin, gather the kClusterSize peer values (one per
// consecutive lane) via DSMEM, sum across the lanes, then scatter back.
cluster.sync();
static_assert(kHistSize == kBlockSize); // we optimize on top of this
constexpr uint32_t kPartition = kHistSize / kClusterSize;
const auto start = this_rank * kPartition;
const auto which = start + tx / kClusterSize;
const auto peer_rank = tx % kClusterSize;
const auto addr = cluster.map_shared_rank(&smem->histogram[which], peer_rank);
const auto value = *addr;
*addr = warp::reduce_sum<kClusterSize>(value);
cluster.sync();
}
// Phase 2: Find the threshold bin (uses global seq_len)
find_threshold(problem.topk, problem.seq_len, smem);
// Phase 3: Collect candidates over this rank's chunk; convert local indices
// back to global by adding chunk_start. Classify by two fp32 boundaries derived
// from the (global) threshold bin instead of recomputing the fp16 bin per
// element -- see TopKStreaming for the rationale. threshold_bin is identical
// across ranks, so v_hi/v_lo are too.
const auto topk = problem.topk;
const auto threshold_bin = smem->threshold_bin;
const float v_hi = coarse_bin_lower_bound<kHistBits>(threshold_bin + 1);
const float v_lo = coarse_bin_lower_bound<kHistBits>(threshold_bin);
const auto cur_out = is_primary ? problem.out : smem->tmp_out;
for_each_input(problem.in, local_seq_len, [&](float val, uint32_t local_idx) {
const auto idx = chunk_start + local_idx;
if (val >= v_hi) {
const auto pos = atomicAdd(&smem->count_gt, 1);
if (pos < topk) [[likely]] {
// rank 0's slots [0, a0) are final; other ranks stage raw indices and
// page-translate them after the cross-rank prefix sum is known.
cur_out[pos] = idx;
}
} else if (val >= v_lo) {
const auto count_eq = atomicAdd(&smem->count_eq, 1);
if (count_eq < kMaxNumTie) [[likely]] {
smem->tie_values[count_eq] = {val, idx};
}
}
});
// Phase 3.5: write tmp out and exit for non-primary blocks
uint32_t start_write = 0;
uint32_t num_write = 0;
if (!is_primary) {
__syncthreads();
const auto local_above_count = smem->count_gt;
const auto local_equal_count = min(smem->count_eq, kMaxNumTie);
const auto smem_0 = cluster.map_shared_rank(smem, 0);
if (tx == 0) {
const auto gt = atomicAdd(&smem_0->count_gt, local_above_count);
const auto eq = atomicAdd(&smem_0->count_eq, local_equal_count);
smem->start_gt_local = gt;
smem->start_eq_local = eq;
}
__syncthreads();
const auto start_gt_local = smem->start_gt_local;
const auto start_eq_local = smem->start_eq_local;
if (tx < local_equal_count && start_eq_local + tx < kMaxNumTie) {
smem_0->tie_values[start_eq_local + tx] = smem->tie_values[tx];
}
start_write = start_gt_local;
num_write = local_above_count;
}
cluster.sync();
if (!is_primary) {
#pragma unroll
for (uint32_t i = 0; i < kTopKItems; ++i) {
if (const auto t = tx + i * kBlockSize; t < num_write && start_write + t < topk) {
problem.emit(start_write + t, smem->tmp_out[t]);
}
}
} else {
// Phase 4: Handle ties.
const auto above_count = smem->count_gt;
const auto equal_count = smem->count_eq;
const auto remain_topk = above_count < topk ? topk - above_count : 0;
const auto tie_count = min(equal_count, kMaxNumTie);
handle_tie(smem->tie_values, problem, above_count, tie_count, remain_topk, &smem->tie_handle_smem);
}
}
};
} // namespace device::topk
python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
+39 -5
View File
@@ -208,6 +208,15 @@ SGL_DEVICE auto offset(const void* ptr, U... offset) -> const void* {
} // namespace pointer
/// PTX pragma that lets the compiler spill registers into otherwise-unused
/// shared memory instead of local memory. The radix kernels run at occupancy 2
/// (32 regs/thread) and rely on this to avoid local-memory traffic.
SGL_DEVICE void enable_smem_spilling() {
#if defined(__CUDA_ARCH__) && CUDART_VERSION >= 13000
asm(".pragma \"enable_smem_spilling\";");
#endif
}
} // namespace device
namespace host {
@@ -233,15 +242,21 @@ inline void RuntimeDeviceCheck(DebugInfo location = {}) {
* Usage:
* \code
* host::LaunchKernel(grid, block, device)
* .enable_pdl(true)
* (my_kernel, arg1, arg2);
* .enable_pdl(true)(my_kernel, arg0, arg1);
* host::LaunchKernel(grid, block, stream)
* .config({.use_pdl = true, .cluster_dim = cluster_dim})(my_kernel, arg0);
* \endcode
*
* The constructor resolves the CUDA stream from a `DLDevice` (via
* `TVMFFIEnvGetStream`) or accepts a raw `cudaStream_t`. The call
* operator launches the kernel and checks for errors.
* The constructor resolves the CUDA stream from a `DLDevice` (via `TVMFFIEnvGetStream`)
* or accepts a raw `cudaStream_t`. The call operator launches the kernel and checks for errors.
*/
struct LaunchKernel {
private:
struct KernelConfig {
bool use_pdl = false;
std::optional<dim3> cluster_dim = std::nullopt;
};
public:
explicit LaunchKernel(
dim3 grid_dim,
@@ -294,6 +309,20 @@ struct LaunchKernel {
return *this;
}
/**
* \brief Configure the kernel launch with the given options.
* \param config The kernel configuration options.
* \return A reference to this `LaunchKernel` for chaining.
* \note This is a convenience method that applies multiple configurations at once.
* We are in favor of this instead of `enable_pdl` and `enable_cluster`.
* We enforce use of designated initializers for better readability.
*/
auto config(const KernelConfig& config) -> LaunchKernel& {
if (config.use_pdl) this->enable_pdl(true);
if (config.cluster_dim) this->enable_cluster(*config.cluster_dim);
return *this;
}
template <typename T, typename... Args>
auto operator()(T&& kernel, Args&&... args) const -> void {
#ifdef USE_ROCM
@@ -310,6 +339,11 @@ struct LaunchKernel {
#endif
}
template <typename T, typename... Args>
auto launch(T&& kernel, Args&&... args) const -> void {
return (*this)(std::forward<T>(kernel), std::forward<Args>(args)...);
}
private:
static auto s_make_config( // Make a config for kernel launch
dim3 grid_dim,