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:
@@ -395,7 +395,7 @@ class MooncakeKVManager(CommonKVManager):
|
||||
self.engine, mooncake_session_id, get_blocks, num_layers=n_layers
|
||||
)
|
||||
|
||||
def register_per_layer_transfer(self, room, page_indices) -> bool:
|
||||
def register_per_layer_transfer(self, room, page_indices, chunk_key=0) -> bool:
|
||||
"""Lever A: before the forward, build + register a per-layer transfer context
|
||||
for `room` so the per-layer notifier overlaps its main-KV transfer with the
|
||||
forward. Reuses send()'s CP filter exactly (no re-derivation). Scoped to
|
||||
@@ -428,10 +428,11 @@ class MooncakeKVManager(CommonKVManager):
|
||||
info.mooncake_session_id, owned_pages, dst_indices
|
||||
)
|
||||
if ctx is not None:
|
||||
mgr.register(room, ctx)
|
||||
mgr.register(room, ctx, chunk_key=chunk_key)
|
||||
logger.info(
|
||||
"[CP_PER_LAYER_TRANSFER] registered room=%s owned_pages=%d",
|
||||
"[CP_PER_LAYER_TRANSFER] registered room=%s chunk=%s owned_pages=%d",
|
||||
room,
|
||||
chunk_key,
|
||||
len(owned_pages),
|
||||
)
|
||||
return True
|
||||
|
||||
Reference in New Issue
Block a user