Reuse IPC descriptors across CP shared-KV layers
CP shared-KV slot remaps already have forward-batch lifetime, but the IPC materialize path rebuilt owner/source/dense descriptor tensors on every layer. Cache prefix/current IPC descriptors on the token and paged slot-remap objects, keyed by layout, spans, device, descriptor kind, and prefix capacity, so all model layers can reuse the same request/batch-plan descriptors. Constraint: Small-extend cache-hit workloads showed descriptor setup could exceed the all-reduce baseline before any IPC kernel work ran. Rejected: Global descriptor cache | slot-remap lifetime is safer and avoids stale entries across request/batch-plan changes. Rejected: Cache without physical page capacity | prefix descriptors encode capacity-invalid pages and must miss when capacity changes. Confidence: high Scope-risk: moderate Directive: Do not reuse descriptors across different slot_logical_pages identity, CP layout, spans, device, or prefix capacity; stale descriptors can alias dense slots across requests. Tested: Local py_compile; local git diff --check; remote g0034 cjy-glm5-new targeted descriptor tests 2 passed; remote full test_cp_shared_kv_runtime.py 146 passed, 21 warnings, 2 subtests passed. Not-tested: Full ETE throughput/accuracy after descriptor cache; CUDA service benchmark still needed to quantify speedup. (cherry picked from commit addd1ca1571e41458315d15304a0e841682fe8fa)
This commit is contained in:
@@ -3832,6 +3832,92 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
],
|
||||
)
|
||||
|
||||
def test_ipc_prefix_descriptors_are_cached_on_token_slot_remap(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0)
|
||||
kv_cache = torch.zeros((64, 1, 1), dtype=torch.float32)
|
||||
remap = runtime.build_shared_token_kv_slot_remap(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=None,
|
||||
remap_logical_pages=torch.tensor([[1, 2, 3, 4]], dtype=torch.int64),
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
first = runtime._get_or_build_prefix_ipc_slot_descriptors(
|
||||
slot_remap=remap,
|
||||
layout=layout,
|
||||
spans=[(0, 4)],
|
||||
device=torch.device("cpu"),
|
||||
physical_page_capacity=16,
|
||||
cache_kind="token",
|
||||
)
|
||||
second = runtime._get_or_build_prefix_ipc_slot_descriptors(
|
||||
slot_remap=remap,
|
||||
layout=layout,
|
||||
spans=[(0, 4)],
|
||||
device=torch.device("cpu"),
|
||||
physical_page_capacity=16,
|
||||
cache_kind="token",
|
||||
)
|
||||
capacity_miss = runtime._get_or_build_prefix_ipc_slot_descriptors(
|
||||
slot_remap=remap,
|
||||
layout=layout,
|
||||
spans=[(0, 4)],
|
||||
device=torch.device("cpu"),
|
||||
physical_page_capacity=2,
|
||||
cache_kind="token",
|
||||
)
|
||||
|
||||
self.assertIs(first.slot_indices, second.slot_indices)
|
||||
self.assertIs(first.owner_ranks, second.owner_ranks)
|
||||
self.assertIs(first.src_page_indices, second.src_page_indices)
|
||||
self.assertIs(first.dense_page_indices, second.dense_page_indices)
|
||||
self.assertIsNot(first.owner_ranks, capacity_miss.owner_ranks)
|
||||
self.assertEqual(first.owner_ranks.tolist(), [0, 1, 0, 1])
|
||||
self.assertEqual(first.src_page_indices.tolist(), [1, 1, 2, 2])
|
||||
self.assertEqual(first.dense_page_indices.tolist(), [1, 2, 3, 4])
|
||||
|
||||
def test_ipc_current_descriptors_are_cached_on_token_slot_remap(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0)
|
||||
kv_cache = torch.zeros((64, 1, 1), dtype=torch.float32)
|
||||
remap = runtime.build_shared_token_kv_slot_remap(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=None,
|
||||
remap_logical_pages=torch.tensor([[1, 2, 3, 4]], dtype=torch.int64),
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
first = runtime._get_or_build_current_ipc_slot_descriptors(
|
||||
slot_remap=remap,
|
||||
layout=layout,
|
||||
spans=[(2, 4)],
|
||||
device=torch.device("cpu"),
|
||||
cache_kind="token",
|
||||
)
|
||||
second = runtime._get_or_build_current_ipc_slot_descriptors(
|
||||
slot_remap=remap,
|
||||
layout=layout,
|
||||
spans=[(2, 4)],
|
||||
device=torch.device("cpu"),
|
||||
cache_kind="token",
|
||||
)
|
||||
|
||||
self.assertIs(first.slot_indices, second.slot_indices)
|
||||
self.assertIs(first.owner_ranks, second.owner_ranks)
|
||||
self.assertIs(first.compact_src_page_indices, second.compact_src_page_indices)
|
||||
self.assertIs(first.dense_page_indices, second.dense_page_indices)
|
||||
self.assertEqual(first.slot_indices.tolist(), [2, 3])
|
||||
self.assertEqual(first.owner_ranks.tolist(), [0, 1])
|
||||
self.assertEqual(first.compact_src_page_indices.tolist(), [0, 1])
|
||||
self.assertEqual(first.dense_page_indices.tolist(), [3, 4])
|
||||
|
||||
def test_forward_batch_token_slot_remap_is_cached_across_layers(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
Reference in New Issue
Block a user