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
@@ -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