Fold the symm compose into one barrier-gated mega gather

The publish-variant staging exchange measured ~equal to the default
compact-current AR (90.4 vs 86.5 ms/batch on the traced scenario): the
publish copy and the barrier serialized behind the 0.65 ms prefix
gather ate the transport win the isolated current exchange showed
(0.196 vs 0.354 ms).  Fix the structure instead of the copy: current
rows are now written straight INTO the staging — the fill kernels take
their write destinations solely from page_inverse, so a per-batch
staging-remapped page inverse on the plan retargets them with zero
kernel changes — then cp_symm_barrier, then ONE slot-dense gather
covers prefix pages (pool pointers) and ALL current pages (staging
pointers, including this rank's own) through a concatenated 2*cp
pointer table where current slots carry owner = cp_size + writer and
src = staging slot.  No publish copy, no prefix pre-gather, no second
gather.

The fused fill's loc outputs are dense-geometry-bound, so the token-KV
path computes mixed_locs/staging row indices once per batch (they are
layer-invariant) and the per-layer fill collapses to a single
index_copy_ into the zeroed staging span.

Benchmark (g0033 8xH200, byte-exact, idle-checked): 62.8 ms/batch vs
84.4 default Step A (-26%) and 60.8 ideal; publish variant was 88.0.
151 unit tests; 8-rank GPU byte-exactness vs v2 across 8 layers,
arena on and off.

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

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()

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)