[Fix] Handle chunked paged MQA metadata in DSV4.1 eager forwards (#40637)

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

@@ -3486,6 +3486,31 @@ class DeepseekV4AttnBackend(
                 self.candidate_indexer.publish_decode(inputs, page_indices, raw_indices)
             )
             return
+        if isinstance(metadata.deep_gemm_metadata, list):
+            topk_plans = metadata.topk_metadata_chunks
+            assert not metadata.use_topk_v2 or topk_plans is not None
+            for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()):
+                logits = deep_gemm_fp4_paged_mqa_logits(
+                    (q_fp4[rows], q_sf[rows]),
+                    k_cache,
+                    weights[rows],
+                    metadata.compressed_seq_lens[rows],
+                    metadata.page_table[rows],
+                    plan,
+                    metadata.max_compressed_seq_len,
+                )
+                # TODO(dark): add bf16 topk
+                topk_transform_paged_from_metadata(
+                    logits,
+                    metadata,
+                    page_indices,
+                    raw_indices,
+                    rows=rows,
+                    topk_metadata=(
+                        topk_plans[chunk_idx] if topk_plans is not None else None
+                    ),
+                )
+        else:
             logits = deep_gemm_fp4_paged_mqa_logits(
                 (q_fp4, q_sf),
                 k_cache,
@@ -3496,7 +3521,9 @@ class DeepseekV4AttnBackend(
                 metadata.max_compressed_seq_len,
             )
             # TODO(dark): add bf16 topk
-        topk_transform_paged_from_metadata(logits, metadata, page_indices, raw_indices)
+            topk_transform_paged_from_metadata(
+                logits, metadata, page_indices, raw_indices
+            )
 
     # TODO(candidate): Hopper decode still publishes / consumes masks inline (torch
     # top-k); move into the candidate indexer with the prefill paths.

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

@@ -21,6 +21,7 @@ from sglang.srt.layers.attention.dsv4.candidate_indexer import (
 )
 from sglang.srt.layers.attention.dsv4.indexer import (
     deep_gemm_fp4_paged_mqa_logits,
+    topk_transform_paged_from_metadata,
 )
 
 CANDIDATE_BLOCK_SIZE = 8  # positions per block; DeepGEMM accepts 8 or 16
@@ -176,6 +177,10 @@ class DeepGemmCandidateIndexer:
         metadata."""
         metadata = inputs.metadata
         seq_lens = metadata.compressed_seq_lens.reshape(-1)
+        if isinstance(metadata.deep_gemm_metadata, list):
+            return self._publish_decode_chunked(
+                inputs, page_indices, raw_indices, seq_lens
+            )
         logits = deep_gemm_fp4_paged_mqa_logits(
             (inputs.q_fp4, inputs.q_sf),
             inputs.k_cache,
@@ -231,6 +236,83 @@ class DeepGemmCandidateIndexer:
                 ready=ready,
             )
 
+    def _publish_decode_chunked(
+        self,
+        inputs: IndexerInputs,
+        page_indices: torch.Tensor,
+        raw_indices: Optional[torch.Tensor],
+        seq_lens: torch.Tensor,
+    ) -> SparseBlockTable:
+        """Publish an eager forward whose dense logits are bounded by row chunks.
+
+        CUDA-graph metadata always carries one tensor schedule and keeps using the
+        asynchronous fast path above.  The exceptional eager path stays on the
+        current stream so each chunk's full logits can be released before the next.
+        """
+        metadata = inputs.metadata
+        block_chunks = []
+        phys_block_chunks = []
+        valid_len_chunks = []
+        topk_plans = metadata.topk_metadata_chunks
+        assert not metadata.use_topk_v2 or topk_plans is not None
+
+        for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()):
+            logits = deep_gemm_fp4_paged_mqa_logits(
+                (inputs.q_fp4[rows], inputs.q_sf[rows]),
+                inputs.k_cache,
+                inputs.weights[rows],
+                metadata.compressed_seq_lens[rows],
+                metadata.page_table[rows],
+                plan,
+                metadata.max_compressed_seq_len,
+            )
+            topk_transform_paged_from_metadata(
+                logits,
+                metadata,
+                page_indices,
+                raw_indices,
+                rows=rows,
+                topk_metadata=(
+                    topk_plans[chunk_idx] if topk_plans is not None else None
+                ),
+            )
+
+            chunk_seq_lens = seq_lens[rows]
+            nblocks, row_valid_lens = candidate_row_lens(
+                chunk_seq_lens, self.topk_blocks
+            )
+            blocks = amax_topk_blocks(logits, chunk_seq_lens, nblocks, self.topk_blocks)
+            phys_blocks = sort_candidate_blocks(
+                blocks,
+                chunk_seq_lens,
+                metadata.page_table[rows],
+                metadata.compressed_page_size,
+            )
+            block_chunks.append(blocks)
+            phys_block_chunks.append(phys_blocks)
+            valid_len_chunks.append(row_valid_lens)
+
+        blocks = torch.cat(block_chunks)
+        phys_blocks = torch.cat(phys_block_chunks)
+        row_valid_lens = torch.cat(valid_len_chunks)
+        schedule = build_sparse_indexer_schedule(
+            blocks,
+            seq_lens,
+            metadata.page_table,
+            metadata.compressed_page_size,
+            inputs.q_fp4.dtype,
+            self._request_ids(inputs.request_ids, inputs.num_rows, blocks.device),
+        )
+        ready = torch.cuda.Event()
+        ready.record(torch.cuda.current_stream())
+        return SparseBlockTable(
+            blocks=blocks,
+            schedule=schedule,
+            phys_blocks=phys_blocks,
+            valid_lens=row_valid_lens,
+            ready=ready,
+        )
+
     def _scores(self, table: SparseBlockTable, inputs: IndexerInputs) -> torch.Tensor:
         return sparse_logits(
             inputs.q_fp4,

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

@@ -475,27 +475,40 @@ def topk_transform_paged_from_metadata(
     metadata,
     page_indices: torch.Tensor,
     raw_indices: Optional[torch.Tensor] = None,
+    *,
+    rows: Optional[slice] = None,
+    topk_metadata: Optional[torch.Tensor] = None,
 ) -> None:
     """Pool slots into ``page_indices`` (``-1`` past the valid count) and, when given,
     positions into ``raw_indices``; ``metadata`` is a ``PagedIndexerMetadata``."""
+    if rows is None:
+        seq_lens = metadata.compressed_seq_lens
+        page_table = metadata.page_table
+        out_page_indices = page_indices
+        out_raw_indices = raw_indices
+    else:
+        seq_lens = metadata.compressed_seq_lens[rows]
+        page_table = metadata.page_table[rows]
+        out_page_indices = page_indices[rows]
+        out_raw_indices = raw_indices[rows] if raw_indices is not None else None
     if metadata.use_topk_v2:
         topk_transform_paged_v2(
             logits,
-            metadata.compressed_seq_lens,
-            metadata.page_table,
-            page_indices,
+            seq_lens,
+            page_table,
+            out_page_indices,
             metadata.compressed_page_size,
-            metadata.topk_metadata,
-            raw_indices,
+            metadata.topk_metadata if topk_metadata is None else topk_metadata,
+            out_raw_indices,
         )
     else:
         topk_transform_paged(
             logits,
-            metadata.compressed_seq_lens,
-            metadata.page_table,
-            page_indices,
+            seq_lens,
+            page_table,
+            out_page_indices,
             metadata.compressed_page_size,
-            raw_indices,
+            out_raw_indices,
         )
 
 

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

@@ -21,6 +21,7 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     is_in_tc_piecewise_cuda_graph,
 )
+from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
 from sglang.srt.utils import is_hip, is_sm120_supported, is_xpu
 
 logger = logging.getLogger(__name__)
@@ -306,7 +307,8 @@ class PagedIndexerMetadata:
         ):
             return None
         if (
-            torch.cuda.is_current_stream_capturing()
+            get_is_capture_mode()
+            or torch.cuda.is_current_stream_capturing()
             or is_in_breakable_cuda_graph()
             or is_in_tc_piecewise_cuda_graph()
         ):
@@ -325,14 +327,27 @@ class PagedIndexerMetadata:
 
     def row_chunks(self):
         num_rows = self.compressed_seq_lens.shape[0]
-        if self.row_chunk <= 0:
+        if self.row_chunk > 0:
+            rows_per_chunk = self.row_chunk
+        elif isinstance(self.deep_gemm_metadata, list):
+            assert self.rows_per_chunk is not None, (
+                "chunked DeepGEMM metadata requires rows_per_chunk"
+            )
+            rows_per_chunk = self.rows_per_chunk
+        else:
             return [(slice(0, num_rows), self.deep_gemm_metadata)]
-        return [
-            (slice(start, min(start + self.row_chunk, num_rows)), plan)
+
+        chunks = [
+            (slice(start, min(start + rows_per_chunk, num_rows)), plan)
             for start, plan in zip(
-                range(0, num_rows, self.row_chunk), self.deep_gemm_metadata
+                range(0, num_rows, rows_per_chunk), self.deep_gemm_metadata
             )
         ]
+        assert chunks and chunks[-1][0].stop == num_rows, (
+            f"chunk schedules do not cover all rows: {num_rows=} {rows_per_chunk=} "
+            f"{len(chunks)=}"
+        )
+        return chunks
 
     def copy_(self, other: PagedIndexerMetadata):
         # A chunked schedule list has no in-place copy; rebind it instead.

test/registered/unit/layers/test_dsv4_nonpaged_indexer.py

@@ -752,6 +752,33 @@ class TestPagedIndexerMetadataChunking(CustomTestCase):
     row chunks the indexer loops over; a mismatch would silently score rows
     with another chunk's schedule."""
 
+    def test_capture_warmup_skips_dynamic_budget_but_eager_forward_uses_it(self):
+        metadata = SimpleNamespace(
+            use_prefill_cuda_graph=False,
+            compressed_seq_lens=SimpleNamespace(
+                is_cuda=True, device=SimpleNamespace(index=0)
+            ),
+            max_compressed_seq_len=65536,
+        )
+        for capture_mode in (True, False):
+            with (
+                self.subTest(capture_mode=capture_mode),
+                patch(f"{_METADATA}.get_is_capture_mode", return_value=capture_mode),
+                patch("torch.cuda.is_current_stream_capturing", return_value=False),
+                patch(f"{_METADATA}.is_in_breakable_cuda_graph", return_value=False),
+                patch(f"{_METADATA}.is_in_tc_piecewise_cuda_graph", return_value=False),
+                patch(
+                    f"{_METADATA}.mqa_logits_budget_bytes", return_value=4096
+                ) as budget,
+            ):
+                result = PagedIndexerMetadata._mqa_logits_budget(metadata, num_rows=256)
+                if capture_mode:
+                    self.assertIsNone(result)
+                    budget.assert_not_called()
+                else:
+                    self.assertEqual(result, 4096)
+                    budget.assert_called_once_with(device_index=0, allow_sync=True)
+
     def _build(self, *, num_rows: int, budget, use_topk_v2: bool):
         deep_gemm = SimpleNamespace(
             get_num_sms=MagicMock(return_value=1),
@@ -806,6 +833,12 @@ class TestPagedIndexerMetadataChunking(CustomTestCase):
 
         self.assertIsInstance(metadata.deep_gemm_metadata, list)
         self.assertEqual(len(metadata.deep_gemm_metadata), len(chunks))
+        metadata_chunks = metadata.row_chunks()
+        self.assertEqual([rows for rows, _ in metadata_chunks], chunks)
+        for (_, actual_plan), expected_plan in zip(
+            metadata_chunks, metadata.deep_gemm_metadata
+        ):
+            self.assertIs(actual_plan, expected_plan)
         schedule_rows = [
             call.args[0]
             for call in deep_gemm.get_paged_mqa_logits_metadata.call_args_list
@@ -875,6 +908,92 @@ class TestChunkedTopKMatchesUnchunked(CustomTestCase):
                 self.assertTrue(torch.equal(run(rows_per_chunk), expected))
 
 
+class TestChunkedCandidatePublisher(CustomTestCase):
+    def test_each_deep_gemm_call_receives_one_tensor_schedule(self):
+        from sglang.srt.layers.attention.dsv4 import candidate_indexer_deep_gemm as mod
+
+        num_rows, width = 5, 16
+        chunks = [slice(0, 2), slice(2, 4), slice(4, 5)]
+        plans = [torch.tensor([i], dtype=torch.uint8) for i in range(len(chunks))]
+        topk_plans = [torch.tensor([i], dtype=torch.int32) for i in range(len(chunks))]
+        metadata = SimpleNamespace(
+            compressed_seq_lens=torch.full((num_rows, 1), width, dtype=torch.int32),
+            page_table=torch.zeros((num_rows, 1), dtype=torch.int32),
+            deep_gemm_metadata=plans,
+            max_compressed_seq_len=width,
+            compressed_page_size=64,
+            topk_metadata_chunks=topk_plans,
+            use_topk_v2=True,
+            row_chunks=lambda: list(zip(chunks, plans)),
+        )
+        inputs = SimpleNamespace(
+            q_fp4=torch.zeros((num_rows, 1, 2, 64), dtype=torch.int8),
+            q_sf=torch.zeros((num_rows, 1, 2), dtype=torch.int32),
+            k_cache=torch.zeros((1, 64, 1, 68), dtype=torch.uint8),
+            weights=torch.zeros((num_rows, 2), dtype=torch.float32),
+            metadata=metadata,
+            request_ids=torch.arange(num_rows),
+            num_rows=num_rows,
+        )
+        page_indices = torch.full((num_rows, 4), -1, dtype=torch.int32)
+        raw_indices = torch.full_like(page_indices, -1)
+        indexer = object.__new__(mod.DeepGemmCandidateIndexer)
+        indexer.topk_blocks = 2
+
+        deep_gemm = MagicMock(
+            side_effect=lambda q, *_args: torch.zeros(
+                (q[0].shape[0], width), dtype=torch.float32
+            )
+        )
+        topk = MagicMock()
+        event = MagicMock()
+        stream = MagicMock()
+        with (
+            patch.object(mod, "deep_gemm_fp4_paged_mqa_logits", deep_gemm),
+            patch.object(mod, "topk_transform_paged_from_metadata", topk),
+            patch.object(
+                mod,
+                "candidate_row_lens",
+                side_effect=lambda lens, _topk: (
+                    torch.ones_like(lens),
+                    lens.clone(),
+                ),
+            ),
+            patch.object(
+                mod,
+                "amax_topk_blocks",
+                side_effect=lambda _logits, lens, _nblocks, topk_blocks: torch.zeros(
+                    (lens.shape[0], topk_blocks), dtype=torch.int32
+                ),
+            ),
+            patch.object(
+                mod,
+                "sort_candidate_blocks",
+                side_effect=lambda blocks, *_args: blocks + 1,
+            ),
+            patch.object(
+                mod,
+                "build_sparse_indexer_schedule",
+                return_value=torch.tensor([7], dtype=torch.uint8),
+            ),
+            patch.object(mod.torch.cuda, "Event", return_value=event),
+            patch.object(mod.torch.cuda, "current_stream", return_value=stream),
+        ):
+            table = indexer.publish_decode(inputs, page_indices, raw_indices)
+
+        self.assertEqual(deep_gemm.call_count, len(chunks))
+        for call, plan in zip(deep_gemm.call_args_list, plans):
+            self.assertIs(call.args[5], plan)
+            self.assertIsInstance(call.args[5], torch.Tensor)
+        self.assertEqual([call.kwargs["rows"] for call in topk.call_args_list], chunks)
+        for call, plan in zip(topk.call_args_list, topk_plans):
+            self.assertIs(call.kwargs["topk_metadata"], plan)
+        self.assertEqual(table.blocks.shape, (num_rows, indexer.topk_blocks))
+        self.assertEqual(table.phys_blocks.shape, table.blocks.shape)
+        self.assertEqual(table.valid_lens.shape, (num_rows,))
+        event.record.assert_called_once_with(stream)
+
+
 class TestCandidateIndexerGating(CustomTestCase):
     def test_candidate_indexer_gating(self):
         from sglang.srt.layers.attention.dsv4 import candidate_indexer