minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[DSA] Trim top-k v2 output modes and tighten its PDL waits (#35041)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
2026-08-19 13:09:21 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent 1c82955861
commit 746418a1ec
5 changed files with 156 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh
+111 -79
View File
@@ -28,6 +28,11 @@ namespace sglang {
namespace impl = device::topk;
using impl::TopKProblem;
enum class TopKMode {
INDICES, ///< raw selected indices into `out`; `page_table` unused
PAGE_TABLE, ///< page-table-transformed indices into `out`
};
using Register2 = impl::TopKRegister<2>; // <= 8192, register-resident, 1 read
using Register4 = impl::TopKRegister<4>; // <= 16384, register-resident, 1 read
using Streaming = impl::TopKStreaming;
@@ -64,7 +69,6 @@ struct TopKLaunchParams {
const int32_t* __restrict__ seq_lens;
const int32_t* __restrict__ page_table;
int32_t* __restrict__ page_indices;
int32_t* __restrict__ raw_indices; // optional raw (pre-transform) indices output; nullptr if unused
const PlanItem* __restrict__ metadata; // [0]=GlobalMetadata, [1+i]=PlanItem
int64_t score_stride;
int64_t page_table_stride;
@@ -89,7 +93,6 @@ struct TopKLaunchParams {
return TopKProblem{
.in = scores + batch_id * score_stride,
.out = page_indices + batch_id * k,
.raw_out = raw_indices != nullptr ? raw_indices + batch_id * k : nullptr,
.page_table = page_table + batch_id * page_table_stride,
.topk = topk,
.seq_len = seq_len,
@@ -133,12 +136,17 @@ SGL_DEVICE void for_each_item(uint32_t topk, const F& f) {
}
}
template <bool kPDL>
template <bool kPDL, TopKMode kMode>
SGL_DEVICE void trivial_transform(const TopKProblem& problem) {
device::PDLWaitPrimary<kPDL>();
device::PDLTriggerSecondary<kPDL>();
for_each_item(problem.topk, [&](uint32_t tx, uint32_t) {
problem.transform_output(tx, tx < problem.seq_len ? static_cast<int32_t>(tx) : -1);
const auto idx = tx < problem.seq_len ? static_cast<int32_t>(tx) : -1;
if constexpr (kMode == TopKMode::INDICES) {
problem.emit(tx, idx);
} else {
problem.transform_output(tx, idx);
}
});
}
@@ -160,17 +168,25 @@ SGL_DEVICE void problem_transform(TopKProblem& problem, int32_t* output_ptr) {
* - Level 2: max_seq_len <= cluster_floor -> trivial + register<4> + streaming
* - Level 3: max_seq_len > cluster_floor -> + epilogue process of cluster path
*/
template <bool kPDL, int kLevel>
template <bool kPDL, int kLevel, TopKMode kMode>
TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams params) {
device::enable_smem_spilling();
auto problem = params.problem(blockIdx.x);
constexpr uint32_t kU32Max = std::numeric_limits<uint32_t>::max();
__shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
__shared__ int32_t topk_indices[kMaxTopK];
problem.out = topk_indices;
constexpr bool kHandleCluster = (kLevel == 3);
// Only the cluster path consumes the cluster kernel's output, so only it waits
// on that kernel (kPDLFinal). Every other path waits at most on the indexer
// (kPDLEarly) and must not be held on an SM slot until the long-running
// persistent pool retires -- that would serialize the short items behind it.
constexpr bool kPDLEarly = kPDL && !kHandleCluster;
constexpr bool kPDLFinal = kPDL && kHandleCluster;
__shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
if (problem.seq_len <= problem.topk) return trivial_transform<kPDLEarly, kMode>(problem);
constexpr bool kNeedStaging = kMode != TopKMode::INDICES;
__shared__ int32_t s_topk_indices[kNeedStaging ? kMaxTopK : 1];
if constexpr (kNeedStaging) problem.out = s_topk_indices;
// non-trivial path: dispatch based on level and seq_len
const auto cluster_threshold = kHandleCluster ? params.cluster_threshold() : kU32Max;
if constexpr (kLevel == 0) {
@@ -181,34 +197,35 @@ TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams param
Register4::forward<kPDL>(problem, &smem); // max_seq_len <= 16384 guarantees seq <= 16384
} else {
static_assert(kLevel == 2 || kLevel == 3, "we only support level = 0,1,2,3 now");
// if using cluster, we can delay the PDL wait
constexpr bool kPDLEarly = kPDL && !kHandleCluster;
constexpr bool kPDLFinal = kPDL && kHandleCluster;
if (problem.seq_len <= kReg4MaxSeqLen) {
Register4::forward<kPDLEarly>(problem, &smem);
} else if (problem.seq_len <= cluster_threshold) {
Streaming::forward<kPDLEarly>(problem, &smem);
} else { // cluster path do nothing here
} else {
// Cluster path: the pool already selected into our output row; the only
// work left is the epilogue, so this is the one path that waits for it.
problem.out = params.get_output_ptr(blockIdx.x);
device::PDLWaitPrimary<kPDLFinal>();
}
device::PDLWaitPrimary<kPDLFinal>();
}
// page-table transform pass (gathers kept out of the hot scatter loop),
// then trigger the dependent kernel only after the full output is written.
device::PDLTriggerSecondary<kPDL>();
__syncthreads();
problem_transform(problem, params.get_output_ptr(blockIdx.x));
if constexpr (kNeedStaging) {
__syncthreads();
problem_transform(problem, params.get_output_ptr(blockIdx.x));
}
}
template <bool kPDL>
template <bool kPDL, TopKMode kMode>
CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLaunchParams params) {
device::enable_smem_spilling();
auto problem = params.problem(blockIdx.x);
__shared__ impl::MaxSmem<Streaming::Smem, Cluster::Smem> smem;
if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
__shared__ int32_t topk_indices[kMaxTopK];
problem.out = topk_indices;
if (problem.seq_len <= problem.topk) return trivial_transform<kPDL, kMode>(problem);
constexpr bool kNeedStaging = kMode != TopKMode::INDICES;
__shared__ int32_t s_topk_indices[kNeedStaging ? kMaxTopK : 1];
if constexpr (kNeedStaging) problem.out = s_topk_indices;
// randomly elect one worker rank to avoid workload imbalance
const auto worker_rank = blockIdx.x % kClusterSize;
@@ -217,30 +234,34 @@ CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLau
if (problem.seq_len <= kReg4MaxSeqLen) {
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) return;
Streaming::forward<kPDL>(problem, &smem);
device::PDLWaitPrimary<kPDL>();
__syncthreads();
} else {
auto cluster = cooperative_groups::this_cluster();
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;
if constexpr (kNeedStaging) {
problem.out = cluster.map_shared_rank(s_topk_indices, worker_rank);
}
Cluster::forward<kPDL>(problem, &smem);
if constexpr (kNeedStaging) {
cluster.sync();
if (blockIdx.y != worker_rank) return;
}
}
// 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));
device::PDLTriggerSecondary<kPDL>();
if constexpr (kNeedStaging) {
// 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 == s_topk_indices);
problem_transform(problem, params.get_output_ptr(blockIdx.x));
}
}
// --- Plan: choose cluster_threshold from the seq_len distribution -----------
@@ -360,11 +381,10 @@ struct TopKKernel {
static void transform(
const tvm::ffi::TensorView scores,
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::TensorView page_table,
const tvm::ffi::Optional<tvm::ffi::TensorView> page_table,
const tvm::ffi::TensorView page_indices,
const uint32_t page_size,
const tvm::ffi::TensorView metadata,
const tvm::ffi::Optional<tvm::ffi::TensorView> raw_indices) {
const tvm::ffi::TensorView metadata) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto Bp1 = SymbolicSize{"batch_size_plus_1"};
@@ -384,11 +404,19 @@ struct TopKKernel {
.with_dtype<int32_t>()
.with_device(device_)
.verify(seq_lens);
TensorMatcher({B, -1}) // page_table
.with_strides({P, 1})
.with_dtype<int32_t>()
.with_device(device_)
.verify(page_table);
// Absent means "no page transform": `page_indices` then receives the raw
// selected indices and nothing dereferences a page table.
const int32_t* page_table_ptr = nullptr;
int64_t page_table_stride = 0;
if (page_table.has_value()) {
TensorMatcher({B, -1}) // page_table
.with_strides({P, 1})
.with_dtype<int32_t>()
.with_device(device_)
.verify(page_table.value());
page_table_ptr = static_cast<const int32_t*>(page_table.value().data_ptr());
page_table_stride = P.unwrap();
}
TensorMatcher({B, K}) // page_indices
.with_dtype<int32_t>()
.with_device(device_)
@@ -398,12 +426,6 @@ struct TopKKernel {
.with_device(device_)
.verify(metadata);
int32_t* raw_indices_ptr = nullptr;
if (raw_indices.has_value()) {
TensorMatcher({B, K}).with_dtype<int32_t>().with_device(device_).verify(raw_indices.value());
raw_indices_ptr = static_cast<int32_t*>(raw_indices.value().data_ptr());
}
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)");
RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
@@ -425,12 +447,11 @@ struct TopKKernel {
const auto params = TopKLaunchParams{
.scores = static_cast<const float*>(scores.data_ptr()),
.seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
.page_table = static_cast<const int32_t*>(page_table.data_ptr()),
.page_table = page_table_ptr,
.page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
.raw_indices = raw_indices_ptr,
.metadata = static_cast<const PlanItem*>(metadata.data_ptr()),
.score_stride = S.unwrap(),
.page_table_stride = P.unwrap(),
.page_table_stride = page_table_stride,
.topk = topk,
.page_bits = page_bits,
.cluster_floor = (batch_size <= kSmallBatchLowFloor) ? kClusterFloorSmall : kClusterFloor,
@@ -438,33 +459,44 @@ struct TopKKernel {
const bool use_cluster = (max_seq_len > params.cluster_floor) && (batch_size <= kClusterMaxBatch);
constexpr bool kUsePDL = true;
if (use_cluster) {
if (batch_size <= kNumPersistentClusters) {
LaunchKernel({batch_size, kClusterSize}, kBlockSize, device)
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
.launch(topk_small_batch_kernel<kUsePDL>, params);
} else {
const uint32_t num_clusters = std::min(batch_size, kNumPersistentClusters);
LaunchKernel({num_clusters, kClusterSize}, kBlockSize, device)
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
.launch(topk_persistent_cluster_kernel<kUsePDL>, params);
const auto mode = page_table.has_value() ? TopKMode::PAGE_TABLE : TopKMode::INDICES;
const auto dispatch = [&]<typename F>(F&& f) {
switch (mode) {
case TopKMode::INDICES:
return f.template operator()<TopKMode::INDICES>();
default:
return f.template operator()<TopKMode::PAGE_TABLE>();
}
};
dispatch([&]<TopKMode kMode>() {
if (use_cluster) {
if (batch_size <= kNumPersistentClusters) {
LaunchKernel({batch_size, kClusterSize}, kBlockSize, device)
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
.launch(topk_small_batch_kernel<kUsePDL, kMode>, params);
} else {
const uint32_t num_clusters = std::min(batch_size, kNumPersistentClusters);
LaunchKernel({num_clusters, kClusterSize}, kBlockSize, device)
.config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
.launch(topk_persistent_cluster_kernel<kUsePDL>, params);
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/3, kMode>, params);
}
} else if (max_seq_len <= kReg2MaxSeqLen) {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/3>, params);
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/0, kMode>, params);
} else if (max_seq_len <= kReg4MaxSeqLen) {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/1, kMode>, params);
} else {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/2, kMode>, params);
}
} else if (max_seq_len <= kReg2MaxSeqLen) {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/0>, params);
} else if (max_seq_len <= kReg4MaxSeqLen) {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/1>, params);
} else {
LaunchKernel(batch_size, kBlockSize, device)
.config({.use_pdl = kUsePDL})
.launch(topk_main_kernel<kUsePDL, /*kLevel=*/2>, params);
}
});
}
};
python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh
+2 -8
View File
@@ -168,25 +168,19 @@ SGL_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, uint3
/// One batch element's worth of work. `emit(pos, raw_idx)` writes the selected raw
/// index to output slot `pos`; `transform_output` then applies the page-table
/// transform in a separate pass (and records the raw index in `raw_out` if set).
/// transform in a separate pass.
struct TopKProblem {
const float* __restrict__ in;
int32_t* __restrict__ out; // page_indices [topk]
int32_t* __restrict__ raw_out; // optional raw (pre-transform) indices [topk]; nullptr if unused
int32_t* __restrict__ out; // page_indices [topk]
const int32_t* __restrict__ page_table;
uint32_t topk;
uint32_t seq_len;
uint32_t page_bits;
// Write the raw selected index; the page-table transform is applied afterwards
// by transform_output() in a separate, pipelined pass. Keeping the per-element
// page_table gather off the atomic-serialized scatter loop is measurably faster
// for both short and long context.
SGL_DEVICE void emit(uint32_t pos, uint32_t raw_idx) const {
out[pos] = static_cast<int32_t>(raw_idx);
}
SGL_DEVICE void transform_output(uint32_t t, int32_t raw) const {
if (raw_out != nullptr) raw_out[t] = raw;
out[t] = raw < 0 ? -1 : page_to_indices(page_table, raw, page_bits);
}
};
python/sglang/kernels/ops/attention/dsv4/topk.py
+11 -4
View File
@@ -89,13 +89,21 @@ def plan_topk_v2(seq_lens: torch.Tensor, static_threshold: int = 0) -> torch.Ten
def topk_transform_512_v2(
scores: torch.Tensor,
seq_lens: torch.Tensor,
page_tables: torch.Tensor,
page_tables: Optional[torch.Tensor],
out_page_indices: torch.Tensor,
page_size: int,
metadata: torch.Tensor,
out_raw_indices: Optional[torch.Tensor] = None,
) -> None:
"""Fused top-k + page-table transform (DeepSeek-V4 top-k v2 kernel).
"""Fused top-k + optional page-table transform (DeepSeek-V4 top-k v2 kernel).
Two output modes, chosen by whether ``page_tables`` is given and resolved to
a device-side template parameter, so an unused page-table gather is compiled
out rather than skipped at runtime:
* ``page_tables=None`` -- ``out_page_indices`` receives the raw selected
indices and no page table is read.
* ``page_tables`` given -- ``out_page_indices`` receives the page-table
transform of them.
IMPORTANT: every entry of ``seq_lens`` must be NON-NEGATIVE, and
``metadata`` must come from :func:`plan_topk_v2` over the same ``seq_lens``
@@ -114,5 +122,4 @@ def topk_transform_512_v2(
out_page_indices,
page_size,
metadata,
out_raw_indices,
)
python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py
+2 -4
View File
@@ -289,8 +289,6 @@ def _topk_transform_v2_paged(
assert 0 < topk <= 2048, f"v2 top-k supports 0 < topk <= 2048, got {topk=}"
page_table = attn_metadata.real_page_table
assert page_table.dtype == torch.int32
lengths_i32 = lengths.to(torch.int32)
# The plan is preprocessed once per forward (DSAMetadata.topk_v2_plan,
# refreshed in-place under CUDA graph) and reused across layers. A missing or
@@ -302,8 +300,8 @@ def _topk_transform_v2_paged(
), "topk_v2_plan must be preprocessed per forward (see DSAMetadata.topk_v2_plan)"
page_size = attn_metadata.page_size
out = logits.new_full((num_rows, topk), -1, dtype=torch.int32)
topk_transform_512_v2(logits, lengths_i32, page_table, out, page_size, plan)
out = logits.new_empty((num_rows, topk), dtype=torch.int32)
topk_transform_512_v2(logits, lengths, page_table, out, page_size, plan)
return out