diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v1.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v1.cuh index 7d3c10e93..121a31e67 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v1.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v1.cuh @@ -11,12 +11,19 @@ namespace sglang { -#ifndef SGL_TOPK -#define SGL_TOPK 512 -#endif - -constexpr uint32_t kTopK = SGL_TOPK; -constexpr uint32_t kTopKBlockSize = SGL_TOPK; +// `topk` is a *runtime* value (<= kMaxTopK), so one module serves every k. It +// used to be baked in via -DSGL_TOPK, which built a separate module per k -- +// and because `kTopK` came from a macro rather than a template parameter, both +// modules exported identically mangled symbols. The function-local static in +// setup_kernel_smem_once() is emitted as STB_GNU_UNIQUE, which the loader +// merges across every loaded object, so whichever module was used second +// skipped its cudaFuncSetAttribute opt-in and then failed to launch with 64 KB +// of dynamic shared memory ("invalid argument"). +constexpr uint32_t kMaxTopK = 1024; +// Fixed, and deliberately not tied to `topk`: run_cumsum() and the histogram +// init below index up to RADIX + 1 == 257 threads, so a block sized after a +// small topk would silently skip part of the histogram. +constexpr uint32_t kTopKBlockSize = kMaxTopK; constexpr uint32_t kSMEM = 16 * 1024 * sizeof(uint32_t); // 64KB (bytes) struct TopKParams { @@ -28,6 +35,7 @@ struct TopKParams { const int64_t score_stride; const int64_t page_table_stride; uint32_t page_bits; + uint32_t topk; }; SGL_DEVICE uint8_t convert_to_uint8(float x) { @@ -54,14 +62,14 @@ SGL_DEVICE void naive_transform( int32_t* __restrict__ indices, int32_t* __restrict__ raw_indices, // optional: output raw abs position indices const uint32_t length, - const uint32_t page_bits) { - static_assert(kTopK <= kTopKBlockSize); + const uint32_t page_bits, + const uint32_t topk) { if (const auto tx = threadIdx.x; tx < length) { indices[tx] = page_to_indices(page_table, tx, page_bits); if (raw_indices != nullptr) { raw_indices[tx] = tx; } - } else if (kTopK == kTopKBlockSize || tx < kTopK) { + } else if (tx < topk) { indices[tx] = -1; // fill invalid indices to -1 if (raw_indices != nullptr) { raw_indices[tx] = -1; @@ -70,7 +78,8 @@ SGL_DEVICE void naive_transform( } [[maybe_unused]] -SGL_DEVICE void radix_topk(const float* __restrict__ input, int32_t* __restrict__ output, const uint32_t length) { +SGL_DEVICE void +radix_topk(const float* __restrict__ input, int32_t* __restrict__ output, const uint32_t length, const uint32_t topk) { constexpr uint32_t RADIX = 256; constexpr uint32_t BLOCK_SIZE = kTopKBlockSize; constexpr uint32_t SMEM_INPUT_SIZE = kSMEM / (2 * sizeof(int32_t)); @@ -84,7 +93,7 @@ SGL_DEVICE void radix_topk(const float* __restrict__ input, int32_t* __restrict_ extern __shared__ uint32_t s_input_idx[][kSMEM / (2 * sizeof(int32_t))]; const uint32_t tx = threadIdx.x; - uint32_t remain_topk = kTopK; + uint32_t remain_topk = topk; auto& s_histogram = _s_histogram_buf[0]; const auto run_cumsum = [&] { @@ -208,7 +217,7 @@ SGL_DEVICE void radix_topk(const float* __restrict__ input, int32_t* __restrict_ if (round == 3) { const auto pos = ::atomicAdd(&s_last_remain, -1); if (pos > 0) { - output[kTopK - pos] = idx; + output[topk - pos] = idx; } } else { const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1); @@ -231,7 +240,7 @@ template __global__ void topk_transform_kernel(const __grid_constant__ TopKParams params) { const auto &[ scores, seq_lens, page_table, page_indices, raw_indices, // pointers - score_stride, page_table_stride, page_bits // sizes + score_stride, page_table_stride, page_bits, topk // sizes ] = params; const uint32_t work_id = blockIdx.x; @@ -239,19 +248,18 @@ __global__ void topk_transform_kernel(const __grid_constant__ TopKParams params) const uint32_t seq_len = seq_lens[work_id]; const auto score_ptr = scores + work_id * score_stride; const auto page_ptr = page_table + work_id * page_table_stride; - const auto indices_ptr = page_indices + work_id * kTopK; - const auto raw_indices_ptr = raw_indices != nullptr ? raw_indices + work_id * kTopK : nullptr; + const auto indices_ptr = page_indices + work_id * topk; + const auto raw_indices_ptr = raw_indices != nullptr ? raw_indices + work_id * topk : nullptr; device::PDLWaitPrimary(); - if (seq_len <= kTopK) { - naive_transform(score_ptr, page_ptr, indices_ptr, raw_indices_ptr, seq_len, page_bits); + if (seq_len <= topk) { + naive_transform(score_ptr, page_ptr, indices_ptr, raw_indices_ptr, seq_len, page_bits, topk); } else { - __shared__ int32_t s_topk_indices[kTopK]; - radix_topk(score_ptr, s_topk_indices, seq_len); - static_assert(kTopK <= kTopKBlockSize); + __shared__ int32_t s_topk_indices[kMaxTopK]; + radix_topk(score_ptr, s_topk_indices, seq_len, topk); const auto tx = threadIdx.x; - if (kTopK == kTopKBlockSize || tx < kTopK) { + if (tx < topk) { indices_ptr[tx] = page_to_indices(page_ptr, s_topk_indices[tx], page_bits); if (raw_indices_ptr != nullptr) { raw_indices_ptr[tx] = s_topk_indices[tx]; @@ -287,6 +295,7 @@ struct TopKKernel { auto B = SymbolicSize{"batch_size"}; auto S = SymbolicSize{"score_stride"}; auto P = SymbolicSize{"page_table_stride"}; + auto K = SymbolicSize{"topk"}; auto device = SymbolicDevice{}; device.set_options(); @@ -304,14 +313,14 @@ struct TopKKernel { .with_dtype() .with_device(device) .verify(page_table); - TensorMatcher({B, kTopK}) // output, must be contiguous + TensorMatcher({B, K}) // output, must be contiguous .with_dtype() .with_device(device) .verify(page_indices); int32_t* raw_indices_ptr = nullptr; if (raw_indices.has_value()) { - TensorMatcher({B, kTopK}) // optional raw indices output, must be contiguous + TensorMatcher({B, K}) // optional raw indices output, must be contiguous .with_dtype() .with_device(device) .verify(raw_indices.value()); @@ -321,6 +330,8 @@ struct TopKKernel { RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2"); const auto page_bits = static_cast(std::countr_zero(page_size)); const auto batch_size = static_cast(B.unwrap()); + const auto topk = static_cast(K.unwrap()); + RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 1024]"); const auto params = TopKParams{ .scores = static_cast(scores.data_ptr()), .seq_lens = static_cast(seq_lens.data_ptr()), @@ -330,6 +341,7 @@ struct TopKKernel { .score_stride = S.unwrap(), .page_table_stride = P.unwrap(), .page_bits = page_bits, + .topk = topk, }; constexpr auto kSMEM_ = kSMEM + sizeof(int32_t); // align up a little setup_kernel_smem_once(); diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh index 0de65a18d..1f8928b2a 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh @@ -215,24 +215,32 @@ CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLau // for small batch, we will fuse in the cluster case if (problem.seq_len <= kReg4MaxSeqLen) { - if (blockIdx.y == worker_rank) Register4::forward(problem, &smem); + if (blockIdx.y != worker_rank) return; + Register4::forward(problem, &smem); + device::PDLWaitPrimary(); + __syncthreads(); } else if (problem.seq_len <= params.cluster_floor) { - if (blockIdx.y == worker_rank) Streaming::forward(problem, &smem); + if (blockIdx.y != worker_rank) return; + Streaming::forward(problem, &smem); + device::PDLWaitPrimary(); + __syncthreads(); } else { auto cluster = cooperative_groups::this_cluster(); - // The mapped alias stays in a copy: the elected rank reads the very same - // bytes back through `topk_indices` below, and letting a shared::cluster - // address reach the `problem.out` that problem_transform loads makes cicc - // segfault on CUDA 13.x (issue #32830). - auto peer_problem = problem; - peer_problem.out = cluster.map_shared_rank(topk_indices, worker_rank); - Cluster::forward(peer_problem, &smem); // write to peer's output shared memory + problem.out = cluster.map_shared_rank(topk_indices, worker_rank); + Cluster::forward(problem, &smem); // write to peer's output shared memory + device::PDLWaitPrimary(); cluster.sync(); + if (blockIdx.y != worker_rank) return; } - device::PDLWaitPrimary(); - __syncthreads(); - if (blockIdx.y == worker_rank) problem_transform(problem, params.get_output_ptr(blockIdx.x)); + // Only the elected worker reaches here, and it mapped `topk_indices` to + // itself, so `problem.out` is this block's own buffer. Stating that keeps the + // shared::cluster address out of the load problem_transform issues -- which is + // load-bearing, not an optimization: without it cicc segfaults on CUDA 13.1+ + // for sm_90a (issue #32830, previously worked around by copying `problem` in + // #32910). Verified: dropping this line reproduces the crash on 13.1/13.2/13.3. + __builtin_assume(problem.out == topk_indices); + problem_transform(problem, params.get_output_ptr(blockIdx.x)); } // --- Plan: choose cluster_threshold from the seq_len distribution ----------- diff --git a/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh index d2672c14b..9a0dd2698 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh @@ -776,28 +776,35 @@ struct TopKCluster : TopKRadixBase<10> { 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; + // Phase 3: collect candidates. The primary scatters straight into + // `problem.out`, the others stage into block-local `smem->tmp_out`. + // + // DO NOT merge these two loops back into one by selecting the destination + // first (`cur_out = is_primary ? problem.out : smem->tmp_out`). `problem.out` + // can be a shared::cluster (DSMEM) alias of the elected rank's buffer while + // `tmp_out` is shared::cta; merging them into a single pointer variable makes + // cicc 13.1+ mis-lower the block-local arm on sm_90a and *silently drop every + // non-primary rank's staged output* -- `tmp_out` stays zero, and phase 3.5 + // then faithfully copies zeros to correct DSMEM addresses. The result is a + // top-k output where only the primary's slots and the handle_tie tail are + // valid, which downstream sparse attention dereferences as garbage KV indices. if (!is_primary) { + // stage to tmp_out first before writing to global/DSMEM + 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]] { + smem->tmp_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}; + } + } + }); __syncthreads(); const auto local_above_count = smem->count_gt; const auto local_equal_count = min(smem->count_eq, kMaxNumTie); @@ -818,12 +825,11 @@ struct TopKCluster : TopKRadixBase<10> { smem_0->tie.values[start_eq_local + t] = smem->tie.values[t]; } } - start_write = start_gt_local; - num_write = local_above_count; - } - cluster.sync(); - if (!is_primary) { + cluster.sync(); + + const auto start_write = start_gt_local; + const auto num_write = local_above_count; #pragma unroll for (uint32_t i = 0; i < kTopKItems; ++i) { if (const auto t = tx + i * kBlockSize; t < num_write && start_write + t < topk) { @@ -831,6 +837,23 @@ struct TopKCluster : TopKRadixBase<10> { } } } else { + 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]] { + 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}; + } + } + }); + + cluster.sync(); + // Phase 4: Handle ties. const auto above_count = smem->count_gt; const auto equal_count = smem->count_eq; diff --git a/python/sglang/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index 53d7ebd52..502c49cf3 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/python/sglang/kernels/ops/attention/dsv4/topk.py @@ -16,15 +16,18 @@ from .utils import make_name @cache_once -def _jit_topk_v1_module(topk: int): +def _jit_topk_v1_module(): + # topk (<= 1024) is a runtime argument, not a compile-time constant, so a + # single module serves every k. Baking it in via -DSGL_TOPK used to build one + # module per k, and since the macro fed a `constexpr` rather than a template + # parameter every module exported identically mangled symbols -- see the + # comment in topk_v1.cuh for how that broke the second module's launch. args = make_cpp_args(is_arch_support_pdl()) - assert topk in (512, 1024), "Only support topk=512 or 1024" return load_jit( - make_name(f"topk_v1_{topk}"), + make_name("topk_v1"), *args, cuda_files=["deepseek_v4/topk_v1.cuh"], cuda_wrappers=[("topk_transform", f"TopKKernel<{args}>::transform")], - extra_cuda_cflags=[f"-DSGL_TOPK={topk}"], ) @@ -55,7 +58,7 @@ def topk_transform_512( scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices ) else: - module = _jit_topk_v1_module(out_page_indices.shape[1]) + module = _jit_topk_v1_module() module.topk_transform( scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices )