diff --git a/python/sglang/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index 24c3bbd96394..003711b06f3e 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/python/sglang/kernels/ops/attention/dsv4/topk.py @@ -28,6 +28,10 @@ def _jit_topk_v1_module(): ) +# Largest ``k`` the v1 kernel selects (kMaxTopK in deepseek_v4/topk_v1.cuh). +TOPK_V1_MAX_K = 1024 + + @cache_once def _jit_topk_v2_module(): from sglang.kernels.jit.utils.occupancy import ( diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py index 883cf0575767..a6b2649b668a 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py @@ -14,7 +14,10 @@ from sglang.kernels.ops.attention.dsv4.index_logits import ( deep_gemm_fp4_paged_mqa_logits, ) -from sglang.kernels.ops.attention.dsv4.topk import topk_transform_paged_torch +from sglang.kernels.ops.attention.dsv4.topk import ( + TOPK_V1_MAX_K, + topk_transform_paged_torch, +) from sglang.srt.layers.attention.dsv4.indexer import topk_transform_paged_from_metadata from .scoring import ( @@ -175,7 +178,12 @@ def _deep_gemm_prefill_captured(self, inputs: CapturedPrefillInputs) -> None: width = metadata.max_compressed_seq_len lens = metadata.compressed_seq_lens page_table = metadata.page_table - topk = min(indexer.index_topk, width) + # With top-k v2 off (the low-ratio prefill metadata), the metadata path runs + # the v1 kernel, which takes k from the output width and caps it at + # TOPK_V1_MAX_K; wider outputs use the torch top-k over min(k, width) columns. + out_width = inputs.out_page_indices.shape[1] + torch_topk = not metadata.use_topk_v2 and out_width > TOPK_V1_MAX_K + topk = min(out_width, width) for rows, plan in metadata.row_chunks(): logits = deep_gemm_fp4_paged_mqa_logits( (q_fp4[rows], q_sf[rows]), @@ -186,11 +194,22 @@ def _deep_gemm_prefill_captured(self, inputs: CapturedPrefillInputs) -> None: plan, width, ) - topk_transform_paged_torch( + if torch_topk: + topk_transform_paged_torch( + logits, + lens[rows], + page_table[rows], + inputs.out_page_indices[rows, :topk], + page_size, + inputs.out_raw_indices[rows, :topk], + ) + continue + # Per row up to its own length: a wide graph then costs only the + # logits buffer, not a full-static-width top-k. + topk_transform_paged_from_metadata( logits, - lens[rows], - page_table[rows], - inputs.out_page_indices[rows, :topk], - page_size, - inputs.out_raw_indices[rows, :topk], + metadata, + inputs.out_page_indices, + inputs.out_raw_indices, + rows=rows, ) diff --git a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py index cbd863837275..994720da8cea 100644 --- a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py +++ b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py @@ -19,6 +19,8 @@ import torch +from sglang.srt.layers.attention.dsv4.v41_indexer import full_topk +from sglang.srt.layers.attention.dsv4.v41_indexer.types import CapturedPrefillInputs from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.test.test_utils import CustomTestCase @@ -831,6 +833,67 @@ def test_sparse_prefill_c128_uses_live_extent(self): ) +class TestDSV4FullTopKCapturedFallback(CustomTestCase): + """CPU-only: the captured full-top-k fallback keeps every configured k.""" + + def _topk_calls(self, index_topk: int, width: int): + rows = 2 + metadata = SimpleNamespace( + compressed_page_size=64, + max_compressed_seq_len=width, + compressed_seq_lens=torch.tensor([width, width // 2], dtype=torch.int32), + page_table=torch.zeros(rows, width // 64, dtype=torch.int32), + use_topk_v2=False, + row_chunks=lambda: [(slice(0, rows), None)], + ) + inputs = CapturedPrefillInputs( + indexer=SimpleNamespace(index_topk=index_topk, n_local_heads=1, n_heads=1), + layer_id=0, + q=torch.zeros(rows, 1, 128, dtype=torch.bfloat16), + weights=torch.ones(rows, 1, dtype=torch.bfloat16), + paged_metadata=metadata, + out_page_indices=torch.full((rows, index_topk), -1, dtype=torch.int32), + out_raw_indices=torch.full((rows, index_topk), -1, dtype=torch.int32), + ) + calls = [] + with ( + mock.patch.object( + full_topk, + "quantize_index_q", + return_value=(torch.zeros(rows, 1, 64), torch.zeros(rows, 1)), + ), + mock.patch.object(full_topk, "get_index_k_cache", return_value=object()), + mock.patch.object( + full_topk, + "deep_gemm_fp4_paged_mqa_logits", + return_value=torch.zeros(rows, width), + ), + mock.patch.object( + full_topk, + "topk_transform_paged_torch", + side_effect=lambda *a, **_: calls.append(("torch", a[3].shape[1])), + ), + mock.patch.object( + full_topk, + "topk_transform_paged_from_metadata", + side_effect=lambda *a, **_: calls.append(("kernel", None)), + ), + ): + full_topk.FullTopKIndexer._deep_gemm_prefill_captured( + SimpleNamespace(token_to_kv_pool=None), inputs + ) + return calls + + def test_v1_limit_uses_output_width(self): + for k, width, expected in ( + (1024, 512, ("kernel", None)), + (2048, 4096, ("torch", 2048)), + (2048, 512, ("torch", 512)), + ): + with self.subTest(k=k, width=width): + self.assertEqual(self._topk_calls(k, width), [expected]) + + class TestDSV4SwaOutCacheLocResolution(CustomTestCase): """SWA writes must translate live locations for idle or missing/mismatched caches. A matching cache on an active forward must be reused.