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