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

Bound CP MQA logits buffers with row chunking

Browse Source 
CP shared-KV bs>1 can build large fp32 MQA-logits temporaries from
DeepGEMM fp8_mqa_logits. The official SGLang path already chunks normal NSA
MQA logits by query rows behind a cached memory budget; carry the same budget
control into our NSA indexer and extend it to CP-ragged topk paths that use
row-wise topk_indices_offset_override.

This keeps the previous one-time cached memory-budget behavior rather than the
recent current-free-mem per-forward variant that regressed performance. A new
optional max-rows env provides an explicit hard cap for debugging or controlled
ETE runs without adding host syncs.

Constraint: DeepGEMM materializes fp32 [q, k] logits internally, so row chunking is the narrowest way to cap temporary memory
Rejected: Restore the reverted syh current-free-mem implementation | it changed hot-path heuristics and showed poor runtime performance
Rejected: Split by K/context dimension | would change topk semantics and require a different transform contract
Confidence: medium
Scope-risk: moderate
Directive: CP-ragged chunking relies on topk_indices_offset_override being row-addressed; do not route non-ragged CP paths through it without separate validation
Tested: Local py_compile for environ.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py
Tested: Remote g0034 cjy-glm5-new py_compile for environ.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py
Tested: Remote pytest TestCpSharedKVTaiMaterializeIntegration, 17 passed
Not-tested: CUDA ETE high-cache-hit bs>1 workload memory/performance after chunking
Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-11 03:15:52 +08:00
co-authored by OmX
parent e0ea8a485c
commit ddc1233955
3 changed files with 195 additions and 68 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 MQA logits are materialized as fp32 [q, k] buffers inside DeepGEMM.
# Lower values split query rows more aggressively to cap peak temporary memory.
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)
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)