diff --git a/omlx/prefill_transient_tracker.py b/omlx/prefill_transient_tracker.py index 965ee647a..7c8e89687 100644 --- a/omlx/prefill_transient_tracker.py +++ b/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,6 +176,11 @@ 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 @@ -172,3 +188,4 @@ def reset(self) -> None: self._last_delta_bytes = 0 self._last_n_tokens = 0 self._observed_max_bytes = 0 + self._recent_reclaim_bytes = 0 diff --git a/omlx/scheduler.py b/omlx/scheduler.py index e5c4b1294..a169a5d52 100644 --- a/omlx/scheduler.py +++ b/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, diff --git a/tests/test_prefill_oom_graceful.py b/tests/test_prefill_oom_graceful.py index 09348e0dd..7ede7e099 100644 --- a/tests/test_prefill_oom_graceful.py +++ b/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(