[AMD] Enable GLM DSA prefill top-k to the v2 kernel (#37889)

Co-authored-by: Thomas Wang <thomawan@amd.com>

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

@@ -250,6 +250,103 @@ TOPK_KERNEL void topk_ragged_kernel(const __grid_constant__ TopKRaggedParams par
   // PDL trigger secondary at the end the block typically has no use, so ignore it
 }
 
+#ifdef USE_ROCM
+// Only the ROCm DSA prefill emits this layout today, so CUDA/XPU builds stay
+// unchanged. Nothing below is AMD-specific; the guard can be dropped later.
+
+/**
+ * \brief Parameters of the packed (DSA extend) layout.
+ *
+ * Same addressing as the ragged layout -- every row's window lives inside one
+ * batch-global score buffer starting at `row_starts[i]` -- but the selected
+ * columns are mapped through a page table before they are written out. Prefill
+ * expands one request into many query-token rows, so several score rows share a
+ * page-table row; `row_to_batch[i]` says which one.
+ */
+struct TopKPackedParams {
+  // NOTE: may write. The head of the window is masked in place, see the kernel.
+  float* __restrict__ scores;
+  const int32_t* __restrict__ seq_lens;      // per-row window length
+  const int32_t* __restrict__ row_starts;    // per-row score column offset
+  const int32_t* __restrict__ row_to_batch;  // per-row page-table row; null => identity
+  const int32_t* __restrict__ page_table;
+  int32_t* __restrict__ page_indices;
+  int64_t score_stride;
+  int64_t page_table_stride;
+  uint32_t topk;
+  uint32_t page_bits;
+
+  SGL_DEVICE PageTransform get_transform(uint32_t bx) const {
+    const auto table_row = row_to_batch == nullptr ? bx : static_cast<uint32_t>(row_to_batch[bx]);
+    return {page_table + static_cast<int64_t>(table_row) * page_table_stride, page_bits, nullptr};
+  }
+};
+
+/**
+ * \brief Top-k over packed rows, emitting page-table indices.
+ * \tparam kPDL whether to use PDL to synchronize with the indexer kernel
+ *
+ * Dispatch mirrors `topk_ragged_kernel`: both are prefill kernels, so the level
+ * is picked per row at runtime and only the register and streaming
+ * implementations are instantiated (no plan, no cluster path).
+ */
+template <bool kPDL>
+TOPK_KERNEL void topk_packed_kernel(const __grid_constant__ TopKPackedParams params) {
+  device::enable_smem_spilling();
+  constexpr uint32_t kVecSize = impl::TopKStreaming::kVecSize;
+  __shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
+  __shared__ int32_t s_topk_indices[kMaxTopK];
+
+  const auto bx = blockIdx.x;
+  const auto seq_len = static_cast<uint32_t>(params.seq_lens[bx]);
+  const auto row_start = static_cast<uint32_t>(params.row_starts[bx]);
+  const auto topk = params.topk;
+  const auto transform = params.get_transform(bx);
+  const auto out = params.page_indices + bx * static_cast<int64_t>(topk);
+  const auto score = params.scores + bx * params.score_stride;
+
+  auto problem = TopKProblem{
+      .in = score + row_start,
+      .out = out,
+      .topk = topk,
+      .seq_len = seq_len,
+  };
+  if (seq_len <= topk) {
+    return trivial_transform<kPDL, TopKMode::PAGE_TABLE>(problem, transform);
+  }
+
+  // Round the window down to a `kVecSize` boundary and mask the <= 3 columns
+  // that pulls in, with the same bias / input_start as `topk_ragged_kernel`.
+  const auto rem = row_start % kVecSize;
+  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] = impl::padding_value();
+    }
+  }
+  using device::topk::broadcast;
+  problem.in -= rem;
+  problem.out = s_topk_indices;  // write into stage buffer in smem first
+  problem.seq_len = seq_len + rem;
+  problem.bias = broadcast(-static_cast<int32_t>(rem));
+  problem.input_start = broadcast(rem);
+
+  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);
+  }
+  device::PDLTriggerSecondary<kPDL>();
+  __syncthreads();
+  paged_transform<TopKMode::PAGE_TABLE>(problem, out, transform);
+}
+#endif  // USE_ROCM
+
 /**
  * \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)
@@ -786,6 +883,94 @@ struct TopKKernel {
         .config({.use_pdl = kUsePDL})
         .launch(topk_ragged_kernel<kUsePDL>, params);
   }
+
+#ifdef USE_ROCM  // see the packed kernel above
+  /**
+   * \brief Packed (DSA extend prefill) variant of `transform_paged`: per-row
+   * window inside one batch-global score buffer, page-table output, no plan.
+   *
+   * `scores` is written in place exactly like `transform_ragged` does, so rows
+   * must not overlap and the buffer must have no consumer after this call.
+   *
+   * `row_to_batch` absent means the page table is indexed by score row; present,
+   * it maps each score row onto the table row of the request it belongs to,
+   * which is what prefill needs (one request expands into many query rows).
+   */
+  static void transform_packed(
+      const tvm::ffi::TensorView scores,
+      const tvm::ffi::TensorView seq_lens,
+      const tvm::ffi::TensorView row_starts,
+      const tvm::ffi::TensorView page_table,
+      const tvm::ffi::TensorView page_indices,
+      const uint32_t page_size,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> row_to_batch) {
+    using namespace host;
+    auto B = SymbolicSize{"batch_size"};
+    auto L = SymbolicSize{"max_seq_len"};
+    auto S = SymbolicSize{"score_stride"};
+    auto R = SymbolicSize{"page_table_rows"};
+    auto K = SymbolicSize{"topk"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLGPU>();
+
+    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})  // row_starts
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(row_starts);
+    TensorMatcher({R, -1})  // page_table
+        .with_strides({-1, 1})
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(page_table);
+    TensorMatcher({B, K})  // page_indices
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(page_indices);
+    const int32_t* row_to_batch_ptr = nullptr;
+    if (row_to_batch.has_value()) {
+      TensorMatcher({B})  // row_to_batch
+          .with_dtype<int32_t>()
+          .with_device(device_)
+          .verify(row_to_batch.value());
+      row_to_batch_ptr = static_cast<const int32_t*>(row_to_batch.value().data_ptr());
+    } else {
+      RuntimeCheck(R.unwrap() == B.unwrap(), "page_table must have one row per score row unless row_to_batch is given");
+    }
+
+    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)");
+    // The kernel masks the head of each window in place, so overlapping rows
+    // would clobber each other.
+    RuntimeCheck(S.unwrap() >= L.unwrap(), "scores rows must not overlap");
+    const auto topk = static_cast<uint32_t>(K.unwrap());
+    RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 2048]");
+
+    const auto params = TopKPackedParams{
+        .scores = static_cast<float*>(scores.data_ptr()),
+        .seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
+        .row_starts = static_cast<const int32_t*>(row_starts.data_ptr()),
+        .row_to_batch = row_to_batch_ptr,
+        .page_table = static_cast<const int32_t*>(page_table.data_ptr()),
+        .page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
+        .score_stride = S.unwrap(),
+        .page_table_stride = page_table.stride(0),
+        .topk = topk,
+        .page_bits = static_cast<uint32_t>(std::countr_zero(page_size)),
+    };
+    LaunchKernel(static_cast<uint32_t>(B.unwrap()), kBlockSize, device_.unwrap())
+        .config({.use_pdl = kUsePDL})
+        .launch(topk_packed_kernel<kUsePDL>, params);
+  }
+#endif  // USE_ROCM
 };
 
 }  // namespace sglang

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

@@ -49,16 +49,20 @@ def _jit_topk_v2_module():
                 f"-DSGL_TOPK_V2_MAX_C{cluster_size}_OCC{occupancy}={max_active_clusters}"
             )
     kernel = f"TopKKernel<{args}>"
+    wrappers = [
+        ("topk_transform_paged", f"{kernel}::transform_paged"),
+        ("topk_transform_ragged", f"{kernel}::transform_ragged"),
+        ("topk_plan", f"{kernel}::plan"),
+    ]
+    if is_hip_runtime():
+        # transform_packed only exists under USE_ROCM, see topk_v2.cuh
+        wrappers.append(("topk_transform_packed", f"{kernel}::transform_packed"))
     return load_jit(
         make_name("topk_v2"),
         *args,
         extra_cuda_cflags=extra_cuda_cflags,
         cuda_files=["deepseek_v4/topk_v2.cuh"],
-        cuda_wrappers=[
-            ("topk_transform_paged", f"{kernel}::transform_paged"),
-            ("topk_transform_ragged", f"{kernel}::transform_ragged"),
-            ("topk_plan", f"{kernel}::plan"),
-        ],
+        cuda_wrappers=wrappers,
     )
 
 
@@ -212,6 +216,9 @@ def topk_transform_paged_v2(
     * Both outputs given -- ``out_page_indices`` receives the page-table
       transform and ``out_raw_indices`` receives the selected raw indices.
 
+    For the packed (DSA extend prefill) layout see
+    :func:`topk_transform_packed_v2`.
+
     NOTE: every entry of `seq_lens` must be NON-NEGATIVE, and `metadata` must
     come from :func:`plan_topk_v2` over the same `seq_lens` values.
     A length of 0 is the valid way to express "no tokens": the row takes the
@@ -247,3 +254,50 @@ def topk_transform_paged_v2(
         metadata,
         out_raw_indices,
     )
+
+
+def topk_transform_packed_v2(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_tables: torch.Tensor,
+    out_page_indices: torch.Tensor,
+    page_size: int,
+    *,
+    row_starts: torch.Tensor,
+    row_to_batch: Optional[torch.Tensor] = None,
+) -> None:
+    """Packed (DSA extend prefill) fused top-k + page-table transform.
+
+    Row ``i`` selects the top-k of ``scores[i, ks : ks + seq_lens[i]]``
+    (``ks = row_starts[i]``) and writes the page-table transform of the selected
+    row-local positions into ``out_page_indices``, ``-1`` padded. Prefill expands
+    one request into many query-token rows, so ``row_to_batch[i]`` (optional,
+    ``(rows,)`` int32) names the ``page_tables`` row of the request row ``i``
+    belongs to; omitting it indexes the table by score row. ``row_to_batch`` is
+    not range-checked.
+
+    This is :func:`topk_transform_ragged_v2` with a page-table output instead of
+    an additive offset. Like ragged, it dispatches the implementation per row at
+    runtime, so it needs no plan and no :func:`plan_topk_v2` metadata.
+
+    NOTE: ``scores`` is MODIFIED 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, so do not
+    pass a view with overlapping rows.
+    ``seq_lens`` entries must be NON-NEGATIVE, as for the paged entry point.
+
+    ROCm only: the kernel is compiled under ``USE_ROCM`` so that CUDA and XPU
+    builds are untouched. Nothing in it is AMD-specific -- no non-ROCm caller
+    produces this layout today.
+    """
+    assert is_hip_runtime(), "topk_transform_packed_v2 is compiled under USE_ROCM only"
+    module = _jit_topk_v2_module()
+    module.topk_transform_packed(
+        scores,
+        seq_lens,
+        row_starts,
+        page_tables,
+        out_page_indices,
+        page_size,
+        row_to_batch,
+    )

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

@@ -7,10 +7,13 @@ import torch
 
 from sglang.srt.environ import envs
 from sglang.srt.runtime_context import get_exec, get_spec
+from sglang.srt.utils import is_hip
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
 
+_is_hip = is_hip()
+
 _FLASHINFER_TIE_BREAK_VALUES = {
     "small": 1,
     "large": 2,
@@ -148,6 +151,39 @@ class DSATopKBackend(Enum):
                 logits, lengths, topk, topk_indices_offset, row_starts
             )
 
+        # Packed PAGED extend (GLM DSA prefill), ROCm-only: CUDA gets the same
+        # fusion from RAGGED above. Unsupported shapes fall back, not raise.
+        # The row -> request map is `token_to_batch_idx` for a whole-forward call
+        # and the chunk's own `batch_idx_list` when the indexer split the logits.
+        if batch_idx_list is None:
+            row_to_batch = attn_metadata.token_to_batch_idx
+        elif isinstance(batch_idx_list, torch.Tensor):
+            row_to_batch = batch_idx_list
+        else:
+            # The prefill-CP list selects requests, not rows: leave it on legacy.
+            row_to_batch = None
+        if (
+            _is_hip
+            and self.should_use_topk_v2()
+            and topk_transform_method == TopkTransformMethod.PAGED
+            and 0 < topk <= 2048
+            and lengths.shape[0] == logits.shape[0]
+            and logits.dtype == torch.float32
+            and logits.stride(1) == 1
+            and logits.stride(0) % 4 == 0
+            and row_starts is not None
+            and row_to_batch is not None
+            and row_to_batch.shape[0] == logits.shape[0]
+        ):
+            return _topk_transform_v2_packed(
+                logits,
+                lengths,
+                topk,
+                attn_metadata,
+                row_starts=row_starts,
+                row_to_batch=row_to_batch,
+            )
+
         # 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.
@@ -291,6 +327,9 @@ def _topk_transform_v2_paged(
     typically 64) yields the same physical slots as gathering the page_size=1
     table, without materializing that wide table.
 
+    For DSA extend's packed batch-global scores see
+    :func:`_topk_transform_v2_packed`.
+
     This is a committed contract, not a best-effort path: ``topk_transform`` routes
     here only for the decode-shaped PAGED case, and the fused-decode CUDA graph
     drops the page_size=1 table for exactly this case (see
@@ -341,6 +380,54 @@ def _topk_transform_v2_paged(
     return out
 
 
+def _topk_transform_v2_packed(
+    logits: torch.Tensor,
+    lengths: torch.Tensor,
+    topk: int,
+    attn_metadata,
+    row_starts: torch.Tensor,
+    row_to_batch: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Fused packed-row top-k + page-table transform (DSA extend prefill).
+
+    Same output contract as :func:`_topk_transform_v2_paged` -- ``(num_rows,
+    topk)`` int32 physical KV slots, ``-1`` padded -- but the scores are packed:
+    row ``i`` owns the window at ``row_starts[i]`` of one batch-global buffer and
+    maps through page-table row ``row_to_batch[i]`` (prefill expands one request
+    into many query-token rows). Selected indices stay row-local.
+
+    Being a prefill-only path it dispatches per row inside the kernel, so unlike
+    the paged entry point it needs no ``topk_v2_plan``.
+
+    NOTE: ``logits`` is MODIFIED IN PLACE (the <= 3 columns ahead of each window
+    are masked); the caller must not reuse it. ``lengths`` must be NON-NEGATIVE,
+    for the same reason as in :func:`_topk_transform_v2_paged`.
+    """
+    from sglang.kernels.ops.attention.dsv4.topk import topk_transform_packed_v2
+
+    num_rows = logits.shape[0]
+    assert (
+        logits.dtype == torch.float32
+        and logits.stride(1) == 1
+        and logits.stride(0) % 4 == 0
+    ), (
+        f"v2 top-k expects fp32 scores with unit row stride and 16B-aligned score_stride, got {logits.dtype=} {logits.stride()=}"
+    )
+    assert 0 < topk <= 2048, f"v2 top-k supports 0 < topk <= 2048, got {topk=}"
+
+    out = logits.new_empty((num_rows, topk), dtype=torch.int32)
+    topk_transform_packed_v2(
+        logits,
+        lengths,
+        attn_metadata.real_page_table,
+        out,
+        attn_metadata.page_size,
+        row_starts=row_starts.to(torch.int32),
+        row_to_batch=(None if row_to_batch is None else row_to_batch.to(torch.int32)),
+    )
+    return out
+
+
 def _topk_transform_v2_ragged(
     logits: torch.Tensor,
     lengths: torch.Tensor,

python/sglang/srt/layers/attention/dsa_backend.py

@@ -747,9 +747,9 @@ class DeepseekSparseAttnBackend(
         # Preprocess the folded top-k v2 plan once per forward (shared across
         # layers), at metadata-build time, from the same seqlens the transform
         # receives as `lengths` (dsa_seqlens_expanded). This must cover EVERY shape
-        # that dispatches to `_topk_transform_v2_paged` -- decode AND MTP
-        # target-verify / draft-extend, whose expanded row count is exactly what v2
-        # sees -- otherwise the helper's plan-present assertion fires. None only
+        # that dispatches to `_topk_transform_v2_paged` -- decode, MTP target-verify
+        # / draft-extend, and packed PAGED extend, whose expanded row count is what
+        # v2 sees -- otherwise the helper's plan-present assertion fires. None only
         # when the SGL v2 path is disabled; such metadata is never dispatched to v2.
         if not self.dsa_topk_backend.should_use_topk_v2():
             return None