-
Notifications
You must be signed in to change notification settings - Fork 105
feat(qwen35): keep the auto prefill budget running under stream overlap #1034
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 3 commits
44ad2ee
8919e81
4f9d087
c2e269e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -67,6 +67,7 @@ pub(super) fn choose_prefill_budget( | |
| base_budget: usize, | ||
| active: &[ActiveDecodeState], | ||
| prefilling: &[PrefillQueueState], | ||
| decode_overlap: bool, | ||
| ) -> usize { | ||
| assert!( | ||
| base_budget > 0, | ||
|
|
@@ -76,9 +77,13 @@ pub(super) fn choose_prefill_budget( | |
| return base_budget; | ||
| } | ||
|
|
||
| if active | ||
| .iter() | ||
| .any(|req| req.remaining_tokens() <= DECODE_FINISH_WINDOW_TOKENS) | ||
| // With stream overlap the finishing request's last tokens keep ticking on | ||
| // the decode stream while the chunk runs aside, so the full prefill | ||
| // deferral would only delay prefill without buying decode latency. | ||
| if !decode_overlap | ||
| && active | ||
| .iter() | ||
| .any(|req| req.remaining_tokens() <= DECODE_FINISH_WINDOW_TOKENS) | ||
| { | ||
| return 0; | ||
| } | ||
|
|
@@ -312,7 +317,13 @@ mod tests { | |
| }]; | ||
|
|
||
| assert_eq!( | ||
| choose_prefill_budget(Qwen35SchedulerPolicy::Off, 1024, &active, &prefilling), | ||
| choose_prefill_budget( | ||
| Qwen35SchedulerPolicy::Off, | ||
| 1024, | ||
| &active, | ||
| &prefilling, | ||
| false | ||
| ), | ||
| 1024, | ||
| "off keeps the fixed chunk budget" | ||
| ); | ||
|
|
@@ -329,7 +340,13 @@ mod tests { | |
| }]; | ||
|
|
||
| assert_eq!( | ||
| choose_prefill_budget(Qwen35SchedulerPolicy::Auto, 1024, &active, &prefilling), | ||
| choose_prefill_budget( | ||
| Qwen35SchedulerPolicy::Auto, | ||
| 1024, | ||
| &active, | ||
| &prefilling, | ||
| false | ||
| ), | ||
| 1024, | ||
| "auto preserves --max-prefill-tokens as a hard per-step cap" | ||
| ); | ||
|
|
@@ -346,7 +363,13 @@ mod tests { | |
| }]; | ||
|
|
||
| assert_eq!( | ||
| choose_prefill_budget(Qwen35SchedulerPolicy::Auto, 1024, &active, &prefilling), | ||
| choose_prefill_budget( | ||
| Qwen35SchedulerPolicy::Auto, | ||
| 1024, | ||
| &active, | ||
| &prefilling, | ||
| false | ||
| ), | ||
| 1024, | ||
| "standard serving cells with long outputs keep the fixed chunk path" | ||
| ); | ||
|
|
@@ -363,7 +386,13 @@ mod tests { | |
| }]; | ||
|
|
||
| assert_eq!( | ||
| choose_prefill_budget(Qwen35SchedulerPolicy::Auto, 1024, &active, &prefilling), | ||
| choose_prefill_budget( | ||
| Qwen35SchedulerPolicy::Auto, | ||
| 1024, | ||
| &active, | ||
| &prefilling, | ||
| false | ||
| ), | ||
| 512, | ||
| "auto may shrink the final chunk but never expands beyond the configured cap" | ||
| ); | ||
|
|
@@ -386,10 +415,45 @@ mod tests { | |
| }]; | ||
|
|
||
| assert_eq!( | ||
| choose_prefill_budget(Qwen35SchedulerPolicy::Auto, 1024, &active, &prefilling), | ||
| choose_prefill_budget( | ||
| Qwen35SchedulerPolicy::Auto, | ||
| 1024, | ||
| &active, | ||
| &prefilling, | ||
| false | ||
| ), | ||
| 0, | ||
| "a near-finished active request gets a decode-priority tick before a long prefill" | ||
| ); | ||
| } | ||
|
|
||
| #[test] | ||
| fn adaptive_prefill_budget_keeps_prefill_running_under_stream_overlap() { | ||
| let active = [ | ||
| ActiveDecodeState { | ||
| generated_count: 252, | ||
| max_tokens: 256, | ||
| }, | ||
| ActiveDecodeState { | ||
| generated_count: 16, | ||
| max_tokens: 4096, | ||
| }, | ||
| ]; | ||
| let prefilling = [PrefillQueueState { | ||
| remaining_tokens: 4096, | ||
| }]; | ||
|
|
||
| assert_eq!( | ||
| choose_prefill_budget( | ||
| Qwen35SchedulerPolicy::Auto, | ||
| 1024, | ||
| &active, | ||
| &prefilling, | ||
| true | ||
| ), | ||
| 1024, | ||
| "with stream overlap the finishing window keeps prefill on the prefill stream instead of deferring it" | ||
| ); | ||
| assert!( | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. no blocking but the first assertion correctly exercises the changed branch: a decoder has four tokens left, overlap is enabled, and the budget remains 1024 instead of returning zero. The final assertion supplies a separate empty queue to build_next_plan(), so it does not use that budget or establish the overlap behaviour, its “zero prefill budget” explanation contradicts this case's expected result. |
||
| matches!( | ||
| build_next_plan::<Pending>(true, vec![]), | ||
|
|
@@ -410,7 +474,13 @@ mod tests { | |
| }]; | ||
|
|
||
| assert_eq!( | ||
| choose_prefill_budget(Qwen35SchedulerPolicy::Auto, 1024, &active, &prefilling), | ||
| choose_prefill_budget( | ||
| Qwen35SchedulerPolicy::Auto, | ||
| 1024, | ||
| &active, | ||
| &prefilling, | ||
| false | ||
| ), | ||
| 0, | ||
| "decode-priority applies before even a final prefill chunk when an active request is finishing" | ||
| ); | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The new
1828 → 1264/873 → 992figures were measured fromfb16fc15plus this patch according to the commit's benchmark contract, but here they remain under the source binding on lines 12–13 that attributes the combination evidence to70a600b7 + this change. Those revisions straddle earlier overlap-policy work, so this prevents a reproducible same-context A/B and can cause readers to attribute the result to the wrong build; record the exact source for the new run separately.AGENTS.md reference: AGENTS.md:L121-L123
Useful? React with 👍 / 👎.