Align MQA logits admission with CP split runtime shape

CP shared-KV batching previously estimated MQA logits from full request
extend/context rows, which overstated memory because CP in-seq split only
computes each rank's two zigzag segments. Add CP-size aware row accounting
that mirrors the fused CP MQA materialization path and take the worst local
rank peak for scheduler admission.

Expose SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB as a more direct cap for one fp32
MQA logits chunk. Runtime and scheduler now both translate this GB cap into
chunk rows from the actual K rows, while keeping the old row cap as a
mutually-exclusive expert override.

Constraint: Scheduler admission must stay CUDA-sync-free and use static budget information only.
Rejected: Keep full-request q*k admission | it over-gates CP bs>1 batches because CP splits q rows per rank.
Rejected: Let rows and GB caps both apply | precedence would be ambiguous during tuning.
Confidence: medium
Scope-risk: moderate
Directive: Keep MQA logits admission tied to the fused CP MQA segment shape; do not revert to full request token counts.
Tested: Local py_compile for touched runtime, scheduler, estimator, and tests.
Tested: Local pytest test_cp_shared_kv_prefill_buffer_estimator.py: 8 passed.
Tested: Remote g0034 cjy-glm5-new py_compile and targeted estimator/runtime tests: 13 passed.
Not-tested: Full ETE high-cache-hit CP bs>1 load with SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.

Co-authored-by: OmX <omx@oh-my-codex.dev>

python/sglang/srt/environ.py

@@ -225,6 +225,9 @@ class Envs:
     SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2)
     # Optional hard cap for rows per MQA-logits chunk. 0 = use memory budget.
     SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS = EnvInt(0)
+    # Optional hard cap for one fp32 MQA-logits chunk, in decimal GB. 0 = use
+    # row cap or memory budget.  Mutually exclusive with *_CHUNK_MAX_ROWS.
+    SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB = EnvFloat(0.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

@@ -1183,10 +1183,21 @@ class Indexer(MultiPlatformOp):
         self, num_q: int, num_k: int, logits_budget_bytes: int
     ) -> int:
         explicit_max_rows = int(envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.get())
+        explicit_max_gb = float(envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.get())
+        if explicit_max_rows > 0 and explicit_max_gb > 0:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][mqa_logits_chunk] "
+                "SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS and "
+                "SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB are mutually exclusive"
+            )
         if explicit_max_rows > 0:
             return min(max(1, explicit_max_rows), max(1, num_q))
 
         bytes_per_row = num_k * self._MQA_LOGITS_BYTES_PER_ELEM
+        if explicit_max_gb > 0:
+            explicit_max_bytes = int(explicit_max_gb * 1_000_000_000)
+            max_rows = max(1, explicit_max_bytes // max(bytes_per_row, 1))
+            return min(max_rows, max(1, num_q))
         max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1)))
         return min(max_rows, max(1, num_q))
 

python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py

@@ -33,6 +33,9 @@ class CPSharedKVPrefillBufferEstimatorContext:
     logprob_chunk_size: int
     bs_gt1_l1_prefetch_enabled: bool = False
     mqa_logits_chunk_max_rows: int = 0
+    mqa_logits_chunk_max_bytes: int = 0
+    mqa_logits_budget_bytes: int = 0
+    cp_size: int = 1
 
 
 @dataclass(frozen=True)
@@ -112,6 +115,120 @@ def _vocab_shard_size(model_config: object | None, tp_size: int) -> int:
     return -(-vocab_size // max(int(tp_size), 1))
 
 
+def _cp_segment_valid_lengths(
+    *,
+    extend_len: int,
+    page_size: int,
+    cp_size: int,
+) -> list[int]:
+    full_pages = int(extend_len) // int(page_size)
+    tail_tokens = int(extend_len) % int(page_size)
+    padded_pages = ceil_pages(extend_len, page_size)
+    cp_segment_num = int(cp_size) * 2
+    if cp_segment_num <= 0:
+        return []
+
+    base_units = padded_pages // cp_segment_num
+    remainder_units = padded_pages % cp_segment_num
+    unit_counts = [
+        base_units + (1 if idx < remainder_units else 0)
+        for idx in range(cp_segment_num)
+    ]
+
+    split_list: list[int] = []
+    page_cursor = 0
+    for unit_count in unit_counts:
+        token_count = 0
+        for _ in range(unit_count):
+            token_count += page_size if page_cursor < full_pages else tail_tokens
+            page_cursor += 1
+        split_list.append(token_count)
+    return split_list
+
+
+def _mqa_logits_row_cap(
+    *,
+    q_rows: int,
+    k_rows: int,
+    explicit_chunk_rows: int,
+    explicit_chunk_bytes: int,
+    budget_bytes: int,
+) -> int:
+    if q_rows <= 0 or k_rows <= 0:
+        return 0
+    if explicit_chunk_rows > 0 and explicit_chunk_bytes > 0:
+        raise ValueError(
+            "mqa_logits_chunk_max_rows and mqa_logits_chunk_max_bytes are "
+            "mutually exclusive"
+        )
+    if explicit_chunk_rows > 0:
+        return min(q_rows, max(1, explicit_chunk_rows))
+    if explicit_chunk_bytes > 0:
+        bytes_per_row = k_rows * 4
+        return min(q_rows, max(1, int(explicit_chunk_bytes) // max(bytes_per_row, 1)))
+    if budget_bytes > 0:
+        bytes_per_row = k_rows * 4
+        return min(q_rows, max(1, int(budget_bytes) // max(bytes_per_row, 1)))
+    return q_rows
+
+
+def _estimate_cp_split_mqa_logits_peak_bytes(
+    *,
+    prefix_lens: Sequence[int],
+    extend_lens: Sequence[int],
+    page_size: int,
+    cp_size: int,
+    explicit_chunk_rows: int,
+    explicit_chunk_bytes: int,
+    budget_bytes: int,
+) -> int:
+    if cp_size <= 1:
+        q_rows = sum(ceil_paged_tokens(tokens, page_size) for tokens in extend_lens)
+        k_rows = sum(
+            ceil_paged_tokens(prefix, page_size) + ceil_paged_tokens(extend, page_size)
+            for prefix, extend in zip(prefix_lens, extend_lens)
+        )
+        row_cap = _mqa_logits_row_cap(
+            q_rows=q_rows,
+            k_rows=k_rows,
+            explicit_chunk_rows=explicit_chunk_rows,
+            explicit_chunk_bytes=explicit_chunk_bytes,
+            budget_bytes=budget_bytes,
+        )
+        return row_cap * k_rows * 4
+
+    cp_segment_num = cp_size * 2
+    max_peak_bytes = 0
+    for cp_rank in range(cp_size):
+        q_rows = 0
+        k_rows = 0
+        mirror_idx = cp_segment_num - cp_rank - 1
+        for prefix_len, extend_len in zip(prefix_lens, extend_lens):
+            split_list = _cp_segment_valid_lengths(
+                extend_len=int(extend_len),
+                page_size=page_size,
+                cp_size=cp_size,
+            )
+            segment_end = 0
+            for segment_idx, segment_len in enumerate(split_list):
+                segment_end += int(segment_len)
+                if segment_idx not in (cp_rank, mirror_idx) or segment_len <= 0:
+                    continue
+                q_rows += int(segment_len)
+                # Match the fused CP MQA path: each owned segment materializes a
+                # compact K range up to this segment's logical end.
+                k_rows += ceil_paged_tokens(prefix_len, page_size) + segment_end
+        row_cap = _mqa_logits_row_cap(
+            q_rows=q_rows,
+            k_rows=k_rows,
+            explicit_chunk_rows=explicit_chunk_rows,
+            explicit_chunk_bytes=explicit_chunk_bytes,
+            budget_bytes=budget_bytes,
+        )
+        max_peak_bytes = max(max_peak_bytes, row_cap * k_rows * 4)
+    return max_peak_bytes
+
+
 def estimate_cp_shared_kv_prefill_buffer_bytes(
     *,
     page_size: int,
@@ -140,6 +257,9 @@ def estimate_cp_shared_kv_prefill_buffer_bytes(
             logprob_chunk_size=2048,
             bs_gt1_l1_prefetch_enabled=False,
             mqa_logits_chunk_max_rows=0,
+            mqa_logits_chunk_max_bytes=0,
+            mqa_logits_budget_bytes=0,
+            cp_size=1,
         )
 
     if page_size <= 0:
@@ -178,15 +298,19 @@ def estimate_cp_shared_kv_prefill_buffer_bytes(
             logits_rows = min(logits_rows, int(context.logprob_chunk_size))
     vocab_shard = _vocab_shard_size(context.model_config, context.tp_size)
     logits_peak_bytes = logits_rows * vocab_shard * int(logits_dtype_bytes)
-    mqa_q_rows = sum(ceil_paged_tokens(tokens, page_size) for tokens in extend_lens)
-    mqa_k_rows = sum(
-        ceil_paged_tokens(prefix, page_size) + ceil_paged_tokens(extend, page_size)
-        for prefix, extend in zip(prefix_lens, extend_lens)
+    mqa_logits_peak_bytes = _estimate_cp_split_mqa_logits_peak_bytes(
+        prefix_lens=prefix_lens,
+        extend_lens=extend_lens,
+        page_size=page_size,
+        cp_size=int(getattr(context, "cp_size", 1) or 1),
+        explicit_chunk_rows=int(
+            getattr(context, "mqa_logits_chunk_max_rows", 0) or 0
+        ),
+        explicit_chunk_bytes=int(
+            getattr(context, "mqa_logits_chunk_max_bytes", 0) or 0
+        ),
+        budget_bytes=int(getattr(context, "mqa_logits_budget_bytes", 0) or 0),
     )
-    mqa_chunk_rows = int(getattr(context, "mqa_logits_chunk_max_rows", 0) or 0)
-    if mqa_chunk_rows > 0:
-        mqa_q_rows = min(mqa_q_rows, mqa_chunk_rows)
-    mqa_logits_peak_bytes = mqa_q_rows * mqa_k_rows * 4
 
     transfer_descriptor_peak_bytes = total_prefix_pages * int(descriptor_bytes)
     backup_descriptor_peak_bytes = total_extend_pages * int(descriptor_bytes)

python/sglang/srt/managers/scheduler.py

@@ -524,6 +524,48 @@ class Scheduler(
     def make_cp_shared_kv_prefill_buffer_estimator_context(
         self,
     ) -> CPSharedKVPrefillBufferEstimatorContext:
+        mqa_logits_chunk_max_rows = int(
+            envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.get()
+        )
+        mqa_logits_chunk_max_gb = float(
+            envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.get()
+        )
+        if mqa_logits_chunk_max_rows > 0 and mqa_logits_chunk_max_gb > 0:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][mqa_logits_chunk] "
+                "SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS and "
+                "SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB are mutually exclusive"
+            )
+        mqa_logits_chunk_max_bytes = (
+            int(mqa_logits_chunk_max_gb * 1_000_000_000)
+            if mqa_logits_chunk_max_gb > 0
+            else 0
+        )
+        mqa_logits_budget_bytes = 0
+        if str(self.device).startswith("cuda"):
+            try:
+                total_mem = torch.cuda.get_device_properties(self.device).total_memory
+                total_mem_budget = int(total_mem * 0.3)
+                if self.server_args.mem_fraction_static is None:
+                    static_budget = total_mem_budget
+                else:
+                    static_free_mem = int(
+                        total_mem
+                        * max(0.0, 1.0 - self.server_args.mem_fraction_static)
+                    )
+                    static_budget = min(
+                        int(
+                            static_free_mem
+                            * envs.SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION.get()
+                        ),
+                        total_mem_budget,
+                    )
+                mqa_logits_budget_bytes = max(1, static_budget)
+            except Exception:
+                # Keep scheduler admission CUDA-sync-free and non-fatal.  A zero
+                # budget makes the estimator fall back to explicit chunk rows or
+                # the conservative unchunked peak.
+                mqa_logits_budget_bytes = 0
         return CPSharedKVPrefillBufferEstimatorContext(
             kvcache=self.token_to_kv_pool_allocator.get_kvcache(),
             model_config=self.model_config,
@@ -535,9 +577,10 @@ class Scheduler(
             # gates off bs>1. Keep the estimate aligned with runtime until that
             # path is enabled.
             bs_gt1_l1_prefetch_enabled=False,
-            mqa_logits_chunk_max_rows=(
-                envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.get()
-            ),
+            mqa_logits_chunk_max_rows=mqa_logits_chunk_max_rows,
+            mqa_logits_chunk_max_bytes=mqa_logits_chunk_max_bytes,
+            mqa_logits_budget_bytes=mqa_logits_budget_bytes,
+            cp_size=self.attn_cp_size,
         )
 
     def maybe_smoke_check_cp_shared_kv_prefill_buffer(self) -> None: