[DSA] Route the ragged prefill top-k to the v2 kernel (#35175)

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

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

@@ -22,6 +22,7 @@
 #include <bit>
 #include <cstdint>
 #include <iterator>
+#include <limits>
 
 namespace sglang {
 
@@ -64,7 +65,7 @@ struct alignas(8) PlanItem {
 };
 static_assert(sizeof(GlobalMetadata) == 2 * sizeof(int32_t) && sizeof(PlanItem) == sizeof(GlobalMetadata));
 
-struct TopKLaunchParams {
+struct TopKPagedParams {
   const float* __restrict__ scores;
   const int32_t* __restrict__ seq_lens;
   const int32_t* __restrict__ page_table;
@@ -104,12 +105,22 @@ struct TopKLaunchParams {
   }
 };
 
+struct TopKRaggedParams {
+  float* __restrict__ scores;  // NOTE: may write
+  const int32_t* __restrict__ seq_lens;
+  const int32_t* __restrict__ row_starts;
+  const int32_t* __restrict__ out_offsets;
+  int32_t* __restrict__ topk_indices;
+  int64_t score_stride;
+  uint32_t topk;
+};
+
 /**
  * \brief Persistent cluster kernel for the long items. It will handle long inputs.
  * The short items are handled by the separate topk_kernel.
  */
 template <bool kPDL>
-CLUSTER_TOPK_KERNEL void topk_persistent_cluster_kernel(const __grid_constant__ TopKLaunchParams params) {
+CLUSTER_TOPK_KERNEL void topk_persistent_cluster_kernel(const __grid_constant__ TopKPagedParams params) {
   device::enable_smem_spilling();
   __shared__ impl::MaxSmem<Cluster::Smem> smem;
   const uint32_t num_cluster_items = params.global().num_cluster_items;
@@ -159,6 +170,85 @@ SGL_DEVICE void problem_transform(TopKProblem& problem, int32_t* output_ptr) {
   for_each_item(problem.topk, [&](uint32_t tx, uint32_t i) { problem.transform_output(tx, source_index[i]); });
 }
 
+/**
+ * \brief Ragged (prefill) top-k: select inside a per-row window, emit indices
+ * rebased onto the flattened KV.
+ *
+ * Row `b` selects the top-k of `scores[b][ks : ks + seq_lens[b]]` (`ks =
+ * row_starts[b]`) and writes `selected_position + out_offsets[b]`, `-1` padded.
+ * No page table and no plan: the DeepGEMM contiguous-KV indexer emits columns
+ * that are already absolute positions in the batch's flattened KV, so an add is
+ * the whole transform. One block per row -- prefill has thousands of rows, so
+ * the cluster path (which exists to split ONE row across blocks) is never worth
+ * it here.
+ *
+ * The window start is an arbitrary token offset, so the 16-byte vectorized load
+ * needs the row pointer rounded down to a 4-float boundary. The <= 3 elements
+ * that pulls in are columns of a preceding request -- real finite scores that
+ * would otherwise win the selection -- so they are masked in place first. That
+ * write races with nothing and needs no barrier of its own:
+ *   - one block owns the row, and a column of row `b` is read by no other row;
+ *   - the score buffer is dead once the top-k has run;
+ *   - every forward() below opens with its smem init and a `__syncthreads()`
+ *     before it reads any score. That barrier both publishes the mask to
+ *     whichever thread loads the head vector and keeps the compiler from
+ *     hoisting those loads above the store -- store and loads reach the same row
+ *     through two `__restrict__` pointers, which otherwise licenses exactly that
+ *     reordering.
+ * It must however land after the PDL wait, or the indexer overwrites it.
+ */
+template <bool kPDL>
+TOPK_KERNEL void topk_ragged_kernel(const __grid_constant__ TopKRaggedParams params) {
+  device::enable_smem_spilling();
+  constexpr uint32_t kVecSize = impl::TopKStreaming::kVecSize;
+  const auto bx = blockIdx.x;
+  // issue all metadata prefetch ahead of time
+  const auto seq_len = static_cast<uint32_t>(params.seq_lens[bx]);
+  const auto offset = params.out_offsets[bx];
+  const auto row_start = params.row_starts == nullptr ? 0u : params.row_starts[bx];
+  const auto topk = params.topk;
+  const auto out = params.topk_indices + bx * static_cast<int64_t>(topk);
+
+  if (seq_len <= topk) {
+    device::PDLWaitPrimary<kPDL>();
+    for_each_item(topk, [&](uint32_t tx, uint32_t) {
+      out[tx] = tx < seq_len ? static_cast<int32_t>(tx) + offset : -1;  // note: need offset
+    });
+    return;
+  }
+
+  const auto rem = row_start % kVecSize;
+  const auto score = params.scores + bx * params.score_stride;
+  if (rem != 0) {
+    // The mask has to land after the indexer has retired
+    // Otherwise it may be accidentally overwritten by DG upstream
+    device::PDLWaitPrimary<kPDL>();
+    static_assert(kVecSize <= kBlockSize, "not enough threads ");
+    if (const auto tx = threadIdx.x; tx < rem) {
+      score[row_start - rem + tx] = -std::numeric_limits<float>::max();
+    }
+  }
+
+  const auto problem = TopKProblem{
+      .in = score + (row_start - rem),
+      .out = out,
+      .page_table = nullptr,  // unused
+      .topk = topk,
+      .seq_len = seq_len + rem,
+      .page_bits = 1,  // unused
+      .bias = offset - static_cast<int32_t>(rem),
+  };
+  __shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
+  if (problem.seq_len <= Register2::kMaxSeqLen) {
+    Register2::forward<kPDL>(problem, &smem);
+  } else if (problem.seq_len <= Register4::kMaxSeqLen) {
+    Register4::forward<kPDL>(problem, &smem);
+  } else {
+    Streaming::forward<kPDL>(problem, &smem);
+  }
+  // PDL trigger secondary at the end the block typically has no use, so ignore it
+}
+
 /**
  * \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)
@@ -169,7 +259,7 @@ SGL_DEVICE void problem_transform(TopKProblem& problem, int32_t* output_ptr) {
  * - Level 3: max_seq_len > cluster_floor   -> + epilogue process of cluster path
  */
 template <bool kPDL, int kLevel, TopKMode kMode>
-TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams params) {
+TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKPagedParams params) {
   device::enable_smem_spilling();
   auto problem = params.problem(blockIdx.x);
   constexpr uint32_t kU32Max = std::numeric_limits<uint32_t>::max();
@@ -217,7 +307,7 @@ TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams param
 }
 
 template <bool kPDL, TopKMode kMode>
-CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLaunchParams params) {
+CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKPagedParams params) {
   device::enable_smem_spilling();
   auto problem = params.problem(blockIdx.x);
   __shared__ impl::MaxSmem<Streaming::Smem, Cluster::Smem> smem;
@@ -378,7 +468,7 @@ struct TopKKernel {
         static_cluster_threshold);
   }
 
-  static void transform(
+  static void transform_paged(
       const tvm::ffi::TensorView scores,
       const tvm::ffi::TensorView seq_lens,
       const tvm::ffi::Optional<tvm::ffi::TensorView> page_table,
@@ -444,7 +534,7 @@ struct TopKKernel {
     // The floor is chosen on the host per launch.
     constexpr uint32_t kClusterFloorSmall = 32768;
     constexpr uint32_t kSmallBatchLowFloor = 15;
-    const auto params = TopKLaunchParams{
+    const auto params = TopKPagedParams{
         .scores = static_cast<const float*>(scores.data_ptr()),
         .seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
         .page_table = page_table_ptr,
@@ -498,6 +588,78 @@ struct TopKKernel {
       }
     });
   }
+
+  /**
+   * \brief Ragged (prefill) variant of `transform`: per-row window, additive
+   * output transform, no page table and no plan.
+   *
+   * `scores` is written in place: the <= 3 columns the 16-byte-aligned read base
+   * pulls in ahead of each row's window are masked out (see
+   * `topk_ragged_kernel`). They are invalid for that row, and the buffer has no
+   * consumer after this call.
+   *
+   * `row_starts` absent means every window starts at column 0, which is the
+   * single-request case; `out_offsets` is added to every selected position and
+   * is what rebases them onto the flattened KV.
+   */
+  static void transform_ragged(
+      const tvm::ffi::TensorView scores,
+      const tvm::ffi::TensorView seq_lens,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> row_starts,
+      const tvm::ffi::TensorView out_offsets,
+      const tvm::ffi::TensorView topk_indices) {
+    using namespace host;
+    auto B = SymbolicSize{"batch_size"};
+    auto L = SymbolicSize{"max_seq_len"};
+    auto S = SymbolicSize{"score_stride"};
+    auto K = SymbolicSize{"topk"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({B, L})  // score
+        .with_strides({S, 1})
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(scores);
+    TensorMatcher({B})  // seq_lens
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(seq_lens);
+    TensorMatcher({B})  // out_offsets
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(out_offsets);
+    TensorMatcher({B, K})  // topk_indices
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(topk_indices);
+    const int32_t* row_starts_ptr = nullptr;
+    if (row_starts.has_value()) {
+      TensorMatcher({B})  // row_starts
+          .with_dtype<int32_t>()
+          .with_device(device_)
+          .verify(row_starts.value());
+      row_starts_ptr = static_cast<const int32_t*>(row_starts.value().data_ptr());
+    }
+
+    RuntimeCheck(S.unwrap() % 4 == 0, "score_stride must be a multiple of 4 (16-byte vectorized load)");
+    const auto topk = static_cast<uint32_t>(K.unwrap());
+    RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 2048]");
+
+    constexpr bool kUsePDL = true;
+    const auto params = TopKRaggedParams{
+        .scores = static_cast<float*>(scores.data_ptr()),
+        .seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
+        .row_starts = row_starts_ptr,
+        .out_offsets = static_cast<const int32_t*>(out_offsets.data_ptr()),
+        .topk_indices = static_cast<int32_t*>(topk_indices.data_ptr()),
+        .score_stride = S.unwrap(),
+        .topk = topk,
+    };
+    LaunchKernel(static_cast<uint32_t>(B.unwrap()), kBlockSize, device_.unwrap())
+        .config({.use_pdl = kUsePDL})
+        .launch(topk_ragged_kernel<kUsePDL>, params);
+  }
 };
 
 }  // namespace sglang

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

@@ -176,9 +176,10 @@ struct TopKProblem {
   uint32_t topk;
   uint32_t seq_len;
   uint32_t page_bits;
+  int32_t bias = 0;  // needed by ragged mode
 
   SGL_DEVICE void emit(uint32_t pos, uint32_t raw_idx) const {
-    out[pos] = static_cast<int32_t>(raw_idx);
+    out[pos] = static_cast<int32_t>(raw_idx) + bias;
   }
   SGL_DEVICE void transform_output(uint32_t t, int32_t raw) const {
     out[t] = raw < 0 ? -1 : page_to_indices(page_table, raw, page_bits);

python/sglang/kernels/ops/attention/dsv4/topk.py

@@ -39,7 +39,8 @@ def _jit_topk_v2_module():
         make_name("topk_v2"),
         cuda_files=["deepseek_v4/topk_v2.cuh"],
         cuda_wrappers=[
-            ("topk_transform", "TopKKernel::transform"),
+            ("topk_transform_paged", "TopKKernel::transform_paged"),
+            ("topk_transform_ragged", "TopKKernel::transform_ragged"),
             ("topk_plan", "TopKKernel::plan"),
         ],
     )
@@ -86,6 +87,34 @@ def plan_topk_v2(seq_lens: torch.Tensor, static_threshold: int = 0) -> torch.Ten
     return metadata
 
 
+def topk_transform_ragged_v2(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    *,
+    out_offsets: torch.Tensor,
+    out_indices: torch.Tensor,
+    row_starts: Optional[torch.Tensor] = None,
+) -> None:
+    """Ragged (prefill) fused top-k for a contiguous-KV score matrix.
+
+    Row ``i`` selects the top-k of ``scores[i, ks : ks + seq_lens[i]]`` (``ks =
+    row_starts[i]``, 0 when ``row_starts`` is omitted) and writes
+    ``selected_position + out_offsets[i]`` into ``out_indices``, ``-1`` padded.
+    With the production convention ``out_offsets == row_starts`` that is the
+    column index itself, i.e. the token's slot in the batch's flattened KV.
+
+    Unlike :func:`topk_transform_512_v2` this needs no page table and no plan
+    (the cluster path only pays off for very few rows, and prefill has many).
+
+    IMPORTANT: ``scores`` is written in place -- the <= 3 columns ahead of each
+    row's window that the 16-byte-aligned read base pulls in are masked out.
+    They are invalid for that row and the buffer must have no other consumer.
+    ``seq_lens`` entries must be NON-NEGATIVE, as for the paged entry point.
+    """
+    module = _jit_topk_v2_module()
+    module.topk_transform_ragged(scores, seq_lens, row_starts, out_offsets, out_indices)
+
+
 def topk_transform_512_v2(
     scores: torch.Tensor,
     seq_lens: torch.Tensor,
@@ -115,7 +144,7 @@ def topk_transform_512_v2(
     the output is all -1.
     """
     module = _jit_topk_v2_module()
-    module.topk_transform(
+    module.topk_transform_paged(
         scores,
         seq_lens,
         page_tables,

python/sglang/srt/layers/attention/dsa/dsa_indexer_metadata.py

@@ -75,6 +75,7 @@ class BaseIndexerMetadata(ABC):
         self,
         logits: torch.Tensor,
         topk: int,
+        **kwargs,
     ) -> torch.Tensor:
         """
         Perform topk selection on the logits and possibly transform the result.

python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py

@@ -114,6 +114,25 @@ class DSATopKBackend(Enum):
         ):
             return _topk_transform_v2_paged(logits, lengths, topk, attn_metadata)
 
+        # Extend-shaped RAGGED top-k for the SGL backend routes to the same v2
+        # kernel through its ragged entry point: no page table (the columns are
+        # already flattened-KV positions), no plan (prefill has enough rows that
+        # the cluster path never applies), just a per-row window and an additive
+        # output transform. `batch_idx_list` is not None only on the prefill-CP
+        # path, whose `topk_indices_offset` is built from cu_seqlens_q rather
+        # than the KV bases -- leave that one on the legacy kernel.
+        if (
+            self.should_use_topk_v2()
+            and topk_transform_method == TopkTransformMethod.RAGGED
+            and topk_indices_offset is not None
+            and batch_idx_list is None
+            and 0 < topk <= 2048
+            and lengths.shape[0] == logits.shape[0] == topk_indices_offset.shape[0]
+        ):
+            return _topk_transform_v2_ragged(
+                logits, lengths, topk, topk_indices_offset, row_starts
+            )
+
         # The legacy transforms below read attn_metadata.page_table_1 (page_size=1),
         # which is always present here: the fold only drops it for the decode case
         # dispatched to v2 above.
@@ -305,6 +324,37 @@ def _topk_transform_v2_paged(
     return out
 
 
+def _topk_transform_v2_ragged(
+    logits: torch.Tensor,
+    lengths: torch.Tensor,
+    topk: int,
+    topk_indices_offset: torch.Tensor,
+    row_starts: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """Fused ragged top-k via the DeepSeek-V4 v2 JIT kernel.
+
+    ``logits`` is written in place: the kernel reads from a 16-byte-aligned base
+    and masks the <= 3 columns that pulls in ahead of the window. Those columns
+    belong to a preceding request of the same row, and the score buffer is dead
+    after the top-k (see ``DSAIndexer._get_topk_ragged``).
+
+    Preconditions match the paged helper: fp32 scores with unit row stride and a
+    16B-aligned row stride (DeepGEMM's contiguous-KV output satisfies this by
+    construction), int32 non-negative lengths, and ``0 < topk <= 2048``.
+    """
+    from sglang.kernels.ops.attention.dsv4.topk import topk_transform_ragged_v2
+
+    out = logits.new_empty((logits.shape[0], topk), dtype=torch.int32)
+    topk_transform_ragged_v2(
+        logits,
+        lengths,
+        out_offsets=topk_indices_offset,
+        out_indices=out,
+        row_starts=row_starts,
+    )
+    return out
+
+
 def _build_flashinfer_paged_args(
     attn_metadata,
     row_starts: Optional[torch.Tensor],