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>
2026-06-11 03:51:20 +08:00
parent 250fab291d
commit d696039092
6 changed files with 319 additions and 11 deletions
@@ -116,6 +116,101 @@ def test_estimator_counts_mqa_logits_peak_from_extend_and_context_rows():
    assert chunked.mqa_logits_peak_bytes == 1024 * k_rows * 4


+def test_estimator_counts_mqa_logits_peak_after_cp_in_seq_split():
+    context = CPSharedKVPrefillBufferEstimatorContext(
+        kvcache=_fake_kvcache(),
+        model_config=SimpleNamespace(vocab_size=32),
+        tp_size=1,
+        page_size=64,
+        logprob_chunk_enabled=False,
+        logprob_chunk_size=2048,
+        bs_gt1_l1_prefetch_enabled=False,
+        cp_size=8,
+    )
+    estimate = estimate_cp_shared_kv_prefill_buffer_bytes(
+        page_size=64,
+        batch_size=1,
+        prefix_lens=[160_000],
+        extend_lens=[65_536],
+        context=context,
+    )
+
+    # 65,536 tokens = 1024 pages.  With CP=8 the in-seq split has 16
+    # page-aligned segments, 4096 valid rows each.  A rank owns two zigzag
+    # segments, so local q rows are 8192, not the full 65,536.
+    q_rows = 4096 + 4096
+    # The fused CP MQA path materializes K per owned segment.  For rank 0 this
+    # is prefix+segment0_end plus prefix+segment15_end.
+    k_rows = (160_000 + 4096) + (160_000 + 65_536)
+    assert estimate.mqa_logits_peak_bytes == q_rows * k_rows * 4
+
+    chunked = estimate_cp_shared_kv_prefill_buffer_bytes(
+        page_size=64,
+        batch_size=1,
+        prefix_lens=[160_000],
+        extend_lens=[65_536],
+        context=replace(context, mqa_logits_chunk_max_rows=4096),
+    )
+    assert chunked.mqa_logits_peak_bytes == 4096 * k_rows * 4
+
+
+def test_estimator_derives_mqa_chunk_rows_from_static_budget():
+    prefix_len = 160_000
+    extend_len = 65_536
+    q_rows = 4096 + 4096
+    k_rows = (prefix_len + 4096) + (prefix_len + extend_len)
+    context = CPSharedKVPrefillBufferEstimatorContext(
+        kvcache=_fake_kvcache(),
+        model_config=SimpleNamespace(vocab_size=32),
+        tp_size=1,
+        page_size=64,
+        logprob_chunk_enabled=False,
+        logprob_chunk_size=2048,
+        bs_gt1_l1_prefetch_enabled=False,
+        cp_size=8,
+        mqa_logits_budget_bytes=1024 * k_rows * 4,
+    )
+
+    estimate = estimate_cp_shared_kv_prefill_buffer_bytes(
+        page_size=64,
+        batch_size=1,
+        prefix_lens=[prefix_len],
+        extend_lens=[extend_len],
+        context=context,
+    )
+
+    assert q_rows > 1024
+    assert estimate.mqa_logits_peak_bytes == 1024 * k_rows * 4
+
+
+def test_estimator_caps_mqa_logits_with_explicit_gb_budget():
+    prefix_len = 160_000
+    extend_len = 65_536
+    k_rows = (prefix_len + 4096) + (prefix_len + extend_len)
+    context = CPSharedKVPrefillBufferEstimatorContext(
+        kvcache=_fake_kvcache(),
+        model_config=SimpleNamespace(vocab_size=32),
+        tp_size=1,
+        page_size=64,
+        logprob_chunk_enabled=False,
+        logprob_chunk_size=2048,
+        bs_gt1_l1_prefetch_enabled=False,
+        cp_size=8,
+        mqa_logits_chunk_max_bytes=2_000_000_000,
+    )
+
+    estimate = estimate_cp_shared_kv_prefill_buffer_bytes(
+        page_size=64,
+        batch_size=1,
+        prefix_lens=[prefix_len],
+        extend_lens=[extend_len],
+        context=context,
+    )
+
+    expected_rows = 2_000_000_000 // (k_rows * 4)
+    assert estimate.mqa_logits_peak_bytes == expected_rows * k_rows * 4
+
+
 def test_smoke_check_allocates_and_releases_probe_with_device_module(monkeypatch):
    events = []

@@ -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()