fix: support chunked-prefill prompts in the disaggregated serving handoff#213
Merged
Conversation
3 tasks
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.
Closes #197
Problem
BatchScheduler::prefill_request_for_handoffandprefill_text_request_for_handoffrejected prompts longer than--prefill-chunk-size: the serving-role entry drove a single full prefill and then extracted, so it bailed on long prompts. This limited disaggregated serving to short prompts, while long prompts are exactly where prefill/decode disaggregation pays off.Implementation
prefill_request_for_handoffnow mirrors theexecute_prefilldispatch: prompts longer than the chunk size go through the standard chunked machinery (start_chunked_prefillfollowed bycontinue_chunked_prefilldriven to completion) instead of bailing. The final chunk samples the first token viafinish_prefillexactly like a single-node chunked prefill (including the terminal-first-chunk handling from Crash ([squeeze] Cannot squeeze axis 1 with size 0) on 3rd turn of multi-turn chat when prompt-prefix cache is enabled — Qwen3-Coder-30B-A3B-Instruct #179), so the full prompt's KV is in the pool before the handoff frame is serialized. Short prompts keep the existingexecute_full_prefillpath unchanged.prefill_text_request_for_handoffis removed (it delegates to the now-chunk-capable entry). A defensive guard rejects entry while another chunked prefill is in progress (the serving-role intake is sequential, so this would be a wiring bug).Verification (M1 Ultra, qwen3-0.6b-4bit, real model)
serving_handoff_parity_chunked_prefill_matches_single_node_qwen3: a scheduler withprefill_chunk_size = 16prefills the ~50-token shared prompt in 4 chunks throughprefill_request_for_handoff, ships the frame over the coordinator pair, ingests it, and decodes 16 tokens. The decode stream is byte-identical to an UNCHUNKED single-node run through the same scheduler, which also re-proves chunked == full prefill across the handoff boundary.-D warnings, both feature sets),cargo fmt.