feat(disagg): unify per-layer transfer over chunked prefill (lever A)

Per code review (HiCache load is layer-by-layer & correctly ordered into the compute
stream before the backup hook; current-reuse is a within-forward read that doesn't
rewrite pool pages): unify registration to the EXACT range this forward's
send_kv_chunk transmits — req_to_token[start_send_idx:end_idx], page-floored for a
non-last chunk. Non-chunked = one full range; chunked = one range per chunk. Drop
the is_chunked/start_send_idx skip.

To avoid the review's collision risk (chunk N still finishing when chunk N+1
registers), the manager keys contexts per (room, start_send_idx): _active[room] is a
FIFO deque of (chunk_key, ctx); register dedups the same chunk but appends a new one;
finish(room) pops the FRONT (chunks finish in send order — no chunk key needed in the
transfer_worker); drop drains all the room's chunks; on_layer_end enqueues for all
active chunk contexts (per-ctx note_enqueued dedup keeps each chunk's own events).

28 unit tests pass incl. chunked FIFO + per-chunk dedup.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-07 07:33:46 +00:00
parent e12afe8ced
commit beffd3a82f
4 changed files with 82 additions and 29 deletions

View File

@@ -250,6 +250,25 @@ class TestPerLayerTransferManager(unittest.TestCase):
# boundaries at L=3 (group [0,3]) and L=7 (group [4,7]) -> 2 groups, not 8
self.assertEqual([(s, e) for _, s, e, _ in items], [(0, 3), (4, 7)])
def test_chunked_fifo_separate_contexts_and_dedup(self):
m = _manager(group_size=1)
c1, c2 = _MockCtx(), _MockCtx()
m.register("r1", c1, chunk_key=0) # chunk 1
m.register("r1", c1, chunk_key=0) # SAME chunk re-register -> deduped
m.register("r1", c2, chunk_key=64) # chunk 2 (different start_send_idx)
self.assertTrue(m.has_chunk("r1", 0))
self.assertTrue(m.has_chunk("r1", 64))
self.assertFalse(m.has_chunk("r1", 128))
m.on_layer_end(5) # both chunk contexts active -> both enqueued
self.assertEqual(len(_drain(m._q)), 2)
# finish pops FIFO: chunk 1 (c1) first, then chunk 2 (c2)
self.assertEqual(m.finish("r1"), 0)
self.assertTrue(c1.finished)
self.assertFalse(c2.finished)
self.assertEqual(m.finish("r1"), 0)
self.assertTrue(c2.finished)
self.assertFalse(m.has_room("r1"))
def test_worker_step_calls_submit_group(self):
m = _manager()
c = _MockCtx()