feat(disagg): wire per-layer overlapped transfer end-to-end (lever A, A3-step3)
Complete the lever-A hot-path integration behind SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER: - PerLayerTransferContext: add num_layers + a completion event so finish() waits until ALL layers are processed before wait_batch_transfers (never races ahead of the worker threads and silently drops in-flight layers); times out to FAILURE. - PerLayerTransferManager: has_room (for the swap) + drop (abort/failure drain so outstanding RDMA finishes before pages are reclaimed). - MooncakeKVManager.register_per_layer_transfer: build + register a context before the forward, reusing send()'s exact CP filter (no re-derivation). - transfer_worker: when a room is per-layer-active, wait those transfers (finish) instead of the monolithic send_kvcache -- no double-send; aux/state/completion unchanged. The skip path drops the context on abort/failure. - prefill scheduler: _register_per_layer_transfers(batch) before run_batch, scoped (first impl) to single-forward, no-cached-prefix requests (the notifier transfers forward-written pages; chunked/cached-prefix are HiCache-loaded -> lever B). Unit-tested (25 cases). e2e output-equality + TTFT verification next. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -49,8 +49,8 @@ def _blocks(layer_id):
|
||||
return ([base], [base + 5000], [64])
|
||||
|
||||
|
||||
def _ctx(engine, blocks=_blocks):
|
||||
return PerLayerTransferContext(engine, "sess", blocks)
|
||||
def _ctx(engine, blocks=_blocks, num_layers=8):
|
||||
return PerLayerTransferContext(engine, "sess", blocks, num_layers=num_layers)
|
||||
|
||||
|
||||
class TestPerLayerTransferContext(unittest.TestCase):
|
||||
@@ -81,7 +81,7 @@ class TestPerLayerTransferContext(unittest.TestCase):
|
||||
|
||||
def test_finish_waits_all_and_returns_zero(self):
|
||||
eng = _FakeEngine()
|
||||
ctx = _ctx(eng)
|
||||
ctx = _ctx(eng, num_layers=3)
|
||||
for L in range(3):
|
||||
ctx.submit_layer(L, _FakeEvent())
|
||||
self.assertEqual(ctx.finish(), 0)
|
||||
@@ -89,19 +89,38 @@ class TestPerLayerTransferContext(unittest.TestCase):
|
||||
|
||||
def test_submit_failure_marks_failed_and_stops(self):
|
||||
eng = _FakeEngine(submit_rets=[101, -1, 103])
|
||||
ctx = _ctx(eng)
|
||||
ctx = _ctx(eng, num_layers=3)
|
||||
for L in range(3):
|
||||
ctx.submit_layer(L, None)
|
||||
self.assertTrue(ctx.failed)
|
||||
self.assertEqual(len(eng.submits), 2) # stopped after the -1
|
||||
self.assertEqual(len(eng.submits), 2) # stopped submitting after the -1
|
||||
self.assertEqual(ctx.finish(), -1)
|
||||
|
||||
def test_finish_propagates_wait_failure(self):
|
||||
eng = _FakeEngine(wait_ret=-9)
|
||||
ctx = _ctx(eng)
|
||||
ctx = _ctx(eng, num_layers=1)
|
||||
ctx.submit_layer(0, None)
|
||||
self.assertEqual(ctx.finish(), -9)
|
||||
|
||||
def test_finish_times_out_to_failure_when_layers_missing(self):
|
||||
# only 1 of 3 layers ever processed -> finish must NOT silently succeed
|
||||
eng = _FakeEngine()
|
||||
ctx = _ctx(eng, num_layers=3)
|
||||
ctx.submit_layer(0, None)
|
||||
self.assertEqual(ctx.finish(timeout=0.05), -1) # incomplete -> failure
|
||||
|
||||
def test_skipped_and_failed_layers_still_count_toward_completion(self):
|
||||
# layer 1 has no blocks (skip), layer 2 submit fails; all 3 still complete
|
||||
eng = _FakeEngine(submit_rets=[101, -1])
|
||||
|
||||
def blocks(L):
|
||||
return None if L == 1 else ([L], [L + 10], [8])
|
||||
|
||||
ctx = _ctx(eng, blocks, num_layers=3)
|
||||
for L in range(3):
|
||||
ctx.submit_layer(L, None)
|
||||
self.assertEqual(ctx.finish(timeout=0.5), -1) # completes (no hang), reports failure
|
||||
|
||||
|
||||
class TestBuildLayerBlocks(unittest.TestCase):
|
||||
def test_addresses_and_lengths(self):
|
||||
|
||||
Reference in New Issue
Block a user