Skip to content

fix: support chunked-prefill prompts in the disaggregated serving handoff#213

Merged
inureyes merged 3 commits into
mainfrom
fix/issue-197-handoff-chunked-prefill
Jun 11, 2026
Merged

fix: support chunked-prefill prompts in the disaggregated serving handoff#213
inureyes merged 3 commits into
mainfrom
fix/issue-197-handoff-chunked-prefill

Conversation

@inureyes

Copy link
Copy Markdown
Member

Closes #197

Problem

BatchScheduler::prefill_request_for_handoff and prefill_text_request_for_handoff rejected 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_handoff now mirrors the execute_prefill dispatch: prompts longer than the chunk size go through the standard chunked machinery (start_chunked_prefill followed by continue_chunked_prefill driven to completion) instead of bailing. The final chunk samples the first token via finish_prefill exactly 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 existing execute_full_prefill path unchanged.
  • The chunked-prefill guard in prefill_text_request_for_handoff is 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).
  • The handoff path is text-only, so the VLM-embeddings full-prefill exemption in the normal dispatch cannot apply.

Verification (M1 Ultra, qwen3-0.6b-4bit, real model)

  • New serving_handoff_parity_chunked_prefill_matches_single_node_qwen3: a scheduler with prefill_chunk_size = 16 prefills the ~50-token shared prompt in 4 chunks through prefill_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.
  • All 5 serving-handoff parity tests green (the existing full-prefill handoff, role-loop TCP parity, and the bind/return smoke tests are unaffected).
  • Cold clippy (-D warnings, both feature sets), cargo fmt.

@inureyes inureyes added type:enhancement New features, capabilities, or significant additions priority:medium Medium priority area:architecture Architecture and code structure changes labels Jun 11, 2026
@inureyes
inureyes merged commit e0def6e into main Jun 11, 2026
@inureyes
inureyes deleted the fix/issue-197-handoff-chunked-prefill branch June 11, 2026 00:45
@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

area:architecture Architecture and code structure changes priority:medium Medium priority type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Disaggregated serving handoff: support chunked-prefill prompts

1 participant