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

9
python/sglang/srt/environ.py
View File

@@ -220,6 +220,15 @@ class Envs:
# large bs) but coarser overlap. 1 = per-layer.
SGLANG_CP_SHARED_KV_PER_LAYER_GROUP = EnvInt(8)
SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
# Step A compose restructure for bs>1 partial-current materialize: one IPC
# prefix gather + one compact-current collective per buffer per layer
# (fast path), or one whole-buffer all-reduce (fallback), instead of one
# all-reduce per request span. 0 restores the legacy per-span collectives.
SGLANG_CP_SHARED_KV_COMPOSE_V2 = EnvBool(True)
# Carve compose dense buffers from the tier-S arena (deterministic bump,
# layer-parity halves). Default off in Step A; Step B registers the arena
# as symmetric memory and flips this on.
SGLANG_CP_SHARED_KV_COMPOSE_ARENA = EnvBool(False)
# NSA MQA logits are materialized as fp32 [q, k] buffers inside DeepGEMM.
# Lower values split query rows more aggressively to cap peak temporary memory.
SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2)

293
python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py Normal file
View File

@@ -0,0 +1,293 @@
"""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).
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
one producer — so the legacy one-sum-all-reduce-per-request-span structure
(48 collectives per F-layer at bs=12) is replaced by:
fast path (c): one IPC slot-dense gather for all prefix spans (slots outside
the prefix spans carry -1 sentinels and are zero-filled by the
kernel, which also replaces the dense zero-fill) + ONE
collective over the compact current pages.
fallback (b): local materialize of all prefix spans (no per-span AR) +
current fill + ONE sum-all-reduce over the whole dense buffer.
Exact because every row is still writer-exclusive (prefix rows
are owner-only, current rows are writer-only, the rest are
zero on all ranks) at reduce time.
Measured on g0034 8xH200 (traced 12-request batch, per batch): legacy
per-span 214 ms -> (b) 119 ms -> (c) 84 ms; production trace shows the legacy
path at 419 ms + 49 ms launch gaps, so real savings are larger.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING
import torch
from sglang.srt.environ import envs
if TYPE_CHECKING:
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
CpSharedKVLayout,
)
logger = logging.getLogger(__name__)
def cp_shared_kv_compose_v2_enabled() -> bool:
return bool(envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.get())
def cp_shared_kv_compose_arena_enabled() -> bool:
return bool(envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.get())
# --------------------------------------------------------------------------
# Compose plan: layer-invariant descriptor tensors, built once per batch.
# --------------------------------------------------------------------------
@dataclass(frozen=True)
class ComposePlan:
"""Layer-invariant compose descriptors for one (batch, buffer-kind).
``prefix_owner_ranks`` / ``prefix_src_pages`` cover the FULL slot range:
slots inside a prefix span carry the page owner/source, every other slot
(current pages, invalid pages) carries -1 and is zero-filled by the IPC
slot-dense kernel. ``current_dense_pages`` are dense page ids (slot + 1,
dummy-page convention) of all current slots in merged-span order.
"""
prefix_owner_ranks: torch.Tensor
prefix_src_pages: torch.Tensor
current_dense_pages: torch.Tensor
total_slots: int
num_prefix_slots: int
num_current_pages: int
def _spans_key(spans: list[tuple[int, int]] | None) -> tuple[tuple[int, int], ...]:
if not spans:
return ()
return tuple((int(s), int(e)) for s, e in spans)
def _tensor_identity(t: torch.Tensor) -> tuple[int, int, tuple[int, ...], str, str]:
return (
int(t.untyped_storage().data_ptr()),
int(t.data_ptr()),
tuple(int(d) for d in t.shape),
str(t.dtype),
str(t.device),
)
_PLAN_CACHE: dict[tuple, ComposePlan] = {}
_PLAN_CACHE_MAX = 32
def get_or_build_compose_plan(
*,
slot_logical_pages: torch.Tensor,
layout: "CpSharedKVLayout",
physical_page_capacity: int | None,
prefix_spans: list[tuple[int, int]],
current_spans: list[tuple[int, int]],
kind: str,
) -> ComposePlan:
"""Build (or fetch the per-batch cached) compose plan.
The plan is identical for every layer of a batch: the cache key is the
identity of ``slot_logical_pages`` (per-batch tensor, see
``get_or_build_shared_token_kv_slot_remap``) plus the span lists and
layout parameters.
"""
key = (
_tensor_identity(slot_logical_pages),
int(layout.cp_size),
int(layout.cp_rank),
int(physical_page_capacity or -1),
_spans_key(prefix_spans),
_spans_key(current_spans),
kind,
)
plan = _PLAN_CACHE.get(key)
if plan is not None:
return plan
# Local import to avoid a circular module dependency (runtime imports us).
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
build_cp_shared_kv_ipc_page_descriptors,
)
flat_pages = slot_logical_pages.reshape(-1)
total_slots = int(flat_pages.numel())
device = flat_pages.device
owner_ranks, src_pages = build_cp_shared_kv_ipc_page_descriptors(
flat_pages,
layout,
physical_page_capacity=physical_page_capacity,
)
prefix_mask = torch.zeros(total_slots, dtype=torch.bool, device=device)
num_prefix_slots = 0
for start, end in prefix_spans:
prefix_mask[start:end] = True
num_prefix_slots += int(end) - int(start)
sentinel = torch.full_like(owner_ranks, -1)
prefix_owner_ranks = torch.where(prefix_mask, owner_ranks, sentinel).contiguous()
prefix_src_pages = torch.where(prefix_mask, src_pages, sentinel).contiguous()
if current_spans:
current_dense_pages = torch.cat(
[
torch.arange(start + 1, end + 1, device=device, dtype=torch.long)
for start, end in current_spans
]
).contiguous()
else:
current_dense_pages = torch.empty(0, device=device, dtype=torch.long)
plan = ComposePlan(
prefix_owner_ranks=prefix_owner_ranks,
prefix_src_pages=prefix_src_pages,
current_dense_pages=current_dense_pages,
total_slots=total_slots,
num_prefix_slots=num_prefix_slots,
num_current_pages=int(current_dense_pages.numel()),
)
if len(_PLAN_CACHE) >= _PLAN_CACHE_MAX:
# Plans are per-batch; a bounded FIFO keeps stale batches from piling
# up without adding bookkeeping on the (hot) cache-hit path.
_PLAN_CACHE.pop(next(iter(_PLAN_CACHE)))
_PLAN_CACHE[key] = plan
return plan
def reset_compose_plan_cache() -> None:
_PLAN_CACHE.clear()
# --------------------------------------------------------------------------
# Compose arena: tier-S carve discipline (symm-registration ready).
# --------------------------------------------------------------------------
class CpComposeArena:
"""Transient arena for dense compose buffers with tier-S discipline.
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.
"""
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._parity_layer: list[int | None] = [None, None]
self._high_water = 0
self._registered = False # Step B: set after IPC registration
@staticmethod
def _align(nbytes: int) -> int:
return (nbytes + 255) & ~255
def begin_layer(self, parity: int, layer_token: int) -> None:
"""Reset the parity half when a new layer starts using it."""
parity &= 1
if self._parity_layer[parity] != layer_token:
self._parity_layer[parity] = layer_token
self._offsets[parity] = 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}"
)
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",
new_half / (1 << 20),
2 * new_half / (1 << 20),
)
self._slab = torch.empty(2 * new_half, dtype=torch.uint8, device=self.device)
self._half_bytes = new_half
def acquire(self, *, parity: int, nbytes: int) -> torch.Tensor:
"""Carve ``nbytes`` from the parity half (deterministic bump order)."""
parity &= 1
aligned = self._align(nbytes)
needed = self._offsets[parity] + aligned
self._ensure_capacity(needed)
start = parity * self._half_bytes + self._offsets[parity]
self._offsets[parity] += aligned
self._high_water = max(self._high_water, self._offsets[parity])
return self._slab[start : start + nbytes]
@property
def high_water_bytes(self) -> int:
return self._high_water
_ARENAS: dict[torch.device, CpComposeArena] = {}
def get_compose_arena(device: torch.device) -> CpComposeArena:
arena = _ARENAS.get(device)
if arena is None:
arena = CpComposeArena(device)
_ARENAS[device] = arena
return arena
def acquire_dense_buffer(
*,
device: torch.device,
shape: tuple[int, ...],
dtype: torch.dtype,
layer_id: int | None,
kind: str,
) -> torch.Tensor:
"""Dense-buffer allocation for compose_v2.
With the arena disabled (Step A default) this is a plain ``torch.empty``;
callers must not assume zero contents either way (the IPC gather sentinel-
fills, the fallback path zeroes explicitly).
"""
if not cp_shared_kv_compose_arena_enabled() or layer_id is None:
return torch.empty(shape, dtype=dtype, device=device)
arena = get_compose_arena(device)
parity = int(layer_id) & 1
arena.begin_layer(parity, int(layer_id))
numel = 1
for dim in shape:
numel *= int(dim)
nbytes = numel * dtype.itemsize
flat = arena.acquire(parity=parity, nbytes=nbytes)
return flat.view(dtype).reshape(shape)

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

@@ -9,6 +9,11 @@ from typing import Any
import torch
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
acquire_dense_buffer,
cp_shared_kv_compose_v2_enabled,
get_or_build_compose_plan,
)
from sglang.srt.layers.attention.nsa.utils import (
cp_shared_kv_bs_gt1_timing_start,
get_cp_shared_kv_local_out_cache_loc,
@@ -4327,6 +4332,207 @@ def materialize_local_token_kv_page_slots_into(
dense_range.copy_(torch.where(owned_view, gathered, zero))
def _resolve_partial_current_spans(
*,
current_slot_spans: list[tuple[int, int]] | None,
prefix_slot_span: tuple[int, int] | None,
prefix_slot_spans: list[tuple[int, int]] | None,
prefix_pages: int,
total_slots: int,
) -> list[tuple[int, int]]:
"""Replicate the legacy current-span defaulting (incl. its error)."""
if current_slot_spans is None:
if prefix_slot_span is not None or prefix_slot_spans is not None:
raise ValueError(
"CP shared KV batched current compose requires explicit "
"current_slot_spans to avoid reducing prefix slots twice."
)
current_slot_spans = (
[(int(prefix_pages), total_slots)]
if int(prefix_pages) < total_slots
else []
)
return _merge_slot_spans(current_slot_spans)
def _reduce_current_pages_compact(
dense_buffer: torch.Tensor,
dense_num_pages: int,
current_dense_pages: torch.Tensor,
layout: CpSharedKVLayout,
*,
nvtx_source: str,
layer_id: int | None,
) -> None:
"""One collective over the compact current pages, scattered back in place.
Used when prefix rows must NOT be reduced (after an IPC prefix gather they
hold real values on every rank). Operates on a uint8 byte view: every
byte of a current page is writer-exclusive (current rows are written by
exactly one rank, non-current rows are zero on all ranks), so the byte sum
has no carries and is exact for any element dtype. ``view`` (not
``reshape``) keeps failure loud if the buffer is ever non-viewable —
a silent copy here would drop the scatter-back.
"""
paged = dense_buffer.view(torch.uint8).view(int(dense_num_pages), -1)
compact = paged.index_select(0, current_dense_pages)
_all_reduce_materialized_buffer(
compact,
layout.cp_size,
nvtx_source=nvtx_source,
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
paged.index_copy_(0, current_dense_pages, compact)
def _compose_token_kv_partial_current_v2(
*,
kv_cache: torch.Tensor,
logical_locs: torch.Tensor,
current_kv_cache: torch.Tensor,
current_locs: torch.Tensor,
slot_remap: SharedTokenKVSlotRemap,
layout: CpSharedKVLayout,
page_size: int,
prefix_spans: list[tuple[int, int]],
current_spans: list[tuple[int, int]],
loc_req_id: torch.Tensor,
current_req_id: torch.Tensor,
layer_id: int | None,
nvtx_source: str,
timing_start: object,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Step A compose for the MLA token-KV dense buffer.
Fast path: one IPC slot-dense gather covers ALL prefix spans (slots
outside them are -1 sentinels, zero-filled by the kernel) + one collective
over the compact current pages. Fallback (no peer IPC): local materialize
of all prefix spans (no per-span collectives) + ONE sum-all-reduce over
the whole dense buffer, exact because every row is writer-exclusive at
reduce time.
"""
plan = get_or_build_compose_plan(
slot_logical_pages=slot_remap.slot_logical_pages,
layout=layout,
physical_page_capacity=kv_cache.shape[0] // page_size,
prefix_spans=prefix_spans,
current_spans=current_spans,
kind="token_kv",
)
dense_rows = int(slot_remap.dense_num_pages) * page_size
# Path selection is a capability decision, not error handling: the probe
# (cached per pool tensor) tells us whether peer-IPC is available. After
# a successful probe, a failing gather is a bug — let it raise.
ipc_state = (
_get_or_open_tai_ipc_peer_ptrs(kv_cache, layout)
if plan.num_prefix_slots > 0
else None
)
gathered = ipc_state is not None
if gathered:
kernels, peer_ptrs = ipc_state
dense_kv_cache = acquire_dense_buffer(
device=kv_cache.device,
shape=(dense_rows, *kv_cache.shape[1:]),
dtype=kv_cache.dtype,
layer_id=layer_id,
kind="token_kv",
)
kernels.materialize_cuda_ipc_peer_pages_slot_dense(
peer_ptrs,
dense_kv_cache,
plan.prefix_owner_ranks,
plan.prefix_src_pages,
page_nbytes=_token_kv_page_nbytes(kv_cache, page_size),
)
else:
dense_kv_cache = kv_cache.new_zeros((dense_rows, *kv_cache.shape[1:]))
for prefix_start_slot, prefix_end_slot in prefix_spans:
materialize_local_token_kv_page_slots_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv_cache,
slot_logical_pages=slot_remap.slot_logical_pages,
layout=layout,
page_size=page_size,
start_slot=prefix_start_slot,
end_slot=prefix_end_slot,
)
logical_locs = filter_locs_mappable_to_physical_pool(
logical_locs=logical_locs,
layout=layout,
physical_token_capacity=kv_cache.shape[0],
)
dense_locs = remap_logical_locs_to_shared_token_slot_dense_locs(
logical_locs,
slot_remap=slot_remap,
page_size=page_size,
loc_req_id=loc_req_id,
)
mixed_kv_cache, mixed_locs, _ = fill_current_kv_page_slots_and_remap_locs(
dense_kv_cache=dense_kv_cache,
materialized_dense_locs=dense_locs,
current_kv_cache=current_kv_cache,
logical_locs=logical_locs,
current_locs=current_locs,
page_inverse=slot_remap.page_inverse,
page_size=page_size,
current_req_id=current_req_id,
mask_non_current_in_current_pages=True,
)
if gathered:
if plan.num_current_pages > 0:
_reduce_current_pages_compact(
mixed_kv_cache,
int(slot_remap.dense_num_pages),
plan.current_dense_pages,
layout,
nvtx_source=f"{nvtx_source}.v2_current_compact",
layer_id=layer_id,
)
elif prefix_spans:
# Unreduced prefix rows present: one reduce must cover them too.
_all_reduce_materialized_buffer(
mixed_kv_cache,
layout.cp_size,
nvtx_source=f"{nvtx_source}.v2_full",
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
elif plan.num_current_pages > 0:
_reduce_current_pages_compact(
mixed_kv_cache,
int(slot_remap.dense_num_pages),
plan.current_dense_pages,
layout,
nvtx_source=f"{nvtx_source}.v2_current_compact",
layer_id=layer_id,
)
log_cp_shared_kv_bs_gt1_timing(
"mla_partial_current_compose_v2",
timing_start,
"cp_rank=%s layer=%s cp_size=%s total_slots=%s dense_rows=%s "
"prefix_span_pages=%s current_pages=%s gathered_by_ipc=%s kv_dtype=%s",
layout.cp_rank,
layer_id,
layout.cp_size,
plan.total_slots,
int(mixed_kv_cache.shape[0]),
plan.num_prefix_slots,
plan.num_current_pages,
gathered,
kv_cache.dtype,
)
return mixed_kv_cache, mixed_locs
def materialize_prefix_and_reuse_current_kv_page_slots(
*,
kv_cache: torch.Tensor,
@@ -4415,6 +4621,30 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
else []
)
if cp_shared_kv_compose_v2_enabled() and layout.cp_size > 1:
return _compose_token_kv_partial_current_v2(
kv_cache=kv_cache,
logical_locs=logical_locs,
current_kv_cache=current_kv_cache,
current_locs=current_locs,
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_spans=prefix_spans,
current_spans=_resolve_partial_current_spans(
current_slot_spans=current_slot_spans,
prefix_slot_span=prefix_slot_span,
prefix_slot_spans=prefix_slot_spans,
prefix_pages=prefix_pages,
total_slots=total_slots,
),
loc_req_id=loc_req_id,
current_req_id=current_req_id,
layer_id=layer_id,
nvtx_source=nvtx_source,
timing_start=timing_start,
)
dense_kv_cache = kv_cache.new_zeros(
(slot_remap.dense_num_pages * page_size, *kv_cache.shape[1:])
)
@@ -4534,6 +4764,127 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
return mixed_kv_cache, mixed_locs
def _compose_index_partial_current_v2(
*,
page_buffer: torch.Tensor,
current_index_k: torch.Tensor,
current_index_scale: torch.Tensor,
current_locs: torch.Tensor,
slot_remap: SharedPagedBufferSlotRemap,
layout: CpSharedKVLayout,
page_size: int,
index_head_dim: int,
prefix_spans: list[tuple[int, int]],
current_spans: list[tuple[int, int]],
current_req_id: torch.Tensor,
layer_id: int | None,
nvtx_source: str,
timing_start: object,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Step A compose for the indexer page buffer (see token-KV variant)."""
plan = get_or_build_compose_plan(
slot_logical_pages=slot_remap.slot_logical_pages,
layout=layout,
physical_page_capacity=page_buffer.shape[0],
prefix_spans=prefix_spans,
current_spans=current_spans,
kind="index",
)
dense_num_pages = int(slot_remap.dense_num_pages)
ipc_state = (
_get_or_open_tai_ipc_peer_ptrs(page_buffer, layout)
if plan.num_prefix_slots > 0
else None
)
gathered = ipc_state is not None
if gathered:
kernels, peer_ptrs = ipc_state
dense_page_buffer = acquire_dense_buffer(
device=page_buffer.device,
shape=(dense_num_pages, *page_buffer.shape[1:]),
dtype=page_buffer.dtype,
layer_id=layer_id,
kind="index",
)
kernels.materialize_cuda_ipc_peer_pages_slot_dense(
peer_ptrs,
dense_page_buffer,
plan.prefix_owner_ranks,
plan.prefix_src_pages,
page_nbytes=_page_nbytes_from_page_tensor(page_buffer),
)
else:
dense_page_buffer = page_buffer.new_zeros(
(dense_num_pages, *page_buffer.shape[1:])
)
for prefix_start_slot, prefix_end_slot in prefix_spans:
materialize_local_paged_buffer_page_slots_into(
page_buffer=page_buffer,
dense_page_buffer=dense_page_buffer,
slot_logical_pages=slot_remap.slot_logical_pages,
layout=layout,
start_slot=prefix_start_slot,
end_slot=prefix_end_slot,
)
dense_page_buffer = fill_current_index_page_slots(
dense_page_buffer=dense_page_buffer,
current_index_k=current_index_k,
current_index_scale=current_index_scale,
current_locs=current_locs,
page_inverse=slot_remap.page_inverse,
page_size=page_size,
index_head_dim=index_head_dim,
current_req_id=current_req_id,
)
if gathered:
if plan.num_current_pages > 0:
_reduce_current_pages_compact(
dense_page_buffer,
dense_num_pages,
plan.current_dense_pages,
layout,
nvtx_source=f"{nvtx_source}.v2_current_compact",
layer_id=layer_id,
)
elif prefix_spans:
_all_reduce_materialized_buffer(
dense_page_buffer,
layout.cp_size,
nvtx_source=f"{nvtx_source}.v2_full",
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
elif plan.num_current_pages > 0:
_reduce_current_pages_compact(
dense_page_buffer,
dense_num_pages,
plan.current_dense_pages,
layout,
nvtx_source=f"{nvtx_source}.v2_current_compact",
layer_id=layer_id,
)
log_cp_shared_kv_bs_gt1_timing(
"index_partial_current_compose_v2",
timing_start,
"cp_rank=%s layer=%s cp_size=%s total_slots=%s dense_pages=%s "
"prefix_span_pages=%s current_pages=%s gathered_by_ipc=%s",
layout.cp_rank,
layer_id,
layout.cp_size,
plan.total_slots,
dense_num_pages,
plan.num_prefix_slots,
plan.num_current_pages,
gathered,
)
return dense_page_buffer, slot_remap.dense_pages
def materialize_prefix_and_reuse_current_index_page_slots(
*,
page_buffer: torch.Tensor,
@@ -4608,6 +4959,30 @@ def materialize_prefix_and_reuse_current_index_page_slots(
else []
)
if cp_shared_kv_compose_v2_enabled() and layout.cp_size > 1:
return _compose_index_partial_current_v2(
page_buffer=page_buffer,
current_index_k=current_index_k,
current_index_scale=current_index_scale,
current_locs=current_locs,
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
index_head_dim=index_head_dim,
prefix_spans=prefix_spans,
current_spans=_resolve_partial_current_spans(
current_slot_spans=current_slot_spans,
prefix_slot_span=prefix_slot_span,
prefix_slot_spans=prefix_slot_spans,
prefix_pages=prefix_pages,
total_slots=total_slots,
),
current_req_id=current_req_id,
layer_id=layer_id,
nvtx_source=nvtx_source,
timing_start=timing_start,
)
dense_page_buffer = page_buffer.new_zeros(
(slot_remap.dense_num_pages, *page_buffer.shape[1:])
)

270
test/manual/test_cp_shared_kv_compose_v2_8rank.py Normal file
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()

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

244
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
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