From c016c6f355f72854fccfd54b316f6de880e127b1 Mon Sep 17 00:00:00 2001 From: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Date: Mon, 6 Jul 2026 11:23:39 +0800 Subject: [PATCH] [JIT Kernel] DeepSeek-V4 DSA indexer: faster top-k + page-table transform (runtime k <= 2048) (#26788) Co-authored-by: Claude --- .../jit_kernel/csrc/deepseek_v4/topk_v2.cuh | 615 +++++++------- python/sglang/jit_kernel/dsv4/topk.py | 25 +- .../sgl_kernel/deepseek_v4/topk/cluster.cuh | 257 ------ .../sgl_kernel/deepseek_v4/topk/common.cuh | 176 ---- .../sgl_kernel/deepseek_v4/topk/ptx.cuh | 54 -- .../sgl_kernel/deepseek_v4/topk/register.cuh | 302 ------- .../sgl_kernel/deepseek_v4/topk/streaming.cuh | 213 ----- .../sgl_kernel/deepseek_v4/topk_impl.cuh | 752 ++++++++++++++++++ .../jit_kernel/include/sgl_kernel/utils.cuh | 44 +- test/registered/jit/benchmark/bench_topk.py | 90 +++ .../jit/deepseek_v4/test_topk_v2.py | 299 +++++++ 11 files changed, 1483 insertions(+), 1344 deletions(-) delete mode 100644 python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/cluster.cuh delete mode 100644 python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/common.cuh delete mode 100644 python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/ptx.cuh delete mode 100644 python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/register.cuh delete mode 100644 python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/streaming.cuh create mode 100644 python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh create mode 100644 test/registered/jit/benchmark/bench_topk.py create mode 100644 test/registered/jit/deepseek_v4/test_topk_v2.py diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh index 8c4a52657..2fb2b533a 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh @@ -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 #include -#include #include -#include -#include -#include -#include -#include +#include #include #include -#include -#include +#include #include #include 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 -void setup_kernel_smem_once(host::DebugInfo where = {}) { - [[maybe_unused]] - static const auto result = [] { - const auto fptr = std::bit_cast(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; -using Medium = impl::StreamingTopK; -using Small = impl::RegisterTopK; +#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(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(topk); + } + SGL_DEVICE TopKProblem problem(uint32_t batch_id, uint32_t seq_len) const { + const auto k = static_cast(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(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(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 +CLUSTER_TOPK_KERNEL void topk_persistent_cluster_kernel(const __grid_constant__ TopKLaunchParams params) { + device::enable_smem_spilling(); + __shared__ impl::MaxSmem smem; + const uint32_t num_cluster_items = params.global().num_cluster_items; + device::PDLWaitPrimary(); + device::PDLTriggerSecondary(); +#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(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 +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 +SGL_DEVICE void trivial_transform(const TopKProblem& problem) { + device::PDLWaitPrimary(); + device::PDLTriggerSecondary(); + for_each_item(problem.topk, [&](uint32_t tx, uint32_t) { + problem.transform_output(tx, tx < problem.seq_len ? static_cast(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 +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::max(); + __shared__ impl::MaxSmem smem; + if (problem.seq_len <= problem.topk) return trivial_transform(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(problem, &smem); + } else if constexpr (kLevel == 1) { + __builtin_assume(problem.seq_len <= kReg4MaxSeqLen); + Register4::forward(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(problem, &smem); + } else if (problem.seq_len <= cluster_threshold) { + Streaming::forward(problem, &smem); + } else { // cluster path do nothing here + problem.out = params.get_output_ptr(blockIdx.x); + } + device::PDLWaitPrimary(); + } + + // 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(); + __syncthreads(); + problem_transform(problem, params.get_output_ptr(blockIdx.x)); +} + +template +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 smem; + if (problem.seq_len <= problem.topk) return trivial_transform(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(problem, &smem); + } else if (problem.seq_len <= params.cluster_floor) { + if (blockIdx.y == worker_rank) Streaming::forward(problem, &smem); + } else { + auto cluster = cooperative_groups::this_cluster(); + problem.out = cluster.map_shared_rank(topk_indices, worker_rank); + Cluster::forward(problem, &smem); // write to peer's output shared memory + cluster.sync(); + } + + device::PDLWaitPrimary(); + __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(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(); - 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(); - device::PDLTriggerSecondary(); - - 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(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(); - 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(); - 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(); - TensorMatcher({B}) // + TensorMatcher({B}) // seq_lens .with_dtype() .with_device(device_) .verify(seq_lens); - TensorMatcher({Bp1, 4}) // + TensorMatcher({Bp1, 2}) // metadata: [0]=GlobalMetadata, [1..N]=PlanItem(batch_id, seq_len) .with_dtype() .with_device(device_) .verify(metadata); const auto batch_size = static_cast(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(seq_lens.data_ptr()), - static_cast(metadata.data_ptr()), + topk_plan, + static_cast(seq_lens.data_ptr()), + static_cast(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 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(); - TensorMatcher({B, L}) // + TensorMatcher({B, L}) // score .with_strides({S, 1}) .with_dtype() .with_device(device_) .verify(scores); - TensorMatcher({B}) // + TensorMatcher({B}) // seq_lens .with_dtype() .with_device(device_) .verify(seq_lens); - TensorMatcher({B, -1}) // + TensorMatcher({B, -1}) // page_table .with_strides({P, 1}) .with_dtype() .with_device(device_) .verify(page_table); - TensorMatcher({B, K}) // + TensorMatcher({B, K}) // page_indices .with_dtype() .with_device(device_) .verify(page_indices); - TensorMatcher({B, D}) // - .with_strides({W, 1}) - .with_dtype() - .with_device(device_) - .verify(workspace); - TensorMatcher({Bp1, 4}) // + TensorMatcher({Bp1, 2}) // metadata: [0]=GlobalMetadata, [1..N]=PlanItem(batch_id, seq_len) .with_dtype() .with_device(device_) .verify(metadata); + int32_t* raw_indices_ptr = nullptr; + if (raw_indices.has_value()) { + TensorMatcher({B, K}).with_dtype().with_device(device_).verify(raw_indices.value()); + raw_indices_ptr = static_cast(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(K.unwrap()); + RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 2048]"); + const auto page_bits = static_cast(std::countr_zero(page_size)); const auto batch_size = static_cast(B.unwrap()); const auto max_seq_len = static_cast(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(seq_lens.data_ptr()), - .scores = static_cast(scores.data_ptr()), - .page_table = static_cast(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(scores.data_ptr()), + .seq_lens = static_cast(seq_lens.data_ptr()), + .page_table = static_cast(page_table.data_ptr()), .page_indices = static_cast(page_indices.data_ptr()), + .raw_indices = raw_indices_ptr, + .metadata = static_cast(metadata.data_ptr()), .score_stride = S.unwrap(), .page_table_stride = P.unwrap(), - .workspace = static_cast(workspace.data_ptr()), - .metadata = static_cast(metadata.data_ptr()), - .workspace_stride = W.unwrap() * static_cast(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(); - 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(); - 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, params); } else { - // stage 1 + stage 2 - constexpr auto kernel_stage_1 = topk_combine_preprocess; - setup_kernel_smem_once(); - 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(); - 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, params); + LaunchKernel(batch_size, kBlockSize, device) + .config({.use_pdl = kUsePDL}) + .launch(topk_main_kernel, params); } + } else if (max_seq_len <= kReg2MaxSeqLen) { + LaunchKernel(batch_size, kBlockSize, device) + .config({.use_pdl = kUsePDL}) + .launch(topk_main_kernel, params); + } else if (max_seq_len <= kReg4MaxSeqLen) { + LaunchKernel(batch_size, kBlockSize, device) + .config({.use_pdl = kUsePDL}) + .launch(topk_main_kernel, params); + } else { + LaunchKernel(batch_size, kBlockSize, device) + .config({.use_pdl = kUsePDL}) + .launch(topk_main_kernel, params); } } }; diff --git a/python/sglang/jit_kernel/dsv4/topk.py b/python/sglang/jit_kernel/dsv4/topk.py index a27245186..273d2434c 100644 --- a/python/sglang/jit_kernel/dsv4/topk.py +++ b/python/sglang/jit_kernel/dsv4/topk.py @@ -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, ) diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/cluster.cuh b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/cluster.cuh deleted file mode 100644 index e58214c95..000000000 --- a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/cluster.cuh +++ /dev/null @@ -1,257 +0,0 @@ -#pragma once -#include -#include -#include - -#include "common.cuh" -#include "ptx.cuh" -#include -#include - -namespace device::top512 { - -template -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); - 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); - 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); - 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(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(*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(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); - 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(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(_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(_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 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/common.cuh b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/common.cuh deleted file mode 100644 index d553032d7..000000000 --- a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/common.cuh +++ /dev/null @@ -1,176 +0,0 @@ -#pragma once -#include -#include -#include -#include - -#include - -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; - -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 -SGL_DEVICE uint32_t extract_coarse_bin(float x) { - static_assert(0 < kBits && kBits < 15); - const auto hx = cast(x); - const uint16_t bits = *reinterpret_cast(&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(_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 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/ptx.cuh b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/ptx.cuh deleted file mode 100644 index 73eef555f..000000000 --- a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/ptx.cuh +++ /dev/null @@ -1,54 +0,0 @@ -#pragma once -#include - -#include - -#include - -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 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/register.cuh b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/register.cuh deleted file mode 100644 index 77d7361ee..000000000 --- a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/register.cuh +++ /dev/null @@ -1,302 +0,0 @@ -#pragma once - -#include -#include -#include - -#include "common.cuh" -#include "ptx.cuh" -#include -#include - -namespace device::top512 { - -template -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; - 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 - 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); - 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(); - - // 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(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(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(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(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(__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(__popc(__ballot_sync(0xFFFFFFFF, pred_0))); - const auto rank_1 = static_cast(__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(__popc(__ballot_sync(0xFFFFFFFF, pred_0))); - const auto rank_1 = static_cast(__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 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/streaming.cuh b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/streaming.cuh deleted file mode 100644 index 4462b89a1..000000000 --- a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/streaming.cuh +++ /dev/null @@ -1,213 +0,0 @@ -#pragma once - -#include -#include -#include - -#include "common.cuh" -#include "ptx.cuh" -#include -#include - -namespace device::top512 { - -template -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; - - 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 - 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 - 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(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(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(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); - 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(scores, length, 0, nullptr, smem); - - // Phase B: locate threshold bin & re-init barriers - find_threshold(length, smem); - - // Phase C: scatter pass. - stream_pass(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); - 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 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh new file mode 100644 index 000000000..528d8c6e4 --- /dev/null +++ b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh @@ -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 +#include +#include +#include + +#include +#include +#include +#include + +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 dependency). +template +constexpr T ct_max(T a) { + return a; +} +template +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 +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 +SGL_DEVICE uint32_t extract_coarse_bin(float x) { + static_assert(0 < kBits && kBits < 15); + const auto hx = cast(x); + const uint16_t bits = *reinterpret_cast(&hx); + const uint16_t key = (bits & 0x8000) ? ~bits : bits | 0x8000; + return key >> (16 - kBits); +} + +// Smallest fp32 value `v` for which `extract_coarse_bin(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 +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(okey); + const uint16_t hb = (ob & 0x8000) ? static_cast(ob ^ 0x8000) : static_cast(~ob); + return cast(*reinterpret_cast(&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(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 +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; + + struct Smem { + using kHistVec = AlignedVector; + 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 + 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 +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 + SGL_DEVICE static void forward(const TopKProblem problem, void* _smem) { + const auto tx = threadIdx.x; + const auto smem = static_cast(_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(); + + // A vector `vi` is fully in bounds iff vi < num_full; only full vectors are + // vector-loaded (16B aligned, never straddling seq_len). The = 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(local_vecs[i][j])], 1); + } + if (tx >= kBlockSize - tail) { + const uint32_t idx = tail_start + tx - (kBlockSize - tail); + atomicAdd(&smem->histogram[extract_coarse_bin(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(threshold_bin + 1); + const auto v_lo = coarse_bin_lower_bound(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::max(); + + template + SGL_DEVICE static void forward(const TopKProblem problem, void* _smem) { + const auto tx = threadIdx.x; + const auto smem = static_cast(_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(); + + // Phase 1: Load and build histogram + for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t) { + const auto bin = extract_coarse_bin(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(threshold_bin + 1); + const float v_lo = coarse_bin_lower_bound(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 +struct TopKCluster : TopKRadixBase<10> { + public: + static constexpr uint32_t kClusterSize = kClusterSize_; + static constexpr uint32_t kMaxSeqLen = std::numeric_limits::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 + SGL_DEVICE static void forward(TopKProblem problem, void* _smem) { + const auto tx = threadIdx.x; + const auto smem = static_cast(_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(); + + // 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(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(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(threshold_bin + 1); + const float v_lo = coarse_bin_lower_bound(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 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index 2dd6f3dc9..25cc89aab 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -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 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 auto operator()(T&& kernel, Args&&... args) const -> void { #ifdef USE_ROCM @@ -310,6 +339,11 @@ struct LaunchKernel { #endif } + template + auto launch(T&& kernel, Args&&... args) const -> void { + return (*this)(std::forward(kernel), std::forward(args)...); + } + private: static auto s_make_config( // Make a config for kernel launch dim3 grid_dim, diff --git a/test/registered/jit/benchmark/bench_topk.py b/test/registered/jit/benchmark/bench_topk.py new file mode 100644 index 000000000..c5d53fff2 --- /dev/null +++ b/test/registered/jit/benchmark/bench_topk.py @@ -0,0 +1,90 @@ +import torch + +from sglang.jit_kernel.benchmark import marker +from sglang.jit_kernel.dsv4.topk import ( + plan_topk_v2, + topk_transform_512, + topk_transform_512_v2, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=120, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + +# Compressed page size used by the DSA indexer (real value is 256 // 4 = 64). +PAGE_SIZE = 64 + + +def _make_inputs(batch_size: int, seq_len: int, k: int): + torch.random.manual_seed(42) + scores = torch.randn(batch_size, seq_len, dtype=torch.float32, device="cuda") + seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") + num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE + page_table = ( + torch.arange(num_pages, dtype=torch.int32, device="cuda") + .unsqueeze(0) + .expand(batch_size, -1) + .contiguous() + ) + out = torch.empty(batch_size, k, dtype=torch.int32, device="cuda") + return scores, seq_lens, page_table, out + + +def _make_p1_table(batch_size: int, seq_len: int): + # flashinfer / torch do a per-token (page_size=1) gather, so they need a + # (batch, seq) table (one entry per position) rather than the page-size-64 one. + src_page_table = ( + torch.arange(seq_len, dtype=torch.int32, device="cuda") + .unsqueeze(0) + .expand(batch_size, -1) + .contiguous() + ) + lengths = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") + return src_page_table, lengths + + +def _build_fn(provider: str, batch_size: int, seq_len: int, k: int): + scores, seq_lens, page_table, out = _make_inputs(batch_size, seq_len, k) + N = PAGE_SIZE + + def fn(scores, seq_lens, page_table): + if provider == "jit_v1": + topk_transform_512(scores, seq_lens, page_table, out, N) + return out + elif provider == "jit_v2": + topk_transform_512_v2(scores, seq_lens, page_table, out, N, metadata) + return out + elif provider == "flashinfer": + from flashinfer import top_k_page_table_transform + + return top_k_page_table_transform(scores, page_table, seq_lens, k) + elif provider == "torch": + idx = scores.topk(k, dim=-1).indices # (batch, k) int64 + return torch.gather(page_table, 1, idx) + else: + raise ValueError(f"unknown provider {provider}") + + if provider in ("flashinfer", "torch"): + page_table, seq_lens = _make_p1_table(batch_size, seq_len) + if provider == "jit_v2": + metadata = plan_topk_v2(seq_lens) + return fn, (scores, seq_lens, page_table) + + +@marker.parametrize("k", [512, 1024, 2048], [512]) +@marker.parametrize("seq_len", [2**x for x in range(10, 19)], [4096, 65536]) +@marker.parametrize("batch_size", [2**x for x in range(13)], [1, 128, 1024]) +@marker.benchmark("provider", ["jit_v1", "jit_v2", "flashinfer", "torch"]) +def benchmark(seq_len: int, batch_size: int, k: int, provider: str): + if k > seq_len: + marker.skip("k cannot be larger than seq_len") + if k == 2048 and provider == "jit_v1": + marker.skip("jit_v1 does not support k=2048") + + fn, input_args = _build_fn(provider, batch_size, seq_len, k) + return marker.do_bench(fn, input_args=input_args, memory_args=input_args[:2]) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/jit/deepseek_v4/test_topk_v2.py b/test/registered/jit/deepseek_v4/test_topk_v2.py new file mode 100644 index 000000000..be6c3ca3d --- /dev/null +++ b/test/registered/jit/deepseek_v4/test_topk_v2.py @@ -0,0 +1,299 @@ +"""Correctness tests for the DeepSeek-V4 (DSA indexer) JIT top-k transform v2. + +The v2 kernel selects the per-row top-k of ``scores`` (ragged ``seq_lens``) and +writes the page-table transform of the selected raw indices into the output. We +validate against ``torch.topk`` with a small tolerance for boundary ties (the +fp16 coarse histogram can swap elements of equal score). + +Coverage is organized around the kernel's dispatch so every template and its +boundaries are exercised: + + template per-row seq reached when + -------- ---------- ------------ + trivial seq <= k + Register2 k < seq <= 8192 max_seq <= 8192 (level 0) + Register4 8192 < seq <= 16384 max_seq <= 16384 (level 1) + Streaming 16384 < seq <= floor max_seq > 16384, non-cluster (level 2) + Cluster seq > floor(=65536) max_seq > floor and batch <= 128 + +and two cluster dispatch shapes: the fused small-batch kernel (batch <= 30) and +the persistent-pool + main kernel (30 < batch <= 128). Boundary seq lengths +(8192/8193, 16384/16385, 65535/65536/65537) and batch sizes (30/31, 128/129) are +included explicitly, across k in {512,1024,2048} and identity/perm page tables. +""" + +from __future__ import annotations + +import sys + +import pytest +import torch + +from sglang.jit_kernel.dsv4.topk import plan_topk_v2, topk_transform_512_v2 +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +PAGE_SIZE = 64 # c4 page size = 256 // 4 +PAGE_BITS = PAGE_SIZE.bit_length() - 1 +PAGE_MASK = PAGE_SIZE - 1 +MAX_PERMIT_ERROR = 5 +FLOOR = 65536 # kClusterFloor + +# (batch, seq) chosen to land on each template and each dispatch boundary. +FIXED_CONFIGS = [ + # --- trivial (seq <= k) --- + (8, 256), # trivial for every k + (16, 1024), # trivial for k>=1024 + # --- Register2 (level 0: max_seq <= 8192) --- + (8, 4096), + (8, 8192), # reg2 upper boundary + (128, 8192), + (300, 8192), # batch > 128, still level 0 + # --- Register4 (level 1: 8192 < max_seq <= 16384) --- + (8, 8193), # just over reg2 + (64, 16384), # reg4 upper boundary + (256, 16384), # batch > 128 + # --- Streaming (level 2: max_seq > 16384, non-cluster) --- + (8, 16385), # just over reg4 (small batch, seq < floor => non-cluster) + (4, 32768), + (16, 65535), # just under floor + (4, 65536), # at floor (seq == floor => non-cluster) + (100, 65536), + # --- Cluster, fused small-batch kernel (batch <= 30, max_seq > floor) --- + (1, 65537), # single row just over floor + (2, 131072), + (8, 98304), + (30, 131072), # batch == pool boundary + # --- Cluster, persistent pool + main kernel (30 < batch <= 128) --- + (31, 131072), # just over small-batch + (40, 262144), # N > pool of 30 => round-robin + (64, 196608), + (128, 131072), # cluster batch upper boundary + # --- batch > 128 => non-cluster streaming even at long ctx --- + (129, 131072), + (200, 262144), +] + + +def _assert_topk_close(scores_cpu, ref_raw, our_raw, bs, seq_lens, k): + """Set-compare our top-k raw indices vs torch's, tolerating equal-score ties.""" + bad = 0 + for i in range(bs): + L = int(seq_lens[i]) + ref, our = set(ref_raw[i]), set(our_raw[i]) + more, less = our - ref, ref - our + if more or less: + mv = sorted(scores_cpu[i, list(more)].tolist()) + lv = sorted(scores_cpu[i, list(less)].tolist()) + if mv != lv: # not merely a tie swap -> genuine error + bad += len(more) + print( + f"b={i} L={L} k={k}: more={list(more)[:4]} less={list(less)[:4]} mv={mv[:3]} lv={lv[:3]}" + ) + assert len(our) == min( + k, L + ), f"b={i} L={L} k={k}: {len(our)} valid != {min(k, L)}" + assert bad <= MAX_PERMIT_ERROR, f"{bad=} > {MAX_PERMIT_ERROR}" + + +def _make_page_table(batch, num_pages, mode, device, per_row=False): + if mode == "identity": + pt = torch.arange(num_pages, dtype=torch.int32, device=device) + full = pt.unsqueeze(0).expand(batch, -1).contiguous() + inv = pt.unsqueeze(0).expand(batch, -1).cpu() + return full, inv + # permutation (optionally a distinct permutation per row) + rows = batch if per_row else 1 + full = torch.stack( + [torch.randperm(num_pages, device=device) for _ in range(rows)] + ).to(torch.int32) + inv = torch.empty_like(full) + ar = torch.arange(num_pages, dtype=torch.int32, device=device) + for r in range(rows): + inv[r, full[r].long()] = ar + if not per_row: + full = full.expand(batch, -1).contiguous() + inv = inv.expand(batch, -1) + return full, inv.cpu() + + +def _invert(out_row, inv_row): + """Undo page_to_indices for one row's page indices (drop -1 padding).""" + return [ + (int(inv_row[v >> PAGE_BITS]) << PAGE_BITS) | (v & PAGE_MASK) + for v in out_row + if v != -1 + ] + + +def _reference(scores, seq_lens, k): + """torch.topk reference indices per row (trivial rows -> all positions).""" + ref = [] + for i in range(scores.shape[0]): + L = int(seq_lens[i]) + if L <= k: + ref.append(list(range(L))) + else: + ref.append( + torch.topk(scores[i, :L], k, sorted=False).indices.cpu().tolist() + ) + return ref + + +def _run(scores, seq_lens, page_table, inv_cpu, k): + batch = scores.shape[0] + out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device) + metadata = plan_topk_v2(seq_lens) + topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata) + torch.cuda.synchronize() + out_cpu = out.cpu().tolist() + return [_invert(out_cpu[i], inv_cpu[i]) for i in range(batch)] + + +def _run_raw(scores, seq_lens, page_table, k): + """Run the kernel and return its optional raw (pre-transform) top-k index + output per row, dropping -1 padding -- the selected positions themselves, + NOT the page-table transform of them.""" + batch = scores.shape[0] + out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device) + raw = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device) + metadata = plan_topk_v2(seq_lens) + topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata, raw) + torch.cuda.synchronize() + raw_cpu = raw.cpu().tolist() + return [[v for v in raw_cpu[i] if v != -1] for i in range(batch)] + + +@pytest.mark.parametrize("page_mode", ["identity", "perm"]) +@pytest.mark.parametrize("k", [512, 1024, 2048]) +@pytest.mark.parametrize("batch,seq", FIXED_CONFIGS) +@torch.inference_mode() +def test_topk_v2(batch: int, seq: int, k: int, page_mode: str) -> None: + torch.manual_seed(batch * 100003 + seq * 7 + k) + device = "cuda" + # Pad the row stride to a multiple of 4 (16-byte vectorized load) while keeping + # the exact seq_len -- this also exercises the scalar-tail path for odd seq. + width = (seq + 3) & ~3 + scores = torch.randn(batch, width, dtype=torch.float32, device=device)[:, :seq] + seq_lens = torch.full((batch,), seq, dtype=torch.int32, device=device) + num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE + page_table, inv_cpu = _make_page_table(batch, num_pages, page_mode, device) + + our_raw = _run(scores, seq_lens, page_table, inv_cpu, k) + ref_raw = _reference(scores, seq_lens, k) + _assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, seq_lens.cpu(), k) + + +@pytest.mark.parametrize("k", [512, 1024, 2048]) +@pytest.mark.parametrize( + "batch,shape", + [ + (20, "small_batch"), # fused small-batch kernel (<= pool of 30) + (64, "persistent"), # persistent pool + main kernel + (128, "persistent"), # cluster batch boundary + ], +) +@pytest.mark.parametrize("per_row_pt", [False, True]) +@torch.inference_mode() +def test_topk_v2_ragged(batch: int, shape: str, k: int, per_row_pt: bool) -> None: + """Ragged lengths spanning trivial..cluster in one launch, both dispatch shapes. + + ``per_row_pt`` gives each row a distinct page-table permutation, exercising + the per-batch page_table indexing (batch_id stride) rather than a shared one. + """ + torch.manual_seed(7777 + batch + k + int(per_row_pt)) + device = "cuda" + seq = 262144 + scores = torch.randn(batch, seq, dtype=torch.float32, device=device) + # span every path; guarantee at least one > floor row so cluster dispatch fires + buckets = [max(1, k // 2), k, 4096, 12000, 40000, 65536, 98304, 262144] + g = torch.Generator(device="cpu").manual_seed(batch + k) + lengths = torch.tensor( + [ + buckets[int(torch.randint(0, len(buckets), (1,), generator=g))] + for _ in range(batch) + ], + dtype=torch.int32, + device=device, + ) + lengths[0] = max(1, k // 2) # a trivial row + lengths[1] = 262144 # a long (cluster) row + num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE + page_table, inv_cpu = _make_page_table( + batch, num_pages, "perm", device, per_row=per_row_pt + ) + + our_raw = _run(scores, lengths, page_table, inv_cpu, k) + ref_raw = _reference(scores, lengths, k) + _assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, lengths.cpu(), k) + + +@pytest.mark.parametrize("page_mode", ["identity", "perm"]) +@pytest.mark.parametrize( + "batch,seq", + [ + (8, 256), # trivial + (8, 4096), # register + (4, 131072), # fused small-batch cluster + (64, 131072), # persistent cluster + main<3> epilogue + (256, 131072), # non-cluster streaming + ], +) +@torch.inference_mode() +def test_topk_v2_raw_indices(batch: int, seq: int, page_mode: str) -> None: + """The optional raw-index output must be the pre-transform position of each + transformed output slot (out[j] == page_to_indices(raw[j])), and -1 aligns.""" + k = 512 + torch.manual_seed(batch * 131 + seq) + device = "cuda" + width = (seq + 3) & ~3 + scores = torch.randn(batch, width, dtype=torch.float32, device=device)[:, :seq] + seq_lens = torch.full((batch,), seq, dtype=torch.int32, device=device) + num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE + page_table, inv_cpu = _make_page_table(batch, num_pages, page_mode, device) + out = torch.full((batch, k), -1, dtype=torch.int32, device=device) + raw = torch.full((batch, k), -1, dtype=torch.int32, device=device) + + metadata = plan_topk_v2(seq_lens) + topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata, raw) + torch.cuda.synchronize() + + out_cpu, raw_cpu = out.cpu().tolist(), raw.cpu().tolist() + for i in range(batch): + for j in range(k): + o, r = out_cpu[i][j], raw_cpu[i][j] + if o == -1: + assert r == -1, f"b={i} j={j}: out=-1 but raw={r}" + else: + inv = (int(inv_cpu[i][o >> PAGE_BITS]) << PAGE_BITS) | (o & PAGE_MASK) + assert r == inv, f"b={i} j={j}: raw={r} != inverse(out)={inv}" + + +@pytest.mark.parametrize("k", [512, 1024, 2048]) +@pytest.mark.parametrize("batch,seq", FIXED_CONFIGS) +@torch.inference_mode() +def test_topk_v2_output_indices(batch: int, seq: int, k: int) -> None: + """Validate the raw (pre-transform) index output DIRECTLY against torch.topk. + + Unlike ``test_topk_v2`` -- which checks the page-transformed output and inverts + it through the page table -- this exercises the selected indices themselves, so + it isolates the top-k selection from the page-table transform. A permuted page + table is used so raw != out, catching any bug that leaks transformed page + indices into the raw buffer. Covers every dispatch template/boundary. + """ + torch.manual_seed(batch * 100003 + seq * 7 + k + 1) + device = "cuda" + width = (seq + 3) & ~3 + scores = torch.randn(batch, width, dtype=torch.float32, device=device)[:, :seq] + seq_lens = torch.full((batch,), seq, dtype=torch.int32, device=device) + num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE + page_table, _ = _make_page_table(batch, num_pages, "perm", device) + + our_raw = _run_raw(scores, seq_lens, page_table, k) + ref_raw = _reference(scores, seq_lens, k) + _assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, seq_lens.cpu(), k) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))