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:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user