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>
This commit is contained in:
2026-06-11 22:09:29 +00:00
co-authored by Claude Fable 5
parent e8e5edb230
commit d63fbd4d79
4 changed files with 362 additions and 147 deletions
@@ -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()