[DSV4] Chunk the indexer MQA logits by query rows under a free-memory budget (#39095)

Signed-off-by: Shiki Wu <shikiw@nvidia.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -32,6 +32,16 @@ from sglang.srt.layers.attention.dsa.utils import (
     is_graph_dsa_split_op_surface,
 )
 from sglang.srt.layers.attention.graph_variants import DSA_DENSE
+from sglang.srt.layers.attention.mqa_logits_utils import (
+    MQA_LOGITS_BYTES_PER_ELEM,
+    MQA_LOGITS_MAX_BYTES_ROCM,
+    MQA_LOGITS_STATIC_SKIP_ELEMS,
+    MQA_LOGITS_TOTAL_MEM_FRACTION,
+    mqa_logits_budget_bytes,
+    mqa_logits_free_mem_fraction,
+    mqa_logits_should_chunk,
+    mqa_logits_static_budget_bytes,
+)
 from sglang.srt.layers.layernorm import LayerNorm, RMSNorm
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
     is_in_breakable_cuda_graph,
@@ -43,7 +53,6 @@ from sglang.srt.runtime_context import (
     get_device,
     get_exec,
     get_parallel,
-    get_schedule,
 )
 from sglang.srt.state_capturer.indexer_topk import (
     maybe_capture_indexer_topk,
@@ -52,7 +61,6 @@ from sglang.srt.utils import (
     add_prefix,
     ceil_align,
     get_bool_env_var,
-    get_device_module,
     is_cuda,
     is_gfx95_supported,
     is_hip,
@@ -202,16 +210,16 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor:
 
 
 class Indexer(DSANPUIndexerMixin, BaseFusedOp):
-    _MQA_LOGITS_BYTES_PER_ELEM = 4
+    _MQA_LOGITS_BYTES_PER_ELEM = MQA_LOGITS_BYTES_PER_ELEM
-    _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000
+    _MQA_LOGITS_STATIC_SKIP_ELEMS = MQA_LOGITS_STATIC_SKIP_ELEMS
-    _MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3
+    _MQA_LOGITS_TOTAL_MEM_FRACTION = MQA_LOGITS_TOTAL_MEM_FRACTION
-    # aiter's fp8_mqa_logits only compiles below 2 GiB of logits (buffer_store).
+    _MQA_LOGITS_MAX_BYTES_ROCM = MQA_LOGITS_MAX_BYTES_ROCM
-    _MQA_LOGITS_MAX_BYTES_ROCM = 2**31 - 1
+    # One measured budget per device for the process lifetime.
     _mqa_logits_budget_bytes: Dict[int, int] = {}
 
     @staticmethod
     def _mqa_logits_free_mem_fraction() -> float:
-        return envs.SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION.get()
+        return mqa_logits_free_mem_fraction()
 
     def __init__(
         self,
@@ -988,68 +996,21 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
         return topk_result
 
     def _get_mqa_logits_budget_bytes(self, device_index: int) -> int:
-        free_mem_fraction = self._mqa_logits_free_mem_fraction()
         cached_budget = self._mqa_logits_budget_bytes.get(device_index)
         if cached_budget is not None:
             return cached_budget
 
-        total_mem = get_device_module().get_device_properties(device_index).total_memory
+        # Graph capture cannot sync the host; use the static guard and do not
-
+        # cache it, so the first eager prefill still measures the real budget.
-        total_mem_budget = int(total_mem * self._MQA_LOGITS_TOTAL_MEM_FRACTION)
-        mem_fraction_static = get_schedule().mem_fraction_static
-        if mem_fraction_static is None:
-            static_budget = total_mem_budget
-        else:
-            static_free_mem = int(total_mem * max(0.0, 1.0 - mem_fraction_static))
-            static_budget = min(
-                int(static_free_mem * free_mem_fraction),
-                total_mem_budget,
-            )
-        static_budget = max(1, static_budget)
-
-        # Keep the static serving-memory guard during CUDA graph capture without
-        # caching it. The first non-capture prefill path will cache the real
-        # free-memory budget below.
         if get_is_capture_mode():
-            return static_budget
+            return mqa_logits_static_budget_bytes(device_index=device_index)
 
-        # Match the original free-memory guard: logits_bytes * 2 > free_mem.
+        budget_bytes = mqa_logits_budget_bytes(
-        # Synchronizes the host; cache the result capped by serving-memory headroom.
+            device_index=device_index, allow_sync=True
-        if _is_xpu:
+        )
-            # On XPU, use total_mem budget as the free-memory estimate;
-            # dynamic free-memory query is not supported the same way as CUDA.
-            # TODO Use torch.xpu.mem_get_info() when available (planned end of 2026).
-            budget_bytes = static_budget
-        else:
-            free_mem, _ = torch.cuda.mem_get_info(device_index)
-            budget_bytes = min(int(free_mem * free_mem_fraction), static_budget)
-
-        budget_bytes = max(1, budget_bytes)
         self._mqa_logits_budget_bytes[device_index] = budget_bytes
         return budget_bytes
 
-    def _should_chunk_mqa_logits(
-        self, num_q: int, num_k: int, device_index: int
-    ) -> Tuple[bool, int]:
-        """
-        Detect whether we need to chunk the MQA logits computation to avoid OOM,
-        and on ROCm to stay under aiter's 2 GiB logits limit
-        Return: (need_chunk, logits_budget_bytes)
-        """
-        # Quick static check for normal batches
-        if num_q * num_k < self._MQA_LOGITS_STATIC_SKIP_ELEMS:
-            return False, 0
-
-        logits_bytes = num_q * num_k * self._MQA_LOGITS_BYTES_PER_ELEM
-        logits_budget_bytes = self._get_mqa_logits_budget_bytes(device_index)
-        if _is_hip:
-            logits_budget_bytes = min(
-                logits_budget_bytes, self._MQA_LOGITS_MAX_BYTES_ROCM
-            )
-
-        need_chunk = logits_bytes > logits_budget_bytes
-        return need_chunk, logits_budget_bytes
-
     def _get_topk_ragged(
         self,
         enable_dual_stream: bool,
@@ -1139,8 +1100,11 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
         token_to_batch_idx = metadata.get_token_to_batch_idx()
         q_offset = ks.shape[0]
         k_offset = k_fp8.shape[0]
-        need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits(
+        need_chunk, logits_budget_bytes = mqa_logits_should_chunk(
-            q_offset, k_offset, device_index
+            num_rows=q_offset,
+            num_cols=k_offset,
+            get_budget_bytes=lambda: self._get_mqa_logits_budget_bytes(device_index),
+            rocm=_is_hip,
         )
 
         if not need_chunk:

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

@@ -35,9 +35,13 @@ from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend
 from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
 from sglang.srt.layers.attention.dsv4.compressor import Compressor
 from sglang.srt.layers.attention.dsv4.metadata import (
-    _SM120_INDEXER_M_CHUNK,
     NonPagedIndexerPlan,
     PagedIndexerMetadata,
+    iter_row_chunks,
+)
+from sglang.srt.layers.attention.mqa_logits_utils import (
+    mqa_logits_row_bytes,
+    mqa_logits_rows_per_chunk,
 )
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
@@ -690,6 +694,14 @@ class C4IndexerBackendMixin:
             ke = torch.where(ke - ks > c4_indexer.index_topk, ke, ks)
         c4_page_size = indexer_metadata.compressed_page_size
         max_seqlen_k = (final_c4_len + c4_page_size - 1) // c4_page_size * c4_page_size
+        rows_per_chunk = None
+        if indexer_metadata.mqa_logits_budget_bytes is not None:
+            # fp8_mqa_logits allocates [query_rows, align256(max_seqlen_k)] fp32.
+            rows_per_chunk = mqa_logits_rows_per_chunk(
+                num_rows=query_rows,
+                row_bytes=mqa_logits_row_bytes(max_seqlen_k),
+                budget_bytes=indexer_metadata.mqa_logits_budget_bytes,
+            )
         plan = NonPagedIndexerPlan(
             page_table=request_page_table,
             gather_seq_lens=gather_seq_lens,
@@ -699,21 +711,18 @@ class C4IndexerBackendMixin:
             max_seq_len=final_c4_len,
             max_seqlen_k=max_seqlen_k,
             query_rows=query_rows,
+            rows_per_chunk=rows_per_chunk,
         )
         indexer_metadata.nonpaged_plan = plan
         return plan
 
     @staticmethod
-    def _forward_nonpaged_indexer(
+    def _gather_nonpaged_index_k(
         *,
-        q_indexer: torch.Tensor,
-        weights: torch.Tensor,
         c4_indexer: C4Indexer,
         token_to_kv_pool: DeepSeekV4TokenToKVPool,
         plan: NonPagedIndexerPlan,
-    ) -> torch.Tensor:
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        import deep_gemm
-
         k_u8, scale_u8 = token_to_kv_pool.get_index_k_scale_buffer(
             layer_id=c4_indexer.layer_id,
             seq_len_tensor=plan.gather_seq_lens,
@@ -721,14 +730,25 @@ class C4IndexerBackendMixin:
             seq_len_sum=plan.seq_len_sum,
             max_seq_len=plan.max_seq_len,
         )
-        k_fp8 = k_u8.view(FP8_DTYPE)
+        return k_u8.view(FP8_DTYPE), scale_u8.view(torch.float32).squeeze(-1)
-        k_scale = scale_u8.view(torch.float32).squeeze(-1)
+
+    @staticmethod
+    def _nonpaged_mqa_logits(
+        *,
+        q_indexer: torch.Tensor,
+        weights: torch.Tensor,
+        kv: Tuple[torch.Tensor, torch.Tensor],
+        plan: NonPagedIndexerPlan,
+        rows: slice,
+    ) -> torch.Tensor:
+        import deep_gemm
+
         return deep_gemm.fp8_mqa_logits(
-            q_indexer[: plan.query_rows],
+            q_indexer[rows],
-            (k_fp8, k_scale),
+            kv,
-            weights[: plan.query_rows],
+            weights[rows],
-            plan.ks,
+            plan.ks[rows],
-            plan.ke,
+            plan.ke[rows],
             clean_logits=False,
             max_seqlen_k=plan.max_seqlen_k,
         )
@@ -904,7 +924,11 @@ class C4IndexerBackendMixin:
 
         all_rows = slice(0, _c4sl.shape[0])
 
-        def run_topk_transform(rows: slice, logits: torch.Tensor) -> None:
+        def run_topk_transform(
+            rows: slice,
+            logits: torch.Tensor,
+            topk_plan: Optional[torch.Tensor] = None,
+        ) -> None:
             row_raw_indices = raw_indices[rows] if raw_indices is not None else None
             if self.dsa_topk_backend.is_torch():
                 topk_transform_pytorch_vectorized(
@@ -925,19 +949,21 @@ class C4IndexerBackendMixin:
                     row_raw_indices,
                 )
             elif self.dsa_topk_backend.should_use_topk_v2():
+                if topk_plan is None:
+                    # The cached plan routes rows by their index in the full
+                    # range, so a chunk needs one built over its own rows.
+                    topk_plan = (
+                        indexer_metadata.topk_metadata
+                        if rows == all_rows
+                        else plan_topk_v2(c4_seq_lens[rows])
+                    )
                 topk_transform_paged_v2(
                     logits,
                     c4_seq_lens[rows],
                     page_table[rows],
                     c4_sparse_page_indices[rows],
                     indexer_metadata.compressed_page_size,
-                    # The cached plan routes rows by their index in the full
+                    topk_plan,
-                    # range, so a chunk needs one built over its own rows.
-                    (
-                        indexer_metadata.topk_metadata
-                        if rows == all_rows or not is_hip()
-                        else plan_topk_v2(c4_seq_lens[rows])
-                    ),
                     row_raw_indices,
                 )
             else:
@@ -952,14 +978,27 @@ class C4IndexerBackendMixin:
 
         if nonpaged_plan is not None:
             assert isinstance(q_indexer, torch.Tensor)
-            logits = self._forward_nonpaged_indexer(
+            # K is gathered once; each row chunk's logits are reduced to top-k
-                q_indexer=q_indexer,
+            # and dropped before the next chunk allocates, so only one chunk
-                weights=weights,
+            # of logits is live at a time.
+            kv = self._gather_nonpaged_index_k(
                 c4_indexer=c4_indexer,
                 token_to_kv_pool=token_to_kv_pool,
                 plan=nonpaged_plan,
             )
-            run_topk_transform(all_rows, logits)
+            for rows in iter_row_chunks(
+                num_rows=nonpaged_plan.query_rows,
+                rows_per_chunk=nonpaged_plan.rows_per_chunk,
+            ):
+                logits = self._nonpaged_mqa_logits(
+                    q_indexer=q_indexer,
+                    weights=weights,
+                    kv=kv,
+                    plan=nonpaged_plan,
+                    rows=rows,
+                )
+                run_topk_transform(rows, logits)
+                del logits
         elif use_aiter_fp4:
             q_fp4, q_scale = q
             is_decode = forward_batch.forward_mode.is_decode()
@@ -1012,7 +1051,11 @@ class C4IndexerBackendMixin:
                 c4_indexer_kv_cache.shape[0], 64, 1, head_dim_with_sf
             )
 
-            def run_paged_indexer(rows: slice, metadata: torch.Tensor) -> None:
+            def run_paged_indexer(
+                rows: slice,
+                metadata: torch.Tensor,
+                topk_plan: Optional[torch.Tensor] = None,
+            ) -> None:
                 row_q = (q[0][rows], q[1][rows]) if isinstance(q, tuple) else q[rows]
                 logits = fn(
                     row_q,
@@ -1024,18 +1067,30 @@ class C4IndexerBackendMixin:
                     indexer_metadata.max_compressed_seq_len,
                     False,
                 )
-                run_topk_transform(rows, logits)
+                run_topk_transform(rows, logits, topk_plan)
 
             deep_gemm_metadata = indexer_metadata.deep_gemm_metadata
             if isinstance(deep_gemm_metadata, list):
-                # SM120 only: DeepGEMM's metadata kernel caps the row count, so
+                # PagedIndexerMetadata split this forward into row chunks (SM120
-                # PagedIndexerMetadata split it; run indexer + topk per chunk.
+                # kernel cap and/or logits memory budget), one schedule each.
                 num_rows = _c4sl.shape[0]
-                for chunk_idx, start in enumerate(
+                assert num_rows == indexer_metadata.compressed_seq_lens.shape[0], (
-                    range(0, num_rows, _SM120_INDEXER_M_CHUNK)
+                    f"chunk schedules were built for "
+                    f"{indexer_metadata.compressed_seq_lens.shape[0]} rows, "
+                    f"got {num_rows}"
+                )
+                topk_plans = indexer_metadata.topk_metadata_chunks
+                for chunk_idx, rows in enumerate(
+                    iter_row_chunks(
+                        num_rows=num_rows,
+                        rows_per_chunk=indexer_metadata.rows_per_chunk,
+                    )
                 ):
-                    rows = slice(start, min(start + _SM120_INDEXER_M_CHUNK, num_rows))
+                    run_paged_indexer(
-                    run_paged_indexer(rows, deep_gemm_metadata[chunk_idx])
+                        rows,
+                        deep_gemm_metadata[chunk_idx],
+                        topk_plans[chunk_idx] if topk_plans is not None else None,
+                    )
             else:
                 run_paged_indexer(all_rows, deep_gemm_metadata)
 

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

@@ -1,14 +1,30 @@
 from __future__ import annotations
 
+import logging
 import warnings
 from dataclasses import dataclass, field, fields
-from typing import Any, List, Optional
+from typing import Any, Iterator, List, Optional
 
 import torch
 
 from sglang.srt.environ import envs
+from sglang.srt.layers.attention.mqa_logits_utils import (
+    mqa_logits_budget_bytes,
+    mqa_logits_needs_budget_check,
+    mqa_logits_row_bytes,
+    mqa_logits_rows_per_chunk,
+    mqa_logits_should_chunk,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    is_in_breakable_cuda_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.utils import is_hip, is_sm120_supported, is_xpu
 
+logger = logging.getLogger(__name__)
+
 _IS_SM120 = is_sm120_supported()
 
 """
@@ -48,9 +64,57 @@ Some other notes:
 """
 _LARGE_INDEXER_QUERY_THRESHOLD = 11673
 
+# DeepGEMM's paged-MQA metadata kernel cannot schedule more rows than this on
+# SM120 (shared-memory cap), so SM120 always splits larger batches.
 _SM120_INDEXER_M_CHUNK = 4096
 
 
+def plan_indexer_row_chunks(
+    *,
+    num_rows: int,
+    num_cols: int,
+    budget_bytes: Optional[int],
+    sm120_row_cap: Optional[int],
+) -> Optional[int]:
+    """Query rows per paged-indexer chunk; None runs the whole batch in one call.
+
+    The fp32 logits are [num_rows, num_cols] per layer, so the chunk is the
+    smaller of the SM120 kernel cap and what the memory budget allows.
+    """
+    rows_per_chunk = None
+    if sm120_row_cap is not None and num_rows > sm120_row_cap:
+        rows_per_chunk = sm120_row_cap
+    if budget_bytes is not None:
+        need_chunk, budget_bytes = mqa_logits_should_chunk(
+            num_rows=num_rows,
+            num_cols=num_cols,
+            get_budget_bytes=lambda: budget_bytes,
+            rocm=is_hip(),
+        )
+        by_budget = (
+            mqa_logits_rows_per_chunk(
+                num_rows=num_rows,
+                row_bytes=mqa_logits_row_bytes(num_cols),
+                budget_bytes=budget_bytes,
+            )
+            if need_chunk
+            else None
+        )
+        if by_budget is not None:
+            rows_per_chunk = (
+                by_budget if rows_per_chunk is None else min(rows_per_chunk, by_budget)
+            )
+    return rows_per_chunk
+
+
+def iter_row_chunks(*, num_rows: int, rows_per_chunk: Optional[int]) -> Iterator[slice]:
+    if rows_per_chunk is None or rows_per_chunk >= num_rows:
+        yield slice(0, num_rows)
+        return
+    for start in range(0, num_rows, rows_per_chunk):
+        yield slice(start, min(start + rows_per_chunk, num_rows))
+
+
 def copy_metadata(
     *,
     src,
@@ -105,6 +169,8 @@ class NonPagedIndexerPlan:
     max_seq_len: int
     max_seqlen_k: int
     query_rows: int
+    # None runs all query rows in one fp8_mqa_logits call.
+    rows_per_chunk: Optional[int] = None
 
 
 @dataclass
@@ -121,11 +187,20 @@ class PagedIndexerMetadata:
     # Rows per logits chunk for the prefill CUDA graph low-ratio indexer; 0 plans
     # all rows at once.
     row_chunk: int = 0
+    # A list when the dynamic-budget forward is row-chunked: one schedule per chunk.
     deep_gemm_metadata: Any = field(init=False, repr=False)
     topk_metadata: torch.Tensor = field(init=False, repr=False)
     nonpaged_plan: Optional[NonPagedIndexerPlan] = field(
         init=False, repr=False, default=None
     )
+    # Decided once per forward and shared by every layer's indexer call.
+    rows_per_chunk: Optional[int] = field(init=False, repr=False, default=None)
+    mqa_logits_budget_bytes: Optional[int] = field(init=False, repr=False, default=None)
+    # The top-k v2 plan routes rows by index within the batch it was built for,
+    # so a row-chunked forward needs one plan per chunk.
+    topk_metadata_chunks: Optional[List[torch.Tensor]] = field(
+        init=False, repr=False, default=None
+    )
 
     def __post_init__(self):
         if (
@@ -149,6 +224,23 @@ class PagedIndexerMetadata:
             compressed_seq_lens = self.compressed_seq_lens.to(torch.int32)
             if compressed_seq_lens.dim() == 1:
                 compressed_seq_lens = compressed_seq_lens.unsqueeze(-1)
+            num_rows = compressed_seq_lens.shape[0]
+            self.mqa_logits_budget_bytes = self._mqa_logits_budget(num_rows=num_rows)
+            self.rows_per_chunk = plan_indexer_row_chunks(
+                num_rows=num_rows,
+                num_cols=self.max_compressed_seq_len,
+                budget_bytes=self.mqa_logits_budget_bytes,
+                sm120_row_cap=_SM120_INDEXER_M_CHUNK if _IS_SM120 else None,
+            )
+            if self.rows_per_chunk is not None:
+                logger.debug(
+                    "DSV4 indexer chunks %d query rows x %d compressed cols into "
+                    "%d-row chunks (logits budget %s bytes)",
+                    num_rows,
+                    self.max_compressed_seq_len,
+                    self.rows_per_chunk,
+                    self.mqa_logits_budget_bytes,
+                )
             if self.row_chunk > 0:
                 self.deep_gemm_metadata = torch.stack(
                     [
@@ -160,16 +252,16 @@ class PagedIndexerMetadata:
                         for _s in range(0, compressed_seq_lens.shape[0], self.row_chunk)
                     ]
                 )
-            elif _IS_SM120 and compressed_seq_lens.shape[0] > _SM120_INDEXER_M_CHUNK:
+            elif self.rows_per_chunk is not None:
                 # Chunk metadata is shared by all indexer layers in this forward.
                 self.deep_gemm_metadata = [
                     get_paged_mqa_logits_metadata(
-                        compressed_seq_lens[_s : _s + _SM120_INDEXER_M_CHUNK],
+                        compressed_seq_lens[rows],
                         self.compressed_page_size,
                         deep_gemm.get_num_sms(),
                     )
-                    for _s in range(
+                    for rows in iter_row_chunks(
-                        0, compressed_seq_lens.shape[0], _SM120_INDEXER_M_CHUNK
+                        num_rows=num_rows, rows_per_chunk=self.rows_per_chunk
                     )
                 ]
             else:
@@ -185,6 +277,14 @@ class PagedIndexerMetadata:
             from sglang.kernels.ops.attention.dsv4 import plan_topk_v2
 
             self.topk_metadata = plan_topk_v2(self.compressed_seq_lens)
+            if self.rows_per_chunk is not None:
+                self.topk_metadata_chunks = [
+                    plan_topk_v2(self.compressed_seq_lens[rows])
+                    for rows in iter_row_chunks(
+                        num_rows=self.compressed_seq_lens.shape[0],
+                        rows_per_chunk=self.rows_per_chunk,
+                    )
+                ]
         else:
             self.topk_metadata = torch.empty((0,))
 
@@ -193,6 +293,28 @@ class PagedIndexerMetadata:
             f"compress_ratio {self.compress_ratio} must divide page_size {self.page_size}"
         )
 
+    def _mqa_logits_budget(self, *, num_rows: int) -> Optional[int]:
+        """Free-memory budget for this forward's logits; None disables chunking.
+
+        Graph-backed forwards keep a single call: their shapes are fixed at
+        capture and the free-memory read would sync the host mid-capture.
+        """
+        if self.use_prefill_cuda_graph or not self.compressed_seq_lens.is_cuda:
+            return None
+        if not mqa_logits_needs_budget_check(
+            num_rows=num_rows, num_cols=self.max_compressed_seq_len
+        ):
+            return None
+        if (
+            torch.cuda.is_current_stream_capturing()
+            or is_in_breakable_cuda_graph()
+            or is_in_tc_piecewise_cuda_graph()
+        ):
+            return None
+        return mqa_logits_budget_bytes(
+            device_index=self.compressed_seq_lens.device.index, allow_sync=True
+        )
+
     @property
     def max_seq_len(self) -> int:
         return self.page_table.shape[1] * self.page_size
@@ -213,13 +335,22 @@ class PagedIndexerMetadata:
         ]
 
     def copy_(self, other: PagedIndexerMetadata):
-        if is_hip():
+        # A chunked schedule list has no in-place copy; rebind it instead.
+        chunked = isinstance(self.deep_gemm_metadata, list) or isinstance(
+            other.deep_gemm_metadata, list
+        )
+        if is_hip() or chunked:
             copy_fields = ["page_table", "compressed_seq_lens"]
             assign_fields = ["deep_gemm_metadata", "nonpaged_plan"]
         else:
             copy_fields = ["page_table", "compressed_seq_lens", "deep_gemm_metadata"]
             assign_fields = ["nonpaged_plan"]
         copy_fields += ["topk_metadata"]
+        assign_fields += [
+            "rows_per_chunk",
+            "mqa_logits_budget_bytes",
+            "topk_metadata_chunks",
+        ]
         copy_metadata(
             src=other,
             dst=self,

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

@@ -0,0 +1,101 @@
+"""Memory budget for the fp32 MQA-logits scratch of the DeepSeek sparse
+attention indexers (DSA, DSV4, DSV4.1).
+
+The indexer scores every query row against every key column into one fp32
+logits matrix that no pool sized by mem_fraction_static accounts for. These
+helpers bound that matrix and slice it by query rows.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Optional, Tuple
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.runtime_context import get_schedule
+from sglang.srt.utils import get_device_module, is_hip, is_xpu
+from sglang.srt.utils.common import ceil_div
+
+MQA_LOGITS_BYTES_PER_ELEM = 4
+MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000
+MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3
+# aiter's fp8_mqa_logits only compiles below 2 GiB of logits (buffer_store).
+MQA_LOGITS_MAX_BYTES_ROCM = 2**31 - 1
+# DeepGEMM pads the logits row stride to 1024 bytes, i.e. 256 fp32 columns.
+MQA_LOGITS_ROW_ALIGN_ELEMS = 256
+
+
+def mqa_logits_free_mem_fraction() -> float:
+    return envs.SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION.get()
+
+
+def mqa_logits_needs_budget_check(*, num_rows: int, num_cols: int) -> bool:
+    return num_rows * num_cols >= MQA_LOGITS_STATIC_SKIP_ELEMS
+
+
+def mqa_logits_row_bytes(num_cols: int) -> int:
+    aligned_cols = (
+        ceil_div(num_cols, MQA_LOGITS_ROW_ALIGN_ELEMS) * MQA_LOGITS_ROW_ALIGN_ELEMS
+    )
+    return aligned_cols * MQA_LOGITS_BYTES_PER_ELEM
+
+
+def mqa_logits_static_budget_bytes(*, device_index: int) -> int:
+    """Budget from configuration alone (no device query); safe during graph capture."""
+    total_mem = get_device_module().get_device_properties(device_index).total_memory
+    total_mem_budget = int(total_mem * MQA_LOGITS_TOTAL_MEM_FRACTION)
+    mem_fraction_static = get_schedule().mem_fraction_static
+    if mem_fraction_static is None:
+        budget = total_mem_budget
+    else:
+        static_free_mem = int(total_mem * max(0.0, 1.0 - mem_fraction_static))
+        budget = min(
+            int(static_free_mem * mqa_logits_free_mem_fraction()), total_mem_budget
+        )
+    return max(1, budget)
+
+
+def mqa_logits_budget_bytes(*, device_index: int, allow_sync: bool) -> int:
+    """Static budget capped by current free memory; mem_get_info syncs, so
+    callers pass allow_sync=False under CUDA graph capture."""
+    budget = mqa_logits_static_budget_bytes(device_index=device_index)
+    if allow_sync and not is_xpu():
+        free_mem, _ = torch.cuda.mem_get_info(device_index)
+        budget = min(int(free_mem * mqa_logits_free_mem_fraction()), budget)
+    if is_hip():
+        budget = min(budget, MQA_LOGITS_MAX_BYTES_ROCM)
+    return max(1, budget)
+
+
+def mqa_logits_should_chunk(
+    *,
+    num_rows: int,
+    num_cols: int,
+    get_budget_bytes: Callable[[], int],
+    rocm: bool,
+) -> Tuple[bool, int]:
+    """Whether a [num_rows, num_cols] fp32 logits matrix must be row-chunked.
+
+    Returns (need_chunk, effective_budget_bytes). Matrices below
+    MQA_LOGITS_STATIC_SKIP_ELEMS never chunk and the budget is not queried
+    (it may synchronize the host), hence the callable. On ROCm the budget is
+    also capped by aiter's logits ceiling.
+    """
+    if not mqa_logits_needs_budget_check(num_rows=num_rows, num_cols=num_cols):
+        return False, 0
+    budget_bytes = get_budget_bytes()
+    if rocm:
+        budget_bytes = min(budget_bytes, MQA_LOGITS_MAX_BYTES_ROCM)
+    logits_bytes = num_rows * num_cols * MQA_LOGITS_BYTES_PER_ELEM
+    return logits_bytes > budget_bytes, budget_bytes
+
+
+def mqa_logits_rows_per_chunk(
+    *, num_rows: int, row_bytes: int, budget_bytes: int
+) -> Optional[int]:
+    """Query rows per chunk so one logits chunk fits the budget; None if all rows fit."""
+    if num_rows * row_bytes <= budget_bytes:
+        return None
+    rows = max(budget_bytes // max(row_bytes, 1), 1)
+    return int(rows) if rows < num_rows else None

test/registered/unit/layers/attention/test_dsa_mqa_logits_chunking.py

@@ -1,4 +1,5 @@
-"""Contract tests for the DSA indexer's MQA-logits chunk budget.
+"""Contract tests for the MQA-logits chunk decision shared by the DSA and DSV4
+indexers.
 
 On ROCm the `[num_q x num_k]` fp32 logits tensor goes to aiter's
 `fp8_mqa_logits`, which only compiles below 2 GiB, so the budget that decides
@@ -8,34 +9,30 @@ The measured memory budget is stubbed: it is the only input the limit has to
 beat, and stubbing it keeps these tests on CPU.
 """
 
-from unittest import mock
-
 import pytest
 
 torch = pytest.importorskip("torch")
 
-from sglang.srt.layers.attention.dsa import dsa_indexer  # noqa: E402
+from sglang.srt.layers.attention.mqa_logits_utils import (  # noqa: E402
+    MQA_LOGITS_MAX_BYTES_ROCM,
+    mqa_logits_should_chunk,
+)
 from sglang.test.ci.ci_register import register_cpu_ci  # noqa: E402
 
 register_cpu_ci(est_time=9, suite="base-a-test-cpu")
 
-CEILING = dsa_indexer.Indexer._MQA_LOGITS_MAX_BYTES_ROCM
+CEILING = MQA_LOGITS_MAX_BYTES_ROCM
 # More than any single logits tensor here needs, so it never decides a case.
 HUGE_MEM_BUDGET = 64 * 2**30
 
 
 def _decide(num_q, num_k, mem_budget=HUGE_MEM_BUDGET, is_hip=True):
-    # __new__ skips an __init__ that needs a model config and a device.
+    return mqa_logits_should_chunk(
-    indexer = dsa_indexer.Indexer.__new__(dsa_indexer.Indexer)
+        num_rows=num_q,
-    with (
+        num_cols=num_k,
-        mock.patch.object(dsa_indexer, "_is_hip", is_hip),
+        get_budget_bytes=lambda: mem_budget,
-        mock.patch.object(
+        rocm=is_hip,
-            dsa_indexer.Indexer,
+    )
-            "_get_mqa_logits_budget_bytes",
-            return_value=mem_budget,
-        ),
-    ):
-        return indexer._should_chunk_mqa_logits(num_q, num_k, 0)
 
 
 def test_the_ceiling_is_the_largest_logits_aiter_still_takes():

test/registered/unit/layers/test_dsv4_nonpaged_indexer.py

@@ -11,19 +11,35 @@ from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend
 from sglang.srt.layers.attention.dsv4.indexer import (
     FP8_DTYPE,
     C4IndexerBackendMixin,
+    topk_transform_pytorch_vectorized,
 )
 from sglang.srt.layers.attention.dsv4.metadata import (
     NonPagedIndexerPlan,
     PagedIndexerMetadata,
+    iter_row_chunks,
+    plan_indexer_row_chunks,
+)
+from sglang.srt.layers.attention.mqa_logits_utils import (
+    MQA_LOGITS_MAX_BYTES_ROCM,
+    mqa_logits_budget_bytes,
+    mqa_logits_row_bytes,
+    mqa_logits_rows_per_chunk,
+    mqa_logits_should_chunk,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_cpu_ci(est_time=11, suite="base-a-test-cpu")
+register_cpu_ci(est_time=17, suite="base-a-test-cpu")
 
 _INDEXER = "sglang.srt.layers.attention.dsv4.indexer"
+_METADATA = "sglang.srt.layers.attention.dsv4.metadata"
+_MQA_UTILS = "sglang.srt.layers.attention.mqa_logits_utils"
+
+# issue #35201: 372K raw tokens -> 92992 c4 columns, padded to 93184 by DeepGEMM.
+_ISSUE_C4_COLS = 92992
+_ISSUE_ALIGNED_COLS = 93184
 
 
 class TestDSV4PagedIndexerMetadata(CustomTestCase):
@@ -247,8 +263,11 @@ class TestDSV4TopKDispatch(CustomTestCase):
                 torch.empty((1, 1, 1)),
             )
         )
-        backend._get_nonpaged_indexer_plan = MagicMock(return_value=object())
+        backend._get_nonpaged_indexer_plan = MagicMock(
-        backend._forward_nonpaged_indexer = MagicMock(return_value=logits)
+            return_value=SimpleNamespace(query_rows=1, rows_per_chunk=None)
+        )
+        backend._gather_nonpaged_index_k = MagicMock(return_value=(object(), object()))
+        backend._nonpaged_mqa_logits = MagicMock(return_value=logits)
 
         indexer_capturer = MagicMock()
         with (
@@ -353,7 +372,9 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
             extend_start_loc=torch.tensor([0], dtype=torch.int32),
             extend_num_tokens=query_rows,
         )
-        metadata = SimpleNamespace(nonpaged_plan=None, compressed_page_size=64)
+        metadata = SimpleNamespace(
+            nonpaged_plan=None, compressed_page_size=64, mqa_logits_budget_bytes=None
+        )
         page_table = torch.tensor([[3, 1]], dtype=torch.int32).repeat(query_rows, 1)
         c4_seq_lens = torch.tensor([62, 63, 64, 65], dtype=torch.int32)
 
@@ -388,6 +409,50 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
         with threshold.override(0):
             self.assertIsNone(build_plan())
 
+    def test_plan_row_chunking_follows_the_forward_budget(self):
+        backend = SimpleNamespace(_can_use_nonpaged_indexer=lambda **_: True)
+        backend.dsa_topk_backend = SimpleNamespace(is_sgl_kernel=lambda: True)
+        c4_indexer = SimpleNamespace(use_fp4_indexer=False, index_topk=512)
+        query_rows = 8192
+        seq_len = 372_000
+        batch = SimpleNamespace(
+            seq_lens=torch.tensor([seq_len], dtype=torch.int32),
+            seq_lens_cpu=[seq_len],
+            extend_seq_lens_cpu=[query_rows],
+            extend_seq_lens=torch.tensor([query_rows], dtype=torch.int32),
+            extend_start_loc=torch.tensor([0], dtype=torch.int32),
+            extend_num_tokens=query_rows,
+        )
+        c4_seq_lens = torch.full((query_rows,), seq_len // 4, dtype=torch.int32)
+        page_table = torch.zeros((query_rows, 1), dtype=torch.int32)
+
+        def build_plan(budget):
+            metadata = SimpleNamespace(
+                nonpaged_plan=None,
+                compressed_page_size=64,
+                mqa_logits_budget_bytes=budget,
+            )
+            return C4IndexerBackendMixin._get_nonpaged_indexer_plan(
+                backend,
+                c4_indexer=c4_indexer,
+                forward_batch=batch,
+                indexer_metadata=metadata,
+                page_table=page_table,
+                c4_seq_lens=c4_seq_lens,
+                query_rows=query_rows,
+            )
+
+        # No budget measured this forward (small batch or graph): one call.
+        self.assertIsNone(build_plan(None).rows_per_chunk)
+        budget = 512 << 20
+        plan = build_plan(budget)
+        # 8192 rows x align256(93056) cols x 4 B is far over 512 MiB.
+        self.assertIsNotNone(plan.rows_per_chunk)
+        self.assertLess(plan.rows_per_chunk, query_rows)
+        self.assertLessEqual(
+            plan.rows_per_chunk * mqa_logits_row_bytes(plan.max_seqlen_k), budget
+        )
+
     def test_extreme_plan_metadata_is_bounded_and_fail_closed(self):
         backend = SimpleNamespace(_can_use_nonpaged_indexer=lambda **_: True)
         backend.dsa_topk_backend = SimpleNamespace(is_sgl_kernel=lambda: True)
@@ -401,7 +466,9 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
             extend_start_loc=torch.tensor([0], dtype=torch.int32),
             extend_num_tokens=query_rows,
         )
-        metadata = SimpleNamespace(nonpaged_plan=None, compressed_page_size=64)
+        metadata = SimpleNamespace(
+            nonpaged_plan=None, compressed_page_size=64, mqa_logits_budget_bytes=None
+        )
         page_table = torch.zeros((query_rows, 1), dtype=torch.int32)
         c4_seq_lens = torch.tensor(
             [124_997, 124_998, 124_999, 125_000], dtype=torch.int32
@@ -439,7 +506,9 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
         backend = SimpleNamespace(_can_use_nonpaged_indexer=can_use_nonpaged_indexer)
         backend.dsa_topk_backend = SimpleNamespace(is_sgl_kernel=lambda: True)
         c4_indexer = SimpleNamespace(use_fp4_indexer=False, index_topk=512)
-        metadata = SimpleNamespace(nonpaged_plan=None, compressed_page_size=64)
+        metadata = SimpleNamespace(
+            nonpaged_plan=None, compressed_page_size=64, mqa_logits_budget_bytes=None
+        )
 
         def build_plan(query_rows):
             batch = SimpleNamespace(
@@ -503,13 +572,18 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
         deep_gemm = SimpleNamespace(fp8_mqa_logits=MagicMock(return_value=expected))
 
         with patch.dict(sys.modules, {"deep_gemm": deep_gemm}):
-            actual = C4IndexerBackendMixin._forward_nonpaged_indexer(
+            kv = C4IndexerBackendMixin._gather_nonpaged_index_k(
-                q_indexer=q_indexer,
-                weights=weights,
                 c4_indexer=c4_indexer,
                 token_to_kv_pool=token_to_kv_pool,
                 plan=plan,
             )
+            actual = C4IndexerBackendMixin._nonpaged_mqa_logits(
+                q_indexer=q_indexer,
+                weights=weights,
+                kv=kv,
+                plan=plan,
+                rows=slice(0, plan.query_rows),
+            )
 
         self.assertIs(actual, expected)
         token_to_kv_pool.get_index_k_scale_buffer.assert_called_once_with(
@@ -531,6 +605,276 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
         self.assertEqual(call.kwargs, {"clean_logits": False, "max_seqlen_k": 128})
 
 
+class TestMqaLogitsBudgetArithmetic(CustomTestCase):
+    def test_row_bytes_follow_deepgemm_stride_alignment(self):
+        # DeepGEMM pads the fp32 logits row stride to 1024 B, i.e. 256 columns.
+        self.assertEqual(mqa_logits_row_bytes(1), 256 * 4)
+        self.assertEqual(mqa_logits_row_bytes(256), 256 * 4)
+        self.assertEqual(mqa_logits_row_bytes(257), 512 * 4)
+        self.assertEqual(mqa_logits_row_bytes(_ISSUE_C4_COLS), _ISSUE_ALIGNED_COLS * 4)
+
+    def test_rows_per_chunk_keeps_one_chunk_inside_budget(self):
+        row_bytes = mqa_logits_row_bytes(_ISSUE_C4_COLS)
+        budget = 512 << 20
+        # 4096 rows x 93184 cols x 4 B = 1.42 GiB > 512 MiB: must slice.
+        rows = mqa_logits_rows_per_chunk(
+            num_rows=4096, row_bytes=row_bytes, budget_bytes=budget
+        )
+        self.assertIsNotNone(rows)
+        self.assertLess(rows, 4096)
+        self.assertLessEqual(rows * row_bytes, budget)
+        # The whole matrix fits: single call.
+        self.assertIsNone(
+            mqa_logits_rows_per_chunk(
+                num_rows=64, row_bytes=row_bytes, budget_bytes=budget
+            )
+        )
+        # A tight budget never yields a chunk larger than the budget, even when
+        # that means fewer rows than a full page of queries.
+        tight = 64 << 20
+        rows = mqa_logits_rows_per_chunk(
+            num_rows=4096, row_bytes=row_bytes, budget_bytes=tight
+        )
+        self.assertEqual(rows, tight // row_bytes)
+        self.assertLessEqual(rows * row_bytes, tight)
+        # Few query rows still chunk when they do not fit.
+        self.assertEqual(
+            mqa_logits_rows_per_chunk(
+                num_rows=100, row_bytes=row_bytes, budget_bytes=40 * row_bytes
+            ),
+            40,
+        )
+        # A budget below one row degrades to single-row chunks, never None.
+        self.assertEqual(
+            mqa_logits_rows_per_chunk(
+                num_rows=4096, row_bytes=row_bytes, budget_bytes=1
+            ),
+            1,
+        )
+
+    def test_should_chunk_caps_the_budget_only_on_rocm(self):
+        huge = 64 << 30
+        # 16384 x 32768 x 4 B is exactly 2 GiB, aiter's compile-time ceiling.
+        self.assertEqual(
+            mqa_logits_should_chunk(
+                num_rows=16384, num_cols=32768, get_budget_bytes=lambda: huge, rocm=True
+            ),
+            (True, MQA_LOGITS_MAX_BYTES_ROCM),
+        )
+        self.assertEqual(
+            mqa_logits_should_chunk(
+                num_rows=16384,
+                num_cols=32768,
+                get_budget_bytes=lambda: huge,
+                rocm=False,
+            ),
+            (False, huge),
+        )
+
+    def test_should_chunk_skips_small_matrices_without_querying_the_budget(self):
+        get_budget = MagicMock(return_value=1)
+        # 64 decode rows x 100K columns is far below the 8M-element threshold.
+        self.assertEqual(
+            mqa_logits_should_chunk(
+                num_rows=64, num_cols=100_000, get_budget_bytes=get_budget, rocm=False
+            ),
+            (False, 0),
+        )
+        get_budget.assert_not_called()
+
+    def test_plan_combines_sm120_cap_with_budget(self):
+        budget = 512 << 20
+        by_budget = mqa_logits_rows_per_chunk(
+            num_rows=8192,
+            row_bytes=mqa_logits_row_bytes(_ISSUE_C4_COLS),
+            budget_bytes=budget,
+        )
+        cases = (
+            # (num_rows, num_cols, budget_bytes, sm120_row_cap) -> rows_per_chunk
+            ((8192, 1024, None, None), None),
+            ((8192, 1024, None, 4096), 4096),
+            ((4096, 1024, None, 4096), None),
+            ((8192, _ISSUE_C4_COLS, budget, None), by_budget),
+            ((8192, _ISSUE_C4_COLS, budget, 4096), min(4096, by_budget)),
+        )
+        for (num_rows, num_cols, budget_bytes, cap), expected in cases:
+            with self.subTest(num_rows=num_rows, num_cols=num_cols, cap=cap):
+                self.assertEqual(
+                    plan_indexer_row_chunks(
+                        num_rows=num_rows,
+                        num_cols=num_cols,
+                        budget_bytes=budget_bytes,
+                        sm120_row_cap=cap,
+                    ),
+                    expected,
+                )
+
+    def test_iter_row_chunks_covers_rows_exactly_once(self):
+        self.assertEqual(
+            list(iter_row_chunks(num_rows=10, rows_per_chunk=4)),
+            [slice(0, 4), slice(4, 8), slice(8, 10)],
+        )
+        self.assertEqual(
+            list(iter_row_chunks(num_rows=10, rows_per_chunk=None)), [slice(0, 10)]
+        )
+        self.assertEqual(
+            list(iter_row_chunks(num_rows=10, rows_per_chunk=64)), [slice(0, 10)]
+        )
+
+    def test_static_budget_never_queries_free_memory(self):
+        total = 80 << 30
+        props = SimpleNamespace(total_memory=total)
+        device_module = SimpleNamespace(
+            get_device_properties=MagicMock(return_value=props)
+        )
+        schedule = SimpleNamespace(mem_fraction_static=0.9)
+        with (
+            envs.SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION.override(0.2),
+            patch(f"{_MQA_UTILS}.get_device_module", return_value=device_module),
+            patch(f"{_MQA_UTILS}.get_schedule", return_value=schedule),
+            patch(f"{_MQA_UTILS}.is_hip", return_value=False),
+            patch(f"{_MQA_UTILS}.is_xpu", return_value=False),
+            patch(
+                "torch.cuda.mem_get_info", return_value=(6 << 30, total)
+            ) as mem_get_info,
+        ):
+            static = mqa_logits_budget_bytes(device_index=0, allow_sync=False)
+            mem_get_info.assert_not_called()
+            live = mqa_logits_budget_bytes(device_index=0, allow_sync=True)
+            mem_get_info.assert_called_once_with(0)
+        # static: 80 GiB x (1 - 0.9) x 0.2; live is further capped by 6 GiB free x 0.2.
+        self.assertEqual(static, int(int(total * 0.1) * 0.2))
+        self.assertEqual(live, int((6 << 30) * 0.2))
+
+
+class TestPagedIndexerMetadataChunking(CustomTestCase):
+    """The schedule list and the top-k plan list must be built over the exact
+    row chunks the indexer loops over; a mismatch would silently score rows
+    with another chunk's schedule."""
+
+    def _build(self, *, num_rows: int, budget, use_topk_v2: bool):
+        deep_gemm = SimpleNamespace(
+            get_num_sms=MagicMock(return_value=1),
+            get_paged_mqa_logits_metadata=MagicMock(
+                side_effect=lambda c4, *_: torch.zeros((2, 2), dtype=torch.int32)
+            ),
+        )
+        c4_seq_lens = torch.arange(1, num_rows + 1, dtype=torch.int32)
+        page_table = torch.zeros(
+            (num_rows, _ISSUE_ALIGNED_COLS // 64), dtype=torch.int32
+        )
+        with (
+            patch.dict(sys.modules, {"deep_gemm": deep_gemm}),
+            envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.override(False),
+            envs.SGLANG_OPT_USE_JIT_INDEXER_METADATA.override(False),
+            patch(f"{_METADATA}.is_hip", return_value=False),
+            patch(f"{_METADATA}.is_xpu", return_value=False),
+            patch(f"{_METADATA}._IS_SM120", False),
+            patch.object(
+                PagedIndexerMetadata, "_mqa_logits_budget", return_value=budget
+            ),
+            patch(
+                "sglang.kernels.ops.attention.dsv4.plan_topk_v2",
+                side_effect=lambda seq_lens: seq_lens.new_zeros(
+                    (seq_lens.shape[0] + 1, 2)
+                ),
+            ) as plan_topk_v2,
+        ):
+            metadata = PagedIndexerMetadata(
+                page_size=256,
+                compressed_page_size=64,
+                page_table=page_table,
+                compressed_seq_lens=c4_seq_lens,
+                use_topk_v2=use_topk_v2,
+            )
+        return metadata, deep_gemm, plan_topk_v2
+
+    def test_budget_splits_schedules_and_topk_plans_over_the_same_rows(self):
+        num_rows, budget = 4096, 512 << 20
+        metadata, deep_gemm, plan_topk_v2 = self._build(
+            num_rows=num_rows, budget=budget, use_topk_v2=True
+        )
+        expected_rows = mqa_logits_rows_per_chunk(
+            num_rows=num_rows,
+            row_bytes=mqa_logits_row_bytes(_ISSUE_ALIGNED_COLS),
+            budget_bytes=budget,
+        )
+        self.assertEqual(metadata.rows_per_chunk, expected_rows)
+        self.assertEqual(metadata.mqa_logits_budget_bytes, budget)
+        chunks = list(iter_row_chunks(num_rows=num_rows, rows_per_chunk=expected_rows))
+        self.assertGreater(len(chunks), 1)
+
+        self.assertIsInstance(metadata.deep_gemm_metadata, list)
+        self.assertEqual(len(metadata.deep_gemm_metadata), len(chunks))
+        schedule_rows = [
+            call.args[0]
+            for call in deep_gemm.get_paged_mqa_logits_metadata.call_args_list
+        ]
+        torch.testing.assert_close(
+            torch.cat(schedule_rows), metadata.compressed_seq_lens.unsqueeze(-1)
+        )
+        self.assertEqual(
+            [r.shape[0] for r in schedule_rows], [c.stop - c.start for c in chunks]
+        )
+
+        self.assertEqual(len(metadata.topk_metadata_chunks), len(chunks))
+        # First call is the full-batch plan; the rest are one per chunk.
+        plan_rows = [call.args[0] for call in plan_topk_v2.call_args_list]
+        torch.testing.assert_close(plan_rows[0], metadata.compressed_seq_lens)
+        torch.testing.assert_close(
+            torch.cat(plan_rows[1:]), metadata.compressed_seq_lens
+        )
+        self.assertEqual(
+            [r.shape[0] for r in plan_rows[1:]], [c.stop - c.start for c in chunks]
+        )
+
+    def test_no_budget_keeps_the_single_call_shape(self):
+        metadata, deep_gemm, plan_topk_v2 = self._build(
+            num_rows=4096, budget=None, use_topk_v2=True
+        )
+        self.assertIsNone(metadata.rows_per_chunk)
+        self.assertIsNone(metadata.topk_metadata_chunks)
+        self.assertIsInstance(metadata.deep_gemm_metadata, torch.Tensor)
+        deep_gemm.get_paged_mqa_logits_metadata.assert_called_once()
+        plan_topk_v2.assert_called_once()
+
+
+class TestChunkedTopKMatchesUnchunked(CustomTestCase):
+    """Each row's top-k depends only on its own logits row, sequence length and
+    page-table row, so scoring the batch in row chunks must select the same
+    pages as one pass. This is the property the chunk loop relies on."""
+
+    def test_row_chunks_select_the_same_pages(self):
+        torch.manual_seed(0)
+        rows, width, topk, page_size = 37, 2048, 64, 64
+        logits = torch.randn(rows, width, dtype=torch.float32)
+        seq_lens = torch.randint(1, width, (rows,), dtype=torch.int32)
+        page_table = torch.randint(
+            0, 4096, (rows, width // page_size), dtype=torch.int32
+        )
+
+        def run(rows_per_chunk):
+            out = torch.full((rows, topk), -1, dtype=torch.int32)
+            for rows_slice in iter_row_chunks(
+                num_rows=rows, rows_per_chunk=rows_per_chunk
+            ):
+                topk_transform_pytorch_vectorized(
+                    logits[rows_slice],
+                    seq_lens[rows_slice],
+                    page_table[rows_slice],
+                    out[rows_slice],
+                    page_size,
+                    None,
+                )
+            # Unsorted top-k: compare the selected sets row by row.
+            return out.sort(dim=1).values
+
+        expected = run(None)
+        for rows_per_chunk in (1, 7, 16, rows - 1):
+            with self.subTest(rows_per_chunk=rows_per_chunk):
+                self.assertTrue(torch.equal(run(rows_per_chunk), expected))
+
+
 class TestCandidateIndexerGating(CustomTestCase):
     def test_candidate_indexer_gating(self):
         from sglang.srt.layers.attention.dsv4 import candidate_indexer