diff --git a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py index ce57f8925..c80315fd8 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py @@ -34,6 +34,9 @@ class DSATopKBackend(Enum): def is_flashinfer(self) -> bool: return self == DSATopKBackend.FLASHINFER + def should_use_topk_v2(self) -> bool: + return self.is_sgl_kernel() and envs.SGLANG_OPT_USE_TOPK_V2.get() + def topk_func( self, score: torch.Tensor, @@ -88,18 +91,19 @@ class DSATopKBackend(Enum): if not envs.SGLANG_DSA_FUSE_TOPK.get() or force_unfused_topk: return self.topk_func(logits, lengths, topk, row_starts=row_starts) - # Decode-shaped PAGED top-k (plain decode AND spec verify / draft-extend, - # whose expanded rows match the same shape) routes to the DeepSeek-V4 top-k - # v2 JIT kernel, which fuses top-k selection and the page-table transform in - # one launch and consumes the indexer's own page_size>=1 table directly, so - # no page_size=1 table is materialized. Shared by DeepSeek-V3.2 and GLM DSA. + # Decode-shaped PAGED top-k for the SGL backend (plain decode AND spec + # verify / draft-extend, whose expanded rows match the same shape) routes + # to the DeepSeek-V4 top-k v2 JIT kernel. It fuses top-k selection and the + # page-table transform in one launch and consumes the indexer's own + # page_size>=1 table directly, so no page_size=1 table is materialized. + # Shared by DeepSeek-V3.2 and GLM DSA. # This is a deterministic dispatch on the work shape, not a best-effort # attempt: the fused-decode CUDA graph drops the page_size=1 table for # exactly this case (see dsa_drop_wide_page_table), so once the shape # matches we commit to v2 and never silently fall back to the legacy # page_size=1 path from here. if ( - envs.SGLANG_OPT_USE_TOPK_V2.get() + self.should_use_topk_v2() and topk_transform_method == TopkTransformMethod.PAGED and row_starts is None and batch_idx_list is None @@ -154,6 +158,26 @@ class DSATopKBackend(Enum): import flashinfer if topk_transform_method == TopkTransformMethod.PAGED: + if row_starts is not None: + # Packed PAGED extend uses batch-global score offsets with + # request-local page tables. FlashInfer applies row_starts + # to both, so reuse the SGL transform. + from sgl_kernel import fast_topk_transform_fused + + page_table_size_1 = ( + attn_metadata.page_table_1[batch_idx_list] + if batch_idx_list is not None + else attn_metadata.page_table_1 + ) + return fast_topk_transform_fused( + score=logits, + lengths=lengths, + page_table_size_1=page_table_size_1, + cu_seqlens_q=cu_seqlens_q_topk, + topk=topk, + row_starts=row_starts, + ) + row_to_batch, local_row_starts = _build_flashinfer_paged_args( attn_metadata=attn_metadata, row_starts=row_starts, @@ -268,7 +292,6 @@ def _topk_transform_v2_paged( ``seqlens_expand_kernel``); 0 takes the trivial all-(-1) output path. """ from sglang.kernels.ops.attention.dsv4.topk import topk_transform_512_v2 - from sglang.srt.model_executor.forward_context import get_token_to_kv_pool num_rows = logits.shape[0] @@ -297,7 +320,7 @@ def _topk_transform_v2_paged( plan is not None and plan.shape[0] == num_rows + 1 ), "topk_v2_plan must be preprocessed per forward (see DSAMetadata.topk_v2_plan)" - page_size = get_token_to_kv_pool().page_size + 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) return out @@ -317,6 +340,8 @@ def _build_flashinfer_paged_args( else None ) + # Both dynamic mappings contain one entry per logit row. Supplying the known + # size avoids synchronizing CUDA to infer the sum of the repeat counts. if ( row_to_batch is not None and cu_seqlens_q_topk is not None @@ -325,7 +350,9 @@ def _build_flashinfer_paged_args( q_lens = (cu_seqlens_q_topk[1:] - cu_seqlens_q_topk[:-1]).to( dtype=torch.int32, device=device ) - row_to_batch = torch.repeat_interleave(row_to_batch, q_lens) + row_to_batch = torch.repeat_interleave( + row_to_batch, q_lens, output_size=num_rows + ) if row_to_batch is None and cu_seqlens_q_topk is not None: # Decode-like case (one query row per batch) does not need an explicit mapping. @@ -338,6 +365,7 @@ def _build_flashinfer_paged_args( row_to_batch = torch.repeat_interleave( torch.arange(q_lens.shape[0], dtype=torch.int32, device=device), q_lens, + output_size=num_rows, ) if row_starts is not None and row_to_batch is None: diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 1223de526..df3a0ed61 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -229,7 +229,7 @@ class DSAMetadata: paged_mqa_ctx_lens_2d: Optional[torch.Tensor] = None # Precomputed once per forward batch and reused across layers: the # DeepSeek-V4 top-k v2 plan (cluster-threshold metadata) for the folded - # decode top-k transform. None unless SGLANG_OPT_USE_TOPK_V2 and decode. + # decode top-k transform. None unless the SGL top-k v2 path is enabled. topk_v2_plan: Optional[torch.Tensor] = None # The sum of sequence lengths for key, prefill only seq_lens_sum: Optional[int] = None @@ -329,6 +329,8 @@ class DSAIndexerMetadata(BaseIndexerMetadata): cu_topk_indices_offset = torch.repeat_interleave( cu_seqlens_q_topk[:-1], cu_seqlens_q, + # Avoid reading sum(cu_seqlens_q) back to the host. + output_size=logits.shape[0], ) else: cu_seqlens_q_topk = self.attn_metadata.cu_seqlens_q @@ -778,8 +780,8 @@ class DeepseekSparseAttnBackend( # 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 - # when the fold is disabled; such metadata is never dispatched to v2. - if not envs.SGLANG_OPT_USE_TOPK_V2.get(): + # 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 from sglang.kernels.ops.attention.dsv4.topk import plan_topk_v2 @@ -1250,12 +1252,12 @@ class DeepseekSparseAttnBackend( # page_size=1 table. This MUST match the exact condition under which # `DSATopKBackend.topk_transform` dispatches decode PAGED to # `_topk_transform_v2_paged` -- otherwise the legacy transform would read a - # dropped (None) table. Hence: fused top-k AND v2 enabled AND index_topk in - # the kernel's supported range, on CUDA with page_size>1. Excludes HIP (its - # indexer reads page_table_1), hisparse (needs page_size=1 loc translation), - # and spec decoding (MTP precompute fast-path + target-verify/draft-extend - # still consume the wide table). Computed once from stable config; the graph - # is captured once per process. + # dropped (None) table. Hence: SGL top-k backend AND fused top-k AND v2 + # enabled AND index_topk in the kernel's supported range, on CUDA with + # page_size>1. Excludes HIP (its indexer reads page_table_1), hisparse + # (needs page_size=1 loc translation), and spec decoding (MTP precompute + # fast-path + target-verify/draft-extend still consume the wide table). + # Computed once from stable config; the graph is captured once per process. self.dsa_drop_wide_page_table = ( is_cuda() and not _is_hip @@ -1263,9 +1265,9 @@ class DeepseekSparseAttnBackend( and self.hisparse_coordinator is None and not self.speculative_num_draft_tokens and self.use_fused_topk - and envs.SGLANG_OPT_USE_TOPK_V2.get() + and self.dsa_topk_backend.should_use_topk_v2() and self.dsa_index_topk is not None - and self.dsa_index_topk <= 2048 + and 0 < self.dsa_index_topk <= 2048 ) max_ctx_len = self.req_to_token.shape[1] diff --git a/test/registered/kernels/ops/attention/test_dsa_indexer.py b/test/registered/kernels/ops/attention/test_dsa_indexer.py index f032c1a76..2cbd4dd62 100644 --- a/test/registered/kernels/ops/attention/test_dsa_indexer.py +++ b/test/registered/kernels/ops/attention/test_dsa_indexer.py @@ -679,7 +679,14 @@ class TestDSAIndexer(CustomTestCase): topk_backend=DSATopKBackend.FLASHINFER, ) - with envs.SGLANG_DSA_FUSE_TOPK.override(True): + repeat_interleave = torch.repeat_interleave + with ( + envs.SGLANG_DSA_FUSE_TOPK.override(True), + patch( + "sglang.srt.layers.attention.dsa_backend.torch.repeat_interleave", + wraps=repeat_interleave, + ) as mock_repeat_interleave, + ): out_sgl = metadata_sgl.topk_transform( logits, topk, @@ -695,6 +702,15 @@ class TestDSAIndexer(CustomTestCase): batch_idx_list=batch_idx_list, ) + if query_lens is not None: + self.assertTrue(mock_repeat_interleave.call_args_list) + self.assertTrue( + all( + call.kwargs.get("output_size") == num_rows + for call in mock_repeat_interleave.call_args_list + ) + ) + self.assertEqual(out_sgl.shape, out_flashinfer.shape) self.assertEqual(out_sgl.dtype, out_flashinfer.dtype) self.assertEqual(out_sgl.dtype, torch.int32) @@ -957,13 +973,6 @@ class TestDSAIndexer(CustomTestCase): TopkTransformMethod.RAGGED, ]: for with_row_starts in [False, True]: - if ( - topk_transform_method == TopkTransformMethod.PAGED - and with_row_starts - ): - # The synthetic paged fixture uses the decode-like row mapping. - # Ragged fused and unfused cases cover shifted row windows. - continue with self.subTest( tie_break=tie_break, topk_transform_method=topk_transform_method.name, @@ -994,6 +1003,71 @@ class TestDSAIndexer(CustomTestCase): with_row_starts=False, query_lens=[1, 2, 3, 1, 2, 1, 3, 2], ) + with self.subTest( + tie_break=tie_break, + topk_transform_method=TopkTransformMethod.PAGED.name, + with_row_starts=True, + query_lens="multi", + ): + with envs.SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK.override(tie_break): + self._run_fused_topk_backend_equivalence_test( + batch_size=batch_size, + max_score_len=max_score_len, + topk=topk, + topk_transform_method=TopkTransformMethod.PAGED, + with_row_starts=True, + query_lens=[1, 2, 3, 1, 2, 1, 3, 2], + ) + + def test_topk_v2_respects_topk_backend(self): + seq_lens = torch.tensor([2048, 4096], dtype=torch.int32, device=self.device) + expected_plan = torch.empty(3, dtype=torch.int32, device=self.device) + + for topk_backend, should_use_topk_v2 in [ + (DSATopKBackend.SGL_KERNEL, True), + (DSATopKBackend.FLASHINFER, False), + ]: + with self.subTest(topk_backend=topk_backend.value): + backend = object.__new__(DeepseekSparseAttnBackend) + backend.device = self.device + backend.real_page_size = 64 + backend.hisparse_coordinator = None + backend.speculative_num_draft_tokens = 0 + backend.use_fused_topk = True + backend.dsa_topk_backend = topk_backend + backend.dsa_index_topk = 2048 + backend.dsa_decode_impl = "fa3" + backend.req_to_token = torch.empty( + 2, 4096, dtype=torch.int32, device=self.device + ) + + with ( + envs.SGLANG_OPT_USE_TOPK_V2.override(True), + patch( + "sglang.kernels.ops.attention.dsv4.topk.plan_topk_v2", + return_value=expected_plan, + ) as mock_plan_topk_v2, + ): + self.assertEqual( + topk_backend.should_use_topk_v2(), should_use_topk_v2 + ) + actual_plan = backend._build_topk_v2_plan(seq_lens) + backend.init_cuda_graph_state(max_bs=2, max_num_tokens=2) + + if should_use_topk_v2: + self.assertIs(actual_plan, expected_plan) + mock_plan_topk_v2.assert_called_once_with(seq_lens) + else: + self.assertIsNone(actual_plan) + mock_plan_topk_v2.assert_not_called() + self.assertEqual( + backend.dsa_drop_wide_page_table, + should_use_topk_v2, + ) + self.assertEqual( + backend.decode_cuda_graph_metadata["page_table"] is None, + should_use_topk_v2, + ) # TODO: enable this test after indexer accuracy aligned # @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm")