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 index f74754e4d..931b7e625 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py @@ -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] = {} diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index a1c221c23..0cf56dbf0 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -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, 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 4937d4581..67c888939 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 @@ -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, diff --git a/test/manual/test_cp_shared_kv_compose_v2_8rank.py b/test/manual/test_cp_shared_kv_compose_v2_8rank.py index fbe9ddc6a..afea40bd5 100644 --- a/test/manual/test_cp_shared_kv_compose_v2_8rank.py +++ b/test/manual/test_cp_shared_kv_compose_v2_8rank.py @@ -183,6 +183,7 @@ def _build_scenario(rank: int, cp_size: int, device: torch.device): return dict( kv_cache=kv_cache, + logical_pages=logical_pages, logical_locs=logical_locs.reshape(-1), loc_req_id=loc_req_id, current_kv=current_kv, @@ -197,6 +198,325 @@ def _build_scenario(rank: int, cp_size: int, device: torch.device): ) +def _build_index_scenario(rank: int, cp_size: int, device: torch.device, s): + """Index page buffer + current K/scale rows over the bs=4 geometry.""" + page_size = s["page_size"] + index_page_nbytes = page_size * (128 + 4) + physical_pages = int(s["kv_cache"].shape[0]) // page_size + page_buffer = torch.zeros( + (physical_pages, index_page_nbytes), dtype=torch.uint8, device=device + ) + # Owned PREFIX pages get deterministic payloads (mirrors the KV pool). + slot_pages = s["slot_remap"].slot_logical_pages.reshape(-1).cpu() + prefix_slots = set() + for start, end in s["prefix_slot_spans"]: + prefix_slots.update(range(int(start), int(end))) + for slot in sorted(prefix_slots): + logical_page = int(slot_pages[slot]) + if logical_page <= 0 or (logical_page - 1) % cp_size != rank: + continue + phys_page = (logical_page - 1) // cp_size + 1 + page_buffer[phys_page] = ( + (torch.arange(index_page_nbytes, dtype=torch.int64) + logical_page * 37) + .remainder_(251) + .to(torch.uint8) + .to(device) + ) + rows = int(s["current_locs"].numel()) + current_k = ( + ( + torch.arange(rows, dtype=torch.int64).view(-1, 1) * 13 + + torch.arange(128, dtype=torch.int64).view(1, -1) + + rank * 17 + ) + .remainder_(249) + .to(torch.uint8) + .to(device) + ) + current_scale = ( + torch.arange(rows, dtype=torch.float32, device=device).view(-1, 1) * 0.25 + + 1.0 + + rank + ) + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as rt + + slot_remap = rt.build_shared_paged_buffer_slot_remap( + page_buffer, + s["logical_pages"], + s["layout"], + ) + return dict( + page_buffer=page_buffer, + current_k=current_k, + current_scale=current_scale, + current_locs=s["current_locs"], + current_req_id=s["current_req_id"], + slot_remap=slot_remap, + ) + + +def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None: + """bs=1 prefetch-path symm consume vs the legacy sync compose (MLA+index).""" + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as rt + from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import ( + CpSharedKVIndexPrefetcher, + CpSharedKVIndexPrefetchHandle, + CpSharedKVMlaPrefetcher, + CpSharedKVMlaPrefetchHandle, + ) + from sglang.srt.mem_cache.cp_shared_kv_compute_owner import ( + build_in_seq_page_compute_owners, + ) + + page_size = 64 + kv_dim = 656 + index_page_nbytes = page_size * (128 + 4) + prefix_len, extend_len = 640, 200 + prefix_pages = prefix_len // page_size + writers = [ + int(o) + for o in build_in_seq_page_compute_owners( + extend_len=extend_len, + extend_prefix_len=prefix_len, + page_size=page_size, + cp_size=cp_size, + ) + ] + pages = list(range(1, prefix_pages + 1)) + next_page = prefix_pages + 1 + for owner in writers: + while (next_page - 1) % cp_size != owner: + next_page += 1 + pages.append(next_page) + next_page += 1 + total_slots = len(pages) + logical_pages = torch.tensor([pages], dtype=torch.int64, device=device) + total_tokens = prefix_len + extend_len + locs = torch.tensor( + [ + pages[i // page_size] * page_size + i % page_size + for i in range(total_tokens) + ], + dtype=torch.int64, + device=device, + ) + physical_pages = (max(pages) - 1) // cp_size + 3 + layout = CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=rank) + + kv_cache = torch.zeros( + (physical_pages * page_size, 1, kv_dim), dtype=torch.uint8, device=device + ) + page_buffer = torch.zeros( + (physical_pages, index_page_nbytes), dtype=torch.uint8, device=device + ) + for logical_page in pages[:prefix_pages]: + if (logical_page - 1) % cp_size != rank: + continue + phys = (logical_page - 1) // cp_size + 1 + kv_cache[phys * page_size : (phys + 1) * page_size] = ( + (torch.arange(page_size * kv_dim, dtype=torch.int64) + logical_page * 53) + .remainder_(251) + .to(torch.uint8) + .view(page_size, 1, kv_dim) + .to(device) + ) + page_buffer[phys] = ( + (torch.arange(index_page_nbytes, dtype=torch.int64) + logical_page * 29) + .remainder_(251) + .to(torch.uint8) + .to(device) + ) + + cur_locs_all = [ + pages[prefix_pages + i // page_size] * page_size + i % page_size + for i in range(extend_len) + ] + mine = [ + i for i in range(extend_len) if writers[i // page_size] == rank + ] + current_locs = torch.tensor( + [cur_locs_all[i] for i in mine], dtype=torch.int64, device=device + ) + rows = len(mine) + current_kv = ( + ( + torch.arange(rows, dtype=torch.int64).view(-1, 1, 1) * 11 + + torch.arange(kv_dim, dtype=torch.int64).view(1, 1, -1) + + rank * 19 + ) + .remainder_(247) + .to(torch.uint8) + .to(device) + ) + current_k = ( + ( + torch.arange(rows, dtype=torch.int64).view(-1, 1) * 23 + + torch.arange(128, dtype=torch.int64).view(1, -1) + + rank * 7 + ) + .remainder_(245) + .to(torch.uint8) + .to(device) + ) + current_scale = ( + torch.arange(rows, dtype=torch.float32, device=device).view(-1, 1) * 0.5 + + 2.0 + + rank + ) + zeros_req = torch.zeros(rows, dtype=torch.long, device=device) + loc_req_id = torch.zeros(total_tokens, dtype=torch.long, device=device) + + kv_remap = rt.build_shared_token_kv_slot_remap( + kv_cache, locs.view(1, -1), logical_pages, layout, page_size + ) + idx_remap = rt.build_shared_paged_buffer_slot_remap( + page_buffer, logical_pages, layout + ) + spans = dict( + prefix_slot_spans=[(0, prefix_pages)], + current_slot_spans=[(prefix_pages, total_slots)], + ) + + # --- references via the legacy sync compose (pure NCCL, no barriers) --- + with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False): + ref_kv, ref_locs = rt.materialize_prefix_and_reuse_current_kv_page_slots( + kv_cache=kv_cache, + logical_locs=locs, + current_kv_cache=current_kv, + current_locs=current_locs, + slot_remap=kv_remap, + layout=layout, + page_size=page_size, + prefix_pages=0, + loc_req_id=loc_req_id, + current_req_id=zeros_req, + layer_id=60, + **spans, + ) + ref_idx, ref_pages = rt.materialize_prefix_and_reuse_current_index_page_slots( + page_buffer=page_buffer, + current_index_k=current_k, + current_index_scale=current_scale, + current_locs=current_locs, + slot_remap=idx_remap, + layout=layout, + page_size=page_size, + index_head_dim=128, + prefix_pages=0, + current_req_id=zeros_req, + layer_id=60, + **spans, + ) + torch.cuda.synchronize() + dist.barrier() + + # --- prefetchers with a hand-built prefix handle (prefix already + # materialized + reduced, exactly what start_next_layer_prefix yields) --- + mla = CpSharedKVMlaPrefetcher( + layout=layout, + page_size=page_size, + prefix_pages=prefix_pages, + slot_logical_pages=kv_remap.slot_logical_pages, + page_inverse=kv_remap.page_inverse, + dense_num_pages=kv_remap.dense_num_pages, + slot_remap=kv_remap, + symm_writers=writers, + ) + dense_kv = kv_cache.new_zeros( + (kv_remap.dense_num_pages * page_size, 1, kv_dim) + ) + rt.materialize_local_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv, + slot_logical_pages=kv_remap.slot_logical_pages, + layout=layout, + page_size=page_size, + start_slot=0, + end_slot=prefix_pages, + ) + prefix_rows = rt.slot_range_to_token_slice(page_size, 0, prefix_pages) + dist.all_reduce(dense_kv[prefix_rows]) + mla_event = torch.cuda.Event() + mla_event.record() + mla.handles[61] = CpSharedKVMlaPrefetchHandle( + layer_id=61, dense_kv_cache=dense_kv, prefix_rows=prefix_rows, + event=mla_event, + ) + + idx = CpSharedKVIndexPrefetcher( + layout=layout, + prefix_pages=prefix_pages, + slot_logical_pages=idx_remap.slot_logical_pages, + page_inverse=idx_remap.page_inverse, + dense_num_pages=idx_remap.dense_num_pages, + slot_remap=idx_remap, + symm_writers=writers, + ) + dense_pb = page_buffer.new_zeros( + (idx_remap.dense_num_pages, index_page_nbytes) + ) + rt.materialize_local_paged_buffer_page_slots_into( + page_buffer=page_buffer, + dense_page_buffer=dense_pb, + slot_logical_pages=idx_remap.slot_logical_pages, + layout=layout, + start_slot=0, + end_slot=prefix_pages, + ) + dist.all_reduce(dense_pb[1 : prefix_pages + 1]) + idx_event = torch.cuda.Event() + idx_event.record() + idx.handles[61] = CpSharedKVIndexPrefetchHandle( + layer_id=61, dense_page_buffer=dense_pb, + prefix_rows=slice(1, prefix_pages + 1), event=idx_event, + ) + torch.cuda.synchronize() + dist.barrier() + + with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override( + True + ), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True): + out_idx = idx.consume_prefix_with_current( + layer_id=61, + logical_pages=logical_pages, + current_index_k=current_k, + current_index_scale=current_scale, + current_locs=current_locs, + page_size=page_size, + index_head_dim=128, + current_req_id=zeros_req, + ) + out_mla = mla.consume_prefix_with_current( + layer_id=61, + kv_cache=kv_cache, + logical_locs=locs, + current_kv_cache=current_kv, + current_locs=current_locs, + loc_req_id=loc_req_id, + current_req_id=zeros_req, + ) + torch.cuda.synchronize() + assert out_idx is not None and out_mla is not None, ( + f"rank{rank}: prefetch consume unexpectedly missed" + ) + sym_idx, sym_pages = out_idx + sym_kv, sym_locs = out_mla + assert torch.equal(ref_pages, sym_pages), f"rank{rank}: prefetch index pages" + assert torch.equal(ref_locs, sym_locs), f"rank{rank}: prefetch mla locs" + if not torch.equal(ref_idx, sym_idx): + bad = (ref_idx != sym_idx).any(dim=-1).nonzero().reshape(-1)[:8] + raise AssertionError( + f"rank{rank}: prefetch index symm mismatch at pages {bad.cpu().tolist()}" + ) + if not torch.equal(ref_kv, sym_kv): + diff = (ref_kv != sym_kv).any(dim=-1).any(dim=-1) + bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist() + raise AssertionError( + f"rank{rank}: prefetch mla symm mismatch at rows {bad} " + f"(of {int(diff.sum())})" + ) + + def _compose(s, layer_id: int, *, writers: list[int] | None = None): return runtime.materialize_prefix_and_reuse_current_kv_page_slots( kv_cache=s["kv_cache"], @@ -316,6 +636,65 @@ def main() -> None: "MiB, fill-to-staging + barrier + mega gather engaged; arena on/off)", flush=True, ) + + # ---- Phase 3: index sync compose, legacy vs symm (staging already + # registered by the token phase — the index compose never registers). ---- + si = _build_index_scenario(rank, world, device, s) + + def _compose_index(layer_id: int, *, writers=None): + return runtime.materialize_prefix_and_reuse_current_index_page_slots( + page_buffer=si["page_buffer"], + current_index_k=si["current_k"], + current_index_scale=si["current_scale"], + current_locs=si["current_locs"], + slot_remap=si["slot_remap"], + layout=s["layout"], + page_size=s["page_size"], + index_head_dim=128, + prefix_pages=0, + current_req_id=si["current_req_id"], + prefix_slot_spans=s["prefix_slot_spans"], + current_slot_spans=s["current_slot_spans"], + layer_id=layer_id, + current_page_writer_ranks=writers, + ) + + with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False): + ref_idx, ref_dense_pages = _compose_index(40) + torch.cuda.synchronize() + dist.barrier() + with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override( + True + ), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True): + for layer_id in range(41, 45): + symm_idx, symm_dense_pages = _compose_index( + layer_id, writers=s["current_page_writer_ranks"] + ) + torch.cuda.synchronize() + assert torch.equal(ref_dense_pages, symm_dense_pages), ( + f"rank{rank} layer{layer_id}: index dense_pages mismatch" + ) + if not torch.equal(ref_idx, symm_idx): + bad = ( + (ref_idx != symm_idx).any(dim=-1).nonzero().reshape(-1)[:8] + ) + raise AssertionError( + f"rank{rank} layer{layer_id}: index symm mismatch at pages " + f"{bad.cpu().tolist()}" + ) + dist.barrier() + if rank == 0: + print("PASS: index sync symm compose byte-identical (4 layers)", flush=True) + + # ---- Phase 4+5: prefetch-path symm consume (MLA + index), bs=1. ---- + _check_prefetch_symm(rank, world, device) + dist.barrier() + if rank == 0: + print( + "PASS: prefetch consume_prefix_with_current symm (MLA + index) " + "byte-identical to the sync compose", + flush=True, + ) dist.destroy_process_group() 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 index 8bf3fed65..26bb1e32a 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_compose.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_compose.py @@ -134,11 +134,11 @@ class TestComposePlanSymm(unittest.TestCase): self.assertEqual(plan.current_dense_pages.tolist(), [3, 4, 5]) self.assertEqual(plan.symm_all_owner_ranks[:2].tolist(), [0, 1]) self.assertEqual(plan.symm_all_owner_ranks[2:].tolist(), [11, 8, 11]) - self.assertEqual(plan.symm_all_src_pages[2:].tolist(), [0, 1, 2]) + self.assertEqual(plan.symm_all_src_pages[2:].tolist(), [1, 2, 3]) # staging page inverse: current dense pages 3,4,5 -> staging slots - # 0,1,2; everything else (dummy, prefix pages) -> -1. + # 1,2,3 (slot 0 = dummy row); everything else -> -1. self.assertEqual( - plan.staging_page_inverse.tolist(), [[-1, -1, -1, 0, 1, 2]] + plan.staging_page_inverse.tolist(), [[-1, -1, -1, 1, 2, 3]] ) def test_writer_count_mismatch_fails_fast(self): @@ -310,7 +310,9 @@ class TestComposeStaging(unittest.TestCase): region_in_half = {} for kind in CpComposeStaging.KINDS: region_in_half[kind] = half - half += staging._align(capacity_pages * staging._page_nbytes[kind]) + half += staging._align( + (capacity_pages + 1) * staging._page_nbytes[kind] + ) for parity in (0, 1): for kind in CpComposeStaging.KINDS: staging._region_offsets[(kind, parity)] = ( @@ -318,6 +320,8 @@ class TestComposeStaging(unittest.TestCase): ) staging._slab = torch.zeros(2 * half, dtype=torch.uint8) staging.peer_slab_bases = torch.full((8,), 1 << 20, dtype=torch.int64) + for key, offset in staging._region_offsets.items(): + staging._region_peer_ptrs[key] = staging.peer_slab_bases + offset return staging, half def test_regions_are_disjoint_and_parity_halves_do_not_overlap(self): @@ -326,12 +330,12 @@ class TestComposeStaging(unittest.TestCase): idx0 = staging.buffer("index", 0) kv1 = staging.buffer("token_kv", 1) # kv and index regions of one half are adjacent but disjoint - # (index starts at the 256-aligned end of kv). - self.assertEqual(idx0.data_ptr() - kv0.data_ptr(), staging._align(4 * 1000)) + # (index starts at the 256-aligned end of kv; +1 = dummy row 0). + self.assertEqual(idx0.data_ptr() - kv0.data_ptr(), staging._align(5 * 1000)) # The parity-1 half starts exactly one half past parity 0. self.assertEqual(kv1.data_ptr() - kv0.data_ptr(), half) - self.assertEqual(kv0.numel(), 4 * 1000) - self.assertEqual(idx0.numel(), 4 * 300) + self.assertEqual(kv0.numel(), 5 * 1000) + self.assertEqual(idx0.numel(), 5 * 300) def test_peer_region_ptrs_offset_matches_local_layout(self): staging, half = self._staging() @@ -342,7 +346,7 @@ class TestComposeStaging(unittest.TestCase): self.assertTrue( torch.equal( staging.peer_region_ptrs("index", 1), - base + half + staging._align(4 * 1000), + base + half + staging._align(5 * 1000), ) )