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

Bound NSA MQA logits peak memory

Browse Source 
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>
This commit is contained in:
laoyao0822
2026-06-11 03:03:01 +08:00
co-authored by OmX
parent cc908dd556
commit 4e49751406
2 changed files with 266 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+5
View File
@@ -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)