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

Wire the symm current exchange into the CP shared-KV prefetchers

Browse Source 
The bs=1 MLA/index prefetchers replaced their consume-side trailing-range
NCCL all-reduce with the staging exchange: fill current rows straight into
this round's staging span (token-KV collapses to one cached index_copy;
the index fill kernel is just pointed at the staging page inverse),
cp_symm_barrier, then gather ALL current pages — this rank's own included
— from the stagings into the prefetched dense buffer.  The symm+prefetcher
FAIL_FAST is gone.

Rank-uniformity moves with it: staging registration now also happens in
maybe_create (batch-logical gates, before any per-rank miss can diverge),
because with a prefetcher active the sync compose runs only on per-rank
misses and its lazy collective registration would hang.  A hit/miss
divergence itself stays barrier-safe — both the prefetch consume and the
sync-compose fallback execute exactly one begin_round + barrier per
(layer, kind), and the counting barrier is shape-free (unlike the AR pair
it replaces, which would shape-mismatch).

Found by the new index test phase: the fill/remap kernel family skips
page id 0 as the SGLang dummy page, so a 0-based first staging slot was
never written.  The staging layout now reserves row 0 (slot of current
page i = i + 1) for every kind, matching the convention instead of
depending on per-kernel behavior.

Launch-path cost: per-(kind,parity) peer pointer tables and the
[pool|staging] concatenations are precomputed/cached (identity pinned by
holding the pool-table reference); all prefetch descriptors, staging row
indices, and mixed_locs are built once per batch.

Validated on g0033 8xH200: 151 unit tests; 8-rank byte-exactness for
token sync symm (8 layers), index sync symm (4 layers, new phase), and
MLA + index prefetch consume_prefix_with_current vs the legacy sync
compose.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
leavelet
2026-06-11 22:39:47 +00:00
co-authored by Claude Fable 5
parent d63fbd4d79
commit 1923fcd67e
5 changed files with 828 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_shared_kv_compose.py
+13 -9
View File
@@ -134,11 +134,11 @@ class TestComposePlanSymm(unittest.TestCase):
self.assertEqual(plan.current_dense_pages.tolist(), [3, 4, 5])
self.assertEqual(plan.symm_all_owner_ranks[:2].tolist(), [0, 1])
self.assertEqual(plan.symm_all_owner_ranks[2:].tolist(), [11, 8, 11])
self.assertEqual(plan.symm_all_src_pages[2:].tolist(), [0, 1, 2])
self.assertEqual(plan.symm_all_src_pages[2:].tolist(), [1, 2, 3])
# staging page inverse: current dense pages 3,4,5 -> staging slots
# 0,1,2; everything else (dummy, prefix pages) -> -1.
# 1,2,3 (slot 0 = dummy row); everything else -> -1.
self.assertEqual(
plan.staging_page_inverse.tolist(), [[-1, -1, -1, 0, 1, 2]]
plan.staging_page_inverse.tolist(), [[-1, -1, -1, 1, 2, 3]]
)
def test_writer_count_mismatch_fails_fast(self):
@@ -310,7 +310,9 @@ class TestComposeStaging(unittest.TestCase):
region_in_half = {}
for kind in CpComposeStaging.KINDS:
region_in_half[kind] = half
half += staging._align(capacity_pages * staging._page_nbytes[kind])
half += staging._align(
(capacity_pages + 1) * staging._page_nbytes[kind]
)
for parity in (0, 1):
for kind in CpComposeStaging.KINDS:
staging._region_offsets[(kind, parity)] = (
@@ -318,6 +320,8 @@ class TestComposeStaging(unittest.TestCase):
)
staging._slab = torch.zeros(2 * half, dtype=torch.uint8)
staging.peer_slab_bases = torch.full((8,), 1 << 20, dtype=torch.int64)
for key, offset in staging._region_offsets.items():
staging._region_peer_ptrs[key] = staging.peer_slab_bases + offset
return staging, half
def test_regions_are_disjoint_and_parity_halves_do_not_overlap(self):
@@ -326,12 +330,12 @@ class TestComposeStaging(unittest.TestCase):
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))
# (index starts at the 256-aligned end of kv; +1 = dummy row 0).
self.assertEqual(idx0.data_ptr() - kv0.data_ptr(), staging._align(5 * 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)
self.assertEqual(kv0.numel(), 5 * 1000)
self.assertEqual(idx0.numel(), 5 * 300)
def test_peer_region_ptrs_offset_matches_local_layout(self):
staging, half = self._staging()
@@ -342,7 +346,7 @@ class TestComposeStaging(unittest.TestCase):
self.assertTrue(
torch.equal(
staging.peer_region_ptrs("index", 1),
base + half + staging._align(4 * 1000),
base + half + staging._align(5 * 1000),
)
)