From 250fab291d7193020715bb1b314a8bd1ee710305 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Thu, 11 Jun 2026 03:25:22 +0800 Subject: [PATCH] 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 --- .../cp_shared_kv_prefill_buffer_estimator.py | 14 ++++++++ python/sglang/srt/managers/schedule_policy.py | 3 +- python/sglang/srt/managers/scheduler.py | 3 ++ ...t_cp_shared_kv_prefill_buffer_estimator.py | 35 +++++++++++++++++++ 4 files changed, 54 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py b/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py index 5a3f92427..93228262a 100644 --- a/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py +++ b/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, diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 99a696b92..2752c6b0e 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 29fa1dc4d..49d42348b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py b/test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py index 6a09378a7..e631867d3 100644 --- a/test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py +++ b/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 = []