[DSA] Fix top-k v2 dropping non-primary ranks' output on CUDA 13.1+ (root cause for #33835) (#34167)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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 <bool kUsePDL>
 __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<kUsePDL>();
 
-  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<kDLCUDA>();
 
@@ -304,14 +313,14 @@ struct TopKKernel {
         .with_dtype<int32_t>()
         .with_device(device)
         .verify(page_table);
-    TensorMatcher({B, kTopK})  // output, must be contiguous
+    TensorMatcher({B, K})  // output, must be contiguous
         .with_dtype<int32_t>()
         .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<int32_t>()
           .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<uint32_t>(std::countr_zero(page_size));
     const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto topk = static_cast<uint32_t>(K.unwrap());
+    RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 1024]");
     const auto params = TopKParams{
         .scores = static_cast<float*>(scores.data_ptr()),
         .seq_lens = static_cast<int32_t*>(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<kernel, kSMEM_>();

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<kPDL>(problem, &smem);
+    if (blockIdx.y != worker_rank) return;
+    Register4::forward<kPDL>(problem, &smem);
+    device::PDLWaitPrimary<kPDL>();
+    __syncthreads();
   } else if (problem.seq_len <= params.cluster_floor) {
-    if (blockIdx.y == worker_rank) Streaming::forward<kPDL>(problem, &smem);
+    if (blockIdx.y != worker_rank) return;
+    Streaming::forward<kPDL>(problem, &smem);
+    device::PDLWaitPrimary<kPDL>();
+    __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<kPDL>(peer_problem, &smem);  // write to peer's output shared memory
+    problem.out = cluster.map_shared_rank(topk_indices, worker_rank);
+    Cluster::forward<kPDL>(problem, &smem);  // write to peer's output shared memory
+    device::PDLWaitPrimary<kPDL>();
     cluster.sync();
+    if (blockIdx.y != worker_rank) return;
   }
 
-  device::PDLWaitPrimary<kPDL>();
-  __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 -----------

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<kHistBits>(threshold_bin + 1);
     const float v_lo = coarse_bin_lower_bound<kHistBits>(threshold_bin);
-    const auto cur_out = is_primary ? problem.out : smem->tmp_out;
-    for_each_input(problem.in, local_seq_len, [&](float val, uint32_t local_idx) {
-      const auto idx = chunk_start + local_idx;
-      if (val >= v_hi) {
-        const auto pos = atomicAdd(&smem->count_gt, 1);
-        if (pos < topk) [[likely]] {
-          // rank 0's slots [0, a0) are final; other ranks stage raw indices and
-          // page-translate them after the cross-rank prefix sum is known.
-          cur_out[pos] = idx;
-        }
-      } else if (val >= v_lo) {
-        const auto count_eq = atomicAdd(&smem->count_eq, 1);
-        if (count_eq < kMaxNumTie) [[likely]] {
-          smem->tie.values[count_eq] = {val, idx};
-        }
-      }
-    });
 
-    // Phase 3.5: write tmp out and exit for non-primary blocks
-    uint32_t start_write = 0;
-    uint32_t num_write = 0;
+    // 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;

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
         )