Add symm-heap current-page exchange for CP shared-KV compose (Step B)

Replaces the last remaining NCCL collective in the bs>1 compose layer loop (the compact-current all-reduce) with a DeepEP-style counting barrier plus a CUDA-IPC gather from peers' symmetric dense buffers, behind SGLANG_CP_SHARED_KV_COMPOSE_SYMM (+ COMPOSE_ARENA, both default off). - CpComposeArena.register_symm: fixes capacity at the pool-derived bound (logical pages x dense page unit, overridable via SGLANG_CP_SHARED_KV_SYMM_HEAP_MB), allocates slab + flags in CUDA-IPC memory, exchanges handles once over the CP group; deterministic bump carve means peer_base + my_offset addresses any peer's dense buffer with no per-layer handshakes. Registration happens only in the token-KV compose (uniform first-use point); growth after registration raises. - Current pages are single-writer at page granularity under the page-aligned in-seq split, so the exchange is the existing gather_cuda_ipc_peer_pages with src==dst page ids and writer (compute owner) ranks; writers are built per batch by build_batch_current_page_writer_ranks and gated by maybe_build_current_page_writer_ranks (page_aligned metadata required). The barrier runs even with zero remote pages (counts must match). - _agreed_tai_ipc_peer_ptrs: the per-rank IPC capability probe is now agreed across the CP group (one-time MIN all-reduce per pool tensor) so ranks can never split between gather and collective paths and deadlock on mismatched NCCL shapes. - ComposePlan cache re-anchored ON the slot_remap object (forward-batch lifetime) instead of a module-level tensor-identity key, which could go stale when a freed tensor's address is reused by the next batch. Validation (g0034 syh-dev-new): tai-kernel cp_symm_barrier correctness (200 adversarial iterations, rotating 10ms producer delays, phase-safety, flags drained at quiescence) and perf (7.7us max-rank latency) both pass; 8-rank GPU test extends to the symm path - byte-identical to compose_v2 across 4 layers (parity halves exercised, slab registered); mem_cache suite 464 passed with the only failures being a documented pre-existing sys.modules stub pollution pair, reproduced identically with SGLANG_CP_SHARED_KV_COMPOSE_V2=0. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-11 20:13:04 +00:00
parent f47b739e30
commit c55e406176
8 changed files with 626 additions and 68 deletions
@@ -21,9 +21,6 @@ import torch.distributed as dist

 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
-from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
-    reset_compose_plan_cache,
-)
 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout


@@ -132,7 +129,12 @@ def _build_scenario(rank: int, cp_size: int, device: torch.device):
                phys_page * page_size : (phys_page + 1) * page_size
            ] = payload.view(page_size, 1, kv_dim).to(device)

-    # Current rows owned by this rank (writer = page owner in this scenario).
+    # Current rows owned by this rank (writer = page owner in this scenario,
+    # so the symm writer list is the storage owner per current page).
+    current_page_writer_ranks = []
+    for req_pages, prefix_pages in zip(logical_rows, prefix_pages_by_req):
+        for page in req_pages[prefix_pages:]:
+            current_page_writer_ranks.append((page - 1) % cp_size)
    cur_locs = [
        loc for req_id, loc in current_locs_all
        if (loc // page_size - 1) % cp_size == rank
@@ -191,10 +193,11 @@ def _build_scenario(rank: int, cp_size: int, device: torch.device):
        page_size=page_size,
        prefix_slot_spans=prefix_slot_spans,
        current_slot_spans=current_slot_spans,
+        current_page_writer_ranks=current_page_writer_ranks,
    )


-def _compose(s, layer_id: int):
+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"],
        logical_locs=s["logical_locs"],
@@ -209,6 +212,7 @@ def _compose(s, layer_id: int):
        prefix_slot_spans=s["prefix_slot_spans"],
        current_slot_spans=s["current_slot_spans"],
        layer_id=layer_id,
+        current_page_writer_ranks=writers,
    )


@@ -230,7 +234,6 @@ def main() -> None:
    torch.cuda.synchronize()
    dist.barrier()

-    reset_compose_plan_cache()
    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(True):
        v2_kv, v2_locs = _compose(s, layer_id=0)
    torch.cuda.synchronize()
@@ -263,6 +266,41 @@ def main() -> None:
            f"(dense rows={int(ref_kv.shape[0])}, world={world})",
            flush=True,
        )
+
+    # ---- Step B: symm exchange (arena + barrier + peer gather, zero NCCL
+    # in the current-page phase). Multiple layers exercise parity halves. ----
+    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(
+        True
+    ), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(
+        True
+    ), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True):
+        for layer_id in range(4):
+            symm_kv, symm_locs = _compose(
+                s, layer_id=layer_id, writers=s["current_page_writer_ranks"]
+            )
+            torch.cuda.synchronize()
+            assert torch.equal(ref_locs, symm_locs), (
+                f"rank{rank} layer{layer_id}: symm locs mismatch"
+            )
+            if not torch.equal(ref_kv, symm_kv):
+                diff = (ref_kv != symm_kv).any(dim=-1).any(dim=-1)
+                bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
+                raise AssertionError(
+                    f"rank{rank} layer{layer_id}: symm dense kv mismatch at "
+                    f"rows {bad} (of {int(diff.sum())})"
+                )
+    dist.barrier()
+    from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
+        get_compose_arena,
+    )
+
+    assert get_compose_arena(device).symm_ready, "symm slab was not registered"
+    if rank == 0:
+        print(
+            "PASS: symm compose byte-identical to v2 across 4 layers "
+            "(arena registered, barrier + peer gather engaged)",
+            flush=True,
+        )
    dist.destroy_process_group()


@@ -8,11 +8,12 @@ import unittest
 import torch

 from sglang.srt.environ import envs
+from types import SimpleNamespace
+
 from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
    CpComposeArena,
    acquire_dense_buffer,
    get_or_build_compose_plan,
-    reset_compose_plan_cache,
 )
 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout

@@ -22,9 +23,6 @@ register_cpu_ci(est_time=1, suite="stage-a-test-cpu")


 class TestComposePlan(unittest.TestCase):
-    def setUp(self):
-        reset_compose_plan_cache()
-
    def _layout(self, cp_size=8, cp_rank=0, page_size=64):
        return CpSharedKVLayout(
            page_size=page_size, cp_size=cp_size, cp_rank=cp_rank
@@ -39,7 +37,7 @@ class TestComposePlan(unittest.TestCase):
        prefix_spans = [(0, 3), (5, 7)]
        current_spans = [(3, 5), (7, 10)]
        plan = get_or_build_compose_plan(
-            slot_logical_pages=slot_logical_pages,
+            slot_remap=SimpleNamespace(slot_logical_pages=slot_logical_pages),
            layout=layout,
            physical_page_capacity=None,
            prefix_spans=prefix_spans,
@@ -77,7 +75,7 @@ class TestComposePlan(unittest.TestCase):
        layout = self._layout()
        slot_logical_pages = torch.tensor([1, 0, 3], dtype=torch.int64)
        plan = get_or_build_compose_plan(
-            slot_logical_pages=slot_logical_pages,
+            slot_remap=SimpleNamespace(slot_logical_pages=slot_logical_pages),
            layout=layout,
            physical_page_capacity=None,
            prefix_spans=[(0, 3)],
@@ -91,7 +89,7 @@ class TestComposePlan(unittest.TestCase):
        layout = self._layout()
        slot_logical_pages = torch.tensor([1, 2, 3, 4], dtype=torch.int64)
        kwargs = dict(
-            slot_logical_pages=slot_logical_pages,
+            slot_remap=SimpleNamespace(slot_logical_pages=slot_logical_pages),
            layout=layout,
            physical_page_capacity=None,
            prefix_spans=[(0, 2)],
@@ -109,6 +107,110 @@ class TestComposePlan(unittest.TestCase):
        self.assertIsNot(plan_a, plan_d)


+class TestComposePlanSymm(unittest.TestCase):
+    def test_remote_current_lists_exclude_self_written_pages(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)
+        plan = get_or_build_compose_plan(
+            slot_remap=SimpleNamespace(slot_logical_pages=slot_logical_pages),
+            layout=layout,
+            physical_page_capacity=None,
+            prefix_spans=[(0, 2)],
+            current_spans=[(2, 5)],
+            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.
+        self.assertEqual(plan.current_dense_pages.tolist(), [3, 4, 5])
+        self.assertEqual(plan.remote_current_writer_ranks.tolist(), [0])
+        self.assertEqual(plan.remote_current_dense_pages.tolist(), [4])
+
+    def test_writer_count_mismatch_fails_fast(self):
+        layout = CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=0)
+        with self.assertRaises(ValueError):
+            get_or_build_compose_plan(
+                slot_remap=SimpleNamespace(
+                    slot_logical_pages=torch.tensor([1, 2, 3], dtype=torch.int64)
+                ),
+                layout=layout,
+                physical_page_capacity=None,
+                prefix_spans=[(0, 1)],
+                current_spans=[(1, 3)],
+                kind="token_kv",
+                current_page_writer_ranks=[0],  # 2 current pages
+            )
+
+    def test_writer_ranks_helper_matches_compute_owner_per_request(self):
+        from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+            build_batch_current_page_writer_ranks,
+        )
+        from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
+            build_in_seq_page_compute_owners,
+        )
+
+        prefix_lens = [640, 320]
+        extend_lens = [95, 130]
+        writers = build_batch_current_page_writer_ranks(
+            prefix_lens_cpu=prefix_lens,
+            extend_lens_cpu=extend_lens,
+            page_size=64,
+            cp_size=8,
+        )
+        expected = []
+        for p, e in zip(prefix_lens, extend_lens):
+            expected.extend(
+                int(o)
+                for o in build_in_seq_page_compute_owners(
+                    extend_len=e,
+                    extend_prefix_len=p,
+                    page_size=64,
+                    cp_size=8,
+                )
+            )
+        self.assertEqual(writers, expected)
+
+    def test_maybe_build_gates_on_env_and_page_alignment(self):
+        from types import SimpleNamespace
+
+        from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+            maybe_build_current_page_writer_ranks,
+        )
+
+        layout = CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=0)
+        fb_aligned = SimpleNamespace(
+            nsa_cp_metadata=SimpleNamespace(page_aligned=True)
+        )
+        fb_unaligned = SimpleNamespace(
+            nsa_cp_metadata=SimpleNamespace(page_aligned=False)
+        )
+        kwargs = dict(
+            prefix_lens_cpu=[640],
+            extend_lens_cpu=[95],
+            page_size=64,
+            layout=layout,
+        )
+        with envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(
+            True
+        ), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(True):
+            self.assertIsNotNone(
+                maybe_build_current_page_writer_ranks(
+                    forward_batch=fb_aligned, **kwargs
+                )
+            )
+            self.assertIsNone(
+                maybe_build_current_page_writer_ranks(
+                    forward_batch=fb_unaligned, **kwargs
+                )
+            )
+        with envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(False):
+            self.assertIsNone(
+                maybe_build_current_page_writer_ranks(
+                    forward_batch=fb_aligned, **kwargs
+                )
+            )
+
+
 class TestComposeArena(unittest.TestCase):
    def test_parity_halves_reset_per_layer_and_offsets_are_deterministic(self):
        arena = CpComposeArena(torch.device("cpu"))
@@ -2166,9 +2166,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
        per-span range reduce, never a whole-buffer reduce)."""
        from sglang.srt.environ import envs
        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
-        from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
-            reset_compose_plan_cache,
-        )
        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout

        page_size = 4
@@ -2185,7 +2182,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
            layout=layout,
            page_size=page_size,
        )
-        reset_compose_plan_cache()

        captured = {}
        test_case = self
@@ -2695,9 +2691,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
        """compose_v2 contract for the index buffer (see token-KV twin)."""
        from sglang.srt.environ import envs
        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
-        from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
-            reset_compose_plan_cache,
-        )
        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout

        page_size = 4
@@ -2717,7 +2710,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
            dtype=torch.uint8,
        )
        current_scale = torch.tensor([[1.25], [2.5]], dtype=torch.float32)
-        reset_compose_plan_cache()

        captured = {}
        test_case = self