[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_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
+    _MQA_LOGITS_BYTES_PER_ELEM = MQA_LOGITS_BYTES_PER_ELEM
+    _MQA_LOGITS_STATIC_SKIP_ELEMS = MQA_LOGITS_STATIC_SKIP_ELEMS
+    _MQA_LOGITS_TOTAL_MEM_FRACTION = MQA_LOGITS_TOTAL_MEM_FRACTION
+    _MQA_LOGITS_MAX_BYTES_ROCM = MQA_LOGITS_MAX_BYTES_ROCM
+    # 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
-
-        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.
+        # 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.
         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.
-        # Synchronizes the host; cache the result capped by serving-memory headroom.
-        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)
+        budget_bytes = mqa_logits_budget_bytes(
+            device_index=device_index, allow_sync=True
+        )
         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(
-            q_offset, k_offset, device_index
+        need_chunk, logits_budget_bytes = mqa_logits_should_chunk(
+            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:
-        import deep_gemm
-
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         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)
-        k_scale = scale_u8.view(torch.float32).squeeze(-1)
+        return k_u8.view(FP8_DTYPE), 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],
-            (k_fp8, k_scale),
-            weights[: plan.query_rows],
-            plan.ks,
-            plan.ke,
+            q_indexer[rows],
+            kv,
+            weights[rows],
+            plan.ks[rows],
+            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
-                    # 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])
-                    ),
+                    topk_plan,
                     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(
-                q_indexer=q_indexer,
-                weights=weights,
+            # K is gathered once; each row chunk's logits are reduced to top-k
+            # and dropped before the next chunk allocates, so only one chunk
+            # 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 it; run indexer + topk per chunk.
+                # PagedIndexerMetadata split this forward into row chunks (SM120
+                # kernel cap and/or logits memory budget), one schedule each.
                 num_rows = _c4sl.shape[0]
-                for chunk_idx, start in enumerate(
-                    range(0, num_rows, _SM120_INDEXER_M_CHUNK)
+                assert num_rows == indexer_metadata.compressed_seq_lens.shape[0], (
+                    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(rows, deep_gemm_metadata[chunk_idx])
+                    run_paged_indexer(
+                        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(
-                        0, compressed_seq_lens.shape[0], _SM120_INDEXER_M_CHUNK
+                    for rows in iter_row_chunks(
+                        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