chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Align MQA logits admission with CP split runtime shape

Browse Source 
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>
This commit is contained in:
laoyao0822
2026-06-11 03:51:20 +08:00
co-authored by OmX
parent 250fab291d
commit d696039092
6 changed files with 319 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+3
View File
@@ -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)