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

Shrink the symm compose region to a compact current-page staging

Browse Source 
Peers only ever read the CURRENT pages of a rank's compose output — the
prefix comes straight from the IPC-registered KV pool — so the symm
region does not need to hold the whole dense buffer (pool-bound ~2.5 GB
double-buffered slab).  It now holds one round of current pages in
merged-span order (extend-cap-bound, ~58-100 MB), and dense buffers
become purely rank-local (plain allocations or the optional local arena;
COMPOSE_SYMM no longer requires COMPOSE_ARENA).

Exchange per compose call: publish my written current pages
dense[page] -> staging[slot i] (slot = the page's batch current index,
identical on every rank, so peers address each other's staging with no
per-batch handshake), cp_symm_barrier, gather peers'
staging[writer][slot] -> dense[page] via the existing src!=dst page
gather.  Reuse safety keeps the parity-half distance-2 argument, now on
the staging.  Capacity sizing comes from the admission caps
(max_total_extend_tokens / max_batch_requests) with a pool-derived
fallback and the SYMM_HEAP_MB override; overflow fails fast
(batch-logical, hence rank-uniform).

Idea credit: laoyao0822's touched-pages-proportional staging
(906ecbe5d4), rebound onto our barrier-gated, group-agreed transport.

Validated on g0033 8xH200: 151 unit tests; 8-rank GPU byte-exactness
vs compose_v2 across 8 layers (arena on and off, parity halves
exercised); benchmark path e (real protocol) byte-exact, current-page
exchange 0.196 ms vs 0.354 ms compact-AR isolated.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
leavelet
2026-06-11 21:50:34 +00:00
co-authored by Claude Fable 5
parent f74f0a557f
commit e8e5edb230
4 changed files with 422 additions and 188 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_shared_kv_compose.py
+97 -13
View File
@@ -1,6 +1,6 @@
"""Unit tests for cp_shared_kv_compose (Step A compose plan + arena).
"""Unit tests for cp_shared_kv_compose (compose plan + arena + symm staging).
Registered: CPU CI (no CUDA needed for plan/arena logic).
Registered: CPU CI (no CUDA needed for plan/arena/staging-layout logic).
"""
import unittest
@@ -12,7 +12,9 @@ from types import SimpleNamespace
from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
CpComposeArena,
CpComposeStaging,
acquire_dense_buffer,
compute_staging_capacity_pages,
get_or_build_compose_plan,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -121,10 +123,15 @@ class TestComposePlanSymm(unittest.TestCase):
current_page_writer_ranks=[3, 0, 3],
)
# current dense pages are slots 2,3,4 -> dense 3,4,5; writers 3,0,3;
# self rank 3 -> only the page written by rank 0 is remote.
# self rank 3 -> only the page written by rank 0 is remote. The
# staging slot of current page i is i (merged-span order).
self.assertEqual(plan.current_dense_pages.tolist(), [3, 4, 5])
self.assertEqual(plan.remote_current_writer_ranks.tolist(), [0])
self.assertEqual(plan.remote_current_slot_indices.tolist(), [1])
self.assertEqual(plan.remote_current_dense_pages.tolist(), [4])
self.assertEqual(plan.local_current_writer_ranks.tolist(), [3, 3])
self.assertEqual(plan.local_current_slot_indices.tolist(), [0, 2])
self.assertEqual(plan.local_current_dense_pages.tolist(), [3, 5])
def test_writer_count_mismatch_fails_fast(self):
layout = CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=0)
@@ -190,9 +197,10 @@ class TestComposePlanSymm(unittest.TestCase):
page_size=64,
layout=layout,
)
# Symm is independent of the arena: gating must work with ARENA off.
with envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(
True
), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(True):
), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(False):
self.assertIsNotNone(
maybe_build_current_page_writer_ranks(
forward_batch=fb_aligned, **kwargs
@@ -251,15 +259,6 @@ class TestComposeArena(unittest.TestCase):
t0 = arena.acquire(parity=p_third, nbytes=256)
self.assertEqual(t0.data_ptr(), d0.data_ptr()) # back to half A
def test_growth_is_forbidden_after_registration(self):
arena = CpComposeArena(torch.device("cpu"))
parity = arena.begin_round(0, "kv")
arena.acquire(parity=parity, nbytes=64)
arena._registered = True
parity = arena.begin_round(1, "kv")
with self.assertRaises(RuntimeError):
arena.acquire(parity=parity, nbytes=arena._half_bytes + 1)
def test_acquire_dense_buffer_plain_alloc_when_arena_disabled(self):
with envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(False):
buf = acquire_dense_buffer(
@@ -286,5 +285,90 @@ class TestComposeArena(unittest.TestCase):
self.assertTrue(buf.is_contiguous())
class TestComposeStaging(unittest.TestCase):
"""Layout/parity logic of the compact symm staging (no IPC needed: the
region offsets and round parity are pure functions of the registration
parameters and the call sequence)."""
def _staging(self, capacity_pages=4, kv_nbytes=1000, index_nbytes=300):
staging = CpComposeStaging(torch.device("cpu"))
# Bypass the IPC allocation: install the layout exactly as register()
# computes it, backed by a plain CPU slab.
staging.capacity_pages = capacity_pages
staging.cp_size = 8
staging.cp_rank = 0
staging._page_nbytes = {"token_kv": kv_nbytes, "index": index_nbytes}
half = 0
region_in_half = {}
for kind in CpComposeStaging.KINDS:
region_in_half[kind] = half
half += staging._align(capacity_pages * staging._page_nbytes[kind])
for parity in (0, 1):
for kind in CpComposeStaging.KINDS:
staging._region_offsets[(kind, parity)] = (
parity * half + region_in_half[kind]
)
staging._slab = torch.zeros(2 * half, dtype=torch.uint8)
staging.peer_slab_bases = torch.full((8,), 1 << 20, dtype=torch.int64)
return staging, half
def test_regions_are_disjoint_and_parity_halves_do_not_overlap(self):
staging, half = self._staging()
kv0 = staging.buffer("token_kv", 0)
idx0 = staging.buffer("index", 0)
kv1 = staging.buffer("token_kv", 1)
# kv and index regions of one half are adjacent but disjoint
# (index starts at the 256-aligned end of kv).
self.assertEqual(idx0.data_ptr() - kv0.data_ptr(), staging._align(4 * 1000))
# The parity-1 half starts exactly one half past parity 0.
self.assertEqual(kv1.data_ptr() - kv0.data_ptr(), half)
self.assertEqual(kv0.numel(), 4 * 1000)
self.assertEqual(idx0.numel(), 4 * 300)
def test_peer_region_ptrs_offset_matches_local_layout(self):
staging, half = self._staging()
base = staging.peer_slab_bases
self.assertTrue(
torch.equal(staging.peer_region_ptrs("token_kv", 0), base)
)
self.assertTrue(
torch.equal(
staging.peer_region_ptrs("index", 1),
base + half + staging._align(4 * 1000),
)
)
def test_round_parity_alternates_and_shares_round_across_kinds(self):
staging, _ = self._staging()
p0 = staging.begin_round(0, "index")
self.assertEqual(staging.begin_round(0, "token_kv"), p0)
p1 = staging.begin_round(1, "index")
self.assertNotEqual(p1, p0)
# Repeated (id, kind) -> new round (EAGLE draft id reuse).
p_again = staging.begin_round(1, "index")
self.assertNotEqual(p_again, p1)
def test_capacity_pages_from_env_override_caps_and_pool(self):
kwargs = dict(
kv_pool_tokens=64 * 1000,
page_size=64,
cp_size=8,
kv_page_nbytes=41_984,
index_page_nbytes=8_448,
)
with envs.SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.override(101):
pages = compute_staging_capacity_pages(**kwargs)
self.assertEqual(
pages, (101 << 20) // (2 * (41_984 + 8_448))
)
# No caps set in unit tests -> pool-derived fallback.
with envs.SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.override(0):
import sglang.srt.server_args as server_args_module
if getattr(server_args_module, "_global_server_args", None) is None:
pages = compute_staging_capacity_pages(**kwargs)
self.assertEqual(pages, 1000 * 8 + 512)
if __name__ == "__main__":
unittest.main()