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

Restructure bs>1 CP shared-KV compose to one gather + one collective

Browse Source 
The bs>1 partial-current materialize issued one sum-all-reduce per request
span per buffer per layer (48 collectives per F-layer at bs=12, 2880 per
batch, 419ms + 49ms launch gaps in the production trace), all inline on
the compute stream. The data is a partitioned gather, not a reduction:
every byte has exactly one producer.

compose_v2 (SGLANG_CP_SHARED_KV_COMPOSE_V2, default on) replaces this with:
- fast path: one tai-kernel CUDA-IPC slot-dense gather covering ALL prefix
  spans (full-range descriptors, -1 sentinels zero-fill current slots and
  replace the dense zero-fill) + ONE collective over the compact current
  pages (uint8 byte view; exact because every byte is writer-exclusive).
- fallback (no peer IPC): local materialize of all prefix spans + ONE
  whole-buffer sum-all-reduce (rows are still writer-exclusive pre-reduce).
IPC capability is decided once by the cached peer-pointer probe; after a
successful probe a failing gather raises (no per-call try/except).

cp_shared_kv_compose.py adds the per-batch ComposePlan descriptor cache
(layer-invariant, keyed on the slot_logical_pages identity) and the
CpComposeArena with tier-S carve discipline (deterministic bump, layer-
parity halves; default off) so the Step B symmetric-memory conversion is
a registration flip.

Microbenchmark (g0034 8xH200, traced 12-req batch, per batch): per-span
214ms -> fused AR 119ms -> IPC prefix + compact current 84ms; symm target
61ms. Validation: 143 unit tests incl. v2-contract twins, legacy siblings
and rank-merged simulations under both paths; mem_cache dir 432 passed;
8-rank GPU byte-exactness vs legacy with real NCCL + IPC (test/manual/
test_cp_shared_kv_compose_v2_8rank.py) passed with no fallback markers.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
leavelet
2026-06-11 19:16:28 +00:00
co-authored by Claude Fable 5
parent ecf7998c80
commit f47b739e30
6 changed files with 1357 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/test_cp_shared_kv_compose_v2_8rank.py
+270
View File
@@ -0,0 +1,270 @@
#!/usr/bin/env python3
"""8-rank GPU byte-exactness test: compose legacy vs v2 (Step A).
Runs the REAL materialize_prefix_and_reuse_current_{kv,index}_page_slots on
every rank with real NCCL collectives and (for v2) real tai-kernel CUDA-IPC
gathers, and asserts the composed dense buffers and locs are byte-identical
between the legacy per-span path and compose_v2.
Run inside the g0034 cjy-glm5-new container:
cd /mnt/beegfs/syh/sglang-stable && \
SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 \
PYTHONPATH=python:/mnt/beegfs/syh/tai-kernel/python \
torchrun --nproc-per-node=8 test/manual/test_cp_shared_kv_compose_v2_8rank.py
"""
from __future__ import annotations
import os
import torch
import torch.distributed as dist
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
reset_compose_plan_cache,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
class _CpGroupShim:
"""Minimal stand-in for the attention CP GroupCoordinator."""
def __init__(self, group: dist.ProcessGroup) -> None:
self.device_group = group
self.world_size = dist.get_world_size(group)
self.unique_name = "compose_v2_test"
self.pynccl_comm = None
def all_gather_into_tensor(self, out: torch.Tensor, t: torch.Tensor) -> None:
dist.all_gather_into_tensor(out, t, group=self.device_group)
def all_reduce(self, t: torch.Tensor) -> torch.Tensor:
dist.all_reduce(t, group=self.device_group)
return t
def _build_scenario(rank: int, cp_size: int, device: torch.device):
"""bs=4 with mixed prefix/extend lengths, extends crossing page bounds."""
page_size = 64
batch_size = 4
prefix_lens = [640, 1280, 320, 640]
extend_lens = [95, 130, 64, 200]
kv_dim = 656 # fp8 layout bytes/token
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_in_seq_page_compute_owners,
)
g = torch.Generator().manual_seed(20260612)
logical_rows = []
current_locs_all = []
next_page = 1
prefix_pages_by_req = []
for req_id, (prefix_len, extend_len) in enumerate(zip(prefix_lens, extend_lens)):
prefix_pages = prefix_len // page_size
prefix_pages_by_req.append(prefix_pages)
req_pages = list(range(next_page, next_page + prefix_pages))
next_page += prefix_pages
owners = build_in_seq_page_compute_owners(
extend_len=extend_len,
extend_prefix_len=prefix_len,
page_size=page_size,
cp_size=cp_size,
)
current_pages = []
for owner in owners:
# pick a logical page owned by `owner`: owner = (page-1) % cp_size
candidate = next_page
while (candidate - 1) % cp_size != int(owner):
candidate += 1
current_pages.append(candidate)
next_page = candidate + 1
req_pages.extend(current_pages)
remaining = extend_len
for page_offset, logical_page in enumerate(current_pages):
valid = min(page_size, remaining)
for off in range(valid):
current_locs_all.append(
(req_id, logical_page * page_size + off)
)
remaining -= valid
logical_rows.append(req_pages)
max_pages = max(len(r) for r in logical_rows)
logical_pages = torch.zeros((batch_size, max_pages), dtype=torch.int64)
for req_id, pages in enumerate(logical_rows):
logical_pages[req_id, : len(pages)] = torch.tensor(pages, dtype=torch.int64)
logical_pages = logical_pages.to(device)
# logical locs per request row (token granularity, -1 padded)
max_tokens = max(
p + e for p, e in zip(prefix_lens, extend_lens)
)
logical_locs = torch.full((batch_size, max_tokens), -1, dtype=torch.int64)
loc_req_id_rows = []
for req_id, pages in enumerate(logical_rows):
total = prefix_lens[req_id] + extend_lens[req_id]
locs = []
for token_idx in range(total):
page = pages[token_idx // page_size]
locs.append(page * page_size + token_idx % page_size)
logical_locs[req_id, : len(locs)] = torch.tensor(locs, dtype=torch.int64)
logical_locs = logical_locs.to(device)
physical_pages = (next_page // cp_size + 3)
kv_cache = torch.zeros(
(physical_pages * page_size, 1, kv_dim), dtype=torch.uint8, device=device
)
layout = CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=rank)
# Fill this rank's owned PREFIX pages with deterministic payloads.
for req_id, pages in enumerate(logical_rows):
for slot, logical_page in enumerate(pages[: prefix_pages_by_req[req_id]]):
if (logical_page - 1) % cp_size != rank:
continue
phys_page = (logical_page - 1) // cp_size + 1
payload = (
torch.arange(page_size * kv_dim, dtype=torch.int64)
+ logical_page * 131
).remainder_(251).to(torch.uint8)
kv_cache[
phys_page * page_size : (phys_page + 1) * page_size
] = payload.view(page_size, 1, kv_dim).to(device)
# Current rows owned by this rank (writer = page owner in this scenario).
cur_locs = [
loc for req_id, loc in current_locs_all
if (loc // page_size - 1) % cp_size == rank
]
cur_req = [
req_id for req_id, loc in current_locs_all
if (loc // page_size - 1) % cp_size == rank
]
current_locs = torch.tensor(cur_locs, dtype=torch.int64, device=device)
current_req_id = torch.tensor(cur_req, dtype=torch.int64, device=device)
current_kv = (
(
torch.arange(len(cur_locs), dtype=torch.int64).view(-1, 1, 1) * 7
+ torch.arange(kv_dim, dtype=torch.int64).view(1, 1, -1)
+ rank * 31
)
.remainder_(249)
.to(torch.uint8)
.to(device)
)
prefix_slot_spans = runtime.build_batch_prefix_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=prefix_lens,
page_size=page_size,
)
current_slot_spans = runtime.build_batch_current_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=prefix_lens,
extend_lens_cpu=extend_lens,
page_size=page_size,
)
slot_remap = runtime.build_shared_token_kv_slot_remap(
kv_cache,
logical_locs,
logical_pages,
layout,
page_size,
)
loc_req_id = torch.repeat_interleave(
torch.arange(batch_size, dtype=torch.int64, device=device),
torch.tensor(
[max_tokens] * batch_size, dtype=torch.int64, device=device
),
)
return dict(
kv_cache=kv_cache,
logical_locs=logical_locs.reshape(-1),
loc_req_id=loc_req_id,
current_kv=current_kv,
current_locs=current_locs,
current_req_id=current_req_id,
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
)
def _compose(s, layer_id: int):
return runtime.materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=s["kv_cache"],
logical_locs=s["logical_locs"],
current_kv_cache=s["current_kv"],
current_locs=s["current_locs"],
slot_remap=s["slot_remap"],
layout=s["layout"],
page_size=s["page_size"],
prefix_pages=0,
loc_req_id=s["loc_req_id"],
current_req_id=s["current_req_id"],
prefix_slot_spans=s["prefix_slot_spans"],
current_slot_spans=s["current_slot_spans"],
layer_id=layer_id,
)
def main() -> None:
dist.init_process_group("nccl")
rank = dist.get_rank()
world = dist.get_world_size()
local_rank = int(os.environ.get("LOCAL_RANK", rank))
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
shim = _CpGroupShim(dist.group.WORLD)
runtime.get_attention_cp_group = lambda: shim # type: ignore[assignment]
s = _build_scenario(rank, world, device)
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
ref_kv, ref_locs = _compose(s, layer_id=0)
torch.cuda.synchronize()
dist.barrier()
reset_compose_plan_cache()
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(True):
v2_kv, v2_locs = _compose(s, layer_id=0)
torch.cuda.synchronize()
dist.barrier()
assert torch.equal(ref_locs, v2_locs), (
f"rank{rank}: locs mismatch legacy vs v2"
)
if not torch.equal(ref_kv, v2_kv):
diff = (ref_kv != v2_kv).any(dim=-1).any(dim=-1)
bad_rows = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
raise AssertionError(
f"rank{rank}: dense kv mismatch at rows {bad_rows} "
f"(of {int(diff.sum())} differing rows)"
)
# Cross-rank: every rank must hold the SAME composed buffer.
ref_sum = ref_kv.to(torch.float64).sum()
sums = torch.zeros(world, dtype=torch.float64, device=device)
dist.all_gather_into_tensor(
sums, ref_sum.reshape(1).to(device)
)
assert torch.allclose(sums, sums[0].expand_as(sums)), (
f"rank{rank}: composed buffers differ across ranks: {sums.tolist()}"
)
if rank == 0:
print(
"PASS: legacy vs compose_v2 byte-identical "
f"(dense rows={int(ref_kv.shape[0])}, world={world})",
flush=True,
)
dist.destroy_process_group()
if __name__ == "__main__":
main()
test/registered/unit/mem_cache/test_cp_shared_kv_compose.py
+170
View File
@@ -0,0 +1,170 @@
"""Unit tests for cp_shared_kv_compose (Step A compose plan + arena).
Registered: CPU CI (no CUDA needed for plan/arena logic).
"""
import unittest
import torch
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
CpComposeArena,
acquire_dense_buffer,
get_or_build_compose_plan,
reset_compose_plan_cache,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
class TestComposePlan(unittest.TestCase):
def setUp(self):
reset_compose_plan_cache()
def _layout(self, cp_size=8, cp_rank=0, page_size=64):
return CpSharedKVLayout(
page_size=page_size, cp_size=cp_size, cp_rank=cp_rank
)
def test_sentinels_cover_exactly_the_non_prefix_slots(self):
layout = self._layout()
# 10 slots: pages 1..8 valid, slot 4 invalid (0), slot 9 valid.
slot_logical_pages = torch.tensor(
[1, 2, 3, 4, 0, 5, 6, 7, 8, 9], dtype=torch.int64
)
prefix_spans = [(0, 3), (5, 7)]
current_spans = [(3, 5), (7, 10)]
plan = get_or_build_compose_plan(
slot_logical_pages=slot_logical_pages,
layout=layout,
physical_page_capacity=None,
prefix_spans=prefix_spans,
current_spans=current_spans,
kind="token_kv",
)
in_prefix = set(range(0, 3)) | set(range(5, 7))
for slot in range(10):
owner = int(plan.prefix_owner_ranks[slot].item())
src = int(plan.prefix_src_pages[slot].item())
if slot in in_prefix:
logical = int(slot_logical_pages[slot].item())
self.assertGreaterEqual(owner, 0, f"slot {slot}")
self.assertEqual(
owner,
int(
layout.owner_for_logical_pages(
torch.tensor([logical])
)[0].item()
),
)
self.assertGreaterEqual(src, 0)
else:
self.assertEqual(owner, -1, f"slot {slot}")
self.assertEqual(src, -1, f"slot {slot}")
self.assertEqual(plan.num_prefix_slots, 5)
# Current dense pages follow the dummy-page (+1) convention, in
# merged-span order.
self.assertEqual(
plan.current_dense_pages.tolist(), [4, 5, 8, 9, 10]
)
self.assertEqual(plan.num_current_pages, 5)
def test_invalid_logical_page_is_sentinel_even_inside_prefix_span(self):
layout = self._layout()
slot_logical_pages = torch.tensor([1, 0, 3], dtype=torch.int64)
plan = get_or_build_compose_plan(
slot_logical_pages=slot_logical_pages,
layout=layout,
physical_page_capacity=None,
prefix_spans=[(0, 3)],
current_spans=[],
kind="token_kv",
)
self.assertEqual(int(plan.prefix_owner_ranks[1].item()), -1)
self.assertEqual(plan.num_current_pages, 0)
def test_cache_hits_for_same_batch_and_misses_for_different_spans(self):
layout = self._layout()
slot_logical_pages = torch.tensor([1, 2, 3, 4], dtype=torch.int64)
kwargs = dict(
slot_logical_pages=slot_logical_pages,
layout=layout,
physical_page_capacity=None,
prefix_spans=[(0, 2)],
current_spans=[(2, 4)],
kind="token_kv",
)
plan_a = get_or_build_compose_plan(**kwargs)
plan_b = get_or_build_compose_plan(**kwargs)
self.assertIs(plan_a, plan_b)
plan_c = get_or_build_compose_plan(
**{**kwargs, "prefix_spans": [(0, 1)]}
)
self.assertIsNot(plan_a, plan_c)
plan_d = get_or_build_compose_plan(**{**kwargs, "kind": "index"})
self.assertIsNot(plan_a, plan_d)
class TestComposeArena(unittest.TestCase):
def test_parity_halves_reset_per_layer_and_offsets_are_deterministic(self):
arena = CpComposeArena(torch.device("cpu"))
arena.begin_layer(0, 0)
a0 = arena.acquire(parity=0, nbytes=1000)
a1 = arena.acquire(parity=0, nbytes=500)
# Aligned carve: second buffer starts at align(1000)=1024 past the
# first (offsets are slab-relative; absolute alignment is the
# allocator's business).
self.assertEqual(a1.data_ptr() - a0.data_ptr(), 1024)
arena.begin_layer(1, 1)
b0 = arena.acquire(parity=1, nbytes=1000)
self.assertNotEqual(b0.data_ptr(), a0.data_ptr())
# Layer 2 reuses parity 0 at the SAME offsets (determinism).
arena.begin_layer(0, 2)
c0 = arena.acquire(parity=0, nbytes=1000)
c1 = arena.acquire(parity=0, nbytes=500)
self.assertEqual(c0.data_ptr(), a0.data_ptr())
self.assertEqual(c1.data_ptr(), a1.data_ptr())
def test_growth_is_forbidden_after_registration(self):
arena = CpComposeArena(torch.device("cpu"))
arena.begin_layer(0, 0)
arena.acquire(parity=0, nbytes=64)
arena._registered = True
arena.begin_layer(0, 1)
with self.assertRaises(RuntimeError):
arena.acquire(parity=0, 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(
device=torch.device("cpu"),
shape=(4, 3),
dtype=torch.float32,
layer_id=0,
kind="token_kv",
)
self.assertEqual(tuple(buf.shape), (4, 3))
self.assertEqual(buf.dtype, torch.float32)
def test_acquire_dense_buffer_from_arena_views_dtype_and_shape(self):
with envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(True):
buf = acquire_dense_buffer(
device=torch.device("cpu"),
shape=(8, 2),
dtype=torch.float32,
layer_id=3,
kind="token_kv",
)
self.assertEqual(tuple(buf.shape), (8, 2))
self.assertEqual(buf.dtype, torch.float32)
self.assertTrue(buf.is_contiguous())
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+240 -4
View File
@@ -2130,7 +2130,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
range_calls.append((start_row, end_row, kwargs.get("nvtx_source")))
return buffer
with patch.object(
from sglang.srt.environ import envs
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False), patch.object(
runtime,
"_try_tai_ipc_materialize_token_kv_page_slots_into",
side_effect=fake_ipc_into,
@@ -2158,6 +2160,106 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertTrue(torch.equal(mixed_kv[12:14], current_kv))
self.assertEqual(range_calls, [(12, 16, "mla.partial_current_sync.current")])
def test_materialize_prefix_current_token_kv_compose_v2_ipc_and_compact_reduce(self):
"""compose_v2 contract: ONE full-range sentinel IPC gather for the
prefix + ONE collective over the compact current pages (never a
per-span range reduce, never a whole-buffer reduce)."""
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
reset_compose_plan_cache,
)
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,
)
reset_compose_plan_cache()
captured = {}
test_case = self
class FakeKernels:
@staticmethod
def materialize_cuda_ipc_peer_pages_slot_dense(
peer_ptrs, dense, owner_ranks, src_pages, *, page_nbytes
):
captured["owner_ranks"] = owner_ranks.clone()
captured["src_pages"] = src_pages.clone()
captured["page_nbytes"] = page_nbytes
# Kernel contract: every dst page is either gathered or
# zero-filled (dummy page 0 + sentinel slots).
dense.zero_()
dense[4:12] = torch.arange(10, 18, dtype=torch.float32).view(
8, 1, 1
)
reduce_calls = []
def record_reduce(buffer, cp_size, **kwargs):
reduce_calls.append(
(tuple(buffer.shape), kwargs.get("nvtx_source"))
)
return buffer
def fail_range_reduce(*args, **kwargs):
test_case.fail("compose_v2 must not issue per-span range reduces")
with patch.object(
runtime,
"_get_or_open_tai_ipc_peer_ptrs",
return_value=(FakeKernels, object()),
), patch.object(
runtime,
"_all_reduce_materialized_buffer",
side_effect=record_reduce,
), patch.object(
runtime,
"_all_reduce_materialized_buffer_range",
side_effect=fail_range_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,
)
)
# Descriptors cover the full slot range with -1 sentinels on the
# current slot (slot 2, logical page 5).
self.assertEqual(captured["owner_ranks"].numel(), 3)
self.assertEqual(int(captured["owner_ranks"][2].item()), -1)
self.assertEqual(int(captured["src_pages"][2].item()), -1)
self.assertGreaterEqual(int(captured["owner_ranks"][0].item()), 0)
self.assertGreaterEqual(int(captured["owner_ranks"][1].item()), 0)
# Exactly one collective: the compact current pages (1 page, uint8
# byte view of page_size * 1 * fp32 = 16 bytes).
self.assertEqual(len(reduce_calls), 1)
self.assertEqual(reduce_calls[0][0], (1, 16))
self.assertIn("v2_current_compact", reduce_calls[0][1])
# Composed result matches the legacy contract.
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
@@ -2557,7 +2659,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
range_calls.append((start_row, end_row, kwargs.get("nvtx_source")))
return buffer
with patch.object(
from sglang.srt.environ import envs
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False), patch.object(
runtime,
"_try_tai_ipc_materialize_paged_buffer_page_slots_into",
side_effect=fake_ipc_into,
@@ -2587,6 +2691,101 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertTrue(torch.equal(dense_page_buffer[2, 4:8], current_k[1]))
self.assertEqual(range_calls, [(2, 3, "index.partial_current_sync.current")])
def test_materialize_prefix_current_index_compose_v2_ipc_and_compact_reduce(self):
"""compose_v2 contract for the index buffer (see token-KV twin)."""
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
reset_compose_plan_cache,
)
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)
reset_compose_plan_cache()
captured = {}
test_case = self
class FakeKernels:
@staticmethod
def materialize_cuda_ipc_peer_pages_slot_dense(
peer_ptrs, dense, owner_ranks, src_pages, *, page_nbytes
):
captured["owner_ranks"] = owner_ranks.clone()
captured["src_pages"] = src_pages.clone()
dense.zero_()
dense[1] = torch.arange(dense.shape[1], dtype=torch.uint8)
reduce_calls = []
def record_reduce(buffer, cp_size, **kwargs):
reduce_calls.append(
(tuple(buffer.shape), kwargs.get("nvtx_source"))
)
return buffer
def fail_range_reduce(*args, **kwargs):
test_case.fail("compose_v2 must not issue per-span range reduces")
with patch.object(
runtime,
"_get_or_open_tai_ipc_peer_ptrs",
return_value=(FakeKernels, object()),
), patch.object(
runtime,
"_all_reduce_materialized_buffer",
side_effect=record_reduce,
), patch.object(
runtime,
"_all_reduce_materialized_buffer_range",
side_effect=fail_range_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,
)
)
# Full-range descriptors: slot 0 = prefix (real owner), slot 1 =
# current (-1 sentinel).
self.assertEqual(captured["owner_ranks"].numel(), 2)
self.assertGreaterEqual(int(captured["owner_ranks"][0].item()), 0)
self.assertEqual(int(captured["owner_ranks"][1].item()), -1)
self.assertEqual(int(captured["src_pages"][1].item()), -1)
# One compact collective over the single current page.
self.assertEqual(len(reduce_calls), 1)
self.assertEqual(reduce_calls[0][0], (1, page_bytes))
self.assertIn("v2_current_compact", reduce_calls[0][1])
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
@@ -3926,6 +4125,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
with patch.object(
runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
), patch.object(
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
), patch.object(
runtime, "_try_tai_ipc_materialize_token_kv_page_slots_into",
return_value=False,
@@ -3958,6 +4159,13 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
actual_row = merged_dense[dense_page * page_size + token_offset]
torch.testing.assert_close(actual_row, expected_row, atol=0, rtol=0)
@unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
def test_fp8_mla_persistent_pages_bs5_cache_hit_materialize_legacy(self):
from sglang.srt.environ import envs
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
self.test_fp8_mla_persistent_pages_survive_bs5_cache_hit_materialize()
@unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
def test_fp8_index_fused_store_persistent_pages_survive_bs5_materialize(self):
from sglang.jit_kernel.fused_store_index_cache import (
@@ -4306,7 +4514,11 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
else:
current_k = torch.empty((0, index_head_dim), dtype=torch.uint8)
current_scale = torch.empty((0, 1), dtype=torch.float32)
with patch.object(runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce):
with patch.object(
runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
), patch.object(
runtime, "_all_reduce_materialized_buffer", _identity_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,
@@ -4334,6 +4546,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
merged = torch.stack(rank_outputs, dim=0).sum(dim=0).to(torch.uint8)
torch.testing.assert_close(merged, expected, atol=0, rtol=0)
def test_cp8_index_partial_current_rank_merged_reference_bs5_legacy(self):
from sglang.srt.environ import envs
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
self.test_cp8_index_partial_current_compose_matches_rank_merged_reference_for_bs5()
def test_cp8_kv_partial_current_keeps_remote_current_valid_locs_after_reduce(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
@@ -4366,7 +4584,11 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
current_kv_cache = torch.tensor([[10.0, 11.0], [12.0, 13.0]])
with patch.object(runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce):
with patch.object(
runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
), patch.object(
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
):
_mixed_cache, mixed_locs = runtime.materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
logical_locs=logical_locs,
@@ -4386,6 +4608,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
"remote-rank valid current loc must remain visible after current slot all-reduce",
)
def test_cp8_kv_partial_current_remote_current_valid_locs_legacy(self):
from sglang.srt.environ import envs
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
self.test_cp8_kv_partial_current_keeps_remote_current_valid_locs_after_reduce()
def test_cp8_batch_kv_partial_current_keeps_request_packed_layout(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
@@ -4490,6 +4718,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
with patch.object(
runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
), patch.object(
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
):
mixed_kv, mixed_locs = runtime.materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
@@ -4517,6 +4747,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertTrue(torch.equal(actual, expected))
def test_cp8_batch_kv_partial_current_request_packed_layout_legacy(self):
from sglang.srt.environ import envs
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
self.test_cp8_batch_kv_partial_current_keeps_request_packed_layout()
@unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
def test_tai_batched_index_mqa_prepare_matches_getk_gets_reference_gsm8k_bs5(self):
from types import SimpleNamespace