Account for MQA logits in CP batch admission

CP shared-KV bs>1 admission already bounds request count, extend tokens,
cached tokens, and an estimated temporary buffer size. The estimate missed
the fp32 MQA logits temporary, whose peak grows with query rows times
context rows and can dominate high-cache-hit multi-request batches.

Add an MQA logits peak term to the CPU-only estimator and include it in
the layer-forward peak enforced by --cp-shared-kv-prefill-max-buffer-size.
When SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS is set, admission estimates the
post-chunk peak using that row cap; otherwise it remains conservative and
assumes the full extend-row count.

Constraint: Scheduler admission must stay CPU-only and cannot query CUDA free memory.
Rejected: Add a separate scheduler limit for MQA logits | the existing max-buffer-size knob is the right aggregate admission budget.
Rejected: Use SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION in scheduler | that depends on runtime CUDA free memory and would make admission host-sync or stale.
Confidence: medium
Scope-risk: moderate
Directive: Keep the estimator conservative when chunk max rows is unset; do not rely on CUDA free-memory queries in scheduler admission.
Tested: Local py_compile for estimator, scheduler, schedule_policy, and estimator tests.
Tested: Local pytest test_cp_shared_kv_prefill_buffer_estimator.py: 5 passed.
Tested: Remote g0034 cjy-glm5-new py_compile and estimator pytest: 5 passed.
Not-tested: ETE scheduler admission under high-cache-hit bs>1 traffic.

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

python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py

@@ -32,6 +32,7 @@ class CPSharedKVPrefillBufferEstimatorContext:
     logprob_chunk_enabled: bool
     logprob_chunk_size: int
     bs_gt1_l1_prefetch_enabled: bool = False
+    mqa_logits_chunk_max_rows: int = 0
 
 
 @dataclass(frozen=True)
@@ -43,6 +44,7 @@ class CPSharedKVPrefillBufferEstimate:
     materialize_peak_bytes: int
     prefetch_peak_bytes: int
     logits_peak_bytes: int
+    mqa_logits_peak_bytes: int
     remap_peak_bytes: int
     transfer_descriptor_peak_bytes: int
     backup_descriptor_peak_bytes: int
@@ -137,6 +139,7 @@ def estimate_cp_shared_kv_prefill_buffer_bytes(
             logprob_chunk_enabled=False,
             logprob_chunk_size=2048,
             bs_gt1_l1_prefetch_enabled=False,
+            mqa_logits_chunk_max_rows=0,
         )
 
     if page_size <= 0:
@@ -175,6 +178,15 @@ 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_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)
@@ -185,6 +197,7 @@ def estimate_cp_shared_kv_prefill_buffer_bytes(
 
     layer_forward_peak_bytes = (
         materialize_peak_bytes
+        + mqa_logits_peak_bytes
         + remap_peak_bytes
         + prefetch_peak_bytes
         + backup_descriptor_peak_bytes
@@ -211,6 +224,7 @@ def estimate_cp_shared_kv_prefill_buffer_bytes(
         materialize_peak_bytes=materialize_peak_bytes,
         prefetch_peak_bytes=prefetch_peak_bytes,
         logits_peak_bytes=logits_peak_bytes,
+        mqa_logits_peak_bytes=mqa_logits_peak_bytes,
         remap_peak_bytes=remap_peak_bytes,
         transfer_descriptor_peak_bytes=transfer_descriptor_peak_bytes,
         backup_descriptor_peak_bytes=backup_descriptor_peak_bytes,

python/sglang/srt/managers/schedule_policy.py

@@ -698,7 +698,7 @@ class PrefillAdder:
                 "admission_max_buffer_size",
                 "stop rid=%s projected=%s limit=%s batch_size=%s "
                 "layer_forward=%s logits_window=%s load_back_window=%s "
-                "materialize=%s prefetch=%s logits=%s remap=%s "
+                "materialize=%s prefetch=%s logits=%s mqa_logits=%s remap=%s "
                 "transfer_desc=%s backup_desc=%s",
                 rid,
                 estimate.total_peak_bytes,
@@ -710,6 +710,7 @@ class PrefillAdder:
                 estimate.materialize_peak_bytes,
                 estimate.prefetch_peak_bytes,
                 estimate.logits_peak_bytes,
+                estimate.mqa_logits_peak_bytes,
                 estimate.remap_peak_bytes,
                 estimate.transfer_descriptor_peak_bytes,
                 estimate.backup_descriptor_peak_bytes,

python/sglang/srt/managers/scheduler.py

@@ -535,6 +535,9 @@ 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()
+            ),
         )
 
     def maybe_smoke_check_cp_shared_kv_prefill_buffer(self) -> None:

test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py

@@ -1,3 +1,4 @@
+from dataclasses import replace
 from types import SimpleNamespace
 
 import torch
@@ -44,6 +45,7 @@ def test_estimator_uses_stream_aware_peak_instead_of_independent_max():
     assert estimate.prefetch_peak_bytes > 0
     assert estimate.layer_forward_peak_bytes == (
         estimate.materialize_peak_bytes
+        + estimate.mqa_logits_peak_bytes
         + estimate.remap_peak_bytes
         + estimate.prefetch_peak_bytes
         + estimate.backup_descriptor_peak_bytes
@@ -81,6 +83,39 @@ def test_estimator_keeps_bs_gt1_prefetch_zero_until_enabled():
     assert estimate.total_peak_bytes >= estimate.materialize_peak_bytes
 
 
+def test_estimator_counts_mqa_logits_peak_from_extend_and_context_rows():
+    context = CPSharedKVPrefillBufferEstimatorContext(
+        kvcache=_fake_kvcache(),
+        model_config=SimpleNamespace(vocab_size=32),
+        tp_size=1,
+        page_size=64,
+        logprob_chunk_enabled=False,
+        logprob_chunk_size=2048,
+        bs_gt1_l1_prefetch_enabled=False,
+    )
+    estimate = estimate_cp_shared_kv_prefill_buffer_bytes(
+        page_size=64,
+        batch_size=2,
+        prefix_lens=[40384, 8192],
+        extend_lens=[4096, 2048],
+        context=context,
+    )
+
+    q_rows = 4096 + 2048
+    k_rows = 40384 + 4096 + 8192 + 2048
+    assert estimate.mqa_logits_peak_bytes == q_rows * k_rows * 4
+    assert estimate.total_peak_bytes >= estimate.mqa_logits_peak_bytes
+
+    chunked = estimate_cp_shared_kv_prefill_buffer_bytes(
+        page_size=64,
+        batch_size=2,
+        prefix_lens=[40384, 8192],
+        extend_lens=[4096, 2048],
+        context=replace(context, mqa_logits_chunk_max_rows=1024),
+    )
+    assert chunked.mqa_logits_peak_bytes == 1024 * k_rows * 4
+
+
 def test_smoke_check_allocates_and_releases_probe_with_device_module(monkeypatch):
     events = []