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:])
)