E2E diagnostic was precise: "finish TIMEOUT processed=78/79", submit_failed=0 —
the per-layer notifier fires 78x but kv_data_ptrs has 79 layers (the 79th is the
MTP/nextn EAGLE buffer: present in kv_data_ptrs so the monolithic send moves it,
but it doesn't fire the per-layer hook). The old completion required all num_layers,
so it both hung to the timeout AND would silently miss that layer's KV (corruption).
Redesign: gate completion on the ACTUAL enqueued count (note_enqueued), and in
finish() SYNCHRONOUSLY transfer any layers the notifier didn't fire for (KV is fully
written post-forward, no event needed). Net: the per-layer set == kv_data_ptrs,
byte-identical to the monolithic send; robust to any model firing fewer hooks than
KV buffers. The fired layers stay overlapped with the forward.
Unit tests updated (25 pass): fallback transfers the missed layers; submit failures
still report -1.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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>
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>
Add PerLayerTransferManager: owns a worker thread-pool + the active per-request
contexts for the current forward batch, registered as a KV-pool notifier.
on_layer_end (forward thread) records ONE CUDA event on the compute stream and
enqueues (ctx, layer, event) per active context; workers do the event-wait +
async submit OFF the forward thread. finish(room) waits the request's transfers.
event_factory/current_stream injected for unit-testability (no CUDA needed).
Unit-tested (5 manager cases, 11 total in test_cp_per_layer_transfer.py):
per-active-ctx enqueue with the event recorded on the stream, worker-step submit +
mark-failed-on-exception, finish pop + idempotency, no-op when idle. The A3
scheduler wiring (notifier registration + setup-before-forward + finish-after +
no-double-send reconcile) is the remaining hot-path step; plan locked in
lever-a-implementation-plan.md.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add PerLayerTransferContext: the per-request coordinator for overlapped per-layer
KV transfer. submit_layer(layer_id, event) waits the layer's CUDA write event (on
a background thread, never the compute stream) before async-submitting that
layer's RDMA via the G1 path, so the transfer never reads a layer before its
write kernel finished — the core correctness invariant for the forward overlap.
finish() waits all accumulated batch_ids; idempotent per layer; fails closed.
Unit-tested (test_cp_per_layer_transfer.py, 6 cases): event-wait-before-submit
ordering, idempotency, empty-layer skip, finish-waits-all, submit-failure stop,
wait-failure propagation. The scheduler/notifier wiring (A2-wiring + A3) builds
on this.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
leaveletandClaude Opus 4.8