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