fix(scheduler): gate prefill on full batch to protect decode - #1437
Merged
Conversation
Hold new prefills until the waiting queue can fill max_num_batched_tokens, else keep decoding. Prevents fast 补发 from firing under-full prefills that preempt decode and drop it out of cudagraph. Tail-escape and pass-budget valves avoid starvation.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a “dense prefill batch” gate to the scheduler to avoid starting under-full new-prefill steps while decodes are active, helping keep decode on the fast path (e.g., staying in cudagraph) and bounding starvation via tail-escape / pass-budget logic.
Changes:
- Introduces a prefill gate based on whether waiting prefills can fill
max_num_batched_tokens, with tail-escape and a max-hold-pass budget. - Tightens PrefillDelayer negotiation so a rank only reports “prefillable” when it can admit a prefill and has enough waiting tokens to fill a dense batch.
- Adds unit tests covering the new gate behavior and the starvation/tail-escape valves.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| atom/model_engine/scheduler.py | Adds the dense-batch prefill gate, waiting-token accounting, and integrates gating into delayer negotiation + Phase-2 admission. |
| tests/test_scheduler.py | Adds tests validating threshold derivation, readiness/hold behavior, tail-escape, pass-budget forcing, and decode-not-prefill behavior under under-full waiting. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+811
to
+813
| _local_prefillable = self._can_admit_head_prefill() and ( | ||
| self._waiting_prefill_tokens() >= self.prefill_batch_token_threshold | ||
| ) |
Comment on lines
+447
to
+449
| # Max consecutive passes we suppress an under-full prefill before firing | ||
| # it anyway (starvation bound). Reuses the delayer's pass budget knob. | ||
| self._prefill_hold_max_passes = 30 |
Comment on lines
+598
to
+602
| total = 0 | ||
| cap = self.max_num_batched_tokens | ||
| for seq in self.waiting: | ||
| if self._unschedulable_reason(seq) is not None: | ||
| continue |
Comment on lines
+445
to
+449
| # Prefill throughput on V4 is linear in token count, so a full-8k batch | ||
| # is already on the efficient part of the curve; hold prefill until the | ||
| # waiting queue can fill 8192 tokens (not the whole max_num_batched_tokens | ||
| # budget, which may be larger). | ||
| self.prefill_batch_token_threshold = min(8192, self.max_num_batched_tokens) |
Comment on lines
+641
to
+646
| if self._waiting_prefill_tokens() >= self.prefill_batch_token_threshold: | ||
| self._prefill_hold_passes = 0 | ||
| return True | ||
| # Under-full: hold prefill (keep decoding) up to the pass budget. | ||
| self._prefill_hold_passes += 1 | ||
| if self._prefill_hold_passes >= self._prefill_hold_max_passes: |
ZhangLirong-amd
force-pushed
the
fix/prefill-batch-gate-woa-arch
branch
from
July 2, 2026 13:33
4383339 to
92d3d6a
Compare
Comment on lines
+623
to
+626
| """Gate for firing a NEW prefill step (dense-batch requirement). | ||
|
|
||
| The threshold is max_num_batched_tokens (a full prefill batch). | ||
|
|
Comment on lines
+814
to
+820
| # A rank counts as "prefillable" for cross-DP alignment only if it | ||
| # can admit a prefill AND has a full batch's worth of waiting tokens. | ||
| # This makes all ranks align on firing dense prefills together | ||
| # instead of straggling partials. | ||
| _local_prefillable = self._can_admit_head_prefill() and ( | ||
| self._waiting_prefill_tokens() >= self.prefill_batch_token_threshold | ||
| ) |
Comment on lines
857
to
861
| # ---- Phase 2: new requests from waiting ---- | ||
| while ( | ||
| _delayer_allows_prefill | ||
| _new_prefill_allowed | ||
| and (self.delay_factor <= 0 or self._passed_delay(time.time())) | ||
| and self.waiting |
| assert list(batch.scheduled_tokens) == toks[-(mtp_k + 1) :] | ||
|
|
||
|
|
||
| # ── Prefill dense-batch gate (threshold = max_num_batched_tokens) ─────────── |
ZhangLirong-amd
force-pushed
the
fix/prefill-batch-gate-woa-arch
branch
from
July 2, 2026 13:54
92d3d6a to
5da51de
Compare
The prefill dense-batch gate only helps cross-DP rank alignment. Disable it when data_parallel_size<=1 so single-GPU/TP-only runs keep the legacy prefill-first behavior (no added TTFT).
| continue | ||
| if seq.status == SequenceStatus.WAITING_FOR_REMOTE_KVS: | ||
| continue | ||
| n = seq.num_tokens - seq.num_cached_tokens |
Comment on lines
+636
to
+639
| # Tail escape: no decode work left — never hold, or we'd deadlock. | ||
| if not self.running: | ||
| self._prefill_hold_passes = 0 | ||
| return True |
Comment on lines
+817
to
+819
| _local_prefillable = self._can_admit_head_prefill() and ( | ||
| self._waiting_prefill_tokens() >= self.prefill_batch_token_threshold | ||
| ) |
| def test_ready_when_waiting_fills_threshold(self, seq_factory): | ||
| # chunked prefill on so a 40-token prompt is clamped to the 32 budget | ||
| # (not rejected as oversized) and counts toward the threshold. | ||
| cfg = MockConfig(max_num_batched_tokens=32, enable_chunked_prefill=True) |
| assert sched._prefill_hold_passes == 0 | ||
|
|
||
| def test_tail_escape_when_no_decode_left(self, seq_factory): | ||
| cfg = MockConfig(max_num_batched_tokens=32) # threshold = 32 |
valarLip
approved these changes
Jul 3, 2026
valarLip
pushed a commit
that referenced
this pull request
Jul 9, 2026
1 task
jpy794
pushed a commit
to RadeonFlow/up-atom
that referenced
this pull request
Jul 13, 2026
* fix(scheduler): gate prefill on full batch to protect decode Hold new prefills until the waiting queue can fill max_num_batched_tokens, else keep decoding. Prevents fast 补发 from firing under-full prefills that preempt decode and drop it out of cudagraph. Tail-escape and pass-budget valves avoid starvation. * style: black format * fix(scheduler): gate dense-batch prefill hold to DP>1 only The prefill dense-batch gate only helps cross-DP rank alignment. Disable it when data_parallel_size<=1 so single-GPU/TP-only runs keep the legacy prefill-first behavior (no added TTFT). --------- Co-authored-by: ZhangLirong-amd <ZhangLirong@amd.com>
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.
Hold new prefills until the waiting queue can fill max_num_batched_tokens, else keep decoding. Prevents fast sending from firing under-full prefills that preempt decode and drop it out of cudagraph.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist