Account for recently reclaimed memory in adaptive prefill estimates

omlx/prefill_transient_tracker.py

@@ -57,6 +57,15 @@ def __init__(self, model_id: str = "") -> None:
         # gates, matching the floor-chunk charge they price. Never used
         # for chunk sizing.
         self._observed_max_bytes: int = 0
+        # Net process footprint released by negative post-chunk deltas. MLX may
+        # need to allocate that pool again on the next chunk, so the scheduler
+        # prices it once until a positive measurement confirms reallocation.
+        self._recent_reclaim_bytes: int = 0
+
+    def record_reclaim(self, reclaimed_bytes: int) -> None:
+        """Accumulate footprint released since the last positive sample."""
+        if reclaimed_bytes > 0:
+            self._recent_reclaim_bytes += int(reclaimed_bytes)
 
     def update(
         self, n_tokens: int, transient_bytes: int, *, floor_sample: bool = False
@@ -87,6 +96,8 @@ def update(
         if transient_bytes <= 0:
             return
 
+        self._recent_reclaim_bytes = 0
+
         # The very first sample after a model load carries weight page-fault
         # and load-residue noise, so it seeds the EWMA but is excluded from
         # the running max.
@@ -165,10 +176,16 @@ def observed_max_bytes(self) -> int:
         """Largest accepted chunk transient this session (0 if none yet)."""
         return self._observed_max_bytes
 
+    @property
+    def recent_reclaim_bytes(self) -> int:
+        """Footprint released since the last positive chunk measurement."""
+        return self._recent_reclaim_bytes
+
     def reset(self) -> None:
         """Drop all observations (e.g. on model reload or after a long idle)."""
         self._ewma_per_token = 0.0
         self._samples = 0
         self._last_delta_bytes = 0
         self._last_n_tokens = 0
         self._observed_max_bytes = 0
+        self._recent_reclaim_bytes = 0

omlx/scheduler.py

@@ -3478,6 +3478,8 @@ def _predicted_chunk_transient(self, n_tokens: int, kv_len: int) -> float:
         if n_tokens <= 0:
             return 0.0
         per_token = 0.0
+        static_per_token = 0.0
+        recent_reclaim = 0
         tracker = self._prefill_transient_tracker
         if tracker is not None:
             if tracker.last_n_tokens > 0 and tracker.last_delta_bytes > 0:
@@ -3486,13 +3488,20 @@ def _predicted_chunk_transient(self, n_tokens: int, kv_len: int) -> float:
                 )
             if tracker.bytes_per_token > 0:
                 per_token = max(per_token, tracker.bytes_per_token)
+            recent_reclaim = tracker.recent_reclaim_bytes
         if self.memory_monitor is not None:
             static = self.memory_monitor.estimate_chunk_transient_bytes(
                 n_tokens, kv_len + n_tokens
             )
             static += self.memory_monitor.estimate_prompt_kv_bytes(n_tokens)
-            per_token = max(per_token, float(static) / n_tokens)
-        return per_token * n_tokens * self._PREFILL_TRANSIENT_SAFETY
+            static_per_token = float(static) / n_tokens
+            per_token = max(per_token, static_per_token)
+        base_prediction = per_token * n_tokens * self._PREFILL_TRANSIENT_SAFETY
+        reallocation_prediction = (
+            static_per_token * n_tokens * self._PREFILL_TRANSIENT_SAFETY
+            + recent_reclaim
+        )
+        return max(base_prediction, reallocation_prediction)
 
     def _admission_transient_bound(self, n_tokens: int, kv_len: int) -> float:
         """Transient charge for admission and the guard's pass/abort gates.
@@ -4230,6 +4239,11 @@ def _record_chunk_transient(
         has to stay conservative. Keeping kv_len in the log is what made that
         analysis possible.
 
+        Negative deltas remain excluded from the per-token EWMA, but their
+        released footprint is retained until the next positive sample. The
+        next predictor prices that one-shot reallocation risk without treating
+        it as a negative per-token cost.
+
         Under speed priority, only a complete requested step is representative
         of the full-size chunks used for admission. A shorter tail or
         boundary-alignment chunk must not replace the last full-step sample:
@@ -4251,8 +4265,10 @@ def _record_chunk_transient(
             )
             return
         if delta <= 0:
+            self._prefill_transient_tracker.record_reclaim(-delta)
             logger.debug(
-                "[throttle:%s] measure rid=%s n=%d delta=%dB (skipped: <=0)",
+                "[throttle:%s] measure rid=%s n=%d delta=%dB "
+                "(excluded from EWMA; tracked as reclaim)",
                 loop_label,
                 request_id,
                 n_tokens,

tests/test_prefill_oom_graceful.py

@@ -471,6 +471,85 @@ def test_predicted_transient_zero_without_signals():
     assert ns._predicted_chunk_transient(4, 1000) == 0.0
 
 
+def test_adaptive_throttle_charges_recently_reclaimed_footprint():
+    """A pool drop must remain priced until the next chunk reallocates it."""
+    static_prediction = 11.18 * _GB
+    released = 6.34 * _GB
+    monitor = SimpleNamespace(
+        estimate_chunk_transient_bytes=lambda _n, _kv: (
+            static_prediction / Scheduler._PREFILL_TRANSIENT_SAFETY
+        ),
+        estimate_prompt_kv_bytes=lambda _n: 0,
+    )
+    ns = _throttle_ctx(
+        current=97.23 * _GB,
+        hard=119.17 * _GB,
+        soft_ratio=110.23 / 119.17,
+        monitor=monitor,
+        abort=200 * _GB,
+    )
+    ns._prefill_headroom_safety = 110.23 / 119.17
+    ns._fake_current = 97.23 * _GB
+    ns.requests = {}
+    ns.config = SimpleNamespace(model_name="model-b")
+    ns._raise_prefill_eviction_if_available = (
+        Scheduler._raise_prefill_eviction_if_available.__get__(ns, Scheduler)
+    )
+    ns._record_chunk_transient = Scheduler._record_chunk_transient.__get__(
+        ns, Scheduler
+    )
+
+    assert _call(ns, 2048, kv_len=147_680) == 2048
+
+    ns._record_chunk_transient(
+        512,
+        100 * _GB,
+        100 * _GB - released,
+        request_id="r",
+        loop_label="test",
+        requested_step=512,
+    )
+
+    assert _call(ns, 2048, kv_len=147_680) < 2048
+
+
+def test_predicted_transient_does_not_double_count_reclaim_covered_by_raw():
+    """A conservative raw-last sample may already cover pool reallocation."""
+    raw_prediction = 11.83 * _GB
+    static_prediction = 4.11 * _GB
+    released = 6.86 * _GB
+    raw_per_token = raw_prediction / (
+        512 * Scheduler._PREFILL_TRANSIENT_SAFETY
+    )
+    monitor = SimpleNamespace(
+        estimate_chunk_transient_bytes=lambda _n, _kv: (
+            static_prediction / Scheduler._PREFILL_TRANSIENT_SAFETY
+        ),
+        estimate_prompt_kv_bytes=lambda _n: 0,
+    )
+    ns = _throttle_ctx(
+        current=99.12 * _GB,
+        hard=118.71 * _GB,
+        samples_bpt=raw_per_token,
+        monitor=monitor,
+    )
+    ns._record_chunk_transient = Scheduler._record_chunk_transient.__get__(
+        ns, Scheduler
+    )
+    ns._record_chunk_transient(
+        512,
+        100 * _GB,
+        100 * _GB - released,
+        request_id="r",
+        loop_label="test",
+        requested_step=512,
+    )
+
+    predicted = ns._predicted_chunk_transient(512, 186_368)
+
+    assert predicted == pytest.approx(raw_prediction)
+
+
 def test_record_chunk_transient_skips_tail_samples():
     tracker = PrefillTransientTracker()
     ns = SimpleNamespace(