diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 8a5167d07..b103775b0 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -1698,7 +1698,11 @@ 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 attention_backends_of(resolved_view(cfg)): + if ( + "flashinfer" in attention_backends_of(resolved_view(cfg)) + or cfg.dsa_topk_backend == "flashinfer" + or cfg.speculative_dsa_topk_backend == "flashinfer" + ): assert_pkg_version( "flashinfer_python", "0.6.18", diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 2f1f1f2f3..fcf42aaf0 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -704,6 +704,7 @@ class DeepseekV4AttnBackend( 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() and not _is_xpu, # The SM120 FP4 kernel schedules split_kv=128, while the generic # JIT metadata planner encodes split_kv=256. force_deep_gemm_metadata=( diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index a5cd7ee8a..e7055bb2a 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -492,6 +492,7 @@ class DeepseekV4HipRadixBackend( 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( diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 60b8872d7..014678319 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -404,11 +404,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, @@ -829,7 +862,7 @@ class C4IndexerBackendMixin: raw_indices, ) elif self.dsa_topk_backend.is_flashinfer(): - topk_transform_512_flashinfer_unfused( + self.flashinfer_topk_transform( logits, c4_seq_lens, page_table, @@ -837,7 +870,7 @@ class C4IndexerBackendMixin: 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, diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index 60f57e1d8..0d01451fc 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/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,7 +153,7 @@ class PagedIndexerMetadata: assert isinstance(self.deep_gemm_metadata, torch.Tensor) - if envs.SGLANG_OPT_USE_TOPK_V2.get() and not is_xpu(): + if self.use_topk_v2: from sglang.kernels.ops.attention.dsv4 import plan_topk_v2 self.topk_metadata = plan_topk_v2(self.c4_seq_lens) @@ -188,6 +189,7 @@ class PagedIndexerMetadata: "page_size", "force_deep_gemm_metadata", "use_prefill_cuda_graph", + "use_topk_v2", ], copy_fields=copy_fields, assign_fields=assign_fields, diff --git a/test/registered/kernels/benchmark/attention/bench_topk.py b/test/registered/kernels/benchmark/attention/bench_topk.py index 8a49d05c1..7bf950570 100644 --- a/test/registered/kernels/benchmark/attention/bench_topk.py +++ b/test/registered/kernels/benchmark/attention/bench_topk.py @@ -19,11 +19,11 @@ PAGE_SIZE = 64 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 _build_paged_fn(provider: str, batch_size: int, seq_len: int, k: int): 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 @@ if not DISABLE_TORCH: @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]) diff --git a/test/registered/kernels/ops/attention/test_dsa_indexer.py b/test/registered/kernels/ops/attention/test_dsa_indexer.py index 25f014347..df1c52082 100644 --- a/test/registered/kernels/ops/attention/test_dsa_indexer.py +++ b/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,6 +26,7 @@ from sglang.srt.layers.attention.dsa_backend import ( 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 @@ -32,7 +34,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMo 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 @@ class TestDSAIndexer(CustomTestCase): 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, diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py index e73e9a992..8d3af8edb 100644 --- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py +++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py @@ -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):