feat: disaggregated serving-role loops over a real TCP transport (B3a) - #190
Merged
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 onServingCoordinatorand runs them over a realTcpTransport.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 reusesprefill_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 localhostTcpTransport(decode loop as a concurrentspawn_localtask 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.Gate: cold clippy
-D warnings --all-targetson both feature sets (metal,accelerateand+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|decodeprocess drives a role loop) and the two-processRequestRouter/StreamBridgeHTTP path are deferred to B3b, where--prefill-peers/--decode-peerssupply the peer addresses a single process otherwise lacks. The role-loop futures and channels are shaped for that wiring.Part of #126.