diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh index b09a15364..426f04788 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh @@ -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(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(topk); + } SGL_DEVICE TopKProblem problem(uint32_t batch_id, uint32_t seq_len) const { const auto k = static_cast(topk); return TopKProblem{ @@ -156,7 +161,7 @@ SGL_DEVICE void for_each_item(uint32_t topk, const F& f) { } template -SGL_DEVICE void trivial_transform(const TopKProblem& problem) { +SGL_DEVICE void trivial_transform(const TopKProblem& problem, int32_t* raw_output_ptr) { device::PDLWaitPrimary(); device::PDLTriggerSecondary(); 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 +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 smem; - if (problem.seq_len <= problem.topk) return trivial_transform(problem); + if (problem.seq_len <= problem.topk) + return trivial_transform(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(); if constexpr (kNeedStaging) { __syncthreads(); - problem_transform(problem, params.get_output_ptr(blockIdx.x)); + problem_transform(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 smem; - if (problem.seq_len <= problem.topk) return trivial_transform(problem); + if (problem.seq_len <= problem.topk) + return trivial_transform(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(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 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 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().with_device(device_).verify(raw_indices.value()); + raw_indices_ptr = static_cast(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(seq_lens.data_ptr()), .page_table = page_table_ptr, .page_indices = static_cast(page_indices.data_ptr()), + .raw_indices = raw_indices_ptr, .metadata = static_cast(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 = [&](F&& f) { switch (mode) { case TopKMode::INDICES: return f.template operator()(); + case TopKMode::DUAL_OUTPUT: + return f.template operator()(); default: return f.template operator()(); } diff --git a/python/sglang/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index 7acd8b405..e5f0d0fd1 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/python/sglang/kernels/ops/attention/dsv4/topk.py @@ -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, ) diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 0e4085a91..a778b2307 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -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( diff --git a/test/registered/kernels/ops/attention/test_topk_v2.py b/test/registered/kernels/ops/attention/test_topk_v2.py index cae08c08c..86adf27d9 100644 --- a/test/registered/kernels/ops/attention/test_topk_v2.py +++ b/test/registered/kernels/ops/attention/test_topk_v2.py @@ -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 diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py index baa490a6a..8c0284cf2 100644 --- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py +++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py @@ -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)