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:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user