chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Avoid redundant CP collectives on sync shared-KV materialize

Browse Source 
Synchronous CP shared KV full-hit and partial-current paths can now use the tai-kernel CUDA IPC slot-dense materialize path instead of first copying local owner pages and then running a dense CP all-reduce. This keeps the async prefetch pipeline unchanged while routing the safer synchronous runtime path through descriptor-driven owner-page reads.

The runtime builds descriptors from the existing page-aligned slot contract, caches peer pointer tables for long-lived KV/index buffers, and falls back with explicit warnings if the tai-kernel IPC capability is missing. Unit coverage locks offset handle exchange, descriptor construction, and both MLA/index full and partial-current calls.

Constraint: Async prefetch currently has unresolved scheduling contention with GEMM/MoE and current-layer KV collectives, so this commit intentionally does not wire IPC into prefetch.

Constraint: tai-kernel must provide offset-aware CUDA IPC symbols from af9fb67.

Rejected: Replace prefetch all-reduce in this slice | Nsight shows prefetch timing and communicator contention need a separate scheduling design.

Rejected: Fail-fast on missing tai IPC immediately | remote ETE still needs to validate production deployment capability before removing the warning fallback.

Confidence: medium

Scope-risk: moderate

Directive: Keep prefetch and synchronous materialize decisions separate until prefetch has a low-SM/copy-engine schedule.

Related: tai-kernel af9fb67

Tested: git diff --check; remote runtime/unit evidence recorded in docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md.

Not-tested: Fresh GLM5 ETE after this commit; async prefetch IPC path.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-01 00:00:15 +08:00
co-authored by OmX
parent b4f5c3bc0e
commit 4342de0463
3 changed files with 844 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+221
View File
@@ -885,6 +885,125 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertTrue(torch.equal(mixed_kv[12:14], current_kv))
self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1]])
def test_ipc_page_descriptor_builder_maps_slots_to_owner_physical_pages(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=4, cp_rank=0)
slot_logical_pages = torch.tensor(
[[1, 2, 0, 5], [-1, 8, 9, 17]], dtype=torch.int64
)
owner_ranks, src_page_indices = (
runtime.build_cp_shared_kv_ipc_page_descriptors(
slot_logical_pages,
layout,
physical_page_capacity=3,
)
)
self.assertEqual(owner_ranks.tolist(), [0, 1, -1, 0, -1, 3, -1, -1])
self.assertEqual(src_page_indices.tolist(), [1, 1, -1, 2, -1, 2, -1, -1])
def test_materialize_shared_token_kv_uses_ipc_without_all_reduce(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
page_size = 4
layout = CpSharedKVLayout(page_size=page_size, cp_size=2, cp_rank=0)
kv_cache = torch.zeros((32, 1, 1), dtype=torch.float32)
remap_logical_pages = torch.tensor([[1, 2, 5]], dtype=torch.int64)
slot_remap = runtime.build_shared_token_kv_slot_remap(
kv_cache=kv_cache,
logical_locs=torch.tensor([4, 8, 20], dtype=torch.int64),
remap_logical_pages=remap_logical_pages,
layout=layout,
page_size=page_size,
)
def fake_ipc_into(**kwargs):
self.assertEqual(kwargs["start_slot"], 0)
self.assertEqual(kwargs["end_slot"], 3)
kwargs["dense_kv_cache"][4:16] = torch.arange(
12, dtype=torch.float32
).view(12, 1, 1)
return True
with patch.object(
runtime,
"_try_tai_ipc_materialize_token_kv_page_slots_into",
side_effect=fake_ipc_into,
), patch.object(
runtime,
"_all_reduce_materialized_buffer",
side_effect=AssertionError("IPC path must not all-reduce"),
):
dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
kv_cache=kv_cache,
logical_locs=torch.tensor([4, 8, 20], dtype=torch.int64),
remap_logical_pages=remap_logical_pages,
layout=layout,
page_size=page_size,
slot_remap=slot_remap,
)
self.assertEqual(dense_locs.tolist(), [4, 8, 12])
self.assertEqual(list(dense_kv[4:16].flatten().tolist()), list(range(12)))
def test_materialize_prefix_current_token_kv_uses_ipc_without_all_reduce(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
page_size = 4
layout = CpSharedKVLayout(page_size=page_size, cp_size=2, cp_rank=0)
kv_cache = torch.zeros((32, 1, 1), dtype=torch.float32)
logical_locs = torch.tensor([[4, 8, 20, 21, 22, 23]], dtype=torch.int64)
current_locs = torch.tensor([20, 21], dtype=torch.int64)
current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1)
remap_logical_pages = torch.tensor([[1, 2, 5]], dtype=torch.int64)
slot_remap = runtime.build_shared_token_kv_slot_remap(
kv_cache=kv_cache,
logical_locs=logical_locs,
remap_logical_pages=remap_logical_pages,
layout=layout,
page_size=page_size,
)
def fake_ipc_into(**kwargs):
self.assertEqual(kwargs["start_slot"], 0)
self.assertEqual(kwargs["end_slot"], 2)
kwargs["dense_kv_cache"][4:12] = torch.arange(
10, 18, dtype=torch.float32
).view(8, 1, 1)
return True
with patch.object(
runtime,
"_try_tai_ipc_materialize_token_kv_page_slots_into",
side_effect=fake_ipc_into,
), patch.object(
runtime,
"_all_reduce_materialized_buffer_range",
side_effect=AssertionError("IPC path must not range all-reduce"),
):
mixed_kv, mixed_locs = (
runtime.materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
logical_locs=logical_locs,
current_kv_cache=current_kv,
current_locs=current_locs,
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=2,
)
)
self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1]])
self.assertEqual(mixed_kv[4].item(), 10)
self.assertEqual(mixed_kv[8].item(), 14)
self.assertTrue(torch.equal(mixed_kv[12:14], current_kv))
def test_mla_prefetch_consume_prefix_with_current_skips_suffix_materialize(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -1195,6 +1314,108 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
)
def test_materialize_shared_index_uses_ipc_without_all_reduce(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)
page_buffer = torch.zeros((8, 3), dtype=torch.uint8)
logical_pages = torch.tensor([[1, 2, 5]], dtype=torch.int64)
slot_remap = runtime.build_shared_paged_buffer_slot_remap(
page_buffer,
logical_pages,
layout,
)
def fake_ipc_into(**kwargs):
self.assertEqual(kwargs["start_slot"], 0)
self.assertEqual(kwargs["end_slot"], 3)
kwargs["dense_page_buffer"][1:4] = torch.tensor(
[[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.uint8
)
return True
with patch.object(
runtime,
"_try_tai_ipc_materialize_paged_buffer_page_slots_into",
side_effect=fake_ipc_into,
), patch.object(
runtime,
"_all_reduce_materialized_buffer",
side_effect=AssertionError("IPC path must not all-reduce"),
):
dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer(
page_buffer=page_buffer,
logical_pages=logical_pages,
layout=layout,
slot_remap=slot_remap,
)
self.assertEqual(dense_pages.tolist(), [[1, 2, 3]])
self.assertEqual(
dense_page_buffer[1:4].tolist(),
[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
)
def test_materialize_prefix_current_index_uses_ipc_without_all_reduce(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
page_size = 4
index_head_dim = 4
scale_bytes = 4
page_bytes = page_size * index_head_dim + page_size * scale_bytes
layout = CpSharedKVLayout(page_size=page_size, cp_size=2, cp_rank=0)
page_buffer = torch.zeros((8, page_bytes), dtype=torch.uint8)
logical_pages = torch.tensor([[1, 2]], dtype=torch.int64)
slot_remap = runtime.build_shared_paged_buffer_slot_remap(
page_buffer,
logical_pages,
layout,
)
current_k = torch.tensor(
[[1, 2, 3, 4], [5, 6, 7, 8]],
dtype=torch.uint8,
)
current_scale = torch.tensor([[1.25], [2.5]], dtype=torch.float32)
def fake_ipc_into(**kwargs):
self.assertEqual(kwargs["start_slot"], 0)
self.assertEqual(kwargs["end_slot"], 1)
kwargs["dense_page_buffer"][1] = torch.arange(
page_bytes, dtype=torch.uint8
)
return True
with patch.object(
runtime,
"_try_tai_ipc_materialize_paged_buffer_page_slots_into",
side_effect=fake_ipc_into,
), patch.object(
runtime,
"_all_reduce_materialized_buffer_range",
side_effect=AssertionError("IPC path must not range all-reduce"),
):
dense_page_buffer, dense_pages = (
runtime.materialize_prefix_and_reuse_current_index_page_slots(
page_buffer=page_buffer,
current_index_k=current_k,
current_index_scale=current_scale,
current_locs=torch.tensor([8, 9], dtype=torch.int64),
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
index_head_dim=index_head_dim,
prefix_pages=1,
layer_id=2,
)
)
self.assertEqual(dense_pages.tolist(), [[1, 2]])
self.assertEqual(dense_page_buffer[1].tolist(), list(range(page_bytes)))
self.assertTrue(torch.equal(dense_page_buffer[2, 0:4], current_k[0]))
self.assertTrue(torch.equal(dense_page_buffer[2, 4:8], current_k[1]))
def test_index_prefetch_partial_current_compose_fills_current_page_slots(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch