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:
@@ -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 = []
|
||||
|
||||
Reference in New Issue
Block a user