From 6bd5d5760c802b99b651c231ebb97af596a665fe Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 7 Jun 2026 00:41:28 +0000 Subject: [PATCH] feat(disagg): per-layer block-address computation (lever A, A3-step2 core) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../disaggregation/cp_per_layer_transfer.py | 25 ++++++++++++++++++ .../test_cp_per_layer_transfer.py | 26 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py index 301a1a901..4766ca1bc 100644 --- a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py +++ b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py @@ -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.""" diff --git a/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py b/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py index 1802ea54f..622eab85a 100644 --- a/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py +++ b/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py @@ -103,6 +103,32 @@ class TestPerLayerTransferContext(unittest.TestCase): self.assertEqual(ctx.finish(), -9) +class TestBuildLayerBlocks(unittest.TestCase): + def test_addresses_and_lengths(self): + from sglang.srt.disaggregation.cp_per_layer_transfer import build_layer_blocks + + src, dst, lens = build_layer_blocks( + 1000, 5000, 64, [[0, 1, 2], [5, 6]], [[10, 11, 12], [20, 21]] + ) + self.assertEqual(src, [1000 + 0 * 64, 1000 + 5 * 64]) + self.assertEqual(dst, [5000 + 10 * 64, 5000 + 20 * 64]) + self.assertEqual(lens, [64 * 3, 64 * 2]) # item_len * run length + + def test_empty(self): + from sglang.srt.disaggregation.cp_per_layer_transfer import build_layer_blocks + + self.assertEqual(build_layer_blocks(1, 2, 8, [], []), ([], [], [])) + + def test_only_base_ptr_changes_across_layers(self): + from sglang.srt.disaggregation.cp_per_layer_transfer import build_layer_blocks + + s0, d0, l0 = build_layer_blocks(1000, 2000, 64, [[3]], [[7]]) + s1, d1, l1 = build_layer_blocks(9000, 8000, 64, [[3]], [[7]]) + self.assertEqual(s0, [1000 + 3 * 64]) + self.assertEqual(s1, [9000 + 3 * 64]) + self.assertEqual(l0, l1) # lengths identical across layers (the invariant) + + class _MockCtx: def __init__(self): self.submitted = []