diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 3fbcccf79..4368fff7b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py new file mode 100644 index 000000000..3671aa796 --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py @@ -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) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 1bebb91bb..15b78d3bd 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -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:]) ) diff --git a/test/manual/test_cp_shared_kv_compose_v2_8rank.py b/test/manual/test_cp_shared_kv_compose_v2_8rank.py new file mode 100644 index 000000000..17682dd0d --- /dev/null +++ b/test/manual/test_cp_shared_kv_compose_v2_8rank.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_compose.py b/test/registered/unit/mem_cache/test_cp_shared_kv_compose.py new file mode 100644 index 000000000..3cb6c77f6 --- /dev/null +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_compose.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 937a21c4c..f545d4618 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -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