Bound CP MQA logits buffers with row chunking

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>
2026-06-11 03:15:52 +08:00
parent e0ea8a485c
commit ddc1233955
3 changed files with 195 additions and 68 deletions
@@ -5817,6 +5817,43 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
        self.assertIs(dense_pages, fallback_pages)
        logger.warning.assert_not_called()

+    def test_nsa_mqa_logits_chunk_budget_uses_env_fraction(self):
+        from sglang.srt.environ import envs
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        indexer._mqa_logits_budget_bytes = {}
+
+        with envs.SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION.override(0.25), patch(
+            "sglang.srt.layers.attention.nsa.nsa_indexer.get_is_capture_mode",
+            return_value=False,
+        ), patch(
+            "sglang.srt.layers.attention.nsa.nsa_indexer.get_global_server_args",
+            return_value=SimpleNamespace(mem_fraction_static=0.5),
+        ), patch.object(
+            nsa_indexer.torch.cuda, "get_device_properties"
+        ) as props, patch.object(
+            nsa_indexer.torch.cuda, "mem_get_info", return_value=(80_000, 100_000)
+        ):
+            props.return_value = SimpleNamespace(total_memory=100_000)
+            self.assertEqual(indexer._get_mqa_logits_budget_bytes(0), 12_500)
+
+    def test_nsa_mqa_logits_chunk_max_rows_overrides_budget_rows(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):
+            self.assertEqual(
+                indexer._mqa_logits_chunk_max_rows(
+                    num_q=1024,
+                    num_k=4096,
+                    logits_budget_bytes=4096 * 4 * 512,
+                ),
+                128,
+            )
+

 if __name__ == "__main__":
    unittest.main()