[DSA] Make the fused top-k exact on an overflowing threshold bin

python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh

@@ -45,7 +45,25 @@ constexpr uint32_t kMaxTopK = impl::TopKConfig::kMaxTopK;
 constexpr uint32_t kReg2MaxSeqLen = Register2::kMaxSeqLen;  // 8192
 constexpr uint32_t kReg4MaxSeqLen = Register4::kMaxSeqLen;  // 16384
 
+#ifdef USE_ROCM
+// HIP reads the second __launch_bounds__ argument as waves per SIMD, CUDA as
+// blocks per SM, so kOccupancy asks ROCm for 2 waves/SIMD -- below the 4 a
+// 1024-thread block already forces. Do NOT translate kOccupancy into this:
+// "kOccupancy blocks per CU" is 8 waves/SIMD, which caps the allocator at
+// 512/8 = 64 VGPRs and spills, and a second resident block is unreachable with
+// a grid of batch_size. wave64 is a literal because __AMDGCN_WAVEFRONT_SIZE__
+// is gone in ROCm 7 and warpSize is not constexpr.
+inline constexpr uint32_t kSimdsPerCu = 4;
+inline constexpr uint32_t kWavefrontSize = 64;  // CDNA
+inline constexpr uint32_t kWavesPerBlock = kBlockSize / kWavefrontSize;
+inline constexpr uint32_t kMinWavesPerSimd = kWavesPerBlock / kSimdsPerCu;
+static_assert(kMinWavesPerSimd > 0, "kBlockSize must cover at least one wave per SIMD");
+#define TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, kMinWavesPerSimd)
+#else
+// topk_small_batch_cluster_kernel shadows kOccupancy with its own template
+// parameter and CLUSTER_TOPK_KERNEL relies on that, so leave the CUDA spelling.
 #define TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, kOccupancy)
+#endif
 
 /// 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.

python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh

@@ -14,7 +14,8 @@
 ///  - the dispatcher selects cluster size 8 or 16 from the probed occupancy.
 ///
 /// Algorithm: fp16 coarse histogram -> threshold bin -> fp32-boundary collect ->
-/// exact radix tie-break.
+/// exact radix tie-break, plus an exact-key refinement when the threshold bin
+/// overflows the staging buffer (refine_and_handle_tie).
 
 #pragma once
 
@@ -216,6 +217,17 @@ struct TopKConfig {
   // by downstream sparse attention.
   static constexpr uint32_t kMaxNumTie = 2048;
   static constexpr uint32_t kRadixSize = 1 << 8;
+  // Radix width of the threshold-bin refinement (refine_and_handle_tie). Wider
+  // than kRadixSize so that ceil(32 / kRefineBits) == 3 rounds cover the whole
+  // key: the histogram overlays the tie staging buffer, which that path
+  // re-derives anyway, so the extra bins are free. 12 is the widest that fits.
+  static constexpr uint32_t kRefineBits = 12;
+  static constexpr uint32_t kRefineSize = 1 << kRefineBits;
+  static constexpr uint32_t kRefineItems = kRefineSize / kBlockSize;
+  static_assert(
+      kRefineSize * sizeof(uint32_t) <= kMaxNumTie * sizeof(TieValue),
+      "the refinement histogram must fit the tie staging buffer it overlays");
+  static_assert(kRefineSize % kBlockSize == 0);
   static constexpr uint32_t kTopKItems = (kMaxTopK + kBlockSize - 1) / kBlockSize;
   // tie candidates owned per thread in the strided handle_tie loops
   static constexpr uint32_t kTieItems = kMaxNumTie / kBlockSize;
@@ -324,6 +336,44 @@ struct TopKConfig {
     }
   }
 
+  /// Threshold-digit search over the kRefineSize-bin refinement histogram in
+  /// `hist`, by the same rule radix_tie_select uses. One thread owns kRefineItems
+  /// contiguous bins. Publishes the result in `smem->match`. Block-wide; ends
+  /// with a barrier.
+  SGL_DEVICE static void refine_find_threshold(  //
+      const uint32_t* hist,
+      const uint32_t total_active,
+      const uint32_t topk_remain,
+      TieHandleSmem* smem) {
+    const auto tx = threadIdx.x;
+    const auto lane_id = tx % kWarpSize;
+    const auto warp_id = broadcast(tx / kWarpSize);
+
+    uint32_t counts[kRefineItems];
+    uint32_t local_sum = 0;
+#pragma unroll
+    for (uint32_t i = 0; i < kRefineItems; ++i) {
+      counts[i] = hist[tx * kRefineItems + i];
+      local_sum += counts[i];
+    }
+    const auto warp_inc = warp_inclusive_sum(lane_id, local_sum);
+    if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc;
+    __syncthreads();
+
+    // Inclusive prefix over all bins below this thread's first one.
+    uint32_t prefix = warp::reduce_sum(lane_id < warp_id ? smem->warp_sum[lane_id] : 0);
+    prefix += warp_inc - local_sum;
+#pragma unroll
+    for (uint32_t i = 0; i < kRefineItems; ++i) {
+      prefix += counts[i];
+      const auto above = total_active - prefix;  // active elements in bins ABOVE this one
+      if (above < topk_remain && above + counts[i] >= topk_remain) {
+        smem->match = {tx * kRefineItems + i, above, counts[i]};
+      }
+    }
+    __syncthreads();
+  }
+
   /// Exact radix select over the tie candidates: each thread owns kItems
   /// strided elements (inactive beyond num_ties). Requires
   /// num_ties <= kItems * kBlockSize.
@@ -447,7 +497,13 @@ struct TopKRadixBase : TopKConfig {
     union {
       alignas(16) uint32_t histogram[kHistSize];
       struct {
-        TieValue tie_values[kMaxNumTie];
+        // refine_and_handle_tie re-derives the candidates instead of trusting
+        // what the collect pass staged, so the buffer is free to carry that
+        // path's radix histogram until its staging pass writes values back.
+        union {
+          TieValue tie_values[kMaxNumTie];
+          uint32_t refine_hist[kRefineSize];
+        };
         TieHandleSmem tie_handle;
       };
     };
@@ -540,6 +596,124 @@ struct TopKRadixBase : TopKConfig {
 
     __syncthreads();
   }
+
+  /// Exact-key refinement of a threshold coarse bin that overflows the staging
+  /// buffer, where the collect pass kept only whichever kMaxNumTie arrived
+  /// first. Replaces the plain handle_tie call when `count_eq > kMaxNumTie`.
+  ///
+  /// Up to three radix rounds over the order-preserving key from
+  /// `extract_exact_bin`, kRefineBits wide (12 / 12 / 8). Exits once the refined
+  /// set fits the buffer, or once all 32 bits are consumed -- the key is
+  /// injective on fp32 bit patterns, so the survivors are bit-identical by then
+  /// and any subset is correct. Each round emits its own "above" set while
+  /// building the next round's histogram, so this costs at most four input
+  /// passes. Every round needs `0 < remain <= active`, which the threshold-bin
+  /// invariant gives and each round restores.
+  ///
+  /// `for_each_input` is the iteration the collect pass uses, so the two agree
+  /// on which elements exist: padding is NaN and fails both boundary compares.
+  SGL_DEVICE static void refine_and_handle_tie(  //
+      const TopKProblem& problem,
+      Smem* smem,
+      const float v_lo,
+      const float v_hi,
+      const uint32_t count_eq) {
+    const auto tx = threadIdx.x;
+    const auto topk = problem.topk;
+    const auto handle = &smem->tie_handle;
+    const auto hist = smem->refine_hist;
+    const auto clear_hist = [&] {
+#pragma unroll
+      for (uint32_t i = 0; i < kRefineItems; ++i)
+        hist[tx * kRefineItems + i] = 0;
+      __syncthreads();
+    };
+
+    if (smem->count_gt >= topk) [[unlikely]] {
+      // Off the threshold-bin invariant: the collect pass already filled every
+      // output slot, so no candidate can be selected. refine_find_threshold has
+      // no bin to publish in that state, so stop before entering it.
+      return;
+    }
+    uint32_t cand_count = count_eq;
+    uint32_t remain = topk - smem->count_gt;
+    uint32_t prefix = 0;  // refined key bits agreed on so far
+    uint32_t mask = 0;    // which key bits `prefix` pins down
+    uint32_t width = kRefineBits;
+    uint32_t shift = 32 - kRefineBits;
+
+    // First histogram: no preceding round to emit for, so it stands alone.
+    clear_hist();
+    for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t) {
+      if (val >= v_lo && val < v_hi) atomicAdd(&hist[extract_exact_bin(val) >> shift], 1);
+    });
+    __syncthreads();
+
+    while (true) {
+      refine_find_threshold(hist, cand_count, remain, handle);
+      const auto match = handle->match;
+
+      // Candidate predicate for this round, pinned before `prefix`/`mask` move on.
+      const auto sel_prefix = prefix;
+      const auto sel_mask = mask;
+      const auto sel_shift = shift;
+      const auto sel_digit = (1u << width) - 1u;
+      prefix |= match.bin << sel_shift;
+      mask |= sel_digit << sel_shift;
+      remain -= match.above_count;
+      cand_count = match.equal_count;
+
+      // Stop once the survivors fit the staging buffer, the key is exhausted
+      // (they are bit-identical, so truncating them is exact), or the output is
+      // already full -- this round's "above" set still has to be emitted.
+      if (cand_count <= kMaxNumTie || sel_shift == 0 || remain == 0) {
+        if (tx == 0) smem->count_eq = 0;
+        __syncthreads();
+        // Overwrites the histogram, which refine_find_threshold has consumed.
+        for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) {
+          if (val >= v_lo && val < v_hi) {
+            const auto key = extract_exact_bin(val);
+            if ((key & sel_mask) != sel_prefix) return;
+            const auto digit = (key >> sel_shift) & sel_digit;
+            if (digit > match.bin) {
+              const auto pos = atomicAdd(&smem->count_gt, 1);
+              if (pos < topk) [[likely]]
+                problem.emit(pos, idx);
+            } else if (digit == match.bin) {
+              const auto slot = atomicAdd(&smem->count_eq, 1);
+              if (slot < kMaxNumTie) smem->tie_values[slot] = {val, idx};
+            }
+          }
+        });
+        __syncthreads();
+        break;
+      }
+
+      width = min(sel_shift, kRefineBits);
+      shift = sel_shift - width;
+      clear_hist();
+      for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) {
+        if (val >= v_lo && val < v_hi) {
+          const auto key = extract_exact_bin(val);
+          if ((key & sel_mask) != sel_prefix) return;
+          const auto digit = (key >> sel_shift) & sel_digit;
+          if (digit > match.bin) {
+            const auto pos = atomicAdd(&smem->count_gt, 1);
+            if (pos < topk) [[likely]]
+              problem.emit(pos, idx);
+          } else if (digit == match.bin) {
+            atomicAdd(&hist[(key >> shift) & ((1u << width) - 1u)], 1);
+          }
+        }
+      });
+      __syncthreads();
+    }
+
+    const auto count_gt = smem->count_gt;
+    const auto tie_count = min(smem->count_eq, kMaxNumTie);
+    const auto remain_topk = count_gt < topk ? topk - count_gt : 0;
+    handle_tie(smem->tie_values, problem, count_gt, tie_count, remain_topk, handle);
+  }
 };
 
 // ---------------------------------------------------------------------------
@@ -642,6 +816,10 @@ struct TopKRegister : TopKRadixBase<12> {
     __syncthreads();
     const auto count_gt = smem->count_gt;
     const auto count_eq = smem->count_eq;
+    if (count_eq > kMaxNumTie) [[unlikely]] {
+      // Buffer holds an arrival-order subset; refine on the exact key first.
+      return refine_and_handle_tie(problem, smem, v_lo, v_hi, count_eq);
+    }
     const auto remain_topk = count_gt < topk ? topk - count_gt : 0;
     const auto tie_count = min(count_eq, kMaxNumTie);
     handle_tie(smem->tie_values, problem, count_gt, tie_count, remain_topk, &smem->tie_handle);
@@ -719,6 +897,10 @@ struct TopKStreaming : TopKRadixBase<12> {
     __syncthreads();
     const auto count_gt = smem->count_gt;
     const auto count_eq = smem->count_eq;
+    if (count_eq > kMaxNumTie) [[unlikely]] {
+      // See the register path.
+      return refine_and_handle_tie(problem, smem, v_lo, v_hi, count_eq);
+    }
     const auto remain_topk = count_gt < topk ? topk - count_gt : 0;
     const auto tie_count = min(count_eq, kMaxNumTie);
     handle_tie(smem->tie_values, problem, count_gt, tie_count, remain_topk, &smem->tie_handle);
@@ -731,6 +913,11 @@ struct TopKStreaming : TopKRadixBase<12> {
 //
 // CUDA only: thread-block clusters and distributed shared memory have no CDNA
 // equivalent.
+//
+// Still truncates an overflowing threshold bin, unlike the register and
+// streaming paths: the candidate set is split across kClusterSize ranks, so
+// refining needs cluster-wide histogram and emit counters rather than the
+// block-local ones refine_and_handle_tie uses.
 // ---------------------------------------------------------------------------
 
 #if SUPPORT_CLUSTER