[FlashInfer V0.6.18] feat(dsv4): support --dsa-topk-backend flashinfer with fused top-k

python/sglang/srt/entrypoints/engine.py

@@ -1678,10 +1678,14 @@ def _set_envs_and_config(server_args: ServerArgs):
 
     # Check flashinfer version
     if not get_bool_env_var("SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK"):
-        if "flashinfer" in cfg.get_attention_backends():
+        if (
+            "flashinfer" in cfg.get_attention_backends()
+            or cfg.dsa_topk_backend == "flashinfer"
+            or cfg.speculative_dsa_topk_backend == "flashinfer"
+        ):
             assert_pkg_version(
                 "flashinfer_python",
-                "0.6.17",
+                "0.6.18",
                 "Please uninstall the old version and "
                 "reinstall the latest version by following the instructions "
                 "at https://docs.flashinfer.ai/installation.html.",

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -705,6 +705,7 @@ def init_forward_metadata_indexer(
             page_size=self.page_size,
             page_table=core_attn_metadata.page_table,
             c4_seq_lens=core_attn_metadata.c4_topk_lengths_raw,
+            use_topk_v2=self.dsa_topk_backend.should_use_topk_v2(),
             # The SM120 FP4 kernel schedules split_kv=128, while the generic
             # JIT metadata planner encodes split_kv=256.
             force_deep_gemm_metadata=(

python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py

@@ -492,6 +492,7 @@ def init_forward_metadata_indexer(self, core_attn_metadata: DSV4AttnMetadata):
             page_size=self.page_size,
             page_table=core_attn_metadata.page_table,
             c4_seq_lens=core_attn_metadata.c4_topk_lengths_raw,
+            use_topk_v2=False,
         )
 
     def init_forward_metadata_decode(

python/sglang/srt/layers/attention/dsv4/indexer.py

@@ -401,11 +401,44 @@ def topk_transform_512_flashinfer_unfused(
     )
 
 
+def topk_transform_512_flashinfer_fused(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_tables: torch.Tensor,
+    out_page_indices: torch.Tensor,
+    page_size: int,
+    out_raw_indices: Optional[torch.Tensor] = None,
+) -> None:
+    import flashinfer
+
+    from sglang.srt.layers.attention.dsa.dsa_topk_backend import (
+        _flashinfer_tie_break_value,
+    )
+
+    flashinfer.top_k_page_table_transform(
+        scores,
+        page_tables.contiguous(),
+        seq_lens.contiguous(),
+        out_page_indices.shape[1],
+        deterministic=envs.SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC.get(),
+        tie_break=_flashinfer_tie_break_value(),
+        dsa_graph_safe=True,
+        page_size=page_size,
+        out=out_page_indices,
+        out_raw_indices=out_raw_indices,
+    )
+
+
 class C4IndexerBackendMixin:
     def __init__(self):
         super().__init__()
         self.debug_use_external_c4_sparse_indices: bool = False
         self.dsa_topk_backend: DSATopKBackend = DSATopKBackend.SGL_KERNEL
+        self.flashinfer_topk_transform: Callable[..., None] = (
+            topk_transform_512_flashinfer_fused
+            if envs.SGLANG_DSA_FUSE_TOPK.get()
+            else topk_transform_512_flashinfer_unfused
+        )
 
     def _forward_prepare_multi_stream(
         self,
@@ -818,15 +851,15 @@ def match_num_queries(tensor: torch.Tensor, value: int) -> torch.Tensor:
                 raw_indices,
             )
         elif self.dsa_topk_backend.is_flashinfer():
-            topk_transform_512_flashinfer_unfused(
+            self.flashinfer_topk_transform(
                 logits,
                 c4_seq_lens,
                 page_table,
                 c4_sparse_page_indices,
                 indexer_metadata.c4_page_size,
                 raw_indices,
             )
-        elif envs.SGLANG_OPT_USE_TOPK_V2.get() and raw_indices is None:
+        elif self.dsa_topk_backend.should_use_topk_v2() and raw_indices is None:
             topk_transform_512_v2(
                 logits,
                 c4_seq_lens,

python/sglang/srt/layers/attention/dsv4/metadata.py

@@ -112,6 +112,7 @@ class PagedIndexerMetadata:
     page_size: int
     page_table: torch.Tensor
     c4_seq_lens: torch.Tensor
+    use_topk_v2: bool
     force_deep_gemm_metadata: bool = False
     use_prefill_cuda_graph: bool = False
     deep_gemm_metadata: Any = field(init=False, repr=False)
@@ -152,9 +153,9 @@ def __post_init__(self):
 
             assert isinstance(self.deep_gemm_metadata, torch.Tensor)
 
-        from sglang.kernels.ops.attention.dsv4 import plan_topk_v2
+        if self.use_topk_v2:
+            from sglang.kernels.ops.attention.dsv4 import plan_topk_v2
 
-        if envs.SGLANG_OPT_USE_TOPK_V2.get():
             self.topk_metadata = plan_topk_v2(self.c4_seq_lens)
         else:
             self.topk_metadata = torch.empty((0,))
@@ -188,6 +189,7 @@ def copy_(self, other: PagedIndexerMetadata):
                 "page_size",
                 "force_deep_gemm_metadata",
                 "use_prefill_cuda_graph",
+                "use_topk_v2",
             ],
             copy_fields=copy_fields,
             assign_fields=assign_fields,

test/registered/kernels/benchmark/attention/bench_topk.py

@@ -19,11 +19,11 @@
 DISABLE_TORCH = True
 
 
-def _make_inputs(batch_size: int, seq_len: int, k: int):
+def _make_inputs(batch_size: int, seq_len: int, k: int, page_size: int = PAGE_SIZE):
     torch.random.manual_seed(42)
     scores = torch.randn(batch_size, seq_len, dtype=torch.float32, device="cuda")
     seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda")
-    num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
+    num_pages = (seq_len + page_size - 1) // page_size
     page_table = (
         torch.arange(num_pages, dtype=torch.int32, device="cuda")
         .unsqueeze(0)
@@ -34,22 +34,11 @@ def _make_inputs(batch_size: int, seq_len: int, k: int):
     return scores, seq_lens, page_table, out
 
 
-def _make_p1_table(batch_size: int, seq_len: int):
-    # flashinfer / torch do a per-token (page_size=1) gather, so they need a
-    # (batch, seq) table (one entry per position) rather than the page-size-64 one.
-    src_page_table = (
-        torch.arange(seq_len, dtype=torch.int32, device="cuda")
-        .unsqueeze(0)
-        .expand(batch_size, -1)
-        .contiguous()
-    )
-    lengths = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda")
-    return src_page_table, lengths
-
-
-def _build_paged_fn(provider: str, batch_size: int, seq_len: int, k: int):
-    scores, seq_lens, page_table, out = _make_inputs(batch_size, seq_len, k)
-    N = PAGE_SIZE
+def _build_paged_fn(
+    provider: str, batch_size: int, seq_len: int, k: int, page_size: int
+):
+    scores, seq_lens, page_table, out = _make_inputs(batch_size, seq_len, k, page_size)
+    N = page_size
 
     def fn(scores, seq_lens, page_table):
         if provider == "jit_v1":
@@ -61,15 +50,23 @@ def fn(scores, seq_lens, page_table):
         elif provider == "flashinfer":
             from flashinfer import top_k_page_table_transform
 
-            return top_k_page_table_transform(scores, page_table, seq_lens, k)
+            return top_k_page_table_transform(
+                scores,
+                page_table,
+                seq_lens,
+                k,
+                dsa_graph_safe=True,
+                page_size=N,
+                out=out,
+            )
         elif provider == "torch":
-            idx = scores.topk(k, dim=-1).indices  # (batch, k) int64
-            return torch.gather(page_table, 1, idx)
+            raw_indices = scores.topk(k, dim=-1).indices
+            physical_pages = torch.gather(page_table, 1, raw_indices // N)
+            out.copy_(physical_pages * N + raw_indices % N)
+            return out
         else:
             raise ValueError(f"unknown provider {provider}")
 
-    if provider in ("flashinfer", "torch"):
-        page_table, seq_lens = _make_p1_table(batch_size, seq_len)
     if provider == "jit_v2":
         metadata = plan_topk_v2(seq_lens)
     return fn, (scores, seq_lens, page_table)
@@ -110,14 +107,17 @@ def fn(scores, seq_lens, offsets):
 @marker.parametrize("k", [512, 1024, 2048], [512])
 @marker.parametrize("seq_len", [2**x for x in range(10, 19)], [4096, 65536])
 @marker.parametrize("batch_size", [2**x for x in range(13)], [1, 128, 1024])
+@marker.parametrize("page_size", [1, 64], [1, 64])
 @marker.benchmark("provider", PRROVIDERS)
-def benchmark_paged(seq_len: int, batch_size: int, k: int, provider: str):
+def benchmark_paged(
+    seq_len: int, batch_size: int, k: int, page_size: int, provider: str
+):
     if k > seq_len:
         marker.skip("k cannot be larger than seq_len")
     if k == 2048 and provider == "jit_v1":
         marker.skip("jit_v1 does not support k=2048")
 
-    fn, input_args = _build_paged_fn(provider, batch_size, seq_len, k)
+    fn, input_args = _build_paged_fn(provider, batch_size, seq_len, k, page_size)
     return marker.do_bench(fn, input_args=input_args, memory_args=input_args[:2])
 
 

test/registered/kernels/ops/attention/test_dsa_indexer.py

@@ -1,4 +1,5 @@
 import unittest
+from itertools import product
 from typing import List, Optional, Tuple
 from unittest.mock import MagicMock, patch
 
@@ -25,14 +26,15 @@
     DeepseekSparseAttnBackend,
     DSAMetadata,
 )
+from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin
 from sglang.srt.layers.layernorm import LayerNorm
 from sglang.srt.layers.linear import LinearBase
 from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.test_utils import CustomTestCase
 
-register_cuda_ci(est_time=18, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-large")
 
 # Global configuration for all indexer tests
 DEFAULT_CONFIG = {
@@ -450,6 +452,151 @@ def _make_tie_free_logits(
             index=perm,
         )
 
+    def _assert_dsv4_compact_topk_result(
+        self,
+        scores: torch.Tensor,
+        seq_lens: torch.Tensor,
+        page_table: torch.Tensor,
+        out_page_indices: torch.Tensor,
+        out_raw_indices: Optional[torch.Tensor],
+        page_size: int,
+    ) -> None:
+        topk = out_page_indices.shape[1]
+        for row in range(scores.shape[0]):
+            seq_len = int(seq_lens[row].item())
+            expected_count = min(seq_len, topk)
+            if seq_len <= topk:
+                expected_raw = torch.arange(
+                    seq_len, dtype=torch.int32, device=scores.device
+                )
+            else:
+                expected_raw = torch.topk(
+                    scores[row, :seq_len], topk, sorted=False
+                ).indices.to(torch.int32)
+            expected_page_indices = (
+                page_table[row, expected_raw.long() // page_size] * page_size
+                + expected_raw % page_size
+            )
+
+            actual_page_indices = out_page_indices[row]
+            actual_page_prefix = actual_page_indices[:expected_count]
+            self.assertTrue(
+                torch.equal(
+                    torch.sort(actual_page_prefix).values,
+                    torch.sort(expected_page_indices).values,
+                )
+            )
+            self.assertTrue(torch.all(actual_page_indices[expected_count:] == -1))
+
+            if out_raw_indices is not None:
+                actual_raw = out_raw_indices[row]
+                actual_raw_prefix = actual_raw[:expected_count]
+                self.assertTrue(
+                    torch.equal(
+                        torch.sort(actual_raw_prefix).values,
+                        torch.sort(expected_raw).values,
+                    )
+                )
+                self.assertTrue(torch.all(actual_raw[expected_count:] == -1))
+                translated_raw = (
+                    page_table[row, actual_raw_prefix.long() // page_size] * page_size
+                    + actual_raw_prefix % page_size
+                )
+                self.assertTrue(torch.equal(actual_page_prefix, translated_raw))
+
+    def test_dsv4_flashinfer_compact_topk_cuda_graph(self):
+        num_rows, max_len, page_size = 4, 2048, 64
+        score_storage = torch.empty(
+            (num_rows, max_len + 16), dtype=torch.float32, device=self.device
+        )
+        scores = score_storage[:, 1 : max_len + 1]
+        self.assertFalse(scores.is_contiguous())
+
+        page_table = 17 + torch.arange(
+            num_rows * (max_len // page_size),
+            dtype=torch.int32,
+            device=self.device,
+        ).reshape(num_rows, max_len // page_size)
+
+        for fuse_topk, topk, with_raw_output in product(
+            (False, True), (512, 1024), (False, True)
+        ):
+            with self.subTest(
+                fuse_topk=fuse_topk,
+                topk=topk,
+                with_raw_output=with_raw_output,
+            ):
+                scores.copy_(
+                    torch.arange(max_len, dtype=torch.float32, device=self.device)
+                    .unsqueeze(0)
+                    .expand(num_rows, -1)
+                )
+                seq_lens = torch.full(
+                    (num_rows,), max_len, dtype=torch.int32, device=self.device
+                )
+                out_page_indices = torch.empty(
+                    (num_rows, topk), dtype=torch.int32, device=self.device
+                )
+                out_raw_indices = (
+                    torch.empty_like(out_page_indices) if with_raw_output else None
+                )
+
+                # Capture with every output slot valid so shrinking lengths on
+                # replay leaves meaningful stale values in any unwritten suffix.
+                with (
+                    envs.SGLANG_DSA_FUSE_TOPK.override(fuse_topk),
+                    envs.SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC.override(True),
+                    envs.SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK.override("small"),
+                ):
+                    backend = C4IndexerBackendMixin()
+                    backend.flashinfer_topk_transform(
+                        scores,
+                        seq_lens,
+                        page_table,
+                        out_page_indices,
+                        page_size,
+                        out_raw_indices,
+                    )
+                    torch.cuda.synchronize()
+
+                    graph = torch.cuda.CUDAGraph()
+                    with torch.cuda.graph(graph):
+                        backend.flashinfer_topk_transform(
+                            scores,
+                            seq_lens,
+                            page_table,
+                            out_page_indices,
+                            page_size,
+                            out_raw_indices,
+                        )
+
+                    scores.copy_(
+                        torch.arange(
+                            max_len, 0, -1, dtype=torch.float32, device=self.device
+                        )
+                        .unsqueeze(0)
+                        .expand(num_rows, -1)
+                    )
+                    page_table.add_(100)
+                    seq_lens.copy_(
+                        torch.tensor(
+                            [0, topk - 1, topk + 73, max_len],
+                            dtype=torch.int32,
+                            device=self.device,
+                        )
+                    )
+                    graph.replay()
+                    torch.cuda.synchronize()
+
+                self._assert_dsv4_compact_topk_result(
+                    scores,
+                    seq_lens,
+                    page_table,
+                    out_page_indices,
+                    out_raw_indices,
+                    page_size,
+                )
+
     def _run_unfused_topk_backend_validity_test(
         self,
         batch_size: int,