chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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:
孙远航
2026-06-11 21:50:34 +00:00
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

327
python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py
View File

@@ -1,9 +1,9 @@
"""Step A compose support for bs>1 CP shared-KV partial-current materialize.
This module holds the batch-level compose *plan* (descriptor tensors that are
identical for every layer of a batch) and the transient *arena* (tier-S carve
discipline for dense compose buffers, designed so the Step B symmetric-memory
conversion is a registration flip, not a rewrite).
identical for every layer of a batch), the rank-local dense-buffer *arena*
(Step A carve discipline), and the compact symm *staging* (Step B: the only
IPC-registered region — one round of current pages, double-buffered).
Why this exists (see docs_internal/perf/cp-shared-kv-symm-materialize-design.md):
the dense-buffer materialize is a partitioned gather — every byte has exactly
@@ -52,10 +52,9 @@ def cp_shared_kv_compose_arena_enabled() -> bool:
def cp_shared_kv_compose_symm_enabled() -> bool:
"""Step B: exchange current pages via symm-heap IPC gather (zero NCCL).
Requires the arena (the symm slab IS the arena) — checked at use site.
"""
"""Step B: exchange current pages via the compact symm staging (zero
NCCL): publish my current pages to staging, barrier, gather peers'.
Independent of the arena (dense buffers stay rank-local)."""
return bool(envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.get())
@@ -75,9 +74,16 @@ class ComposePlan:
dummy-page convention) of all current slots in merged-span order.
When the caller provides per-current-page writer (compute-owner) ranks,
``remote_current_*`` hold the gather list for the symm exchange: only the
pages written by OTHER ranks, src page == dst page (peers' dense buffers
are at identical offsets).
the plan splits the current pages by writer for the symm staging exchange.
The staging slot of current page ``i`` is ``i`` (its index in merged-span
order) — identical on every rank, so no per-batch offset handshake:
- ``local_current_*``: pages THIS rank wrote — published from the dense
buffer into staging slots before the barrier. The writer-ranks tensor
is carried (all entries == cp_rank) so the publish kernel can index a
uniform pointer table without a per-layer allocation.
- ``remote_current_*``: pages OTHER ranks wrote — gathered after the
barrier from ``staging[writer][slot]`` into ``dense[page]``.
"""
prefix_owner_ranks: torch.Tensor
@@ -86,7 +92,11 @@ class ComposePlan:
total_slots: int
num_prefix_slots: int
num_current_pages: int
local_current_writer_ranks: torch.Tensor | None = None
local_current_slot_indices: torch.Tensor | None = None
local_current_dense_pages: torch.Tensor | None = None
remote_current_writer_ranks: torch.Tensor | None = None
remote_current_slot_indices: torch.Tensor | None = None
remote_current_dense_pages: torch.Tensor | None = None
@@ -172,7 +182,11 @@ def get_or_build_compose_plan(
else:
current_dense_pages = torch.empty(0, device=device, dtype=torch.long)
local_writer_ranks = None
local_slot_indices = None
local_dense_pages = None
remote_writer_ranks = None
remote_slot_indices = None
remote_dense_pages = None
if current_page_writer_ranks is not None:
num_current = int(current_dense_pages.numel())
@@ -185,8 +199,14 @@ def get_or_build_compose_plan(
writers = torch.tensor(
current_page_writer_ranks, dtype=torch.long, device=device
)
slot_indices = torch.arange(num_current, dtype=torch.long, device=device)
remote = writers != int(layout.cp_rank)
local = ~remote
local_writer_ranks = writers[local].contiguous()
local_slot_indices = slot_indices[local].contiguous()
local_dense_pages = current_dense_pages[local].contiguous()
remote_writer_ranks = writers[remote].contiguous()
remote_slot_indices = slot_indices[remote].contiguous()
remote_dense_pages = current_dense_pages[remote].contiguous()
plan = ComposePlan(
@@ -196,7 +216,11 @@ def get_or_build_compose_plan(
total_slots=total_slots,
num_prefix_slots=num_prefix_slots,
num_current_pages=int(current_dense_pages.numel()),
local_current_writer_ranks=local_writer_ranks,
local_current_slot_indices=local_slot_indices,
local_current_dense_pages=local_dense_pages,
remote_current_writer_ranks=remote_writer_ranks,
remote_current_slot_indices=remote_slot_indices,
remote_current_dense_pages=remote_dense_pages,
)
plans[key] = plan
@@ -208,77 +232,69 @@ def get_or_build_compose_plan(
# --------------------------------------------------------------------------
class CpComposeArena:
"""Transient arena for dense compose buffers with tier-S discipline.
class _RoundParity:
"""Monotonic round-epoch parity shared by the arena and the staging.
Step A behavior: a plain device slab, carved by a bump allocator that is
deterministic given the per-layer acquisition order (fixed ``kind`` order,
same shapes on every rank — shapes derive from the batch-logical slot
layout, which is rank-invariant). Buffers alternate between two halves by
layer parity so a buffer stays untouched for one extra layer.
Step B converts this to true symmetric memory by (1) fixing ``capacity``
at startup from the pool-derived materialize bound, (2) IPC-registering
the slab once and exchanging base pointers, (3) adding the flags page.
The carve discipline is enforced *now* so that conversion does not change
any offsets. Growth (allowed in Step A) is forbidden once registered.
A "round" is one layer-instance's compose calls (index then kv on
F-layers). Parity derives from a monotonic epoch, NOT from layer_id —
EAGLE draft layers reuse decoder layer ids, so raw-id parity would stop
alternating halves. The call sequence is identical on every rank, so
epochs (and any offsets derived from them) stay symmetric.
"""
def __init__(self, device: torch.device) -> None:
self.device = device
self._slab: torch.Tensor | None = None
self._half_bytes = 0
self._offsets = [0, 0]
self._high_water = 0
self._registered = False # Step B: set after IPC registration
# Round/epoch state: a "round" is one layer-instance's compose calls
# (index then kv on F-layers). Parity derives from a monotonic epoch,
# NOT from layer_id — EAGLE draft layers reuse decoder layer ids, so
# raw-id parity would stop alternating halves (and same-id rounds
# would keep appending into one half across forwards).
def __init__(self) -> None:
self._epoch = 0
self._round_layer_id: int | None = None
self._round_kinds: set[str] = set()
# Step B symm state (populated by register_symm):
self.peer_slab_bases: torch.Tensor | None = None # CPU int64 [cp]
self.flag_ptrs: torch.Tensor | None = None # CUDA int64 [cp]
self._flags: torch.Tensor | None = None
self.cp_rank: int = -1
@staticmethod
def _align(nbytes: int) -> int:
return (nbytes + 255) & ~255
def begin_round(self, layer_id: int, kind: str) -> int:
"""Advance to a new round (and the other parity half) when needed.
A new round starts when the layer id changes OR when the same buffer
kind repeats for one layer id (e.g. draft layer 0 followed by the
next forward's target layer 0). Within a round, all kinds carve the
same half at deterministic offsets. The call sequence is identical
on every rank, so epochs (and therefore offsets) stay symmetric.
Returns the parity of the current round.
next forward's target layer 0). Returns the round parity.
"""
if layer_id != self._round_layer_id or kind in self._round_kinds:
self._epoch += 1
self._round_layer_id = layer_id
self._round_kinds = set()
self._offsets[self._epoch & 1] = 0
self._on_new_round()
self._round_kinds.add(kind)
return self._epoch & 1
def _on_new_round(self) -> None:
pass
class CpComposeArena(_RoundParity):
"""Local slab for dense compose buffers (Step A allocation discipline).
A plain device slab, carved by a bump allocator that is deterministic
given the per-layer acquisition order. Buffers alternate between two
halves by round parity so a buffer stays untouched for one extra layer.
Purely rank-local: with the compact staging exchange peers never read a
dense buffer, so the slab is never IPC-registered and may grow freely.
"""
def __init__(self, device: torch.device) -> None:
super().__init__()
self.device = device
self._slab: torch.Tensor | None = None
self._half_bytes = 0
self._offsets = [0, 0]
self._high_water = 0
@staticmethod
def _align(nbytes: int) -> int:
return (nbytes + 255) & ~255
def _on_new_round(self) -> None:
self._offsets[self._epoch & 1] = 0
def _ensure_capacity(self, half_bytes: int) -> None:
if self._slab is not None and half_bytes <= self._half_bytes:
return
if self._registered:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][compose_arena] arena growth requested "
f"after symm registration: need_half={half_bytes} "
f"have_half={self._half_bytes}. Increase "
"SGLANG_CP_SHARED_KV_SYMM_HEAP_MB (total slab MB) or check "
"the pool-derived sizing in compute_symm_half_bytes."
)
new_half = max(half_bytes, self._half_bytes * 2, 64 << 20)
logger.info(
"[CP-Compose-Arena] (re)allocating slab: half=%.1f MiB total=%.1f MiB",
@@ -303,32 +319,65 @@ class CpComposeArena:
def high_water_bytes(self) -> int:
return self._high_water
class CpComposeStaging(_RoundParity):
"""Compact symm staging for the Step B current-page exchange.
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 holds just one round of current pages, not the whole dense buffer
(extend-bound ~100 MB instead of the pool-bound ~2.5 GB dense slab).
Layout (fixed at registration, rank-identical, no per-batch handshake):
two parity halves, each ``[kv region | index region]``, each region
``capacity_pages`` pages of that kind. The staging slot of current page
``i`` is ``i`` (its index in the batch's merged current-span order, the
same on every rank), so ``peer_base + slot * page_nbytes`` addresses any
peer's copy of page ``i`` directly.
Reuse safety is the double-buffer parity argument: my publish into half
H at round R+2 starts only after my round-R+1 barrier completed, which
proves every peer's R+1 barrier kernel ran, which (stream order) proves
every peer's round-R gather of half H finished.
"""
KINDS = ("token_kv", "index")
def __init__(self, device: torch.device) -> None:
super().__init__()
self.device = device
self._slab: torch.Tensor | None = None
self.capacity_pages = 0
self.cp_size = 0
self.cp_rank = -1
self._page_nbytes: dict[str, int] = {}
self._region_offsets: dict[tuple[str, int], int] = {}
self.peer_slab_bases: torch.Tensor | None = None # CPU int64 [cp]
self.flag_ptrs: torch.Tensor | None = None # CUDA int64 [cp]
self._flags: torch.Tensor | None = None
@staticmethod
def _align(nbytes: int) -> int:
return (nbytes + 255) & ~255
@property
def symm_ready(self) -> bool:
return self._registered
def registered(self) -> bool:
return self._slab is not None
def slab_offset_of(self, buffer: torch.Tensor) -> int:
assert self._slab is not None
offset = buffer.data_ptr() - self._slab.data_ptr()
if not (0 <= offset < 2 * self._half_bytes):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][compose_arena] buffer is not from "
f"this arena slab: offset={offset}"
)
return int(offset)
def register_symm(
def register(
self,
*,
cp_group,
cp_rank: int,
cp_size: int,
half_bytes: int,
capacity_pages: int,
kv_page_nbytes: int,
index_page_nbytes: int,
) -> None:
"""Fix capacity, allocate the slab + flags in CUDA-IPC memory, and
"""Allocate the staging slab + barrier flags in CUDA-IPC memory and
exchange handles ONCE across the CP group (collective — every rank
must call this at the same point; the first symm compose of a batch
is such a point)."""
must call this at the same point; the first symm token-KV compose of
a batch is such a point)."""
from tai_kernel.nsa_prefill.ipc import (
allocate_cuda_ipc_buffer,
@@ -336,17 +385,27 @@ class CpComposeArena:
open_cuda_ipc_mem_handles,
)
if self._registered:
if self.registered:
return
if self._slab is not None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][compose_arena] register_symm must "
"run before any non-symm slab allocation (enable "
"SGLANG_CP_SHARED_KV_COMPOSE_SYMM from startup, not mid-run)"
)
half_bytes = self._align(half_bytes)
self.capacity_pages = int(capacity_pages)
self.cp_size = int(cp_size)
self.cp_rank = int(cp_rank)
self._page_nbytes = {
"token_kv": int(kv_page_nbytes),
"index": int(index_page_nbytes),
}
half_bytes = 0
region_in_half = {}
for kind in self.KINDS:
region_in_half[kind] = half_bytes
half_bytes += self._align(self.capacity_pages * self._page_nbytes[kind])
for parity in (0, 1):
for kind in self.KINDS:
self._region_offsets[(kind, parity)] = (
parity * half_bytes + region_in_half[kind]
)
self._slab = allocate_cuda_ipc_buffer(2 * half_bytes, device=self.device)
self._half_bytes = half_bytes
flags = allocate_cuda_ipc_buffer(max(cp_size * 4, 8), device=self.device)
flags.zero_()
self._flags = flags
@@ -369,53 +428,37 @@ class CpComposeArena:
flag_peer_ptrs = _exchange(flags)
torch.cuda.synchronize(self.device)
self.flag_ptrs = flag_peer_ptrs.to(self.device)
self.cp_rank = int(cp_rank)
self._registered = True
logger.info(
"[CP-Compose-Arena] symm slab registered: half=%.1f MiB "
"total=%.1f MiB cp_rank=%s cp_size=%s",
half_bytes / (1 << 20),
"[CP-Compose-Staging] symm staging registered: capacity=%s pages "
"(kv=%s B + index=%s B per page) total=%.1f MiB cp_rank=%s cp_size=%s",
self.capacity_pages,
self._page_nbytes["token_kv"],
self._page_nbytes["index"],
2 * half_bytes / (1 << 20),
cp_rank,
cp_size,
)
def peer_dense_ptrs_for(self, buffer: torch.Tensor) -> torch.Tensor:
"""CPU int64 [cp]: each peer's pointer to ITS copy of ``buffer``.
def page_nbytes(self, kind: str) -> int:
return self._page_nbytes[kind]
Valid because every rank carves identical offsets (deterministic
bump over rank-invariant shapes)."""
def _offset(self, kind: str, parity: int) -> int:
return self._region_offsets[(kind, parity & 1)]
def buffer(self, kind: str, parity: int) -> torch.Tensor:
"""This rank's staging region for ``kind`` at ``parity`` (uint8)."""
assert self._slab is not None
start = self._offset(kind, parity)
return self._slab[start : start + self.capacity_pages * self._page_nbytes[kind]]
def peer_region_ptrs(self, kind: str, parity: int) -> torch.Tensor:
"""CPU int64 [cp]: each peer's pointer to ITS region for this kind/parity."""
assert self.peer_slab_bases is not None
return self.peer_slab_bases + self.slab_offset_of(buffer)
def compute_symm_half_bytes(
*,
kv_pool_tokens: int,
page_size: int,
cp_size: int,
kv_page_nbytes: int,
index_page_nbytes: int,
slack_pages: int = 512,
) -> int:
"""Pool-derived hard bound for one arena half (see design doc §1).
Every logical page a batch can materialize is backed by the shared pool:
logical capacity = per-rank pool pages × cp_size. One half must hold the
dense KV and dense index buffers of one layer plus rounding slack.
"""
explicit_mb = int(envs.SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.get())
if explicit_mb > 0:
return explicit_mb << 19 # MB of TOTAL slab -> half bytes
logical_pages = (int(kv_pool_tokens) // int(page_size)) * int(cp_size)
pages = logical_pages + int(slack_pages)
return pages * (int(kv_page_nbytes) + int(index_page_nbytes))
return self.peer_slab_bases + self._offset(kind, parity)
_ARENAS: dict[torch.device, CpComposeArena] = {}
_STAGINGS: dict[torch.device, CpComposeStaging] = {}
def get_compose_arena(device: torch.device) -> CpComposeArena:
@@ -426,6 +469,54 @@ def get_compose_arena(device: torch.device) -> CpComposeArena:
return arena
def get_compose_staging(device: torch.device) -> CpComposeStaging:
staging = _STAGINGS.get(device)
if staging is None:
staging = CpComposeStaging(device)
_STAGINGS[device] = staging
return staging
def compute_staging_capacity_pages(
*,
kv_pool_tokens: int,
page_size: int,
cp_size: int,
kv_page_nbytes: int,
index_page_nbytes: int,
) -> int:
"""Current-page capacity of the symm staging (one parity half, per kind).
Sized from the admission caps when available: a batch's current pages are
bounded by ceil(max_total_extend_tokens / page) plus one boundary page per
request. Falls back to the pool-derived logical-page bound (always
sufficient, just larger) when the caps are unset, e.g. in tests.
``SGLANG_CP_SHARED_KV_SYMM_HEAP_MB`` (MB of the TOTAL slab) overrides.
"""
explicit_mb = int(envs.SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.get())
if explicit_mb > 0:
per_page = int(kv_page_nbytes) + int(index_page_nbytes)
return max(1, (explicit_mb << 20) // (2 * per_page))
import sglang.srt.server_args as server_args_module
server_args = getattr(server_args_module, "_global_server_args", None)
max_extend_tokens = (
getattr(server_args, "cp_shared_kv_prefill_max_total_extend_tokens", None)
if server_args is not None
else None
)
if max_extend_tokens:
max_requests = (
getattr(server_args, "cp_shared_kv_prefill_max_batch_requests", None)
or 256
)
return -(-int(max_extend_tokens) // int(page_size)) + int(max_requests)
return (int(kv_pool_tokens) // int(page_size)) * int(cp_size) + 512
def acquire_dense_buffer(
*,
device: torch.device,

113
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -11,11 +11,10 @@ import torch
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
acquire_dense_buffer,
compute_symm_half_bytes,
cp_shared_kv_compose_arena_enabled,
compute_staging_capacity_pages,
cp_shared_kv_compose_symm_enabled,
cp_shared_kv_compose_v2_enabled,
get_compose_arena,
get_compose_staging,
get_or_build_compose_plan,
)
from sglang.srt.layers.attention.nsa.utils import (
@@ -4481,9 +4480,7 @@ def maybe_build_current_page_writer_ranks(
layer per call site would sit on the launch-critical path for nothing.
"""
if not (
cp_shared_kv_compose_symm_enabled() and cp_shared_kv_compose_arena_enabled()
):
if not cp_shared_kv_compose_symm_enabled():
return None
# The prefetchers issue their own per-span collectives and never the
# barrier/symm exchange; letting them coexist would make SYMM a silent
@@ -4521,38 +4518,42 @@ def maybe_build_current_page_writer_ranks(
return writers
def _symm_arena_ready_or_register(
def _symm_staging_ready_or_register(
*,
layout: CpSharedKVLayout,
kv_cache: torch.Tensor,
page_size: int,
) -> bool:
"""Register the symm slab on first use (collective; uniform call point).
"""Register the compact symm staging on first use (collective; uniform
call point).
Only the token-KV compose registers (it knows the exact KV page bytes);
the index compose uses symm once registration has happened.
"""
arena = get_compose_arena(kv_cache.device)
if arena.symm_ready:
staging = get_compose_staging(kv_cache.device)
if staging.registered:
return True
cp_group = get_attention_cp_group()
# NSA index page bytes are model constants (index_head_dim 128,
# quant_block 128 -> head + 4 scale bytes per token).
index_page_nbytes = page_size * (128 + 4)
arena.register_symm(
kv_page_nbytes = _token_kv_page_nbytes(kv_cache, page_size)
staging.register(
cp_group=cp_group,
cp_rank=int(layout.cp_rank),
cp_size=int(layout.cp_size),
half_bytes=compute_symm_half_bytes(
capacity_pages=compute_staging_capacity_pages(
kv_pool_tokens=int(kv_cache.shape[0]),
page_size=page_size,
cp_size=int(layout.cp_size),
kv_page_nbytes=_token_kv_page_nbytes(kv_cache, page_size),
kv_page_nbytes=kv_page_nbytes,
index_page_nbytes=index_page_nbytes,
),
kv_page_nbytes=kv_page_nbytes,
index_page_nbytes=index_page_nbytes,
)
return arena.symm_ready
return staging.registered
def _symm_exchange_current_pages(
@@ -4561,16 +4562,25 @@ def _symm_exchange_current_pages(
layout: CpSharedKVLayout,
*,
page_nbytes: int,
kind: str,
layer_id: int,
) -> None:
"""Step B current-page exchange: barrier + IPC gather from peers' symm
dense buffers (identical offsets on every rank). The barrier runs even
with zero remote pages — barrier COUNTS must match across ranks.
"""Step B current-page exchange through the compact symm staging:
publish my written current pages: dense[page] -> staging[slot]
barrier cp_symm_barrier (counting; release/acquire at system scope)
gather peers' pages: peer_staging[writer][slot] -> dense[page]
The staging slot of current page ``i`` is ``i`` (merged-span order),
identical on every rank, so peers address each other's staging without a
per-batch handshake. The barrier runs even with zero local/remote pages
— barrier COUNTS must match across ranks.
INVARIANT: every gate on the path to this call (compose env flags,
page_aligned metadata, symm_ready, the agreed IPC capability, prefetch
absence) MUST be rank-uniform. A per-rank divergence here desyncs the
barrier counting and hangs the CP group — never add a per-rank condition
without routing it through a group agreement first
page_aligned metadata, staging.registered, the agreed IPC capability,
prefetch absence) MUST be rank-uniform. A per-rank divergence here
desyncs the barrier counting and hangs the CP group — never add a
per-rank condition without routing it through a group agreement first
(see _agreed_tai_ipc_peer_ptrs)."""
from tai_kernel.nsa_prefill.ipc import (
@@ -4578,17 +4588,50 @@ def _symm_exchange_current_pages(
gather_cuda_ipc_peer_pages,
)
arena = get_compose_arena(dense_buffer.device)
cp_symm_barrier(arena.flag_ptrs, self_rank=int(layout.cp_rank))
staging = get_compose_staging(dense_buffer.device)
if int(plan.num_current_pages) > staging.capacity_pages:
# Batch-logical quantity -> raises uniformly on every rank.
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][compose_symm] batch current pages "
f"exceed the symm staging capacity: pages={plan.num_current_pages} "
f"capacity={staging.capacity_pages}. Raise "
"SGLANG_CP_SHARED_KV_SYMM_HEAP_MB or lower "
"--cp-shared-kv-prefill-max-total-extend-tokens."
)
if page_nbytes != staging.page_nbytes(kind):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][compose_symm] page bytes diverge from "
f"the registered staging layout: kind={kind} page_nbytes="
f"{page_nbytes} registered={staging.page_nbytes(kind)}"
)
parity = staging.begin_round(int(layer_id), kind)
if (
plan.remote_current_dense_pages is not None
and plan.remote_current_dense_pages.numel() > 0
plan.local_current_slot_indices is not None
and plan.local_current_slot_indices.numel() > 0
):
# Local copy via the gather kernel: a uniform pointer table aimed at
# my own dense buffer turns it into dense[page] -> staging[slot].
self_ptrs = torch.full(
(staging.cp_size,), dense_buffer.data_ptr(), dtype=torch.int64
)
gather_cuda_ipc_peer_pages(
self_ptrs,
staging.buffer(kind, parity),
plan.local_current_writer_ranks,
plan.local_current_dense_pages,
plan.local_current_slot_indices,
page_nbytes=page_nbytes,
)
cp_symm_barrier(staging.flag_ptrs, self_rank=int(layout.cp_rank))
if (
plan.remote_current_slot_indices is not None
and plan.remote_current_slot_indices.numel() > 0
):
gather_cuda_ipc_peer_pages(
arena.peer_dense_ptrs_for(dense_buffer),
staging.peer_region_ptrs(kind, parity),
dense_buffer,
plan.remote_current_writer_ranks,
plan.remote_current_dense_pages,
plan.remote_current_slot_indices,
plan.remote_current_dense_pages,
page_nbytes=page_nbytes,
)
@@ -4681,7 +4724,6 @@ def _compose_token_kv_partial_current_v2(
use_symm = (
cp_shared_kv_compose_symm_enabled()
and cp_shared_kv_compose_arena_enabled()
and current_page_writer_ranks is not None
and layer_id is not None
)
@@ -4709,9 +4751,9 @@ def _compose_token_kv_partial_current_v2(
)
if use_symm:
# Symm requires the prefix IPC capability (same transport) and the
# registered slab; registration is collective and happens here, at
# the uniform first-use point.
use_symm = ipc_state is not None and _symm_arena_ready_or_register(
# registered staging; registration is collective and happens here,
# at the uniform first-use point.
use_symm = ipc_state is not None and _symm_staging_ready_or_register(
layout=layout,
kv_cache=kv_cache,
page_size=page_size,
@@ -4776,6 +4818,8 @@ def _compose_token_kv_partial_current_v2(
plan,
layout,
page_nbytes=_token_kv_page_nbytes(kv_cache, page_size),
kind="token_kv",
layer_id=layer_id,
)
elif plan.num_current_pages > 0:
_reduce_current_pages_compact(
@@ -5076,7 +5120,7 @@ def _compose_index_partial_current_v2(
) -> tuple[torch.Tensor, torch.Tensor]:
"""Step A/B compose for the indexer page buffer (see token-KV variant).
The index compose never registers the symm slab itself (the token-KV
The index compose never registers the symm staging itself (the token-KV
compose does, knowing the exact KV page bytes); it uses symm only once
registration already happened — deterministic across ranks because the
layer order is identical everywhere.
@@ -5084,10 +5128,9 @@ def _compose_index_partial_current_v2(
use_symm = (
cp_shared_kv_compose_symm_enabled()
and cp_shared_kv_compose_arena_enabled()
and current_page_writer_ranks is not None
and layer_id is not None
and get_compose_arena(page_buffer.device).symm_ready
and get_compose_staging(page_buffer.device).registered
)
plan = get_or_build_compose_plan(
slot_remap=slot_remap,
@@ -5157,6 +5200,8 @@ def _compose_index_partial_current_v2(
plan,
layout,
page_nbytes=_page_nbytes_from_page_tensor(page_buffer),
kind="index",
layer_id=layer_id,
)
elif plan.num_current_pages > 0:
_reduce_current_pages_compact(

60
test/manual/test_cp_shared_kv_compose_v2_8rank.py
View File

@@ -6,7 +6,7 @@ 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:
Run inside the g0033 syh-dev-new container:
cd /mnt/beegfs/syh/sglang-stable && \
SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 \
PYTHONPATH=python:/mnt/beegfs/syh/tai-kernel/python \
@@ -267,38 +267,52 @@ def main() -> None:
flush=True,
)
# ---- Step B: symm exchange (arena + barrier + peer gather, zero NCCL
# in the current-page phase). Multiple layers exercise parity halves. ----
# ---- Step B: compact symm staging (publish + barrier + peer gather,
# zero NCCL in the current-page phase). Multiple layers exercise the
# staging parity halves; the arena is exercised in a second pass to
# prove the two knobs are independent. ----
def _check_symm(layer_id: int, tag: str) -> None:
symm_kv, symm_locs = _compose(
s, layer_id=layer_id, writers=s["current_page_writer_ranks"]
)
torch.cuda.synchronize()
assert torch.equal(ref_locs, symm_locs), (
f"rank{rank} layer{layer_id} [{tag}]: symm locs mismatch"
)
if not torch.equal(ref_kv, symm_kv):
diff = (ref_kv != symm_kv).any(dim=-1).any(dim=-1)
bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
raise AssertionError(
f"rank{rank} layer{layer_id} [{tag}]: symm dense kv mismatch "
f"at rows {bad} (of {int(diff.sum())})"
)
with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(
True
), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(
True
), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True):
for layer_id in range(4):
symm_kv, symm_locs = _compose(
s, layer_id=layer_id, writers=s["current_page_writer_ranks"]
)
torch.cuda.synchronize()
assert torch.equal(ref_locs, symm_locs), (
f"rank{rank} layer{layer_id}: symm locs mismatch"
)
if not torch.equal(ref_kv, symm_kv):
diff = (ref_kv != symm_kv).any(dim=-1).any(dim=-1)
bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
raise AssertionError(
f"rank{rank} layer{layer_id}: symm dense kv mismatch at "
f"rows {bad} (of {int(diff.sum())})"
)
_check_symm(layer_id, "no-arena")
with envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(True):
for layer_id in range(4, 8):
_check_symm(layer_id, "arena")
dist.barrier()
from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
get_compose_arena,
get_compose_staging,
)
assert get_compose_arena(device).symm_ready, "symm slab was not registered"
staging = get_compose_staging(device)
assert staging.registered, "symm staging was not registered"
if rank == 0:
staged_mb = (
2
* staging.capacity_pages
* (staging.page_nbytes("token_kv") + staging.page_nbytes("index"))
/ (1 << 20)
)
print(
"PASS: symm compose byte-identical to v2 across 4 layers "
"(arena registered, barrier + peer gather engaged)",
"PASS: compact-staging symm compose byte-identical to v2 across "
f"8 layers (staging {staging.capacity_pages} pages ~{staged_mb:.1f} "
"MiB, publish + barrier + peer gather engaged; arena on and off)",
flush=True,
)
dist.destroy_process_group()

110
test/registered/unit/mem_cache/test_cp_shared_kv_compose.py
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()