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

Wire the symm current exchange into the CP shared-KV prefetchers

Browse Source 
The bs=1 MLA/index prefetchers replaced their consume-side trailing-range
NCCL all-reduce with the staging exchange: fill current rows straight into
this round's staging span (token-KV collapses to one cached index_copy;
the index fill kernel is just pointed at the staging page inverse),
cp_symm_barrier, then gather ALL current pages — this rank's own included
— from the stagings into the prefetched dense buffer.  The symm+prefetcher
FAIL_FAST is gone.

Rank-uniformity moves with it: staging registration now also happens in
maybe_create (batch-logical gates, before any per-rank miss can diverge),
because with a prefetcher active the sync compose runs only on per-rank
misses and its lazy collective registration would hang.  A hit/miss
divergence itself stays barrier-safe — both the prefetch consume and the
sync-compose fallback execute exactly one begin_round + barrier per
(layer, kind), and the counting barrier is shape-free (unlike the AR pair
it replaces, which would shape-mismatch).

Found by the new index test phase: the fill/remap kernel family skips
page id 0 as the SGLang dummy page, so a 0-based first staging slot was
never written.  The staging layout now reserves row 0 (slot of current
page i = i + 1) for every kind, matching the convention instead of
depending on per-kernel behavior.

Launch-path cost: per-(kind,parity) peer pointer tables and the
[pool|staging] concatenations are precomputed/cached (identity pinned by
holding the pool-table reference); all prefetch descriptors, staging row
indices, and mixed_locs are built once per batch.

Validated on g0033 8xH200: 151 unit tests; 8-rank byte-exactness for
token sync symm (8 layers), index sync symm (4 layers, new phase), and
MLA + index prefetch consume_prefix_with_current vs the legacy sync
compose.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-11 22:39:47 +00:00
parent d63fbd4d79
commit 1923fcd67e
5 changed files with 828 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -76,14 +76,14 @@ class ComposePlan:
When the caller provides per-current-page writer (compute-owner) ranks,
the plan additionally carries the symm staging descriptors. The staging
slot of current page ``i`` is ``i`` (its index in merged-span order) —
identical on every rank, so no per-batch offset handshake:
slot of current page ``i`` is ``i + 1`` (merged-span order; row 0 = dummy)
— identical on every rank, so no per-batch offset handshake:
- ``symm_all_owner_ranks`` / ``symm_all_src_pages``: ONE slot-dense
gather descriptor pair covering the whole dense buffer, indexed against
a concatenated ``[pool peer ptrs | staging peer ptrs]`` table — prefix
slots keep the pool owner/page, current slots carry
``cp_size + writer`` / staging slot ``i``, everything else stays -1
``cp_size + writer`` / staging slot ``i + 1``, everything else stays -1
(sentinel zero-fill).
- ``staging_page_inverse``: ``page_inverse`` with current pages remapped
to their staging slot (others -1), so the unchanged fill kernels write
@@ -115,11 +115,13 @@ def _build_staging_page_inverse(
``page_inverse`` maps (request row, logical page) -> dense page id; the
fill kernels derive every current-row write destination from it. In the
staging copy a current dense page becomes its index in the (ascending)
merged-span order — the staging slot — and every other entry becomes -1,
so the unchanged fill kernels write current rows into the compact staging
viewed as a page-major buffer (no dummy page). Current locs only ever
map to current pages, so the -1 entries are never exercised by the fill.
staging copy a current dense page becomes ONE PLUS its index in the
(ascending) merged-span order — the staging slot — and every other entry
becomes -1. The +1 keeps staging row 0 unused: the fill/remap kernel
family treats page id 0 as the SGLang dummy page and silently skips
writes to it, so a 0-based first staging slot would never be filled.
Current locs only ever map to current pages, so the -1 entries are never
exercised by the fill.
"""
if page_inverse.dim() == 1:
@@ -132,7 +134,7 @@ def _build_staging_page_inverse(
current_dense_pages, page_inverse.clamp(min=0).contiguous()
).clamp(max=num_current - 1)
is_current = (page_inverse > 0) & (current_dense_pages[pos] == page_inverse)
return torch.where(is_current, pos, invalid)
return torch.where(is_current, pos + 1, invalid)
def get_or_build_compose_plan(
@@ -230,8 +232,10 @@ def get_or_build_compose_plan(
)
current_slots = current_dense_pages - 1
symm_all_owner_ranks[current_slots] = writers + int(layout.cp_size)
# Staging slot of current page i is i + 1 (row 0 = dummy; see
# _build_staging_page_inverse).
symm_all_src_pages[current_slots] = torch.arange(
num_current, dtype=torch.long, device=device
1, num_current + 1, dtype=torch.long, device=device
)
page_inverse = getattr(slot_remap, "page_inverse", None)
if page_inverse is not None:
@@ -357,10 +361,11 @@ class CpComposeStaging(_RoundParity):
Layout (fixed at registration, rank-identical, no per-batch handshake):
two parity halves, each ``[kv region | index region]``, each region
``capacity_pages`` pages of that kind. The staging slot of current page
``i`` is ``i`` (its index in the batch's merged current-span order, the
same on every rank), so ``peer_base + slot * page_nbytes`` addresses any
peer's copy of page ``i`` directly.
``capacity_pages + 1`` pages of that kind (row 0 unused — the fill/remap
kernel family skips page id 0 as the dummy page). The staging slot of
current page ``i`` is ``i + 1`` (one plus its index in the batch's merged
current-span order, the same on every rank), so ``peer_base + slot *
page_nbytes`` addresses any peer's copy of page ``i`` directly.
Reuse safety is the double-buffer parity argument: my fill into half
H at round R+2 starts only after my round-R+1 barrier completed, which
@@ -382,6 +387,12 @@ class CpComposeStaging(_RoundParity):
self.peer_slab_bases: torch.Tensor | None = None # CPU int64 [cp]
self.flag_ptrs: torch.Tensor | None = None # CUDA int64 [cp]
self._flags: torch.Tensor | None = None
# Launch-path caches: the per-(kind, parity) peer pointer tables and
# the [pool | staging] concatenations are tiny CPU tensors, but
# rebuilding them per layer per kind is avoidable overhead.
self._region_peer_ptrs: dict[tuple[str, int], torch.Tensor] = {}
self._combined_pool_table: torch.Tensor | None = None
self._combined_ptrs: dict[tuple[str, int], torch.Tensor] = {}
@staticmethod
def _align(nbytes: int) -> int:
@@ -425,7 +436,10 @@ class CpComposeStaging(_RoundParity):
region_in_half = {}
for kind in self.KINDS:
region_in_half[kind] = half_bytes
half_bytes += self._align(self.capacity_pages * self._page_nbytes[kind])
# +1: staging row 0 stays unused (dummy-page convention).
half_bytes += self._align(
(self.capacity_pages + 1) * self._page_nbytes[kind]
)
for parity in (0, 1):
for kind in self.KINDS:
self._region_offsets[(kind, parity)] = (
@@ -455,6 +469,10 @@ class CpComposeStaging(_RoundParity):
flag_peer_ptrs = _exchange(flags)
torch.cuda.synchronize(self.device)
self.flag_ptrs = flag_peer_ptrs.to(self.device)
for key, offset in self._region_offsets.items():
self._region_peer_ptrs[key] = (
self.peer_slab_bases + offset
).contiguous()
logger.info(
"[CP-Compose-Staging] symm staging registered: capacity=%s pages "
"(kv=%s B + index=%s B per page) total=%.1f MiB cp_rank=%s cp_size=%s",
@@ -476,12 +494,34 @@ class CpComposeStaging(_RoundParity):
"""This rank's staging region for ``kind`` at ``parity`` (uint8)."""
assert self._slab is not None
start = self._offset(kind, parity)
return self._slab[start : start + self.capacity_pages * self._page_nbytes[kind]]
nbytes = (self.capacity_pages + 1) * self._page_nbytes[kind]
return self._slab[start : start + nbytes]
def peer_region_ptrs(self, kind: str, parity: int) -> torch.Tensor:
"""CPU int64 [cp]: each peer's pointer to ITS region for this kind/parity."""
assert self.peer_slab_bases is not None
return self.peer_slab_bases + self._offset(kind, parity)
return self._region_peer_ptrs[(kind, parity & 1)]
def combined_ptr_table(
self, pool_peer_ptrs: torch.Tensor, kind: str, parity: int
) -> torch.Tensor:
"""CPU int64 [2*cp]: ``[pool peer ptrs | staging peer ptrs]``.
Cached per (kind, parity) against the (long-lived, per-pool) agreed
pointer table — holding the table reference pins its identity, so
the identity check cannot go stale on address reuse.
"""
key = (kind, parity & 1)
if self._combined_pool_table is not pool_peer_ptrs:
self._combined_pool_table = pool_peer_ptrs
self._combined_ptrs = {}
table = self._combined_ptrs.get(key)
if table is None:
table = torch.cat(
[pool_peer_ptrs, self.peer_region_ptrs(kind, parity)]
).contiguous()
self._combined_ptrs[key] = table
return table
_ARENAS: dict[torch.device, CpComposeArena] = {}

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

@@ -7,9 +7,20 @@ from typing import Any, Optional
import torch
from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
cp_shared_kv_compose_symm_enabled,
get_compose_staging,
get_or_build_compose_plan,
)
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
_all_reduce_materialized_buffer_async,
_all_reduce_materialized_buffer_range,
_page_nbytes_from_page_tensor,
_symm_begin_current_staging,
_symm_staging_ready_or_register,
_token_kv_page_nbytes,
build_current_loc_remap,
build_current_page_mask,
cp_shared_kv_debug_enabled,
cp_shared_kv_mla_prefetch_enabled,
cp_shared_kv_mla_prefetch_log,
@@ -26,6 +37,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
get_or_build_shared_token_kv_slot_remap,
materialize_local_paged_buffer_page_slots_into,
materialize_local_token_kv_page_slots_into,
maybe_build_current_page_writer_ranks,
remap_logical_pages_to_slot_dense_pages,
remap_logical_locs_to_slot_dense_locs_optimized,
slot_range_to_page_slice,
@@ -38,6 +50,54 @@ from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
logger = logging.getLogger(__name__)
class _PrefetchSymmState:
"""Layer-invariant symm-exchange descriptors for one prefetch batch.
``num_current_pages`` / ``staging_page_inverse`` make this object a valid
``plan`` argument for ``_symm_begin_current_staging``. The gather covers
ALL current pages (this rank's own come from its own staging — the fill
wrote them there, not into the dense buffer), so unlike the sync mega
gather no prefix descriptors are involved: the prefix was already
materialized by the prefetch stream.
"""
__slots__ = (
"num_current_pages",
"staging_page_inverse",
"writer_ranks",
"staging_slots",
"current_dense_pages",
"page_nbytes",
"staging_current_rows",
"mixed_locs",
"dense_pages",
)
def __init__(self, **kwargs) -> None:
for name in self.__slots__:
setattr(self, name, kwargs.get(name))
def _prefetch_symm_active(prefetcher: Any, device: torch.device) -> bool:
"""Rank-uniform gate for the prefetch-path symm current exchange.
Every condition is batch-logical or group-agreed: ``symm_writers`` comes
from ``maybe_build_current_page_writer_ranks`` (env + batch metadata),
and ``registered`` only flips inside a collective registration. A
prefetch hit/miss divergence across ranks stays barrier-safe because the
sync-compose fallback also runs exactly one begin_round + barrier per
(layer, kind).
"""
return (
prefetcher.symm_writers is not None
and prefetcher.layout.cp_size > 1
and prefetcher.prefix_pages < prefetcher.total_slots
and cp_shared_kv_compose_symm_enabled()
and get_compose_staging(device).registered
)
def _prefetch_log(message: str, *args) -> None:
cp_shared_kv_mla_prefetch_log(message, *args)
@@ -335,6 +395,8 @@ class CpSharedKVMlaPrefetcher:
owned_prefix_pages: int = -1,
owned_total_pages: int = -1,
stream: Optional[torch.cuda.Stream] = None,
slot_remap: Any = None,
symm_writers: Optional[list] = None,
) -> None:
self.layout = layout
self.page_size = page_size
@@ -352,6 +414,9 @@ class CpSharedKVMlaPrefetcher:
self.pending_attention_handle: Optional[CpSharedKVMlaPrefetchHandle] = None
self.disabled = False
self._cpu_timing = _PrefetchCpuTiming()
self.slot_remap = slot_remap
self.symm_writers = symm_writers
self._symm_state: Optional[_PrefetchSymmState] = None
@classmethod
def maybe_create(
@@ -509,6 +574,25 @@ class CpSharedKVMlaPrefetcher:
logger.exception("Failed to initialize CP shared KV MLA prefetcher.")
return None
symm_writers = None
if cp_shared_kv_compose_symm_enabled() and layout.cp_size > 1:
symm_writers = maybe_build_current_page_writer_ranks(
forward_batch=forward_batch,
prefix_lens_cpu=extend_prefix_lens_cpu,
extend_lens_cpu=extend_seq_lens_cpu,
page_size=page_size,
layout=layout,
)
if symm_writers is not None:
# Collective registration at a batch-uniform point: with a
# prefetcher active the sync compose only runs on per-rank
# misses, so its lazy first-compose registration would
# diverge. Must NOT be swallowed — a half-registered group
# is a hang, not a fallback.
_symm_staging_ready_or_register(
layout=layout, kv_cache=kv_cache, page_size=page_size
)
owned_prefix_pages = _debug_owned_pages_count(
layout, remap.slot_logical_pages[:prefix_pages]
)
@@ -556,8 +640,82 @@ class CpSharedKVMlaPrefetcher:
owned_prefix_pages=owned_prefix_pages,
owned_total_pages=owned_total_pages,
stream=prefetch_stream,
slot_remap=remap,
symm_writers=symm_writers,
)
def _get_or_build_symm_state(
self,
*,
kv_cache: torch.Tensor,
logical_locs: torch.Tensor,
current_locs: torch.Tensor,
loc_req_id: torch.Tensor,
current_req_id: torch.Tensor,
) -> _PrefetchSymmState:
state = self._symm_state
if state is not None:
return state
plan = get_or_build_compose_plan(
slot_remap=self.slot_remap,
layout=self.layout,
physical_page_capacity=kv_cache.shape[0] // self.page_size,
prefix_spans=[(0, self.prefix_pages)],
current_spans=[(self.prefix_pages, self.total_slots)],
kind="token_kv",
current_page_writer_ranks=self.symm_writers,
)
num_current = int(plan.num_current_pages)
writer_ranks = (
plan.symm_all_owner_ranks.index_select(0, plan.current_dense_pages - 1)
- int(self.layout.cp_size)
).contiguous()
# Staging slot of current page i is i + 1 (row 0 = dummy page).
staging_slots = torch.arange(
1, num_current + 1, dtype=torch.long, device=writer_ranks.device
)
staging_current_rows = remap_logical_locs_to_slot_dense_locs_optimized(
current_locs.reshape(-1),
page_inverse=plan.staging_page_inverse,
page_size=self.page_size,
loc_req_id=current_req_id.reshape(-1),
).to(torch.long)
# mixed_locs is layer-invariant; the fused fill that used to produce
# it builds masks sized by its target buffer, which is now the
# staging — so compute it once here in logical space instead.
logical_locs = filter_locs_mappable_to_physical_pool(
logical_locs=logical_locs,
layout=self.layout,
physical_token_capacity=kv_cache.shape[0],
)
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
logical_locs,
page_inverse=self.page_inverse,
page_size=self.page_size,
loc_req_id=loc_req_id,
)
current_mask, _ = build_current_loc_remap(logical_locs, current_locs)
current_page_mask = build_current_page_mask(
logical_locs, current_locs, page_size=self.page_size
)
mixed_locs = torch.where(
current_page_mask & (~current_mask),
torch.full_like(dense_locs, -1),
dense_locs,
)
state = _PrefetchSymmState(
num_current_pages=num_current,
staging_page_inverse=plan.staging_page_inverse,
writer_ranks=writer_ranks,
staging_slots=staging_slots,
current_dense_pages=plan.current_dense_pages,
page_nbytes=_token_kv_page_nbytes(kv_cache, self.page_size),
staging_current_rows=staging_current_rows,
mixed_locs=mixed_locs,
)
self._symm_state = state
return state
def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool:
start_layer = int(getattr(token_to_kv_pool, "start_layer", 0))
kv_buffer = getattr(token_to_kv_pool, "kv_buffer", None)
@@ -782,15 +940,81 @@ class CpSharedKVMlaPrefetcher:
dense_kv_cache = handle.dense_kv_cache
remap_cpu = _cpu_timing_start()
if loc_req_id is None:
loc_req_id = torch.zeros_like(logical_locs, dtype=torch.long)
if current_req_id is None:
current_req_id = torch.zeros_like(current_locs, dtype=torch.long)
if _prefetch_symm_active(self, dense_kv_cache.device):
# Symm current exchange: fill current rows straight into this
# round's staging span, barrier, gather ALL current pages from
# the stagings (this rank's own included) into the prefetched
# dense buffer. Zero NCCL; descriptors and mixed_locs are
# layer-invariant and cached on the prefetcher.
from tai_kernel.nsa_prefill.ipc import (
cp_symm_barrier,
gather_cuda_ipc_peer_pages,
)
state = self._get_or_build_symm_state(
kv_cache=kv_cache,
logical_locs=logical_locs,
current_locs=current_locs,
loc_req_id=loc_req_id,
current_req_id=current_req_id,
)
staging = get_compose_staging(dense_kv_cache.device)
parity, staging_span = _symm_begin_current_staging(
staging=staging,
plan=state,
kind="token_kv",
layer_id=layer_id,
page_nbytes=state.page_nbytes,
)
num_rows = int(state.staging_current_rows.numel())
if num_rows > 0:
staging_rows = staging_span.view(kv_cache.dtype).reshape(
(state.num_current_pages + 1) * self.page_size,
*kv_cache.shape[1:],
)
staging_rows.index_copy_(
0,
state.staging_current_rows,
current_kv_cache[:num_rows],
)
cp_symm_barrier(
staging.flag_ptrs, self_rank=int(self.layout.cp_rank)
)
gather_cuda_ipc_peer_pages(
staging.peer_region_ptrs("token_kv", parity),
dense_kv_cache,
state.writer_ranks,
state.staging_slots,
state.current_dense_pages,
page_nbytes=state.page_nbytes,
)
remap_ms = _cpu_timing_ms(remap_cpu)
total_ms = _cpu_timing_ms(consume_cpu)
self._log_layer(
layer_id,
"consume_prefix_current_hit layer=%s prefix_pages=%s "
"dense_rows=%s current_rows=%s symm=1 total_ms=%.3f "
"wait_ms=%.3f remap_ms=%.3f",
layer_id,
self.prefix_pages,
int(dense_kv_cache.shape[0]),
int(current_kv_cache.shape[0]),
total_ms,
wait_ms,
remap_ms,
)
return dense_kv_cache, state.mixed_locs
logical_locs = filter_locs_mappable_to_physical_pool(
logical_locs=logical_locs,
layout=self.layout,
physical_token_capacity=kv_cache.shape[0],
)
if loc_req_id is None:
loc_req_id = torch.zeros_like(logical_locs, dtype=torch.long)
if current_req_id is None:
current_req_id = torch.zeros_like(current_locs, dtype=torch.long)
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
logical_locs,
page_inverse=self.page_inverse,
@@ -1120,6 +1344,8 @@ class CpSharedKVIndexPrefetcher:
owned_prefix_pages: int = -1,
owned_total_pages: int = -1,
stream: Optional[torch.cuda.Stream] = None,
slot_remap: Any = None,
symm_writers: Optional[list] = None,
) -> None:
self.layout = layout
self.prefix_pages = prefix_pages
@@ -1136,6 +1362,9 @@ class CpSharedKVIndexPrefetcher:
self.pending_attention_handle: Optional[CpSharedKVIndexPrefetchHandle] = None
self.disabled = False
self._cpu_timing = _PrefetchCpuTiming()
self.slot_remap = slot_remap
self.symm_writers = symm_writers
self._symm_state: Optional[_PrefetchSymmState] = None
@classmethod
def maybe_create(
@@ -1335,6 +1564,31 @@ class CpSharedKVIndexPrefetcher:
logger.exception("Failed to initialize CP shared KV index prefetcher.")
return None
symm_writers = None
if cp_shared_kv_compose_symm_enabled() and layout.cp_size > 1:
symm_writers = maybe_build_current_page_writer_ranks(
forward_batch=forward_batch,
prefix_lens_cpu=extend_prefix_lens_cpu,
extend_lens_cpu=extend_seq_lens_cpu,
page_size=page_size,
layout=layout,
)
if symm_writers is not None:
# Registration sizing needs the token-KV page bytes, so fetch
# the key buffer; a no-op when the MLA prefetcher (created
# first) already registered. Collective — must not be
# swallowed (see the MLA twin).
_symm_staging_ready_or_register(
layout=layout,
kv_cache=_prefetch_pool_get_key_buffer(
token_to_kv_pool=token_to_kv_pool,
layer_id=first_layer_id,
stream=prefetch_stream,
path="index_symm_register",
),
page_size=page_size,
)
owned_prefix_pages = _debug_owned_pages_count(
layout, remap.slot_logical_pages[:prefix_pages]
)
@@ -1378,8 +1632,55 @@ class CpSharedKVIndexPrefetcher:
owned_prefix_pages=owned_prefix_pages,
owned_total_pages=owned_total_pages,
stream=prefetch_stream,
slot_remap=remap,
symm_writers=symm_writers,
)
def _get_or_build_symm_state(
self,
*,
dense_page_buffer: torch.Tensor,
logical_pages: torch.Tensor,
) -> _PrefetchSymmState:
state = self._symm_state
if state is not None:
return state
plan = get_or_build_compose_plan(
slot_remap=self.slot_remap,
layout=self.layout,
physical_page_capacity=None,
prefix_spans=[(0, self.prefix_pages)],
current_spans=[(self.prefix_pages, self.total_slots)],
kind="index",
current_page_writer_ranks=self.symm_writers,
)
num_current = int(plan.num_current_pages)
writer_ranks = (
plan.symm_all_owner_ranks.index_select(0, plan.current_dense_pages - 1)
- int(self.layout.cp_size)
).contiguous()
# Staging slot of current page i is i + 1 (row 0 = dummy page).
staging_slots = torch.arange(
1, num_current + 1, dtype=torch.long, device=writer_ranks.device
)
# The returned dense-pages remap is layer-invariant too.
dense_pages = remap_logical_pages_to_slot_dense_pages(
logical_pages,
page_inverse=self.page_inverse,
page_req_id=build_page_table_row_req_id(logical_pages),
)
state = _PrefetchSymmState(
num_current_pages=num_current,
staging_page_inverse=plan.staging_page_inverse,
writer_ranks=writer_ranks,
staging_slots=staging_slots,
current_dense_pages=plan.current_dense_pages,
page_nbytes=_page_nbytes_from_page_tensor(dense_page_buffer),
dense_pages=dense_pages,
)
self._symm_state = state
return state
def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool:
start_layer = int(getattr(token_to_kv_pool, "start_layer", 0))
kv_buffer = getattr(token_to_kv_pool, "kv_buffer", None)
@@ -1599,6 +1900,72 @@ class CpSharedKVIndexPrefetcher:
remap_cpu = _cpu_timing_start()
if current_req_id is None:
current_req_id = torch.zeros_like(current_locs, dtype=torch.long)
if _prefetch_symm_active(self, dense_page_buffer.device):
# Symm current exchange (see the MLA twin): fill into the
# staging, barrier, gather all current pages into the prefetched
# dense page buffer. Zero NCCL.
from tai_kernel.nsa_prefill.ipc import (
cp_symm_barrier,
gather_cuda_ipc_peer_pages,
)
state = self._get_or_build_symm_state(
dense_page_buffer=dense_page_buffer,
logical_pages=logical_pages,
)
staging = get_compose_staging(dense_page_buffer.device)
parity, staging_span = _symm_begin_current_staging(
staging=staging,
plan=state,
kind="index",
layer_id=layer_id,
page_nbytes=state.page_nbytes,
)
if state.num_current_pages > 0:
staging_pages = staging_span.view(
dense_page_buffer.dtype
).reshape(
state.num_current_pages + 1, *dense_page_buffer.shape[1:]
)
fill_current_index_page_slots(
dense_page_buffer=staging_pages,
current_index_k=current_index_k,
current_index_scale=current_index_scale,
current_locs=current_locs,
page_inverse=state.staging_page_inverse,
page_size=page_size,
index_head_dim=index_head_dim,
current_req_id=current_req_id,
)
cp_symm_barrier(
staging.flag_ptrs, self_rank=int(self.layout.cp_rank)
)
gather_cuda_ipc_peer_pages(
staging.peer_region_ptrs("index", parity),
dense_page_buffer,
state.writer_ranks,
state.staging_slots,
state.current_dense_pages,
page_nbytes=state.page_nbytes,
)
remap_ms = _cpu_timing_ms(remap_cpu)
total_ms = _cpu_timing_ms(consume_cpu)
self._log_layer(
layer_id,
"index_consume_prefix_current_hit layer=%s prefix_pages=%s "
"dense_pages=%s current_rows=%s symm=1 total_ms=%.3f "
"wait_ms=%.3f remap_ms=%.3f",
layer_id,
self.prefix_pages,
int(dense_page_buffer.shape[0]),
int(current_index_k.shape[0]),
total_ms,
wait_ms,
remap_ms,
)
return dense_page_buffer, state.dense_pages
dense_pages = remap_logical_pages_to_slot_dense_pages(
logical_pages,
page_inverse=self.page_inverse,

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

@@ -4482,22 +4482,6 @@ def maybe_build_current_page_writer_ranks(
if not cp_shared_kv_compose_symm_enabled():
return None
# The prefetchers issue their own per-span collectives and never the
# barrier/symm exchange; letting them coexist would make SYMM a silent
# no-op on every prefetch hit. Fail fast instead of measuring parity.
if (
getattr(forward_batch, "cp_shared_kv_mla_prefetcher", None) is not None
or getattr(forward_batch, "cp_shared_kv_index_prefetcher", None) is not None
or envs.SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH.get()
):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][compose_symm] "
"SGLANG_CP_SHARED_KV_COMPOSE_SYMM requires the CP shared-KV "
"prefetchers to be disabled (the prefetch path bypasses the "
"symm exchange and would keep issuing per-span all-reduces). "
"Unset SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH / index prefetch "
"or disable COMPOSE_SYMM."
)
metadata = getattr(forward_batch, "nsa_cp_metadata", None)
if metadata is None or not getattr(metadata, "page_aligned", False):
return None
@@ -4659,8 +4643,10 @@ def _symm_begin_current_staging(
"staging page inverse (slot_remap.page_inverse missing?)"
)
parity = staging.begin_round(int(layer_id), kind)
# +1: staging row 0 is the (unused) dummy page; zeroing it keeps any
# accidental slot-0 read deterministic.
span = staging.buffer(kind, parity)[
: int(plan.num_current_pages) * page_nbytes
: (int(plan.num_current_pages) + 1) * page_nbytes
]
span.zero_()
return parity, span
@@ -4701,11 +4687,8 @@ def _symm_barrier_and_gather_all(
from tai_kernel.nsa_prefill.ipc import cp_symm_barrier
cp_symm_barrier(staging.flag_ptrs, self_rank=int(layout.cp_rank))
combined_ptrs = torch.cat(
[pool_peer_ptrs, staging.peer_region_ptrs(kind, parity)]
)
kernels.materialize_cuda_ipc_peer_pages_slot_dense(
combined_ptrs,
staging.combined_ptr_table(pool_peer_ptrs, kind, parity),
dense_buffer,
plan.symm_all_owner_ranks,
plan.symm_all_src_pages,
@@ -4894,7 +4877,8 @@ def _compose_token_kv_partial_current_v2(
num_rows = int(fill_state.staging_current_rows.numel())
if num_rows > 0:
staging_rows = staging_span.view(kv_cache.dtype).reshape(
int(plan.num_current_pages) * page_size, *kv_cache.shape[1:]
(int(plan.num_current_pages) + 1) * page_size,
*kv_cache.shape[1:],
)
staging_rows.index_copy_(
0,
@@ -5332,7 +5316,7 @@ def _compose_index_partial_current_v2(
)
if plan.num_current_pages > 0:
staging_pages = staging_span.view(page_buffer.dtype).reshape(
int(plan.num_current_pages), *page_buffer.shape[1:]
int(plan.num_current_pages) + 1, *page_buffer.shape[1:]
)
fill_current_index_page_slots(
dense_page_buffer=staging_pages,