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

[DSV4] Support raw-index output in TopK v2 (#33672)

Browse Source 
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>
This commit is contained in:
weireweire
2026-09-11 07:17:36 -07:00
committed by GitHub
co-authored by weireweire Brayden Zhong Po-Han Huang
parent ab9750fb35
commit 335f6aab27
5 changed files with 167 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh
+37 -11
View File
@@ -30,8 +30,9 @@ 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`
INDICES, ///< raw selected indices into `out`; `page_table` unused
PAGE_TABLE, ///< page-table-transformed indices into `out`
DUAL_OUTPUT, ///< page-table-transformed indices into `out` and raw indices into `raw_indices`
};
using Register2 = impl::TopKRegister<2>; // <= 8192, register-resident, 1 read
@@ -76,6 +77,7 @@ struct TopKPagedParams {
const int32_t* __restrict__ seq_lens;
const int32_t* __restrict__ page_table;
int32_t* __restrict__ page_indices;
int32_t* __restrict__ raw_indices;
const PlanItem* __restrict__ metadata; // [0]=GlobalMetadata, [1+i]=PlanItem
int64_t score_stride;
int64_t page_table_stride;
@@ -95,6 +97,9 @@ struct TopKPagedParams {
SGL_DEVICE int32_t* get_output_ptr(uint32_t batch_id) const {
return page_indices + batch_id * static_cast<int64_t>(topk);
}
SGL_DEVICE int32_t* get_raw_output_ptr(uint32_t batch_id) const {
return raw_indices == nullptr ? nullptr : raw_indices + batch_id * static_cast<int64_t>(topk);
}
SGL_DEVICE TopKProblem problem(uint32_t batch_id, uint32_t seq_len) const {
const auto k = static_cast<int64_t>(topk);
return TopKProblem{
@@ -156,7 +161,7 @@ SGL_DEVICE void for_each_item(uint32_t topk, const F& f) {
}
template <bool kPDL, TopKMode kMode>
SGL_DEVICE void trivial_transform(const TopKProblem& problem) {
SGL_DEVICE void trivial_transform(const TopKProblem& problem, int32_t* raw_output_ptr) {
device::PDLWaitPrimary<kPDL>();
device::PDLTriggerSecondary<kPDL>();
for_each_item(problem.topk, [&](uint32_t tx, uint32_t) {
@@ -165,17 +170,23 @@ SGL_DEVICE void trivial_transform(const TopKProblem& problem) {
problem.emit(tx, idx);
} else {
problem.transform_output(tx, idx);
if constexpr (kMode == TopKMode::DUAL_OUTPUT) raw_output_ptr[tx] = idx;
}
});
}
SGL_DEVICE void problem_transform(TopKProblem& problem, int32_t* output_ptr) {
template <TopKMode kMode>
SGL_DEVICE void problem_transform(TopKProblem& problem, int32_t* output_ptr, int32_t* raw_output_ptr) {
static_assert(kMode != TopKMode::INDICES, "problem_transform requires page-table output");
static_assert(kMaxTopK % kBlockSize == 0);
constexpr uint32_t kNumElems = kMaxTopK / kBlockSize;
int32_t source_index[kNumElems];
for_each_item(problem.topk, [&](uint32_t tx, uint32_t i) { source_index[i] = problem.out[tx]; });
problem.out = output_ptr;
for_each_item(problem.topk, [&](uint32_t tx, uint32_t i) { problem.transform_output(tx, source_index[i]); });
for_each_item(problem.topk, [&](uint32_t tx, uint32_t i) {
problem.transform_output(tx, source_index[i]);
if constexpr (kMode == TopKMode::DUAL_OUTPUT) raw_output_ptr[tx] = source_index[i];
});
}
/**
@@ -279,7 +290,8 @@ TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKPagedParams params
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);
if (problem.seq_len <= problem.topk)
return trivial_transform<kPDLEarly, kMode>(problem, params.get_raw_output_ptr(blockIdx.x));
constexpr bool kNeedStaging = kMode != TopKMode::INDICES;
__shared__ int32_t s_topk_indices[kNeedStaging ? kMaxTopK : 1];
@@ -310,7 +322,7 @@ TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKPagedParams params
device::PDLTriggerSecondary<kPDL>();
if constexpr (kNeedStaging) {
__syncthreads();
problem_transform(problem, params.get_output_ptr(blockIdx.x));
problem_transform<kMode>(problem, params.get_output_ptr(blockIdx.x), params.get_raw_output_ptr(blockIdx.x));
}
}
@@ -320,7 +332,8 @@ CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKPag
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, kMode>(problem);
if (problem.seq_len <= problem.topk)
return trivial_transform<kPDL, kMode>(problem, params.get_raw_output_ptr(blockIdx.x));
constexpr bool kNeedStaging = kMode != TopKMode::INDICES;
__shared__ int32_t s_topk_indices[kNeedStaging ? kMaxTopK : 1];
@@ -359,7 +372,7 @@ CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKPag
// 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));
problem_transform<kMode>(problem, params.get_output_ptr(blockIdx.x), params.get_raw_output_ptr(blockIdx.x));
}
}
#endif // !USE_ROCM
@@ -490,7 +503,8 @@ struct TopKKernel {
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::TensorView metadata,
const tvm::ffi::Optional<tvm::ffi::TensorView> raw_indices) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto Bp1 = SymbolicSize{"batch_size_plus_1"};
@@ -532,6 +546,13 @@ struct TopKKernel {
.with_device(device_)
.verify(metadata);
int32_t* raw_indices_ptr = nullptr;
if (raw_indices.has_value()) {
RuntimeCheck(page_table.has_value(), "raw_indices requires a page table");
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");
@@ -555,6 +576,7 @@ struct TopKKernel {
.seq_lens = static_cast<const int32_t*>(seq_lens.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 = page_table_stride,
@@ -567,11 +589,15 @@ struct TopKKernel {
const bool use_cluster = (max_seq_len > params.cluster_floor) && (batch_size <= kClusterMaxBatch);
#endif
constexpr bool kUsePDL = true;
const auto mode = page_table.has_value() ? TopKMode::PAGE_TABLE : TopKMode::INDICES;
const auto mode = raw_indices.has_value() ? TopKMode::DUAL_OUTPUT
: 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>();
case TopKMode::DUAL_OUTPUT:
return f.template operator()<TopKMode::DUAL_OUTPUT>();
default:
return f.template operator()<TopKMode::PAGE_TABLE>();
}
python/sglang/kernels/ops/attention/dsv4/topk.py
+17 -3
View File
@@ -136,17 +136,20 @@ def topk_transform_paged_v2(
out_page_indices: torch.Tensor,
page_size: int,
metadata: torch.Tensor,
out_raw_indices: Optional[torch.Tensor] = None,
) -> None:
"""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:
Output mode is chosen from ``page_tables`` and ``out_raw_indices`` 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.
* Both outputs given -- ``out_page_indices`` receives the page-table
transform and ``out_raw_indices`` receives the selected raw indices.
IMPORTANT: every entry of ``seq_lens`` must be NON-NEGATIVE, and
``metadata`` must come from :func:`plan_topk_v2` over the same ``seq_lens``
@@ -158,6 +161,16 @@ def topk_transform_paged_v2(
the output is all -1.
"""
if is_xpu():
if out_raw_indices is not None:
topk_transform_paged(
scores,
seq_lens,
page_tables,
out_page_indices,
page_size,
out_raw_indices,
)
return
torch.ops.sgl_kernel.topk_transform_paged(
scores,
seq_lens,
@@ -175,4 +188,5 @@ def topk_transform_paged_v2(
out_page_indices,
page_size,
metadata,
out_raw_indices,
)
python/sglang/srt/layers/attention/dsv4/indexer.py
+5 -4
View File
@@ -837,14 +837,14 @@ class C4IndexerBackendMixin:
)
raw_indices = None
if capture_enabled:
if core_metadata.c4_sparse_raw_indices is not None:
raw_indices = core_metadata.c4_sparse_raw_indices
elif capture_enabled:
raw_indices = torch.empty_like(c4_sparse_page_indices)
elif hisparse_decode:
raw_indices = hisparse_coordinator.raw_indices_buffer[
: c4_sparse_page_indices.size(0)
]
elif core_metadata.c4_sparse_raw_indices is not None:
raw_indices = core_metadata.c4_sparse_raw_indices
all_rows = slice(0, _c4sl.shape[0])
@@ -868,7 +868,7 @@ class C4IndexerBackendMixin:
indexer_metadata.compressed_page_size,
row_raw_indices,
)
elif self.dsa_topk_backend.should_use_topk_v2() and raw_indices is None:
elif self.dsa_topk_backend.should_use_topk_v2():
topk_transform_paged_v2(
logits,
c4_seq_lens[rows],
@@ -882,6 +882,7 @@ class C4IndexerBackendMixin:
if rows == all_rows or not is_hip()
else plan_topk_v2(c4_seq_lens[rows])
),
row_raw_indices,
)
else:
topk_transform_paged(
test/registered/kernels/ops/attention/test_topk_v2.py
+35
View File
@@ -183,6 +183,20 @@ def _run_raw(scores, seq_lens, k):
return [[v for v in out_cpu[i] if v != -1] for i in range(batch)]
def _run_dual(scores, seq_lens, page_table, inv_cpu, k):
batch = scores.shape[0]
metadata = _plan(seq_lens)
out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
raw = torch.full_like(out, -1)
topk_transform_paged_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata, raw)
torch.cuda.synchronize()
out_cpu = out.cpu().tolist()
raw_cpu = raw.cpu().tolist()
transformed_raw = [_invert(out_cpu[i], inv_cpu[i]) for i in range(batch)]
direct_raw = [[v for v in raw_cpu[i] if v != -1] for i in range(batch)]
return transformed_raw, direct_raw
@pytest.mark.parametrize("page_mode", ["identity", "perm"])
@pytest.mark.parametrize("k", [512, 1024, 2048])
@pytest.mark.parametrize("batch,seq", FIXED_CONFIGS)
@@ -270,6 +284,27 @@ def test_topk_v2_output_indices(batch: int, seq: int, k: int) -> None:
_assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, seq_lens.cpu(), k)
@pytest.mark.parametrize(
"batch,seq", [(8, 256), (8, 8192), (4, 32768), (2, 131072), (31, 131072)]
)
@torch.inference_mode()
def test_topk_v2_dual_output(batch: int, seq: int) -> None:
"""The dual mode returns the same selection before and after page transform."""
k = 512
torch.manual_seed(batch * 100003 + seq * 7 + k + 2)
device = "cuda"
scores = torch.randn(batch, seq, dtype=torch.float32, device=device)
seq_lens = torch.full((batch,), seq, dtype=torch.int32, device=device)
num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE
page_table, inv_cpu = _make_page_table(batch, num_pages, "perm", device)
transformed_raw, direct_raw = _run_dual(scores, seq_lens, page_table, inv_cpu, k)
for row in range(batch):
assert sorted(transformed_raw[row]) == sorted(direct_raw[row])
ref_raw = _reference(scores, seq_lens, k)
_assert_topk_close(scores.cpu(), ref_raw, direct_raw, batch, seq_lens.cpu(), k)
# --- ragged entry point ------------------------------------------------------
# Rows select inside `[row_start, row_start + seq_len)` of their score row and
# emit `position + offset`. The window start is an arbitrary token offset, so
test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
+73
View File
@@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch
import torch
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend
from sglang.srt.layers.attention.dsv4.indexer import (
FP8_DTYPE,
C4IndexerBackendMixin,
@@ -209,6 +210,78 @@ class TestDSV4FlashInferTopK(CustomTestCase):
)
class TestDSV4TopKDispatch(CustomTestCase):
def test_v2_raw_output_uses_sparse_prefill_buffer_with_capture(self):
page_table = torch.zeros((1, 1), dtype=torch.int32)
c4_seq_lens = torch.ones(1, dtype=torch.int32)
page_indices = torch.full((1, 512), -1, dtype=torch.int32)
raw_indices = torch.full_like(page_indices, -1)
topk_metadata = torch.zeros((2, 2), dtype=torch.int32)
indexer_metadata = object.__new__(PagedIndexerMetadata)
indexer_metadata.page_size = 256
indexer_metadata.page_table = page_table
indexer_metadata.c4_seq_lens = c4_seq_lens
indexer_metadata.topk_metadata = topk_metadata
logits = torch.empty((1, 65), dtype=torch.float32)
backend = C4IndexerBackendMixin()
backend.dsa_topk_backend = DSATopKBackend.SGL_KERNEL
backend.token_to_kv_pool = SimpleNamespace(
layer_mapping={0: SimpleNamespace(compress_layer_id=7)}
)
backend.forward_metadata = SimpleNamespace(
indexer_metadata=indexer_metadata,
core_metadata=SimpleNamespace(
positions=torch.arange(1, dtype=torch.int64),
page_table=page_table,
c4_sparse_page_indices=page_indices,
c4_sparse_raw_indices=raw_indices,
),
)
backend.hisparse_coordinator = None
backend._forward_prepare_normal = MagicMock(
return_value=(
torch.empty((1, 1, 128)),
torch.empty((1, 1, 1)),
)
)
backend._get_nonpaged_indexer_plan = MagicMock(return_value=object())
backend._forward_nonpaged_indexer = MagicMock(return_value=logits)
indexer_capturer = MagicMock()
with (
envs.SGLANG_OPT_USE_TILELANG_INDEXER.override(False),
envs.SGLANG_OPT_USE_AITER_INDEXER.override(False),
envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.override(True),
envs.SGLANG_OPT_USE_TOPK_V2.override(True),
patch(
f"{_INDEXER}.get_global_indexer_capturer",
return_value=indexer_capturer,
),
patch(f"{_INDEXER}.topk_transform_paged") as topk_v1,
patch(f"{_INDEXER}.topk_transform_paged_v2") as topk_v2,
):
backend.forward_c4_indexer(
x=torch.empty((1, 1)),
q_lora=torch.empty((1, 1)),
c4_indexer=SimpleNamespace(use_fp4_indexer=False, layer_id=0),
forward_batch=SimpleNamespace(forward_mode=ForwardMode.EXTEND),
)
topk_v2.assert_called_once()
args = topk_v2.call_args.args
self.assertIs(args[0], logits)
torch.testing.assert_close(args[1], c4_seq_lens)
torch.testing.assert_close(args[2], page_table)
torch.testing.assert_close(args[3], page_indices)
self.assertEqual(args[4], 64)
self.assertIs(args[5], topk_metadata)
self.assertEqual(args[6].data_ptr(), raw_indices.data_ptr())
topk_v1.assert_not_called()
indexer_capturer.capture.assert_called_once_with(7, raw_indices)
class TestDSV4NonPagedIndexer(CustomTestCase):
def _is_eligible(self, **overrides):
backend = SimpleNamespace(hisparse_coordinator=None)