Harden symm compose per review: prefetch gate, round parity, hot-path CPU

Addresses the opus review of the Step A/B compose stack:

- FAIL FAST on symm + prefetcher coexistence: the MLA/index prefetchers
  issue their own per-span collectives and bypass the symm exchange, so
  letting them coexist would make COMPOSE_SYMM a silent no-op on every
  prefetch hit (measured "parity" that never ran). maybe_build_current_
  page_writer_ranks now raises when a prefetcher is attached or
  SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH is set.
- Arena parity now derives from a monotonic round epoch instead of raw
  layer_id: EAGLE draft layers reuse decoder layer ids, so id-parity would
  stop alternating halves (draft L0 -> next forward target L0 lands on the
  same half) and repeated-id rounds would keep appending into one half.
  begin_round(layer_id, kind) starts a new round (other half, offsets
  reset) when the id changes OR a kind repeats; regression test included.
- Per-batch writer-ranks cache on the forward batch and a presence flag in
  the plan key: previously the ~bs x current-pages writer list was rebuilt
  AND tuple-hashed on every layer per call site, pure launch-path waste.
- Rank-uniformity invariant documented at _symm_exchange_current_pages
  (any per-rank gate must go through a group agreement first) and an
  actionable arena-overflow message naming SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.

Deferred follow-ups (documented in the design doc): factor the two near-
clone v2 helpers, single barrier per F-layer (~308us/batch), persistent
device-side peer-ptr tables for the gather.

182 targeted tests + 8-rank GPU byte-exactness (v2 and symm) re-validated
on g0034 syh-dev-new.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-06-11 20:26:27 +00:00
co-authored by Claude Fable 5
parent c55e406176
commit b356773d2f
3 changed files with 102 additions and 29 deletions
@@ -212,35 +212,53 @@ class TestComposePlanSymm(unittest.TestCase):
class TestComposeArena(unittest.TestCase):
def test_parity_halves_reset_per_layer_and_offsets_are_deterministic(self):
def test_round_halves_alternate_and_offsets_are_deterministic(self):
arena = CpComposeArena(torch.device("cpu"))
arena.begin_layer(0, 0)
a0 = arena.acquire(parity=0, nbytes=1000)
a1 = arena.acquire(parity=0, nbytes=500)
p0 = arena.begin_round(0, "index")
a0 = arena.acquire(parity=p0, nbytes=1000)
self.assertEqual(arena.begin_round(0, "kv"), p0) # same round
a1 = arena.acquire(parity=p0, nbytes=500)
# Aligned carve: second buffer starts at align(1000)=1024 past the
# first (offsets are slab-relative; absolute alignment is the
# allocator's business).
self.assertEqual(a1.data_ptr() - a0.data_ptr(), 1024)
arena.begin_layer(1, 1)
b0 = arena.acquire(parity=1, nbytes=1000)
p1 = arena.begin_round(1, "index")
self.assertNotEqual(p1, p0)
b0 = arena.acquire(parity=p1, nbytes=1000)
self.assertNotEqual(b0.data_ptr(), a0.data_ptr())
# Layer 2 reuses parity 0 at the SAME offsets (determinism).
arena.begin_layer(0, 2)
c0 = arena.acquire(parity=0, nbytes=1000)
c1 = arena.acquire(parity=0, nbytes=500)
# Layer 2 reuses the first half at the SAME offsets (determinism).
p2 = arena.begin_round(2, "index")
self.assertEqual(p2, p0)
c0 = arena.acquire(parity=p2, nbytes=1000)
c1 = arena.acquire(parity=p2, nbytes=500)
self.assertEqual(c0.data_ptr(), a0.data_ptr())
self.assertEqual(c1.data_ptr(), a1.data_ptr())
def test_repeated_layer_id_starts_a_new_round_on_other_half(self):
# EAGLE draft reuses decoder layer ids: target L77 -> draft L0 ->
# (next forward) target L0. Each repeat of (id, kind) must rotate to
# the OTHER half — raw layer_id parity would stop alternating.
arena = CpComposeArena(torch.device("cpu"))
p_draft = arena.begin_round(0, "kv")
d0 = arena.acquire(parity=p_draft, nbytes=256)
p_next = arena.begin_round(0, "kv") # same id, same kind -> new round
n0 = arena.acquire(parity=p_next, nbytes=256)
self.assertNotEqual(p_next, p_draft)
self.assertNotEqual(n0.data_ptr(), d0.data_ptr())
p_third = arena.begin_round(0, "kv")
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"))
arena.begin_layer(0, 0)
arena.acquire(parity=0, nbytes=64)
parity = arena.begin_round(0, "kv")
arena.acquire(parity=parity, nbytes=64)
arena._registered = True
arena.begin_layer(0, 1)
parity = arena.begin_round(1, "kv")
with self.assertRaises(RuntimeError):
arena.acquire(parity=0, nbytes=arena._half_bytes + 1)
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):