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 7c0158ff9..f74754e4d 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 @@ -52,8 +52,9 @@ def cp_shared_kv_compose_arena_enabled() -> bool: def cp_shared_kv_compose_symm_enabled() -> bool: - """Step B: exchange current pages via the compact symm staging (zero - NCCL): publish my current pages to staging, barrier, gather peers'. + """Step B: zero-NCCL compose via the compact symm staging — current rows + are written straight into the staging, then barrier + ONE slot-dense + gather covers prefix (pool) and current (staging) pages together. Independent of the arena (dense buffers stay rank-local).""" return bool(envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.get()) @@ -74,16 +75,19 @@ class ComposePlan: dummy-page convention) of all current slots in merged-span order. When the caller provides per-current-page writer (compute-owner) ranks, - the plan splits the current pages by writer for the symm staging exchange. - 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: + 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: - - ``local_current_*``: pages THIS rank wrote — published from the dense - buffer into staging slots before the barrier. The writer-ranks tensor - is carried (all entries == cp_rank) so the publish kernel can index a - uniform pointer table without a per-layer allocation. - - ``remote_current_*``: pages OTHER ranks wrote — gathered after the - barrier from ``staging[writer][slot]`` into ``dense[page]``. + - ``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 + (sentinel zero-fill). + - ``staging_page_inverse``: ``page_inverse`` with current pages remapped + to their staging slot (others -1), so the unchanged fill kernels write + current rows straight INTO the staging instead of the dense buffer. """ prefix_owner_ranks: torch.Tensor @@ -92,12 +96,9 @@ class ComposePlan: total_slots: int num_prefix_slots: int num_current_pages: int - local_current_writer_ranks: torch.Tensor | None = None - local_current_slot_indices: torch.Tensor | None = None - local_current_dense_pages: torch.Tensor | None = None - remote_current_writer_ranks: torch.Tensor | None = None - remote_current_slot_indices: torch.Tensor | None = None - remote_current_dense_pages: torch.Tensor | None = None + symm_all_owner_ranks: torch.Tensor | None = None + symm_all_src_pages: torch.Tensor | None = None + staging_page_inverse: torch.Tensor | None = None def _spans_key(spans: list[tuple[int, int]] | None) -> tuple[tuple[int, int], ...]: @@ -106,6 +107,34 @@ def _spans_key(spans: list[tuple[int, int]] | None) -> tuple[tuple[int, int], .. return tuple((int(s), int(e)) for s, e in spans) +def _build_staging_page_inverse( + page_inverse: torch.Tensor, + current_dense_pages: torch.Tensor, +) -> torch.Tensor: + """Remap ``page_inverse`` so current pages point at their staging slot. + + ``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. + """ + + if page_inverse.dim() == 1: + page_inverse = page_inverse.unsqueeze(0) + num_current = int(current_dense_pages.numel()) + invalid = torch.full_like(page_inverse, -1) + if num_current == 0: + return invalid + pos = torch.searchsorted( + 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) + + def get_or_build_compose_plan( *, slot_remap, @@ -182,12 +211,9 @@ def get_or_build_compose_plan( else: current_dense_pages = torch.empty(0, device=device, dtype=torch.long) - local_writer_ranks = None - local_slot_indices = None - local_dense_pages = None - remote_writer_ranks = None - remote_slot_indices = None - remote_dense_pages = None + symm_all_owner_ranks = None + symm_all_src_pages = None + staging_page_inverse = None if current_page_writer_ranks is not None: num_current = int(current_dense_pages.numel()) if len(current_page_writer_ranks) != num_current: @@ -196,18 +222,22 @@ def get_or_build_compose_plan( f"count mismatch: writers={len(current_page_writer_ranks)} " f"current_pages={num_current} (page-aligned split required)" ) - writers = torch.tensor( - current_page_writer_ranks, dtype=torch.long, device=device - ) - slot_indices = torch.arange(num_current, dtype=torch.long, device=device) - remote = writers != int(layout.cp_rank) - local = ~remote - local_writer_ranks = writers[local].contiguous() - local_slot_indices = slot_indices[local].contiguous() - local_dense_pages = current_dense_pages[local].contiguous() - remote_writer_ranks = writers[remote].contiguous() - remote_slot_indices = slot_indices[remote].contiguous() - remote_dense_pages = current_dense_pages[remote].contiguous() + symm_all_owner_ranks = prefix_owner_ranks.clone() + symm_all_src_pages = prefix_src_pages.clone() + if num_current > 0: + writers = torch.tensor( + current_page_writer_ranks, dtype=torch.long, device=device + ) + current_slots = current_dense_pages - 1 + symm_all_owner_ranks[current_slots] = writers + int(layout.cp_size) + symm_all_src_pages[current_slots] = torch.arange( + num_current, dtype=torch.long, device=device + ) + page_inverse = getattr(slot_remap, "page_inverse", None) + if page_inverse is not None: + staging_page_inverse = _build_staging_page_inverse( + page_inverse, current_dense_pages + ) plan = ComposePlan( prefix_owner_ranks=prefix_owner_ranks, @@ -216,12 +246,9 @@ def get_or_build_compose_plan( total_slots=total_slots, num_prefix_slots=num_prefix_slots, num_current_pages=int(current_dense_pages.numel()), - local_current_writer_ranks=local_writer_ranks, - local_current_slot_indices=local_slot_indices, - local_current_dense_pages=local_dense_pages, - remote_current_writer_ranks=remote_writer_ranks, - remote_current_slot_indices=remote_slot_indices, - remote_current_dense_pages=remote_dense_pages, + symm_all_owner_ranks=symm_all_owner_ranks, + symm_all_src_pages=symm_all_src_pages, + staging_page_inverse=staging_page_inverse, ) plans[key] = plan return plan @@ -335,7 +362,7 @@ class CpComposeStaging(_RoundParity): 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 publish into half + 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 proves every peer's R+1 barrier kernel ran, which (stream order) proves every peer's round-R gather of half H finished. 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 4847c1a3c..4937d4581 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 @@ -4556,41 +4556,90 @@ def _symm_staging_ready_or_register( return staging.registered -def _symm_exchange_current_pages( - dense_buffer: torch.Tensor, - plan, - layout: CpSharedKVLayout, +class _SymmTokenFillState: + """Layer-invariant pieces of the symm token-KV compose, built once per + batch and anchored on the (batch-lifetime) compose plan. + + ``staging_current_rows`` maps each current KV row to its row in the + compact staging (via the staging page inverse), so the per-layer fill is + a single ``index_copy_``. ``mixed_locs`` is the page-slack-masked loc + remap the fused fill kernel used to recompute every layer — it depends + only on the batch's locs, never on buffer contents. + """ + + __slots__ = ("staging_current_rows", "mixed_locs") + + def __init__(self, staging_current_rows, mixed_locs): + self.staging_current_rows = staging_current_rows + self.mixed_locs = mixed_locs + + +def _get_or_build_symm_token_fill_state( *, - page_nbytes: int, + plan, + slot_remap, + layout: CpSharedKVLayout, + kv_cache: torch.Tensor, + logical_locs: torch.Tensor, + current_locs: torch.Tensor, + loc_req_id: torch.Tensor, + current_req_id: torch.Tensor, + page_size: int, +) -> "_SymmTokenFillState": + state = getattr(plan, "_symm_token_fill_state", None) + if state is not None: + return state + 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, + ) + # Same masking as the fused fill kernel / torch reference: page-slack + # rows inside current pages become -1 (invisible to attention). + current_mask, _ = build_current_loc_remap(logical_locs, current_locs) + current_page_mask = build_current_page_mask( + logical_locs, current_locs, page_size=page_size + ) + mixed_locs = torch.where( + current_page_mask & (~current_mask), + torch.full_like(dense_locs, -1), + dense_locs, + ) + staging_current_rows = remap_logical_locs_to_slot_dense_locs_optimized( + current_locs.reshape(-1), + page_inverse=plan.staging_page_inverse, + page_size=page_size, + loc_req_id=current_req_id.reshape(-1), + ).to(torch.long) + state = _SymmTokenFillState(staging_current_rows, mixed_locs) + object.__setattr__(plan, "_symm_token_fill_state", state) + return state + + +def _symm_begin_current_staging( + *, + staging, + plan, kind: str, layer_id: int, -) -> None: - """Step B current-page exchange through the compact symm staging: + page_nbytes: int, +) -> tuple[int, torch.Tensor]: + """Open this round's staging span for the current-row fill. - publish my written current pages: dense[page] -> staging[slot] - barrier cp_symm_barrier (counting; release/acquire at system scope) - gather peers' pages: peer_staging[writer][slot] -> dense[page] + Validates the batch against the registered staging layout (both checks + are batch-logical, hence rank-uniform), advances the round parity, and + returns the zeroed span — zeroed so that tail-slack rows inside current + pages stay zero in every peer's gathered copy, exactly like the sentinel + zero-fill did on the dense buffer. + """ - The staging slot of current page ``i`` is ``i`` (merged-span order), - identical on every rank, so peers address each other's staging without a - per-batch handshake. The barrier runs even with zero local/remote pages - — barrier COUNTS must match across ranks. - - INVARIANT: every gate on the path to this call (compose env flags, - page_aligned metadata, staging.registered, the agreed IPC capability, - prefetch absence) MUST be rank-uniform. A per-rank divergence here - desyncs the barrier counting and hangs the CP group — never add a - per-rank condition without routing it through a group agreement first - (see _agreed_tai_ipc_peer_ptrs).""" - - from tai_kernel.nsa_prefill.ipc import ( - cp_symm_barrier, - gather_cuda_ipc_peer_pages, - ) - - staging = get_compose_staging(dense_buffer.device) if int(plan.num_current_pages) > staging.capacity_pages: - # Batch-logical quantity -> raises uniformly on every rank. raise RuntimeError( "[CP_SHARED_KV_FAIL_FAST][compose_symm] batch current pages " f"exceed the symm staging capacity: pages={plan.num_current_pages} " @@ -4604,37 +4653,64 @@ def _symm_exchange_current_pages( f"the registered staging layout: kind={kind} page_nbytes=" f"{page_nbytes} registered={staging.page_nbytes(kind)}" ) + if plan.staging_page_inverse is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][compose_symm] compose plan has no " + "staging page inverse (slot_remap.page_inverse missing?)" + ) parity = staging.begin_round(int(layer_id), kind) - if ( - plan.local_current_slot_indices is not None - and plan.local_current_slot_indices.numel() > 0 - ): - # Local copy via the gather kernel: a uniform pointer table aimed at - # my own dense buffer turns it into dense[page] -> staging[slot]. - self_ptrs = torch.full( - (staging.cp_size,), dense_buffer.data_ptr(), dtype=torch.int64 - ) - gather_cuda_ipc_peer_pages( - self_ptrs, - staging.buffer(kind, parity), - plan.local_current_writer_ranks, - plan.local_current_dense_pages, - plan.local_current_slot_indices, - page_nbytes=page_nbytes, - ) + span = staging.buffer(kind, parity)[ + : int(plan.num_current_pages) * page_nbytes + ] + span.zero_() + return parity, span + + +def _symm_barrier_and_gather_all( + *, + kernels, + pool_peer_ptrs: torch.Tensor, + dense_buffer: torch.Tensor, + plan, + layout: CpSharedKVLayout, + staging, + kind: str, + parity: int, + page_nbytes: int, +) -> None: + """Step B compose: barrier, then ONE slot-dense gather for everything. + + The current rows were already written into this rank's staging span (the + fill kernels were pointed there via ``plan.staging_page_inverse``), so + after the barrier a single gather against the concatenated + ``[pool peer ptrs | staging peer ptrs]`` table materializes prefix pages + (owner < cp_size, from the KV pools), current pages (owner = cp_size + + writer, from the stagings — including this rank's own), and zero-fills + the remaining sentinel slots. No publish copy, no second gather. + + The barrier runs even with zero current pages — barrier COUNTS must + match across ranks. + + INVARIANT: every gate on the path to this call (compose env flags, + page_aligned metadata, staging.registered, the agreed IPC capability, + prefetch absence) MUST be rank-uniform. A per-rank divergence here + desyncs the barrier counting and hangs the CP group — never add a + per-rank condition without routing it through a group agreement first + (see _agreed_tai_ipc_peer_ptrs).""" + + from tai_kernel.nsa_prefill.ipc import cp_symm_barrier + cp_symm_barrier(staging.flag_ptrs, self_rank=int(layout.cp_rank)) - if ( - plan.remote_current_slot_indices is not None - and plan.remote_current_slot_indices.numel() > 0 - ): - gather_cuda_ipc_peer_pages( - staging.peer_region_ptrs(kind, parity), - dense_buffer, - plan.remote_current_writer_ranks, - plan.remote_current_slot_indices, - plan.remote_current_dense_pages, - page_nbytes=page_nbytes, - ) + combined_ptrs = torch.cat( + [pool_peer_ptrs, staging.peer_region_ptrs(kind, parity)] + ) + kernels.materialize_cuda_ipc_peer_pages_slot_dense( + combined_ptrs, + dense_buffer, + plan.symm_all_owner_ranks, + plan.symm_all_src_pages, + page_nbytes=page_nbytes, + ) def _resolve_partial_current_spans( @@ -4759,6 +4835,7 @@ def _compose_token_kv_partial_current_v2( page_size=page_size, ) gathered = ipc_state is not None + kv_page_nbytes = _token_kv_page_nbytes(kv_cache, page_size) if gathered: kernels, peer_ptrs = ipc_state dense_kv_cache = acquire_dense_buffer( @@ -4768,13 +4845,14 @@ def _compose_token_kv_partial_current_v2( 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), - ) + if not use_symm: + kernels.materialize_cuda_ipc_peer_pages_slot_dense( + peer_ptrs, + dense_kv_cache, + plan.prefix_owner_ranks, + plan.prefix_src_pages, + page_nbytes=kv_page_nbytes, + ) else: dense_kv_cache = kv_cache.new_zeros((dense_rows, *kv_cache.shape[1:])) for prefix_start_slot, prefix_end_slot in prefix_spans: @@ -4788,6 +4866,70 @@ def _compose_token_kv_partial_current_v2( end_slot=prefix_end_slot, ) + if use_symm: + # Step B: write current rows straight into this round's staging span + # (one index_copy; the loc remap and slack masking are layer-invariant + # and cached on the plan), then barrier + ONE slot-dense gather for + # prefix AND current pages. The prefix-only gather above was skipped + # — the mega gather covers it. + staging = get_compose_staging(kv_cache.device) + parity, staging_span = _symm_begin_current_staging( + staging=staging, + plan=plan, + kind="token_kv", + layer_id=layer_id, + page_nbytes=kv_page_nbytes, + ) + fill_state = _get_or_build_symm_token_fill_state( + plan=plan, + slot_remap=slot_remap, + layout=layout, + kv_cache=kv_cache, + logical_locs=logical_locs, + current_locs=current_locs, + loc_req_id=loc_req_id, + current_req_id=current_req_id, + page_size=page_size, + ) + 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:] + ) + staging_rows.index_copy_( + 0, + fill_state.staging_current_rows, + current_kv_cache[:num_rows], + ) + _symm_barrier_and_gather_all( + kernels=kernels, + pool_peer_ptrs=peer_ptrs, + dense_buffer=dense_kv_cache, + plan=plan, + layout=layout, + staging=staging, + kind="token_kv", + parity=parity, + page_nbytes=kv_page_nbytes, + ) + mixed_kv_cache = dense_kv_cache + mixed_locs = fill_state.mixed_locs + 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 symm=1 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, + kv_cache.dtype, + ) + return mixed_kv_cache, mixed_locs + logical_locs = filter_locs_mappable_to_physical_pool( logical_locs=logical_locs, layout=layout, @@ -4812,16 +4954,7 @@ def _compose_token_kv_partial_current_v2( ) if gathered: - if use_symm: - _symm_exchange_current_pages( - mixed_kv_cache, - plan, - layout, - page_nbytes=_token_kv_page_nbytes(kv_cache, page_size), - kind="token_kv", - layer_id=layer_id, - ) - elif plan.num_current_pages > 0: + if plan.num_current_pages > 0: _reduce_current_pages_compact( mixed_kv_cache, int(slot_remap.dense_num_pages), @@ -5161,13 +5294,14 @@ def _compose_index_partial_current_v2( 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), - ) + if not use_symm: + 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:]) @@ -5182,6 +5316,60 @@ def _compose_index_partial_current_v2( end_slot=prefix_end_slot, ) + if use_symm: + # Step B: fill current index rows straight into this round's staging + # span (the unchanged fill kernel — its write destinations all come + # from the page inverse, which the plan remapped to staging slots), + # then barrier + ONE slot-dense gather for prefix AND current pages. + index_page_nbytes = _page_nbytes_from_page_tensor(page_buffer) + staging = get_compose_staging(page_buffer.device) + parity, staging_span = _symm_begin_current_staging( + staging=staging, + plan=plan, + kind="index", + layer_id=layer_id, + page_nbytes=index_page_nbytes, + ) + if plan.num_current_pages > 0: + staging_pages = staging_span.view(page_buffer.dtype).reshape( + int(plan.num_current_pages), *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=plan.staging_page_inverse, + page_size=page_size, + index_head_dim=index_head_dim, + current_req_id=current_req_id, + ) + _symm_barrier_and_gather_all( + kernels=kernels, + pool_peer_ptrs=peer_ptrs, + dense_buffer=dense_page_buffer, + plan=plan, + layout=layout, + staging=staging, + kind="index", + parity=parity, + page_nbytes=index_page_nbytes, + ) + 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 symm=1", + layout.cp_rank, + layer_id, + layout.cp_size, + plan.total_slots, + dense_num_pages, + plan.num_prefix_slots, + plan.num_current_pages, + ) + return dense_page_buffer, slot_remap.dense_pages + dense_page_buffer = fill_current_index_page_slots( dense_page_buffer=dense_page_buffer, current_index_k=current_index_k, @@ -5194,16 +5382,7 @@ def _compose_index_partial_current_v2( ) if gathered: - if use_symm: - _symm_exchange_current_pages( - dense_page_buffer, - plan, - layout, - page_nbytes=_page_nbytes_from_page_tensor(page_buffer), - kind="index", - layer_id=layer_id, - ) - elif plan.num_current_pages > 0: + if plan.num_current_pages > 0: _reduce_current_pages_compact( dense_page_buffer, dense_num_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 1c0430c5e..fbe9ddc6a 100644 --- a/test/manual/test_cp_shared_kv_compose_v2_8rank.py +++ b/test/manual/test_cp_shared_kv_compose_v2_8rank.py @@ -267,7 +267,8 @@ def main() -> None: flush=True, ) - # ---- Step B: compact symm staging (publish + barrier + peer gather, + # ---- Step B: compact symm staging (fill-to-staging + barrier + one + # mega slot-dense gather, # zero NCCL in the current-page phase). Multiple layers exercise the # staging parity halves; the arena is exercised in a second pass to # prove the two knobs are independent. ---- @@ -312,7 +313,7 @@ def main() -> None: print( "PASS: compact-staging symm compose byte-identical to v2 across " f"8 layers (staging {staging.capacity_pages} pages ~{staged_mb:.1f} " - "MiB, publish + barrier + peer gather engaged; arena on and off)", + "MiB, fill-to-staging + barrier + mega gather engaged; arena on/off)", 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 9d0d1d483..8bf3fed65 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 @@ -110,11 +110,16 @@ class TestComposePlan(unittest.TestCase): class TestComposePlanSymm(unittest.TestCase): - def test_remote_current_lists_exclude_self_written_pages(self): + def test_symm_descriptors_cover_prefix_pool_and_current_staging(self): layout = CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=3) slot_logical_pages = torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64) + # logical page p lives at dense page p (identity batch layout). + page_inverse = torch.tensor([[-1, 1, 2, 3, 4, 5]], dtype=torch.int64) plan = get_or_build_compose_plan( - slot_remap=SimpleNamespace(slot_logical_pages=slot_logical_pages), + slot_remap=SimpleNamespace( + slot_logical_pages=slot_logical_pages, + page_inverse=page_inverse, + ), layout=layout, physical_page_capacity=None, prefix_spans=[(0, 2)], @@ -122,16 +127,19 @@ class TestComposePlanSymm(unittest.TestCase): kind="token_kv", current_page_writer_ranks=[3, 0, 3], ) - # current dense pages are slots 2,3,4 -> dense 3,4,5; writers 3,0,3; - # self rank 3 -> only the page written by rank 0 is remote. The - # staging slot of current page i is i (merged-span order). + # current dense pages are slots 2,3,4 -> dense 3,4,5 with staging + # slots 0,1,2. The mega-gather descriptors keep the pool owner on + # prefix slots and switch current slots to cp_size + writer with the + # staging slot as the source page. self.assertEqual(plan.current_dense_pages.tolist(), [3, 4, 5]) - self.assertEqual(plan.remote_current_writer_ranks.tolist(), [0]) - self.assertEqual(plan.remote_current_slot_indices.tolist(), [1]) - self.assertEqual(plan.remote_current_dense_pages.tolist(), [4]) - self.assertEqual(plan.local_current_writer_ranks.tolist(), [3, 3]) - self.assertEqual(plan.local_current_slot_indices.tolist(), [0, 2]) - self.assertEqual(plan.local_current_dense_pages.tolist(), [3, 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]) + # staging page inverse: current dense pages 3,4,5 -> staging slots + # 0,1,2; everything else (dummy, prefix pages) -> -1. + self.assertEqual( + plan.staging_page_inverse.tolist(), [[-1, -1, -1, 0, 1, 2]] + ) def test_writer_count_mismatch_fails_fast(self): layout = CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=0)