Account for recently reclaimed memory in adaptive prefill estimates

[mvdbos]

Context

This follows fix: estimate DeepSeek V4 prefill memory (#2521) (745d9c22), included in oMLX 0.5.8.dev1. That change added architecture-aware static prefill and KV-memory estimates for DeepSeek V4.

While testing that path at high context, I found a separate gap in the adaptive measurement layer: MLX can release pooled memory after one chunk and allocate it again during the next, but the tracker currently forgets the release. This PR leaves the DeepSeek V4 static estimator unchanged and adds the missing dynamic signal on top of it.

The measurement fix itself is model-agnostic. PrefillTransientTracker is created for every scheduler, and both prefill loops feed it process-footprint deltas without an architecture check. DeepSeek V4 is the model used for the live reproduction; any model showing the same release-then-reallocation pattern can hit this path.

What this changes

This PR makes the adaptive prefill tracker remember process footprint released by a chunk and account for it once when sizing the next chunk.

The change:

• records the absolute value of negative post-chunk footprint deltas as recent_reclaim_bytes;
• keeps reclaimed bytes out of the positive-sample EWMA and raw per-token estimate;
• compares the existing prediction with static prediction + recent reclaim and uses the larger value;
• clears the reclaim charge after the next positive footprint measurement; and
• clears it on tracker reset.

This closes a gap where MLX can release several GiB of pooled memory after a small or boundary chunk, then allocate roughly that memory again during the next larger chunk.

Why

The existing measurement path ignores non-positive process-footprint deltas:

delta = post_bytes - pre_bytes
if delta <= 0:
    return
tracker.update(n_tokens, delta, ...)

That is appropriate for the positive per-token EWMA, but it also means the scheduler forgets a recent pool release completely. The next prediction can only use the latest positive sample, EWMA, and static model estimate.

In one controlled run:

1. A 512-token chunk released 6.34 GiB, leaving process footprint at about 97.23 GiB.
2. The release was ignored.
3. The static estimate for the next 2,048-token chunk was 11.18 GiB.
4. Process footprint grew by about 17.16 GiB before the memory guard aborted the request.
5. 11.18 + 6.34 = 17.52 GiB, which closely matches that growth.

Reproduction

I reproduced this on:

• oMLX 0.5.8.dev1, build 2159;
• source matching commit 745d9c22;
• Jundot/DeepSeek-V4-Flash-0731-oQ2e-mtp;
• MTP and VLM MTP disabled;
• 350,000-token context window;
• aggressive memory guard;
• 120 GB Metal cap;
• prompt cache disabled; and
• an Apple M5 Max with 128 GB unified memory.

The request is synthetic and deterministic. It uses one user message containing "x " repeated 245,243 times:

import json
import urllib.request

base_url = "http://127.0.0.1:58891"
payload = {
    "model": "DeepSeek-V4-Flash-0731-oQ2e-mtp",
    "messages": [
        {"role": "user", "content": "x " * 245_243},
    ],
    "max_tokens": 1,
    "temperature": 0,
    "top_p": 1.0,
    "stream": False,
}

request = urllib.request.Request(
    base_url + "/v1/messages",
    data=json.dumps(payload, separators=(",", ":")).encode(),
    headers={"Content-Type": "application/json"},
    method="POST",
)
with urllib.request.urlopen(request, timeout=1800) as response:
    print(response.read().decode())

/v1/messages/count_tokens reports 245,248 tokens for the message. The completed /v1/messages response reports 245,327 input tokens after the API template is applied.

The serialized request is 490,629 bytes and has SHA-256:

4b9dcdec01ff8d56ea1bb22e3a011bdd90176f8127259e9d3dd69f4cde15ef2d

Restart the server between the unmodified and patched runs so both start cold, and leave prompt caching disabled.

Before and after

Version	Result	Prefill progress	Hard-pressure transitions	Prefill API errors	Wall time
Before this change	prefill_memory_aborted	149,728 / 245,327	1	1	587.56 s
With this change	Valid response	245,327 / 245,327	0	0	1,146.21 s

The successful run returned HTTP 200 with 245,327 input tokens and 1 output token. During that run the tracker saw:

• 89 footprint releases;
• a maximum retained release of 21.33 GiB;
• 227 adaptive chunk reductions;
• three soft-pressure transitions;
• no hard-pressure transition; and
• no abort or prefill API error.

A trace from the patched run shows the new path directly:

1. At 107,712 tokens, a 32-token boundary chunk released 21.33 GiB.
2. The tracker retained that release.
3. The estimate for a requested 2,048-token chunk became 31.91 GiB.
4. The scheduler selected an 832-token chunk instead.
5. That chunk grew process footprint by 8.53 GiB.
6. The tracker then cleared the one-shot reclaim charge.
7. The request continued past the previous abort point and completed all 245,327 input tokens.

Elapsed time to roughly the same prefill position was about 5.1% higher with the change. This is only a rough safety-cost comparison: the patched measurement completed 736 tokens beyond the abort position, while the unmodified run includes hard-pressure abort deferral.

Avoiding double counting

The reclaim charge is not added blindly to the current winning prediction. The predictor uses:

max(
    existing_raw_ewma_static_prediction,
    static_prediction + recent_reclaim_bytes,
)

For example, a second controlled high-context sequence had:

• raw-last prediction: 11.83 GiB;
• static prediction: 4.11 GiB;
• recent release: 6.86 GiB; and
• next observed growth: 7.54 GiB.

The result stays at 11.83 GiB because raw-last already covers the likely reallocation. Adding the release to raw-last would incorrectly raise the estimate to 18.69 GiB.

Tests

This PR adds:

• test_adaptive_throttle_charges_recently_reclaimed_footprint
• test_predicted_transient_does_not_double_count_reclaim_covered_by_raw

Focused suites run:

• tests/test_prefill_transient_tracker.py
• tests/test_prefill_oom_graceful.py
• tests/test_scheduler_prefill_memory_guard.py
• tests/test_memory_monitor.py

Result: 162 passed.

The full fast suite completed with 7,979 passed, 66 skipped, 71 deselected, and 5 failures. The same five unrelated GLM/Inkling MTP numerical-parity tests fail with the same deltas on a clean origin/main worktree, so they are not introduced by this change.

Scope

This change does not treat adaptive headroom or large footprint samples as bugs. In 64 positive high-context measurements (kv_len >= 180k, chunks no larger than 512 tokens), median footprint growth was 6.17 GiB and the maximum was 12.14 GiB. The large pool reallocations are real.

The narrower problem is that a release was forgotten even though the next chunk could allocate that memory again. This PR keeps the existing positive-sample predictor intact and adds a one-shot safety charge for that specific case.

The reproduction uses repeated synthetic text and no prompt cache. It tests the scheduler and memory-guard behavior, not model-output quality or a particular agent workload.

The before-and-after live run has only been performed with DeepSeek V4. The regression tests cover the generic scheduler/tracker path, but I have not run the same high-context live comparison on another model family.

[jundot]

Thanks for the deep investigation, the repro and the double counting analysis made this easy to verify. The approach is right and I'll merge this, with a small follow-up on my side for two gate ordering corners.

I ran a live A/B with DeepSeek V4 Flash on a 512GB machine using a pinned custom ceiling. I could not reproduce the abort itself in context mode there, without real OS memory pressure the enforcer reclaim always wins on this box, but the mechanism shows up clearly:

Run	main	this PR
context, 90k, tight watermarks	completed, hard pressure 2x, 127 releases forgotten	completed, same wall time and peak, hard pressure 1x, 114 releases charged
speed, 175k	aborted at 118,784 tokens	aborted at the exact same step, charge fired only 2x

Speed coming out neutral is expected, that abort comes from a single full-size chunk spike no predictor can see, so no regression on that path.

The two follow-up corners, both from the gate order in _record_chunk_transient:

Corner	Why	Effect
Release on a tail chunk below min_chunk	Tail gate returns before the delta <= 0 branch	That release is still forgotten (probe: 17 token tail releasing 6GB leaves the charge at 0)
Positive realloc on a skipped sample	Charge is only cleared inside tracker.update	Stale charge double counts near the cap until the next full positive sample

Neither blocks the merge, both are fixed by moving the sign handling above the skip gates, and I'll take care of that in the follow-up.