diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index f8a0c9add2..bf12fc42e6 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 + 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 @@ -583,6 +588,65 @@ 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. + """ + # 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 + 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 +810,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 +855,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/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 659d900813..a5b0188160 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -673,3 +673,100 @@ 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): + # 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 + 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_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) + # 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, data_parallel_size=8) + 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, + data_parallel_size=8, # gate only active under DP>1 + ) + 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