Bound NSA MQA logits peak memory

Paged and CP-ragged NSA indexer paths could materialize q x context fp32
MQA-logits buffers large enough to OOM high-cache-hit bs>1 prefill batches.
Port the syh branch chunking logic so paged and ragged paths split logits by
query rows when the estimated logits buffer exceeds the current free-memory
budget.

The free-memory query is cached on forward_batch so the OOM guard uses current
free memory without adding a torch.cuda.mem_get_info host sync on every layer.
The only new env kept from the syh commits is
SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS, which forces chunking for equivalence
validation.

Constraint: DeepGEMM fp8_mqa_logits still materializes fp32 logits internally, so limiting q rows is the least invasive way to cap peak memory
Rejected: Carry unrelated syh envs for page trace/source-fingerprint strictness | not part of the logits peak-memory fix
Rejected: Static mem_fraction-only budget | overestimates logits headroom shared with other forward activations
Confidence: medium
Scope-risk: moderate
Directive: Keep chunking row-split only; changing K/context partitioning needs topk_transform equivalence validation
Related: 40a0389a9c feat(nsa): chunk paged + CP-ragged MQA-logits by current-free-mem budget
Related: 108fa1f538 perf(nsa): cache MQA-logits free-mem budget per-forward
Tested: Local py_compile for environ.py and nsa_indexer.py
Tested: Remote g0034 cjy-glm5-new py_compile for environ.py and nsa_indexer.py
Not-tested: CUDA ETE run with forced SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS equivalence check
Not-tested: Full high-cache-hit bs>1 prefill OOM regression workload
Co-authored-by: OmX <omx@oh-my-codex.dev>

python/sglang/srt/environ.py

@@ -220,6 +220,11 @@ class Envs:
     # large bs) but coarser overlap. 1 = per-layer.
     SGLANG_CP_SHARED_KV_PER_LAYER_GROUP = EnvInt(8)
     SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
+    # NSA paged MQA-logits chunking equivalence test: when >0, force the paged
+    # topk path to chunk at this many query rows AND assert the chunked topk_result
+    # is byte-identical to the unchunked single-call result. For validation only
+    # (run a small batch so the unchunked reference fits); 0 = off (production).
+    SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS = EnvInt(0)
     SGLANG_CP_SHARED_KV_FUSED_MLA_STORE = EnvBool(False)
     SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False)
     SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -1073,6 +1073,7 @@ class Indexer(MultiPlatformOp):
         # When attn_tp_size > 1 or in the MAX_LEN padding mode, padding may exist in the hidden states,
         # and it is necessary to extract the actual q length.
         q_offset = sum(metadata.get_nsa_extend_len_cpu())
+        topk_result = None
         if _is_hip:
             from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits
 
@@ -1098,19 +1099,100 @@ class Indexer(MultiPlatformOp):
                 WavePerEU=5,
             )
         else:
-            logits = deep_gemm.fp8_paged_mqa_logits(
+            device_index = q_fp8.device.index
-                q_fp8[:q_offset],
+            assert device_index is not None
-                kv_cache_fp8,
+            # The kernel allocates logits of width align(max_seq_len, 256) (DeepGEMM
-                weights[:q_offset],
+            # attention.hpp), so budget/chunk against the aligned width.
-                seqlens_32_2d,
+            aligned_ctx = ((max_seq_len + 255) // 256) * 256
-                block_tables,
+            force_rows = int(envs.SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS.get())
-                schedule_metadata,
+            need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits(
-                max_seq_len,
+                q_offset, aligned_ctx, device_index, forward_batch=forward_batch
-                clean_logits=False,
             )
+            if force_rows > 0:
+                need_chunk = True
+            if not need_chunk:
+                logits = deep_gemm.fp8_paged_mqa_logits(
+                    q_fp8[:q_offset],
+                    kv_cache_fp8,
+                    weights[:q_offset],
+                    seqlens_32_2d,
+                    block_tables,
+                    schedule_metadata,
+                    max_seq_len,
+                    clean_logits=False,
+                )
+            else:
+                # Bound the q_offset x align(max_seq_len,256) f32 logits buffer by
+                # chunking over query rows (each paged q-row is its own length-1 entry,
+                # so any row split is valid). Recompute the SM schedule per chunk (it
+                # encodes the work split for this chunk's context_lens). Run the topk
+                # transform per chunk with the per-chunk paged args (ke_offset /
+                # batch_idx_list / cu_seqlens_q override) so we never materialize the
+                # full logits buffer. Mirrors the ragged chunk loop.
+                bytes_per_row = aligned_ctx * self._MQA_LOGITS_BYTES_PER_ELEM
+                if force_rows > 0:
+                    max_rows = force_rows
+                else:
+                    max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1)))
+                max_rows = min(max(1, max_rows), q_offset)
+                seqlens_expanded_full = metadata.get_seqlens_expanded()
+                start = 0
+                while start < q_offset:
+                    end = min(start + max_rows, q_offset)
+                    sched_chunk = deep_gemm.get_paged_mqa_logits_metadata(
+                        seqlens_32_2d[start:end], blocksize, self.sm_count
+                    )
+                    logits_chunk = deep_gemm.fp8_paged_mqa_logits(
+                        q_fp8[start:end],
+                        kv_cache_fp8,
+                        weights[start:end],
+                        seqlens_32_2d[start:end],
+                        block_tables[start:end],
+                        sched_chunk,
+                        max_seq_len,
+                        clean_logits=False,
+                    )
+                    cu_chunk = torch.arange(
+                        0,
+                        (end - start) + 1,
+                        dtype=torch.int32,
+                        device=logits_chunk.device,
+                    )
+                    topk_chunk = metadata.topk_transform(
+                        logits_chunk,
+                        self.index_topk,
+                        ke_offset=seqlens_expanded_full[start:end],
+                        batch_idx_list=list(range(start, end)),
+                        cu_seqlens_q_topk_override=cu_chunk,
+                    )
+                    if topk_result is None:
+                        topk_result = topk_chunk.new_full(
+                            (q_offset, topk_chunk.shape[1]), -1
+                        )
+                    topk_result[start:end] = topk_chunk
+                    start = end
+                if force_rows > 0:
+                    # Equivalence gate: chunked topk_result must be byte-identical to
+                    # the unchunked single-call path (run a small batch so this fits).
+                    ref_logits = deep_gemm.fp8_paged_mqa_logits(
+                        q_fp8[:q_offset],
+                        kv_cache_fp8,
+                        weights[:q_offset],
+                        seqlens_32_2d,
+                        block_tables,
+                        schedule_metadata,
+                        max_seq_len,
+                        clean_logits=False,
+                    )
+                    ref_topk = metadata.topk_transform(ref_logits, self.index_topk)
+                    assert torch.equal(topk_result, ref_topk), (
+                        "[MQA_LOGITS_CHUNK_VERIFY] paged chunked topk_result != "
+                        f"unchunked (q_offset={q_offset}, max_rows={max_rows})"
+                    )
 
-        # NOTE(dark): logits should be cleaned in topk_transform
+        if topk_result is None:
-        topk_result = metadata.topk_transform(logits, self.index_topk)
+            # NOTE(dark): logits should be cleaned in topk_transform
+            topk_result = metadata.topk_transform(logits, self.index_topk)
         # Restore possible padding exist in the hidden states.
         if not _is_hip and q_offset < q_fp8.shape[0]:
             pad_len = q_fp8.shape[0] - q_offset
@@ -1158,8 +1240,35 @@ class Indexer(MultiPlatformOp):
         self._mqa_logits_budget_bytes[device_index] = budget_bytes
         return budget_bytes
 
+    def _current_free_mem_logits_budget(
+        self, device_index: int, forward_batch=None
+    ) -> int:
+        """0.5 x CURRENT free-memory budget for the MQA-logits buffer, queried at
+        most ONCE per forward.
+
+        ``torch.cuda.mem_get_info`` host-syncs; the indexer runs once per layer, so
+        querying it per call serializes the host ~num_layers x per forward and
+        starves the GPU between batches (commit 40a0389a9c introduced the per-call
+        query for OOM safety -- this caches it without losing that safety). Free
+        memory is ~constant across layers within a forward: eager mode frees each
+        layer's activations and the KV pool is pre-reserved, so the driver-level
+        free-mem high-water-mark is set early and stays flat. We snapshot it on the
+        ``forward_batch`` (recreated per forward) and reuse it for later layers; a new
+        forward gets a fresh snapshot. Falls back to a direct query when no
+        ``forward_batch`` is threaded (keeps the call correct, just uncached).
+        """
+        if forward_batch is not None:
+            cached = getattr(forward_batch, "_nsa_mqa_free_budget", None)
+            if cached is not None and cached[0] == device_index:
+                return cached[1]
+        free_mem, _ = torch.cuda.mem_get_info(device_index)
+        budget_bytes = max(1, int(free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION))
+        if forward_batch is not None:
+            forward_batch._nsa_mqa_free_budget = (device_index, budget_bytes)
+        return budget_bytes
+
     def _should_chunk_mqa_logits(
-        self, num_q: int, num_k: int, device_index: int
+        self, num_q: int, num_k: int, device_index: int, forward_batch=None
     ) -> Tuple[bool, int]:
         """
         Detect whether we need to chunk the MQA logits computation to avoid OOM
@@ -1170,11 +1279,122 @@ class Indexer(MultiPlatformOp):
             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)
+        # Budget against CURRENT free memory, not the cached first-prefill / static
+        # estimate. The static headroom (1 - mem_fraction_static) is SHARED with the
+        # rest of the forward's activations, so a cached estimate over-counts what the
+        # logits buffer alone may use and OOMs at large batch (observed: 15.1 GiB
+        # logits "fit" a 24 GiB static budget but only 14.67 GiB was actually free).
+        # Reached only for large logits (post static-skip). The per-forward snapshot
+        # (see _current_free_mem_logits_budget) keeps that current-free-mem safety
+        # while collapsing the per-layer host-sync to once per forward.
+        # Keep the static guard during CUDA-graph capture (mem_get_info unreliable).
+        if get_is_capture_mode():
+            logits_budget_bytes = self._get_mqa_logits_budget_bytes(device_index)
+        else:
+            logits_budget_bytes = self._current_free_mem_logits_budget(
+                device_index, forward_batch
+            )
 
         need_chunk = logits_bytes > logits_budget_bytes
         return need_chunk, logits_budget_bytes
 
+    def _mqa_logits_topk_ragged_chunked(
+        self,
+        metadata,
+        q_fp8,
+        kv_fp8,
+        weights,
+        ks,
+        ke,
+        *,
+        actual_seq_q,
+        ke_offset,
+        batch_idx_list,
+        topk_indices_offset_override,
+        forward_batch=None,
+    ):
+        """RAGGED fp8_mqa_logits + topk_transform, byte-budget-chunked over query rows.
+
+        Mirrors the `_get_topk_ragged` chunk loop so the unbounded `q_offset x kv_len`
+        f32 logits buffer can't OOM when the CP prefill batch grows. Per-row inputs
+        (q/weights/ks/ke/ke_offset/topk_indices_offset) are sliced; the shared `kv_fp8`
+        stays whole. The RAGGED transform keys off per-row `ks` + `ke_offset` +
+        `topk_indices_offset_override`, so per-chunk results are byte-identical
+        (`cu_seqlens_q`/`batch_idx_list` are unused once the override is set --
+        nsa_backend.py:569). SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS>0 forces chunking
+        and asserts equivalence vs the single-call path.
+        """
+        device_index = q_fp8.device.index
+        assert device_index is not None
+        q_offset = int(q_fp8.shape[0])
+        k_offset = int(kv_fp8[0].shape[0])
+        force_rows = int(envs.SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS.get())
+        need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits(
+            q_offset, k_offset, device_index, forward_batch=forward_batch
+        )
+        if force_rows > 0:
+            need_chunk = True
+
+        def _single():
+            with self._with_real_sm_count():
+                logits = deep_gemm.fp8_mqa_logits(
+                    q_fp8, kv_fp8, weights, ks, ke, clean_logits=False
+                )
+            return metadata.topk_transform(
+                logits,
+                self.index_topk,
+                ks=ks,
+                cu_seqlens_q=actual_seq_q,
+                ke_offset=ke_offset,
+                batch_idx_list=batch_idx_list,
+                topk_indices_offset_override=topk_indices_offset_override,
+            )
+
+        if not need_chunk:
+            return _single()
+
+        bytes_per_row = k_offset * self._MQA_LOGITS_BYTES_PER_ELEM
+        if force_rows > 0:
+            max_rows = force_rows
+        else:
+            max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1)))
+        max_rows = min(max(1, max_rows), q_offset)
+
+        topk_result = None
+        start = 0
+        while start < q_offset:
+            end = min(start + max_rows, q_offset)
+            with self._with_real_sm_count():
+                logits_chunk = deep_gemm.fp8_mqa_logits(
+                    q_fp8[start:end],
+                    kv_fp8,
+                    weights[start:end],
+                    ks[start:end],
+                    ke[start:end],
+                    clean_logits=False,
+                )
+            topk_chunk = metadata.topk_transform(
+                logits_chunk,
+                self.index_topk,
+                ks=ks[start:end],
+                ke_offset=ke_offset[start:end],
+                topk_indices_offset_override=topk_indices_offset_override[start:end],
+            )
+            if topk_result is None:
+                topk_result = topk_chunk.new_full(
+                    (q_offset, topk_chunk.shape[1]), -1
+                )
+            topk_result[start:end] = topk_chunk
+            start = end
+
+        if force_rows > 0:
+            ref = _single()
+            assert torch.equal(topk_result, ref), (
+                "[MQA_LOGITS_CHUNK_VERIFY] cp-ragged chunked topk_result != "
+                f"unchunked (q_offset={q_offset}, max_rows={max_rows})"
+            )
+        return topk_result
+
     def _get_topk_ragged(
         self,
         enable_dual_stream: bool,
@@ -1256,7 +1476,7 @@ class Indexer(MultiPlatformOp):
         device_index = device.index
         assert device_index is not None, "q_fp8 must be on an indexed CUDA device"
         need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits(
-            q_offset, k_offset, device_index
+            q_offset, k_offset, device_index, forward_batch=forward_batch
         )
 
         if not need_chunk:
@@ -1588,23 +1808,18 @@ class Indexer(MultiPlatformOp):
                         q_lens_list, dtype=torch.int32, device=q_fp8.device
                     )
                     ke = ks + ke_offset
-                    with self._with_real_sm_count():
+                    topk_result = self._mqa_logits_topk_ragged_chunked(
-                        logits = deep_gemm.fp8_mqa_logits(
+                        metadata,
-                            q_fp8,
+                        q_fp8,
-                            kv_fp8,
+                        kv_fp8,
-                            weights,
+                        weights,
-                            ks,
+                        ks,
-                            ke,
+                        ke,
-                            clean_logits=False,
+                        actual_seq_q=actual_seq_q,
-                        )
-                    topk_result = metadata.topk_transform(
-                        logits,
-                        self.index_topk,
-                        ks=ks,
-                        cu_seqlens_q=actual_seq_q,
                         ke_offset=ke_offset,
                         batch_idx_list=batch_idx_list,
                         topk_indices_offset_override=topk_indices_offset_override,
+                        forward_batch=forward_batch,
                     )
                     return topk_result
             else:
@@ -1653,23 +1868,18 @@ class Indexer(MultiPlatformOp):
                         q_lens_list, dtype=torch.int32, device=q_fp8.device
                     )
                     ke = ks + ke_offset
-                    with self._with_real_sm_count():
+                    topk_result = self._mqa_logits_topk_ragged_chunked(
-                        logits = deep_gemm.fp8_mqa_logits(
+                        metadata,
-                            q_fp8,
+                        q_fp8,
-                            kv_fp8,
+                        kv_fp8,
-                            weights,
+                        weights,
-                            ks,
+                        ks,
-                            ke,
+                        ke,
-                            clean_logits=False,
+                        actual_seq_q=actual_seq_q,
-                        )
-                    topk_result = metadata.topk_transform(
-                        logits,
-                        self.index_topk,
-                        ks=ks,
-                        cu_seqlens_q=actual_seq_q,
                         ke_offset=ke_offset,
                         batch_idx_list=batch_idx_list,
                         topk_indices_offset_override=topk_indices_offset_override,
+                        forward_batch=forward_batch,
                     )
                     return topk_result
 
@@ -1748,23 +1958,18 @@ class Indexer(MultiPlatformOp):
             ke_offset = torch.cat(ke_offset_list, dim=0)
             ke = ks + ke_offset
             actual_seq_q = torch.cat(actual_seq_q_list, dim=0)
-            with self._with_real_sm_count():
+            topk_result = self._mqa_logits_topk_ragged_chunked(
-                logits = deep_gemm.fp8_mqa_logits(
+                metadata,
-                    q_fp8,
+                q_fp8,
-                    kv_fp8,
+                kv_fp8,
-                    weights,
+                weights,
-                    ks,
+                ks,
-                    ke,
+                ke,
-                    clean_logits=False,
+                actual_seq_q=actual_seq_q,
-                )
-            topk_result = metadata.topk_transform(
-                logits,
-                self.index_topk,
-                ks=ks,
-                cu_seqlens_q=actual_seq_q,
                 ke_offset=ke_offset,
                 batch_idx_list=batch_idx_list,
                 topk_indices_offset_override=topk_indices_offset_override,
+                forward_batch=forward_batch,
             )
         else:
             seq_len = int(forward_batch.seq_lens_cpu[batch_idx].item())