Align MQA logits admission with CP split runtime shape

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
@@ -5854,6 +5854,38 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
128,
)
def test_nsa_mqa_logits_chunk_max_gb_caps_rows_by_bytes(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
indexer = object.__new__(Indexer)
with envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.override(
0
), envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.override(1.0):
self.assertEqual(
indexer._mqa_logits_chunk_max_rows(
num_q=8192,
num_k=1_000_000,
logits_budget_bytes=1,
),
250,
)
def test_nsa_mqa_logits_chunk_rows_and_gb_conflict_fail_fast(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
indexer = object.__new__(Indexer)
with envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.override(
128
), envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.override(1.0):
with self.assertRaisesRegex(RuntimeError, "mqa_logits_chunk"):
indexer._mqa_logits_chunk_max_rows(
num_q=8192,
num_k=1_000_000,
logits_budget_bytes=1,
)
if __name__ == "__main__":
unittest.main()