Keep CP batch slot descriptors on IPC-compatible runtime

The syh rebase kept callers that expect reusable batch slot spans, while the restored CUDA IPC runtime lacked the helper and indexer still passed a symm-only writer-rank argument. Restore the batch-scoped span cache and remove the stale symm argument so the production compose path stays on CUDA IPC with exact per-request spans.\n\nConstraint: Production main-stream compose should use CUDA IPC, not symm writer-rank routing.\nRejected: Re-enable symm writer-rank parameters | benchmark showed no main-stream win and callers fail against IPC runtime contracts.\nConfidence: high\nScope-risk: narrow\nDirective: Keep slot-span metadata batch-scoped; do not rebuild Python span descriptors per layer without benchmarking.\nTested: Local py_compile for cp_shared_kv_runtime.py, cp_shared_kv_prefetch.py, nsa_indexer.py, nsa_backend.py.\nTested: Remote cjy-glm5-new pytest test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 157 passed, 2 subtests passed.\nNot-tested: Full ETE prefill/decode runtime after restarting services.
This commit is contained in:
laoyao0822
2026-06-13 01:15:33 +08:00
parent e0c388076e
commit 254d667853
3 changed files with 90 additions and 8 deletions
@@ -2304,6 +2304,40 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
[(0, 2), (4, 5)],
)
def test_get_or_build_batch_slot_spans_caches_exact_request_spans(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
forward_batch = SimpleNamespace()
logical_pages = torch.tensor(
[
[1, 2, 5, 0],
[9, 11, 12, 13],
],
dtype=torch.int64,
)
prefix_spans, current_spans = runtime.get_or_build_batch_slot_spans(
forward_batch,
logical_pages=logical_pages,
prefix_lens_cpu=[8, 4],
extend_lens_cpu=[2, 7],
page_size=4,
want_prefix=True,
)
self.assertEqual(prefix_spans, [(0, 2), (4, 5)])
self.assertEqual(current_spans, [(2, 3), (5, 7)])
cached = runtime.get_or_build_batch_slot_spans(
forward_batch,
logical_pages=logical_pages,
prefix_lens_cpu=[8, 4],
extend_lens_cpu=[2, 7],
page_size=4,
want_prefix=True,
)
self.assertIs(cached[0], prefix_spans)
self.assertIs(cached[1], current_spans)
def test_mla_prefetch_create_batch_uses_exact_prefix_and_current_spans(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch