fix(dsa): correct packed FlashInfer top-k and backend selection semantics (#32490)

Co-authored-by: Parth Chadha <parth@humansand.ai>

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,
+        # Decode-shaped PAGED top-k for the SGL backend (plain decode AND spec
-        # whose expanded rows match the same shape) routes to the DeepSeek-V4 top-k
+        # verify / draft-extend, whose expanded rows match the same shape) routes
-        # v2 JIT kernel, which fuses top-k selection and the page-table transform in
+        # to the DeepSeek-V4 top-k v2 JIT kernel. It fuses top-k selection and the
-        # one launch and consumes the indexer's own page_size>=1 table directly, so
+        # page-table transform in one launch and consumes the indexer's own
-        # no page_size=1 table is materialized. Shared by DeepSeek-V3.2 and GLM DSA.
+        # 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:

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.
+        # when the SGL v2 path is disabled; such metadata is never dispatched to v2.
-        if not envs.SGLANG_OPT_USE_TOPK_V2.get():
+        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
+        # dropped (None) table. Hence: SGL top-k backend AND fused top-k AND v2
-        # the kernel's supported range, on CUDA with page_size>1. Excludes HIP (its
+        # enabled AND index_topk in the kernel's supported range, on CUDA with
-        # indexer reads page_table_1), hisparse (needs page_size=1 loc translation),
+        # page_size>1. Excludes HIP (its indexer reads page_table_1), hisparse
-        # and spec decoding (MTP precompute fast-path + target-verify/draft-extend
+        # (needs page_size=1 loc translation), and spec decoding (MTP precompute
-        # still consume the wide table). Computed once from stable config; the graph
+        # fast-path + target-verify/draft-extend still consume the wide table).
-        # is captured once per process.
+        # 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]

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")