Shrink the symm compose region to a compact current-page staging

Peers only ever read the CURRENT pages of a rank's compose output — the
prefix comes straight from the IPC-registered KV pool — so the symm
region does not need to hold the whole dense buffer (pool-bound ~2.5 GB
double-buffered slab).  It now holds one round of current pages in
merged-span order (extend-cap-bound, ~58-100 MB), and dense buffers
become purely rank-local (plain allocations or the optional local arena;
COMPOSE_SYMM no longer requires COMPOSE_ARENA).

Exchange per compose call: publish my written current pages
dense[page] -> staging[slot i] (slot = the page's batch current index,
identical on every rank, so peers address each other's staging with no
per-batch handshake), cp_symm_barrier, gather peers'
staging[writer][slot] -> dense[page] via the existing src!=dst page
gather.  Reuse safety keeps the parity-half distance-2 argument, now on
the staging.  Capacity sizing comes from the admission caps
(max_total_extend_tokens / max_batch_requests) with a pool-derived
fallback and the SYMM_HEAP_MB override; overflow fails fast
(batch-logical, hence rank-uniform).

Idea credit: laoyao0822's touched-pages-proportional staging
(906ecbe5d4), rebound onto our barrier-gated, group-agreed transport.

Validated on g0033 8xH200: 151 unit tests; 8-rank GPU byte-exactness
vs compose_v2 across 8 layers (arena on and off, parity halves
exercised); benchmark path e (real protocol) byte-exact, current-page
exchange 0.196 ms vs 0.354 ms compact-AR isolated.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

test/manual/test_cp_shared_kv_compose_v2_8rank.py

@@ -6,7 +6,7 @@ every rank with real NCCL collectives and (for v2) real tai-kernel CUDA-IPC
 gathers, and asserts the composed dense buffers and locs are byte-identical
 between the legacy per-span path and compose_v2.
 
-Run inside the g0034 cjy-glm5-new container:
+Run inside the g0033 syh-dev-new container:
   cd /mnt/beegfs/syh/sglang-stable && \
   SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 \
   PYTHONPATH=python:/mnt/beegfs/syh/tai-kernel/python \
@@ -267,38 +267,52 @@ def main() -> None:
             flush=True,
         )
 
-    # ---- Step B: symm exchange (arena + barrier + peer gather, zero NCCL
-    # in the current-page phase). Multiple layers exercise parity halves. ----
+    # ---- Step B: compact symm staging (publish + barrier + peer 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. ----
+    def _check_symm(layer_id: int, tag: str) -> None:
+        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} [{tag}]: 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} [{tag}]: symm dense kv mismatch "
+                f"at rows {bad} (of {int(diff.sum())})"
+            )
+
     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())})"
-                )
+            _check_symm(layer_id, "no-arena")
+        with envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(True):
+            for layer_id in range(4, 8):
+                _check_symm(layer_id, "arena")
     dist.barrier()
     from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
-        get_compose_arena,
+        get_compose_staging,
     )
 
-    assert get_compose_arena(device).symm_ready, "symm slab was not registered"
+    staging = get_compose_staging(device)
+    assert staging.registered, "symm staging was not registered"
     if rank == 0:
+        staged_mb = (
+            2
+            * staging.capacity_pages
+            * (staging.page_nbytes("token_kv") + staging.page_nbytes("index"))
+            / (1 << 20)
+        )
         print(
-            "PASS: symm compose byte-identical to v2 across 4 layers "
-            "(arena registered, barrier + peer gather engaged)",
+            "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)",
             flush=True,
         )
     dist.destroy_process_group()

test/registered/unit/mem_cache/test_cp_shared_kv_compose.py

@@ -1,6 +1,6 @@
-"""Unit tests for cp_shared_kv_compose (Step A compose plan + arena).
+"""Unit tests for cp_shared_kv_compose (compose plan + arena + symm staging).
 
-Registered: CPU CI (no CUDA needed for plan/arena logic).
+Registered: CPU CI (no CUDA needed for plan/arena/staging-layout logic).
 """
 
 import unittest
@@ -12,7 +12,9 @@ from types import SimpleNamespace
 
 from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
     CpComposeArena,
+    CpComposeStaging,
     acquire_dense_buffer,
+    compute_staging_capacity_pages,
     get_or_build_compose_plan,
 )
 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -121,10 +123,15 @@ class TestComposePlanSymm(unittest.TestCase):
             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 rank 3 -> only the page written by rank 0 is remote.  The
+        # staging slot of current page i is i (merged-span order).
         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])
 
     def test_writer_count_mismatch_fails_fast(self):
         layout = CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=0)
@@ -190,9 +197,10 @@ class TestComposePlanSymm(unittest.TestCase):
             page_size=64,
             layout=layout,
         )
+        # Symm is independent of the arena: gating must work with ARENA off.
         with envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(
             True
-        ), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(True):
+        ), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(False):
             self.assertIsNotNone(
                 maybe_build_current_page_writer_ranks(
                     forward_batch=fb_aligned, **kwargs
@@ -251,15 +259,6 @@ class TestComposeArena(unittest.TestCase):
         t0 = arena.acquire(parity=p_third, nbytes=256)
         self.assertEqual(t0.data_ptr(), d0.data_ptr())  # back to half A
 
-    def test_growth_is_forbidden_after_registration(self):
-        arena = CpComposeArena(torch.device("cpu"))
-        parity = arena.begin_round(0, "kv")
-        arena.acquire(parity=parity, nbytes=64)
-        arena._registered = True
-        parity = arena.begin_round(1, "kv")
-        with self.assertRaises(RuntimeError):
-            arena.acquire(parity=parity, nbytes=arena._half_bytes + 1)
-
     def test_acquire_dense_buffer_plain_alloc_when_arena_disabled(self):
         with envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(False):
             buf = acquire_dense_buffer(
@@ -286,5 +285,90 @@ class TestComposeArena(unittest.TestCase):
         self.assertTrue(buf.is_contiguous())
 
 
+class TestComposeStaging(unittest.TestCase):
+    """Layout/parity logic of the compact symm staging (no IPC needed: the
+    region offsets and round parity are pure functions of the registration
+    parameters and the call sequence)."""
+
+    def _staging(self, capacity_pages=4, kv_nbytes=1000, index_nbytes=300):
+        staging = CpComposeStaging(torch.device("cpu"))
+        # Bypass the IPC allocation: install the layout exactly as register()
+        # computes it, backed by a plain CPU slab.
+        staging.capacity_pages = capacity_pages
+        staging.cp_size = 8
+        staging.cp_rank = 0
+        staging._page_nbytes = {"token_kv": kv_nbytes, "index": index_nbytes}
+        half = 0
+        region_in_half = {}
+        for kind in CpComposeStaging.KINDS:
+            region_in_half[kind] = half
+            half += staging._align(capacity_pages * staging._page_nbytes[kind])
+        for parity in (0, 1):
+            for kind in CpComposeStaging.KINDS:
+                staging._region_offsets[(kind, parity)] = (
+                    parity * half + region_in_half[kind]
+                )
+        staging._slab = torch.zeros(2 * half, dtype=torch.uint8)
+        staging.peer_slab_bases = torch.full((8,), 1 << 20, dtype=torch.int64)
+        return staging, half
+
+    def test_regions_are_disjoint_and_parity_halves_do_not_overlap(self):
+        staging, half = self._staging()
+        kv0 = staging.buffer("token_kv", 0)
+        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))
+        # 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)
+
+    def test_peer_region_ptrs_offset_matches_local_layout(self):
+        staging, half = self._staging()
+        base = staging.peer_slab_bases
+        self.assertTrue(
+            torch.equal(staging.peer_region_ptrs("token_kv", 0), base)
+        )
+        self.assertTrue(
+            torch.equal(
+                staging.peer_region_ptrs("index", 1),
+                base + half + staging._align(4 * 1000),
+            )
+        )
+
+    def test_round_parity_alternates_and_shares_round_across_kinds(self):
+        staging, _ = self._staging()
+        p0 = staging.begin_round(0, "index")
+        self.assertEqual(staging.begin_round(0, "token_kv"), p0)
+        p1 = staging.begin_round(1, "index")
+        self.assertNotEqual(p1, p0)
+        # Repeated (id, kind) -> new round (EAGLE draft id reuse).
+        p_again = staging.begin_round(1, "index")
+        self.assertNotEqual(p_again, p1)
+
+    def test_capacity_pages_from_env_override_caps_and_pool(self):
+        kwargs = dict(
+            kv_pool_tokens=64 * 1000,
+            page_size=64,
+            cp_size=8,
+            kv_page_nbytes=41_984,
+            index_page_nbytes=8_448,
+        )
+        with envs.SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.override(101):
+            pages = compute_staging_capacity_pages(**kwargs)
+            self.assertEqual(
+                pages, (101 << 20) // (2 * (41_984 + 8_448))
+            )
+        # No caps set in unit tests -> pool-derived fallback.
+        with envs.SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.override(0):
+            import sglang.srt.server_args as server_args_module
+
+            if getattr(server_args_module, "_global_server_args", None) is None:
+                pages = compute_staging_capacity_pages(**kwargs)
+                self.assertEqual(pages, 1000 * 8 + 512)
+
+
 if __name__ == "__main__":
     unittest.main()