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:
2026-06-07 09:51:05 +00:00
co-authored by Claude Opus 4.8
parent aa6acc9485
commit ae18e3adc8
4 changed files with 179 additions and 28 deletions
@@ -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):