From 4b7b05dda31a34d5f0921f0ef20649a9f5695d69 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Thu, 2 Jul 2026 04:52:06 +0000 Subject: [PATCH 1/3] fix(scheduler): gate prefill on full batch to protect decode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- atom/model_engine/scheduler.py | 71 ++++++++++++++++++++++++++++- tests/test_scheduler.py | 81 ++++++++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+), 2 deletions(-) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index f8a0c9add2..528bab6f75 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -442,6 +442,11 @@ def __init__(self, config: Config): # Latency of the last prompt step self.last_prompt_latency = 0.0 self.delay_factor = config.scheduler_delay_factor + self.prefill_batch_token_threshold = self.max_num_batched_tokens + self._prefill_hold_passes = 0 + # 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 # Speculative decoding self.use_spec = config.speculative_config is not None @@ -583,6 +588,59 @@ def _can_admit_head_prefill(self) -> bool: return True return False + def _waiting_prefill_tokens(self) -> int: + """Sum of admissible new-prefill tokens sitting in the waiting queue, + capped at max_num_batched_tokens (we only care whether a *dense* batch + can be filled). Skips unschedulable / remote-KV-waiting entries and + clamps each seq to the chunked-prefill / budget limit, mirroring the + Phase-2 admission math so the count reflects what would actually pack + into one prefill step.""" + total = 0 + cap = self.max_num_batched_tokens + for seq in self.waiting: + if self._unschedulable_reason(seq) is not None: + continue + if seq.status == SequenceStatus.WAITING_FOR_REMOTE_KVS: + continue + n = seq.num_tokens - seq.num_cached_tokens + if self.enable_chunked_prefill and 0 < self.long_prefill_token_threshold < n: + n = self.long_prefill_token_threshold + if n > cap: + n = cap + total += n + if total >= cap: + return cap + return total + + def _prefill_batch_ready(self) -> bool: + """Gate for firing a NEW prefill step (dense-batch requirement). + + The threshold is max_num_batched_tokens (a full prefill batch). + + Returns True (allow prefill) when: + - the waiting queue can fill a dense batch (>= threshold tokens), or + - there is nothing left to decode (running empty) — tail escape so a + partial final batch still goes out, or + - we've suppressed prefill for too many consecutive passes + (_prefill_hold_max_passes) — starvation bound. + + Otherwise returns False (keep decoding) and advances the hold counter. + The counter resets whenever prefill is allowed to fire. + """ + # Tail escape: no decode work left — never hold, or we'd deadlock. + if not self.running: + self._prefill_hold_passes = 0 + return True + 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: + self._prefill_hold_passes = 0 + return True + return False + def _kv_usage(self) -> float: """Fraction of KV-cache blocks currently in use ∈ [0, 1]. @@ -746,14 +804,23 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: # ─── Cross-DP prefill alignment (PrefillDelayer) ─────────────── _delayer_allows_prefill = True if self.prefill_delayer is not None: + # 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 + ) _delayer_allows_prefill = self.prefill_delayer.should_allow_prefill( - local_prefillable=self._can_admit_head_prefill(), + local_prefillable=_local_prefillable, token_usage=self._kv_usage(), ) if not self.running and not self.waiting: return None + _new_prefill_allowed = _delayer_allows_prefill and self._prefill_batch_ready() + # ---- Phase 1: resume partial prefills from running ---- # Gated by `_delayer_allows_prefill` so cross-DP alignment still # holds when one rank is mid-chunked-prefill: a delayer veto skips @@ -782,7 +849,7 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: # ---- 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 and num_seqs_prefill < self.max_num_seqs diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 659d900813..bb562661e5 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -673,3 +673,84 @@ def test_normal_decode_window_unchanged(self): ) assert list(batch.scheduled_tokens) == toks[-(mtp_k + 1) :] + + +# ── Prefill dense-batch gate (threshold = max_num_batched_tokens) ─────────── + + +class TestPrefillBatchGate: + def test_threshold_defaults_to_batch_budget(self): + cfg = MockConfig(max_num_batched_tokens=32) + sched = Scheduler(cfg) + # Threshold is derived from the batch-token budget, not a separate knob. + assert sched.prefill_batch_token_threshold == 32 + + 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) + sched = Scheduler(cfg) + # Keep a decode running so the tail-escape doesn't trivially allow. + r = seq_factory([1, 2, 3]) + r.status = SequenceStatus.RUNNING + sched.block_manager.allocate(r, 0) + sched.running.append(r) + # 40 waiting tokens, clamped to 32 == threshold -> ready. + sched.waiting.append(seq_factory(list(range(40)))) + assert sched._waiting_prefill_tokens() >= 32 + assert sched._prefill_batch_ready() is True + + def test_holds_when_under_full(self, seq_factory): + cfg = MockConfig(max_num_batched_tokens=32) # threshold = 32 + sched = Scheduler(cfg) + r = seq_factory([1, 2, 3]) + r.status = SequenceStatus.RUNNING + sched.block_manager.allocate(r, 0) + sched.running.append(r) + # Only 8 waiting tokens < 32 -> hold (keep decoding). + sched.waiting.append(seq_factory(list(range(8)))) + assert sched._prefill_batch_ready() is False + assert sched._prefill_hold_passes == 1 + + def test_tail_escape_when_no_decode_left(self, seq_factory): + cfg = MockConfig(max_num_batched_tokens=32) # threshold = 32 + sched = Scheduler(cfg) + # running empty -> even an under-full waiting batch must be allowed, + # else the final partial batch would deadlock. + sched.waiting.append(seq_factory(list(range(8)))) + assert not sched.running + assert sched._prefill_batch_ready() is True + + def test_hold_pass_budget_forces_fire(self, seq_factory): + cfg = MockConfig(max_num_batched_tokens=32) # threshold = 32 + sched = Scheduler(cfg) + sched._prefill_hold_max_passes = 3 + r = seq_factory([1, 2, 3]) + r.status = SequenceStatus.RUNNING + sched.block_manager.allocate(r, 0) + sched.running.append(r) + sched.waiting.append(seq_factory(list(range(8)))) # under-full + # First (max_passes - 1) calls hold, then it force-fires and resets. + assert sched._prefill_batch_ready() is False # pass 1 + assert sched._prefill_batch_ready() is False # pass 2 + assert sched._prefill_batch_ready() is True # pass 3 -> force + assert sched._prefill_hold_passes == 0 + + def test_gate_holds_new_prefill_but_decodes(self, seq_factory): + """End-to-end: under-full waiting + running decode -> schedule() should + NOT start a new prefill; it should return a decode batch instead.""" + cfg = MockConfig( + max_num_batched_tokens=32, # threshold = 32 + enable_chunked_prefill=False, + ) + sched = Scheduler(cfg) + r = seq_factory([1, 2, 3]) + r.status = SequenceStatus.RUNNING + r.type = SequenceType.DECODE + sched.block_manager.allocate(r, 0) + sched.running.append(r) + sched.waiting.append(seq_factory(list(range(8)))) # under-full prefill + batch, _ = sched.schedule() + assert batch is not None + assert batch.total_seqs_num_prefill == 0 + assert batch.total_seqs_num_decode == 1 From 5da51de04fb46dc6563025e35653f1d932331b07 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Thu, 2 Jul 2026 05:00:42 +0000 Subject: [PATCH 2/3] style: black format --- atom/model_engine/scheduler.py | 5 ++++- tests/test_scheduler.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 528bab6f75..5f557ed1b1 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -603,7 +603,10 @@ def _waiting_prefill_tokens(self) -> int: if seq.status == SequenceStatus.WAITING_FOR_REMOTE_KVS: continue n = seq.num_tokens - seq.num_cached_tokens - if self.enable_chunked_prefill and 0 < self.long_prefill_token_threshold < n: + if ( + self.enable_chunked_prefill + and 0 < self.long_prefill_token_threshold < n + ): n = self.long_prefill_token_threshold if n > cap: n = cap diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index bb562661e5..970712aaba 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -733,7 +733,7 @@ def test_hold_pass_budget_forces_fire(self, seq_factory): # First (max_passes - 1) calls hold, then it force-fires and resets. assert sched._prefill_batch_ready() is False # pass 1 assert sched._prefill_batch_ready() is False # pass 2 - assert sched._prefill_batch_ready() is True # pass 3 -> force + assert sched._prefill_batch_ready() is True # pass 3 -> force assert sched._prefill_hold_passes == 0 def test_gate_holds_new_prefill_but_decodes(self, seq_factory): From ac99b69133edb42c6e6f6f8d499b5ab1307af728 Mon Sep 17 00:00:00 2001 From: ZhangLirong-amd Date: Fri, 3 Jul 2026 12:45:08 +0000 Subject: [PATCH 3/3] 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). --- atom/model_engine/scheduler.py | 7 +++++-- tests/conftest.py | 9 ++++++++- tests/test_scheduler.py | 20 ++++++++++++++++++-- 3 files changed, 31 insertions(+), 5 deletions(-) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 5f557ed1b1..bf12fc42e6 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -444,9 +444,9 @@ def __init__(self, config: Config): self.delay_factor = config.scheduler_delay_factor self.prefill_batch_token_threshold = self.max_num_batched_tokens self._prefill_hold_passes = 0 - # 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 + _pc = getattr(config, "parallel_config", None) + self._prefill_gate_enabled = getattr(_pc, "data_parallel_size", 1) > 1 # Speculative decoding self.use_spec = config.speculative_config is not None @@ -630,6 +630,9 @@ def _prefill_batch_ready(self) -> bool: Otherwise returns False (keep decoding) and advances the hold counter. The counter resets whenever prefill is allowed to fire. """ + # Gate only applies under DP (>1); otherwise never hold (legacy path). + if not self._prefill_gate_enabled: + return True # Tail escape: no decode work left — never hold, or we'd deadlock. if not self.running: self._prefill_hold_passes = 0 diff --git a/tests/conftest.py b/tests/conftest.py index f875e6fd9b..7b60e9b9bc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -52,7 +52,8 @@ def __init__(self, **kwargs): class _StubParallelConfig: """Placeholder for ParallelConfig.""" - pass + def __init__(self, data_parallel_size: int = 1): + self.data_parallel_size = data_parallel_size _atom_config.Config = _StubConfig @@ -146,6 +147,12 @@ def __init__(self, **overrides): stop_token_ids=[], scheduler_delay_factor=0.0, speculative_config=None, + # DP size gates the dense-batch prefill hold (see Scheduler). Default + # 1 (gate off) so unrelated tests keep legacy behavior; gate tests + # pass data_parallel_size>1. + parallel_config=_StubParallelConfig( + data_parallel_size=overrides.pop("data_parallel_size", 1) + ), # Scheduler.__init__ reads config.hf_config.architectures for V4 # SWA-warmup detection; a non-V4 stub keeps that path inert. hf_config=_MockHFConfig(), diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 970712aaba..a5b0188160 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -701,7 +701,8 @@ def test_ready_when_waiting_fills_threshold(self, seq_factory): assert sched._prefill_batch_ready() is True def test_holds_when_under_full(self, seq_factory): - cfg = MockConfig(max_num_batched_tokens=32) # threshold = 32 + # gate only active under DP>1 + cfg = MockConfig(max_num_batched_tokens=32, data_parallel_size=8) sched = Scheduler(cfg) r = seq_factory([1, 2, 3]) r.status = SequenceStatus.RUNNING @@ -712,6 +713,20 @@ def test_holds_when_under_full(self, seq_factory): assert sched._prefill_batch_ready() is False assert sched._prefill_hold_passes == 1 + def test_gate_disabled_without_dp(self, seq_factory): + # DP<=1 (default): gate off -> under-full prefill is NOT held. + cfg = MockConfig(max_num_batched_tokens=32) # data_parallel_size=1 + sched = Scheduler(cfg) + assert sched._prefill_gate_enabled is False + r = seq_factory([1, 2, 3]) + r.status = SequenceStatus.RUNNING + sched.block_manager.allocate(r, 0) + sched.running.append(r) + sched.waiting.append(seq_factory(list(range(8)))) # under-full + # Gate off -> ready True (legacy behavior), no hold counter advance. + assert sched._prefill_batch_ready() is 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 sched = Scheduler(cfg) @@ -722,7 +737,7 @@ def test_tail_escape_when_no_decode_left(self, seq_factory): assert sched._prefill_batch_ready() is True def test_hold_pass_budget_forces_fire(self, seq_factory): - cfg = MockConfig(max_num_batched_tokens=32) # threshold = 32 + cfg = MockConfig(max_num_batched_tokens=32, data_parallel_size=8) sched = Scheduler(cfg) sched._prefill_hold_max_passes = 3 r = seq_factory([1, 2, 3]) @@ -742,6 +757,7 @@ def test_gate_holds_new_prefill_but_decodes(self, seq_factory): cfg = MockConfig( max_num_batched_tokens=32, # threshold = 32 enable_chunked_prefill=False, + data_parallel_size=8, # gate only active under DP>1 ) sched = Scheduler(cfg) r = seq_factory([1, 2, 3])