[Perf] deepseek-v4.1: Select the captured prefill full top-k per row

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 (

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

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.