fix(scheduler): gate prefill on full batch to protect decode

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

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(),

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