feat(disagg): per-layer block-address computation (lever A, A3-step2 core)
Add build_layer_blocks: the pure per-layer transfer-address computation (src/dst addrs + lengths for layer L's owned page blocks), the core of the context's get_blocks closure. Mirrors the mooncake set_transfer_blocks math; the page index lists are identical across layers, so only the per-layer base ptr + item_len change. Unit-tested (3 cases incl. the cross-layer invariant). 27 per-layer/async tests green total. The remaining A3 step assembles get_blocks from the scheduler's per-request data (transfer_infos dst indices + decode_kv_args_table dst ptrs + out_cache_loc src + CP owner filter) before run_batch, and reconciles finish() with send_kv_chunk — the hot-path integration, to be verified by the bitwise-equivalence harness. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -25,6 +25,31 @@ from typing import Callable, List, Optional, Tuple
|
||||
BlocksFn = Callable[[int], Optional[Tuple[List[int], List[int], List[int]]]]
|
||||
|
||||
|
||||
def build_layer_blocks(
|
||||
src_ptr: int,
|
||||
dst_ptr: int,
|
||||
item_len: int,
|
||||
prefill_blocks,
|
||||
dst_blocks,
|
||||
) -> Tuple[List[int], List[int], List[int]]:
|
||||
"""Per-layer transfer addresses for one layer, given the request's page blocks
|
||||
(each block a contiguous run of page ids: prefill_blocks[i] <-> dst_blocks[i]).
|
||||
Mirrors the mooncake set_transfer_blocks math exactly: the page index lists are
|
||||
identical across layers, so only the per-layer base ptr + item_len change.
|
||||
|
||||
src_ptr/dst_ptr: layer L's kv_buffer base (prefill src, decode dst).
|
||||
item_len: bytes per page for layer L. Returns (src_addrs, dst_addrs, lengths).
|
||||
"""
|
||||
src_addrs: List[int] = []
|
||||
dst_addrs: List[int] = []
|
||||
lengths: List[int] = []
|
||||
for pblock, dblock in zip(prefill_blocks, dst_blocks):
|
||||
src_addrs.append(src_ptr + int(pblock[0]) * item_len)
|
||||
dst_addrs.append(dst_ptr + int(dblock[0]) * item_len)
|
||||
lengths.append(item_len * len(pblock))
|
||||
return src_addrs, dst_addrs, lengths
|
||||
|
||||
|
||||
class PerLayerTransferContext:
|
||||
"""Tracks the per-layer async transfers for one prefill request."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user