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 09:51:05 +00:00
co-authored by Claude Opus 4.8
parent e12afe8ced
commit beffd3a82f
4 changed files with 82 additions and 29 deletions
@@ -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