Skip to content

feat: disaggregated serving-role loops over a real TCP transport (B3a) - #190

Merged
inureyes merged 1 commit into
mainfrom
feature/issue-126-b3a-role-loops
Jun 9, 2026
Merged

feat: disaggregated serving-role loops over a real TCP transport (B3a)#190
inureyes merged 1 commit into
mainfrom
feature/issue-126-b3a-role-loops

Conversation

@inureyes

@inureyes inureyes commented Jun 9, 2026

Copy link
Copy Markdown
Member

Summary

Epic #116 step #126 B3a. Adds the prefill/decode serving-role loops that drive the disaggregated KV handoff over a transport, and proves them end to end over a real localhost TCP socket.

B2b landed the scheduler handoff entries (prefill_request_for_handoff, ingest_handoff_as_active, decode_handoff_until_idle) and B2c proved them over an in-process MockTransport. B3a wires those entries into role loops on ServingCoordinator and runs them over a real TcpTransport.

Changes

  • BatchScheduler::prefill_text_request_for_handoff (src/server/batch/scheduler.rs): builds a queued text sequence from raw request parts (prompt tokens, sampling, max_tokens, response channel, cancel flag) and reuses prefill_request_for_handoff. The empty-prompt and chunked-prefill guards run before allocation, so a rejected request leaks no pool state.
  • ServingCoordinator::run_prefill_role / run_decode_role (src/distributed/disaggregated/coordinator.rs): async role loops. Prefill drains requests, prefills each, and ships the extracted frame to the decode peer; decode receives frames, ingests each onto a fresh pool slot, and decodes to completion, streaming tokens on the request channel. The scheduler is borrowed (the coordinator stays model-free for its transport unit tests); the futures are !Send (the scheduler owns the per-process MLX pool), so callers drive them on a current-thread runtime.
  • PrefillRoleRequest / DecodeRoleHandoff: the per-request work item and decode coordination metadata. The KV frame carries the cache, prompt history, and generated tokens; the per-request budget, sampling, and output stream stay with the node holding the client connection.

Verification

Real-model two-node parity test (serving_role_loop_parity_matches_single_node_qwen3, #[ignore]): a qwen3-0.6b prefill node and a separate decode node, each its own model and scheduler, connected over a real localhost TcpTransport (decode loop as a concurrent spawn_local task on a current-thread runtime). The prefill node emits the first token and the decode node the continuation, the same split a router merges for a client; the concatenated stream is byte-identical to a single-node run.

reference text: " \n\n###Answer: 4\n\n###Step 1: 2 + "
handoff text:   " \n\n###Answer: 4\n\n###Step 1: 2 + "
  (prefill first token: " \n\n", decode continuation: "###Answer: 4\n\n###Step 1: 2 + ")

Gate: cold clippy -D warnings --all-targets on both feature sets (metal,accelerate and +test-utils), fmt, and the lib unit suites (3081 passed, 0 failed).

Scope / follow-up (B3b)

Additive only; no existing logic changed. The worker flip (a --node-role prefill|decode process drives a role loop) and the two-process RequestRouter/StreamBridge HTTP path are deferred to B3b, where --prefill-peers/--decode-peers supply the peer addresses a single process otherwise lacks. The role-loop futures and channels are shaped for that wiring.

Part of #126.

Wire the prefill and decode serving-role loops that drive the disaggregated
KV handoff over a transport, and prove them end to end over a real localhost
TCP socket.

B2b added the scheduler handoff entries (prefill_request_for_handoff,
ingest_handoff_as_active, decode_handoff_until_idle) and B2c proved them over
an in-process MockTransport. B3a turns those entries into role loops on
ServingCoordinator and runs them over a real TcpTransport.

- BatchScheduler::prefill_text_request_for_handoff builds a queued text
  sequence from raw request parts and reuses prefill_request_for_handoff, so
  the prefill loop can drive a request from prompt to handoff frame. The
  empty-prompt and chunked-prefill guards run before allocation so a rejected
  request leaks no pool state.
- ServingCoordinator::run_prefill_role / run_decode_role are async loops:
  prefill drains requests, prefills each, and ships the extracted frame to the
  decode peer; decode receives frames, ingests each onto a fresh pool slot, and
  decodes to completion, streaming tokens on the request channel. The scheduler
  is borrowed so the coordinator stays model-free for its transport unit tests;
  the futures are !Send (the scheduler owns the MLX pool), so callers drive them
  on a current-thread runtime.
- PrefillRoleRequest / DecodeRoleHandoff carry the per-request parts and the
  decode coordination metadata. The KV frame carries the cache and tokens; the
  budget, sampling, and output stream stay with the node holding the client
  connection.

Verified with a real-model two-node test: a qwen3-0.6b prefill node and a
separate decode node over a real localhost TcpTransport (decode loop as a
concurrent spawn_local task) stream byte-identical text to a single-node run,
the prefill node emitting the first token and the decode node the continuation.

The worker flip and the two-process RequestRouter/StreamBridge HTTP path land
in B3b, where --prefill-peers/--decode-peers supply the peer a single process
otherwise lacks.

Part of #126.
@inureyes
inureyes merged commit 24114fd into main Jun 9, 2026
5 checks passed
@inureyes
inureyes deleted the feature/issue-126-b3a-role-loops branch June 9, 2026 09:20
@inureyes inureyes self-assigned this Jun 10, 2026
@inureyes inureyes added status:done Completed type:enhancement New features, capabilities, or significant additions priority:medium Medium priority labels 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

priority:medium Medium priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant