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

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

@@ -3486,17 +3486,44 @@ class DeepseekV4AttnBackend(
                 self.candidate_indexer.publish_decode(inputs, page_indices, raw_indices)
             )
             return
-        logits = deep_gemm_fp4_paged_mqa_logits(
-            (q_fp4, q_sf),
-            k_cache,
-            weights,
-            metadata.compressed_seq_lens,
-            metadata.page_table,
-            metadata.deep_gemm_metadata,
-            metadata.max_compressed_seq_len,
-        )
-        # TODO(dark): add bf16 topk
-        topk_transform_paged_from_metadata(logits, metadata, page_indices, raw_indices)
+        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,
+                weights,
+                metadata.compressed_seq_lens,
+                metadata.page_table,
+                metadata.deep_gemm_metadata,
+                metadata.max_compressed_seq_len,
+            )
+            # TODO(dark): add bf16 topk
+            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.