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/manual/test_cp_shared_kv_compose_v2_8rank.py
+379
View File
@@ -183,6 +183,7 @@ def _build_scenario(rank: int, cp_size: int, device: torch.device):
return dict(
kv_cache=kv_cache,
logical_pages=logical_pages,
logical_locs=logical_locs.reshape(-1),
loc_req_id=loc_req_id,
current_kv=current_kv,
@@ -197,6 +198,325 @@ def _build_scenario(rank: int, cp_size: int, device: torch.device):
)
def _build_index_scenario(rank: int, cp_size: int, device: torch.device, s):
"""Index page buffer + current K/scale rows over the bs=4 geometry."""
page_size = s["page_size"]
index_page_nbytes = page_size * (128 + 4)
physical_pages = int(s["kv_cache"].shape[0]) // page_size
page_buffer = torch.zeros(
(physical_pages, index_page_nbytes), dtype=torch.uint8, device=device
)
# Owned PREFIX pages get deterministic payloads (mirrors the KV pool).
slot_pages = s["slot_remap"].slot_logical_pages.reshape(-1).cpu()
prefix_slots = set()
for start, end in s["prefix_slot_spans"]:
prefix_slots.update(range(int(start), int(end)))
for slot in sorted(prefix_slots):
logical_page = int(slot_pages[slot])
if logical_page <= 0 or (logical_page - 1) % cp_size != rank:
continue
phys_page = (logical_page - 1) // cp_size + 1
page_buffer[phys_page] = (
(torch.arange(index_page_nbytes, dtype=torch.int64) + logical_page * 37)
.remainder_(251)
.to(torch.uint8)
.to(device)
)
rows = int(s["current_locs"].numel())
current_k = (
(
torch.arange(rows, dtype=torch.int64).view(-1, 1) * 13
+ torch.arange(128, dtype=torch.int64).view(1, -1)
+ rank * 17
)
.remainder_(249)
.to(torch.uint8)
.to(device)
)
current_scale = (
torch.arange(rows, dtype=torch.float32, device=device).view(-1, 1) * 0.25
+ 1.0
+ rank
)
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as rt
slot_remap = rt.build_shared_paged_buffer_slot_remap(
page_buffer,
s["logical_pages"],
s["layout"],
)
return dict(
page_buffer=page_buffer,
current_k=current_k,
current_scale=current_scale,
current_locs=s["current_locs"],
current_req_id=s["current_req_id"],
slot_remap=slot_remap,
)
def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None:
"""bs=1 prefetch-path symm consume vs the legacy sync compose (MLA+index)."""
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as rt
from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import (
CpSharedKVIndexPrefetcher,
CpSharedKVIndexPrefetchHandle,
CpSharedKVMlaPrefetcher,
CpSharedKVMlaPrefetchHandle,
)
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_in_seq_page_compute_owners,
)
page_size = 64
kv_dim = 656
index_page_nbytes = page_size * (128 + 4)
prefix_len, extend_len = 640, 200
prefix_pages = prefix_len // page_size
writers = [
int(o)
for o in build_in_seq_page_compute_owners(
extend_len=extend_len,
extend_prefix_len=prefix_len,
page_size=page_size,
cp_size=cp_size,
)
]
pages = list(range(1, prefix_pages + 1))
next_page = prefix_pages + 1
for owner in writers:
while (next_page - 1) % cp_size != owner:
next_page += 1
pages.append(next_page)
next_page += 1
total_slots = len(pages)
logical_pages = torch.tensor([pages], dtype=torch.int64, device=device)
total_tokens = prefix_len + extend_len
locs = torch.tensor(
[
pages[i // page_size] * page_size + i % page_size
for i in range(total_tokens)
],
dtype=torch.int64,
device=device,
)
physical_pages = (max(pages) - 1) // cp_size + 3
layout = CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=rank)
kv_cache = torch.zeros(
(physical_pages * page_size, 1, kv_dim), dtype=torch.uint8, device=device
)
page_buffer = torch.zeros(
(physical_pages, index_page_nbytes), dtype=torch.uint8, device=device
)
for logical_page in pages[:prefix_pages]:
if (logical_page - 1) % cp_size != rank:
continue
phys = (logical_page - 1) // cp_size + 1
kv_cache[phys * page_size : (phys + 1) * page_size] = (
(torch.arange(page_size * kv_dim, dtype=torch.int64) + logical_page * 53)
.remainder_(251)
.to(torch.uint8)
.view(page_size, 1, kv_dim)
.to(device)
)
page_buffer[phys] = (
(torch.arange(index_page_nbytes, dtype=torch.int64) + logical_page * 29)
.remainder_(251)
.to(torch.uint8)
.to(device)
)
cur_locs_all = [
pages[prefix_pages + i // page_size] * page_size + i % page_size
for i in range(extend_len)
]
mine = [
i for i in range(extend_len) if writers[i // page_size] == rank
]
current_locs = torch.tensor(
[cur_locs_all[i] for i in mine], dtype=torch.int64, device=device
)
rows = len(mine)
current_kv = (
(
torch.arange(rows, dtype=torch.int64).view(-1, 1, 1) * 11
+ torch.arange(kv_dim, dtype=torch.int64).view(1, 1, -1)
+ rank * 19
)
.remainder_(247)
.to(torch.uint8)
.to(device)
)
current_k = (
(
torch.arange(rows, dtype=torch.int64).view(-1, 1) * 23
+ torch.arange(128, dtype=torch.int64).view(1, -1)
+ rank * 7
)
.remainder_(245)
.to(torch.uint8)
.to(device)
)
current_scale = (
torch.arange(rows, dtype=torch.float32, device=device).view(-1, 1) * 0.5
+ 2.0
+ rank
)
zeros_req = torch.zeros(rows, dtype=torch.long, device=device)
loc_req_id = torch.zeros(total_tokens, dtype=torch.long, device=device)
kv_remap = rt.build_shared_token_kv_slot_remap(
kv_cache, locs.view(1, -1), logical_pages, layout, page_size
)
idx_remap = rt.build_shared_paged_buffer_slot_remap(
page_buffer, logical_pages, layout
)
spans = dict(
prefix_slot_spans=[(0, prefix_pages)],
current_slot_spans=[(prefix_pages, total_slots)],
)
# --- references via the legacy sync compose (pure NCCL, no barriers) ---
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
ref_kv, ref_locs = rt.materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
logical_locs=locs,
current_kv_cache=current_kv,
current_locs=current_locs,
slot_remap=kv_remap,
layout=layout,
page_size=page_size,
prefix_pages=0,
loc_req_id=loc_req_id,
current_req_id=zeros_req,
layer_id=60,
**spans,
)
ref_idx, ref_pages = rt.materialize_prefix_and_reuse_current_index_page_slots(
page_buffer=page_buffer,
current_index_k=current_k,
current_index_scale=current_scale,
current_locs=current_locs,
slot_remap=idx_remap,
layout=layout,
page_size=page_size,
index_head_dim=128,
prefix_pages=0,
current_req_id=zeros_req,
layer_id=60,
**spans,
)
torch.cuda.synchronize()
dist.barrier()
# --- prefetchers with a hand-built prefix handle (prefix already
# materialized + reduced, exactly what start_next_layer_prefix yields) ---
mla = CpSharedKVMlaPrefetcher(
layout=layout,
page_size=page_size,
prefix_pages=prefix_pages,
slot_logical_pages=kv_remap.slot_logical_pages,
page_inverse=kv_remap.page_inverse,
dense_num_pages=kv_remap.dense_num_pages,
slot_remap=kv_remap,
symm_writers=writers,
)
dense_kv = kv_cache.new_zeros(
(kv_remap.dense_num_pages * page_size, 1, kv_dim)
)
rt.materialize_local_token_kv_page_slots_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv,
slot_logical_pages=kv_remap.slot_logical_pages,
layout=layout,
page_size=page_size,
start_slot=0,
end_slot=prefix_pages,
)
prefix_rows = rt.slot_range_to_token_slice(page_size, 0, prefix_pages)
dist.all_reduce(dense_kv[prefix_rows])
mla_event = torch.cuda.Event()
mla_event.record()
mla.handles[61] = CpSharedKVMlaPrefetchHandle(
layer_id=61, dense_kv_cache=dense_kv, prefix_rows=prefix_rows,
event=mla_event,
)
idx = CpSharedKVIndexPrefetcher(
layout=layout,
prefix_pages=prefix_pages,
slot_logical_pages=idx_remap.slot_logical_pages,
page_inverse=idx_remap.page_inverse,
dense_num_pages=idx_remap.dense_num_pages,
slot_remap=idx_remap,
symm_writers=writers,
)
dense_pb = page_buffer.new_zeros(
(idx_remap.dense_num_pages, index_page_nbytes)
)
rt.materialize_local_paged_buffer_page_slots_into(
page_buffer=page_buffer,
dense_page_buffer=dense_pb,
slot_logical_pages=idx_remap.slot_logical_pages,
layout=layout,
start_slot=0,
end_slot=prefix_pages,
)
dist.all_reduce(dense_pb[1 : prefix_pages + 1])
idx_event = torch.cuda.Event()
idx_event.record()
idx.handles[61] = CpSharedKVIndexPrefetchHandle(
layer_id=61, dense_page_buffer=dense_pb,
prefix_rows=slice(1, prefix_pages + 1), event=idx_event,
)
torch.cuda.synchronize()
dist.barrier()
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(
True
), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True):
out_idx = idx.consume_prefix_with_current(
layer_id=61,
logical_pages=logical_pages,
current_index_k=current_k,
current_index_scale=current_scale,
current_locs=current_locs,
page_size=page_size,
index_head_dim=128,
current_req_id=zeros_req,
)
out_mla = mla.consume_prefix_with_current(
layer_id=61,
kv_cache=kv_cache,
logical_locs=locs,
current_kv_cache=current_kv,
current_locs=current_locs,
loc_req_id=loc_req_id,
current_req_id=zeros_req,
)
torch.cuda.synchronize()
assert out_idx is not None and out_mla is not None, (
f"rank{rank}: prefetch consume unexpectedly missed"
)
sym_idx, sym_pages = out_idx
sym_kv, sym_locs = out_mla
assert torch.equal(ref_pages, sym_pages), f"rank{rank}: prefetch index pages"
assert torch.equal(ref_locs, sym_locs), f"rank{rank}: prefetch mla locs"
if not torch.equal(ref_idx, sym_idx):
bad = (ref_idx != sym_idx).any(dim=-1).nonzero().reshape(-1)[:8]
raise AssertionError(
f"rank{rank}: prefetch index symm mismatch at pages {bad.cpu().tolist()}"
)
if not torch.equal(ref_kv, sym_kv):
diff = (ref_kv != sym_kv).any(dim=-1).any(dim=-1)
bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
raise AssertionError(
f"rank{rank}: prefetch mla symm mismatch at rows {bad} "
f"(of {int(diff.sum())})"
)
def _compose(s, layer_id: int, *, writers: list[int] | None = None):
return runtime.materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=s["kv_cache"],
@@ -316,6 +636,65 @@ def main() -> None:
"MiB, fill-to-staging + barrier + mega gather engaged; arena on/off)",
flush=True,
)
# ---- Phase 3: index sync compose, legacy vs symm (staging already
# registered by the token phase — the index compose never registers). ----
si = _build_index_scenario(rank, world, device, s)
def _compose_index(layer_id: int, *, writers=None):
return runtime.materialize_prefix_and_reuse_current_index_page_slots(
page_buffer=si["page_buffer"],
current_index_k=si["current_k"],
current_index_scale=si["current_scale"],
current_locs=si["current_locs"],
slot_remap=si["slot_remap"],
layout=s["layout"],
page_size=s["page_size"],
index_head_dim=128,
prefix_pages=0,
current_req_id=si["current_req_id"],
prefix_slot_spans=s["prefix_slot_spans"],
current_slot_spans=s["current_slot_spans"],
layer_id=layer_id,
current_page_writer_ranks=writers,
)
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
ref_idx, ref_dense_pages = _compose_index(40)
torch.cuda.synchronize()
dist.barrier()
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(
True
), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True):
for layer_id in range(41, 45):
symm_idx, symm_dense_pages = _compose_index(
layer_id, writers=s["current_page_writer_ranks"]
)
torch.cuda.synchronize()
assert torch.equal(ref_dense_pages, symm_dense_pages), (
f"rank{rank} layer{layer_id}: index dense_pages mismatch"
)
if not torch.equal(ref_idx, symm_idx):
bad = (
(ref_idx != symm_idx).any(dim=-1).nonzero().reshape(-1)[:8]
)
raise AssertionError(
f"rank{rank} layer{layer_id}: index symm mismatch at pages "
f"{bad.cpu().tolist()}"
)
dist.barrier()
if rank == 0:
print("PASS: index sync symm compose byte-identical (4 layers)", flush=True)
# ---- Phase 4+5: prefetch-path symm consume (MLA + index), bs=1. ----
_check_prefetch_symm(rank, world, device)
dist.barrier()
if rank == 0:
print(
"PASS: prefetch consume_prefix_with_current symm (MLA + index) "
"byte-identical to the sync compose",
flush=True,
)
dist.destroy_process_group()
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),
)
)