Skip to content

feat: in-process paged KV serving-role handoff mechanism#187

Merged
inureyes merged 1 commit into
mainfrom
feature/issue-126-b1-handoff
Jun 9, 2026
Merged

feat: in-process paged KV serving-role handoff mechanism#187
inureyes merged 1 commit into
mainfrom
feature/issue-126-b1-handoff

Conversation

@inureyes

@inureyes inureyes commented Jun 9, 2026

Copy link
Copy Markdown
Member

Summary

In-process building block for disaggregated serving-role KV handoff (#126 step B1). This wires the #125 paged KV serde layer to the transport layer, which previously never referenced each other in production, and proves the mechanism byte-identical against real models.

What it adds

New distributed::disaggregated::handoff_impl:

  • extract_sequence_handoff (prefill side) serializes a finished pool-backed sequence (dense metadata, paged block table, and the referenced pool block contents) into a single wire frame.
  • ingest_sequence_handoff (decode side) validates an incoming frame, allocates a fresh pool-backed sequence, and reconstructs its KV onto local pool blocks anchored to the decode model's real block geometry, releasing the slot on restore failure so a rejected handoff leaks nothing.
  • probe_block_geometry derives the decode model's exact (n_kv_heads, head_dim, dtype) with a one-token probe through a throwaway pool-backed sequence. This captures the runtime KV dtype (bf16 for a 4-bit checkpoint, not fp16), which no static model accessor exposes, with no per-model code.
  • send_handoff_payload / recv_handoff_payload move a frame over any Transport as a tagged message (the async serde/transport bridge).

BatchScheduler gains thin extract_sequence_handoff / ingest_sequence_handoff hooks plus a probe-once geometry cache. These are additive (#[allow(dead_code)]): the live caller, a serving-role serve loop, lands in a later step, matching the additive-then-wired pattern used across this epic.

Verification

Orchestrator-local (Metal cannot run on CI runners):

  • Real-model tests/paged_handoff_parity.rs (requires test-utils): qwen3-0.6b (28 layers) and llama3-1b (16 layers) over the full prefill-extract, MockTransport wire, probe, decode-ingest path produce a byte-identical 16-token decode versus a single-node run.
  • In-crate handoff_impl_tests.rs: transport bridge round-trip plus two rejection paths.
  • Cold clippy -D warnings (metal,accelerate, with and without test-utils) and fmt clean.
cargo test --test paged_handoff_parity --release \
    --features metal,accelerate,test-utils -- --ignored --nocapture --test-threads=1

Correctness rests on #125's proven serialize/restore path plus the Some(geometry) anchor, which only rejects a mismatch and does not alter a valid restore.

Design notes

  • Geometry anchor by runtime probe. The dims are config-derivable but the KV dtype is the runtime activation dtype with no static accessor, so a single one-token probe reads the exact geometry off the gathered window. Model-agnostic and correct for bf16 checkpoints.
  • Async bridge primitive, not the sync trait. The scaffolding HandoffProtocol trait is synchronous and PrefillHandoff-shaped, which fights the async Transport (it would force a block_on) and only covers the send half. The async byte bridge here is the primitive the real serve loop should build on.

Part of #126

Add the in-process building block for disaggregated serving-role KV
handoff (#126 step B1): the seam that joins the #125 paged KV serde
layer to the transport layer, which previously never referenced each
other in production.

New `distributed::disaggregated::handoff_impl`:
- `extract_sequence_handoff` serializes a finished pool-backed sequence
  (dense metadata, paged block table, and the referenced pool block
  contents) into a single wire frame.
- `ingest_sequence_handoff` validates an incoming frame, allocates a
  fresh pool-backed sequence, and reconstructs its KV onto local pool
  blocks anchored to the decode model's real block geometry, releasing
  the slot on restore failure so a rejected handoff leaks nothing.
- `probe_block_geometry` derives the decode model's exact
  (n_kv_heads, head_dim, dtype) with a one-token probe through a
  throwaway pool-backed sequence. This captures the runtime KV dtype
  (e.g. bf16 for a 4-bit checkpoint, not fp16), which no static model
  accessor exposes, with no per-model code.
- `send_handoff_payload` / `recv_handoff_payload` move a frame over any
  `Transport` as a tagged message (the async serde/transport bridge).

`BatchScheduler` gains thin `extract_sequence_handoff` /
`ingest_sequence_handoff` hooks plus a probe-once geometry cache. They
are additive: a live caller (the serving-role serve loop) lands in a
later step, matching the additive-then-wired pattern used across this
epic.

Verified byte-identical to a single-node run for qwen3 and llama3 over
the full prefill-extract, transport, probe, decode-ingest path
(`tests/paged_handoff_parity.rs`, requires the test-utils feature),
plus in-crate transport bridge round-trip and rejection tests.

The synchronous scaffolding `HandoffProtocol` trait is intentionally
left untouched: it predates the async `Transport` and would force a
`block_on`, so the async byte bridge here is the primitive the real
serve loop should build on.

Part of #126
@inureyes inureyes added type:enhancement New features, capabilities, or significant additions status:review Under review labels Jun 9, 2026
@inureyes
inureyes merged commit 704eece into main Jun 9, 2026
5 checks passed
@inureyes
inureyes deleted the feature/issue-126-b1-handoff branch June 9, 2026 03:59
@inureyes inureyes self-assigned this Jun 10, 2026
@inureyes inureyes added this to the 0.2 milestone Jun 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

status:review Under review type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant