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

2026-09-01 01:18:10 -07:00
parent 3484f7f836
commit 5edcd0a445
8 changed files with 333 additions and 31 deletions
@@ -1,12 +1,16 @@
 import sys
 import unittest
+from itertools import product
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch

 import torch

 from sglang.srt.environ import envs
-from sglang.srt.layers.attention.dsv4.indexer import FP8_DTYPE, C4IndexerBackendMixin
+from sglang.srt.layers.attention.dsv4.indexer import (
+    FP8_DTYPE,
+    C4IndexerBackendMixin,
+)
 from sglang.srt.layers.attention.dsv4.metadata import (
    NonPagedIndexerPlan,
    PagedIndexerMetadata,
@@ -43,6 +47,7 @@ class TestDSV4PagedIndexerMetadata(CustomTestCase):
                page_size=256,
                page_table=torch.zeros((1, 1), dtype=torch.int32),
                c4_seq_lens=torch.tensor([65], dtype=torch.int32),
+                use_topk_v2=False,
                force_deep_gemm_metadata=True,
            )

@@ -64,10 +69,119 @@ class TestDSV4PagedIndexerMetadata(CustomTestCase):
                page_size=256,
                page_table=torch.zeros((1, 1), dtype=torch.int32),
                c4_seq_lens=torch.tensor([65], dtype=torch.int32),
+                use_topk_v2=False,
            )

        self.assertIsNone(metadata.deep_gemm_metadata)

+    def test_topk_v2_ineligible_backend_skips_plan(self):
+        with (
+            envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.override(True),
+            envs.SGLANG_OPT_USE_AITER_INDEXER.override(False),
+            envs.SGLANG_OPT_USE_TOPK_V2.override(True),
+            patch("sglang.kernels.ops.attention.dsv4.plan_topk_v2") as plan_topk_v2,
+        ):
+            metadata = PagedIndexerMetadata(
+                page_size=256,
+                page_table=torch.zeros((1, 1), dtype=torch.int32),
+                c4_seq_lens=torch.tensor([65], dtype=torch.int32),
+                use_topk_v2=False,
+            )
+
+        plan_topk_v2.assert_not_called()
+        self.assertEqual(metadata.topk_metadata.numel(), 0)
+
+
+class TestDSV4FlashInferTopK(CustomTestCase):
+    def test_compact_page_transform_respects_fuse_topk(self):
+        score_storage = torch.arange(160, dtype=torch.float32).reshape(2, 80)
+        scores = score_storage[:, 1:65]
+        self.assertFalse(scores.is_contiguous())
+
+        seq_lens = torch.tensor([63, 64], dtype=torch.int32)
+        page_tables = torch.tensor(
+            [[7, 17, 8, 18], [11, 21, 12, 22]], dtype=torch.int32
+        )[:, ::2]
+        self.assertFalse(page_tables.is_contiguous())
+        out_page_indices = torch.empty((2, 8), dtype=torch.int32)
+
+        for fuse_topk, with_raw_output in product((False, True), repeat=2):
+            with self.subTest(fuse_topk=fuse_topk, with_raw_output=with_raw_output):
+                out_raw_indices = (
+                    torch.empty_like(out_page_indices) if with_raw_output else None
+                )
+
+                def fake_top_k(input: torch.Tensor, k: int, **kwargs):
+                    return torch.topk(
+                        input,
+                        k,
+                        dim=-1,
+                        largest=True,
+                        sorted=kwargs["sorted"],
+                    )
+
+                top_k = MagicMock(side_effect=fake_top_k)
+                top_k_page_table_transform = MagicMock()
+                flashinfer = SimpleNamespace(
+                    top_k=top_k,
+                    top_k_page_table_transform=top_k_page_table_transform,
+                )
+
+                with (
+                    patch.dict(sys.modules, {"flashinfer": flashinfer}),
+                    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_tables,
+                        out_page_indices,
+                        page_size=64,
+                        out_raw_indices=out_raw_indices,
+                    )
+
+                if not fuse_topk:
+                    top_k_page_table_transform.assert_not_called()
+                    top_k.assert_called_once()
+                    call = top_k.call_args
+                    self.assertTrue(call.args[0].is_contiguous())
+                    self.assertEqual(call.args[0].shape, scores.shape)
+                    self.assertEqual(call.args[1], out_page_indices.shape[1])
+                    self.assertEqual(
+                        call.kwargs,
+                        {
+                            "sorted": False,
+                            "deterministic": True,
+                            "tie_break": 1,
+                            "dsa_graph_safe": True,
+                        },
+                    )
+                    continue
+
+                top_k.assert_not_called()
+                top_k_page_table_transform.assert_called_once()
+                call = top_k_page_table_transform.call_args
+                self.assertIs(call.args[0], scores)
+                self.assertIsNot(call.args[1], page_tables)
+                self.assertTrue(call.args[1].is_contiguous())
+                self.assertTrue(torch.equal(call.args[1], page_tables))
+                self.assertIs(call.args[2], seq_lens)
+                self.assertEqual(call.args[3], out_page_indices.shape[1])
+                self.assertEqual(
+                    call.kwargs,
+                    {
+                        "deterministic": True,
+                        "tie_break": 1,
+                        "dsa_graph_safe": True,
+                        "page_size": 64,
+                        "out": out_page_indices,
+                        "out_raw_indices": out_raw_indices,
+                    },
+                )
+

 class TestDSV4NonPagedIndexer(CustomTestCase):
    def _is_eligible(self, **overrides):