Running oMLX unattended for coding agents: a TurboQuant batched-decode fix, and pausing requests instead of aborting them under memory pressure

[e-Evolution]

I run oMLX as an unattended server for coding agents on an Apple M5 Max (128 GB). The agents are pi with the gentle-shell harness (package gentle-pi) and subagents, working on a documentation pilot: an ADempiere wiki with images of the new version and an assistant that answers questions. The model is Qwen3.8-27B-oQ4e-mtp: 4-bit, Lightning MTP and TurboQuant 4-bit KV. Prompts are long (median 107k tokens, p90 168k), and several sessions run at once.

While making this setup run unattended I found four things, and I'd like your view on how to contribute them:

1. A padded-row defect in batched TurboQuant decode that is still on main: left-padded rows attend over their padding, and batched decode is slow. The fix is ready as a focused PR: fix(turboquant): attend each left-padded row's own slice in batched decode and MTP verify #4441.
2. A pause-resume cycle. Under memory pressure it pauses requests and resumes them later, where main aborts every in-flight request. It is large, so I'd like to agree on the approach here before opening anything.
3. A second TurboQuant defect: the suffix prefill after a paged-SSD restore of a TurboQuant prefix is about 2.4x slower than a cold prefill. That fix is ready as a small PR too: fix(scheduler): prefill long suffixes over a dense copy of a restored TurboQuant prefix #4442.
4. Two server-side causes of prefix-cache breaks in agent workloads.

Full write-up (method, every number and its limits): https://e-evolution.github.io/omlx/. The numbers come from my own runs, one per configuration unless stated, so read them as evidence rather than an independent reproduction.

1. TurboQuant batched decode with left-padded rows (PR ready)

BatchTurboQuantKVCache left-pads its rows to a common length. In batched decode (B > 1):

• Decode (L = 1). The batched attention path dequantizes the whole padded context and lets the zero padding columns of shorter rows into the softmax. A short row's output is shrunk: its norm ratio against a single-row reference is 0.32.
• Padded MTP verify (L = 2..5). mlx-vlm slices the cache's own mask per layer (mask[..., :kv]), so verify never takes a per-row path.

Reproduction on main (synthetic caches, no download; the tests are in the PR):

Check	main	Expected
Short-row output norm vs single-row reference	0.3196	1.0
Short-row max abs diff	0.0988	< 2e-3
Greedy logits, B=3 vs B=1 (left-padded decode)	2.90	< 0.05

On the real model, I ran a greedy B=3 restored decode of 20k / 45k / 135k prompts on main 25aebb3 (the TurboQuant files are unchanged upstream since) and on the PR, from one venv with main's pinned dependencies, one run per build. Both complete without aborts or degeneration; the difference is speed and memory:

	main	with the fix
Batched decode per row, MTP on	1.5 tok/s	8.3 tok/s
B=3 phase, MTP on	345 s	66 s
Batched decode per row, MTP off	0.66 tok/s	8.1 tok/s
B=3 phase, MTP off	780 s	65 s
Process peak	112-117 GB	101-102 GB

With MTP off the B=1 reference is deterministic. main's batched rows leave it at about token 1-2 on the 20k and 45k rows, and the fix's at about token 8-10. Without logprobs I can't classify those divergences, so this is consistent with the padding defect, not proof of it.

Per-layer attention at 20k / 45k / 135k goes from 86.0 to 3.05 ms (L = 1) and from 88.3 to 8.4 ms (L = 3, verify).

An earlier run on main b7c7929 with the v0.7.0 app runtime aborted this batch (6 memory aborts at a 120 GB peak). That does not reproduce on 25aebb3, and with one run per build I can't say whether upstream's commits or the runtime changed it.

My branch also fixed the lossy verify rollback, which #3127 has since fixed on main. I verified that it matches my rollback on every unmasked column, so the PR is down to 2 commits (+815 / -1, 28 tests). Built on main 013bc3f, it merges cleanly into current main (e02ae51), and CI's command on that merge gives main plus the 28 new tests, with no new failures.

2. Pause and resume instead of abort (proposal; I'd like your view first)

Problem

At hard memory pressure, the enforcer on main calls abort_all_requests on the non-emergency path. Every request loses its prefilled KV and its generated text, and the client has to re-issue, usually cold. I ran a forced-pressure scenario: three concurrent prompts of about 130k, 115k and 100k tokens under a 48 GB custom ceiling, with hard_threshold 0.8. On main (b7c7929), all 3 requests were aborted in each case: text decode, image decode and image prefill. A request that can never fit is also not distinguished from one that only has to wait. A 60k prompt with max_tokens 2,000,000 gets HTTP 200 on main.

This looks related to #4213, where the prefill memory guard stops an agent workflow, and to the direction of #4252 (free hot cache instead of rejecting).

Design

• Pause one request. The server pauses one request, in prefill or in decode. It saves the request's KV to the paged SSD cache as an exact, domain-keyed, pinned chain and frees the memory. The client stream stays open, and the pending token is emitted first.
• Resume. The request resumes from the exact prefix with no prefill and no RNG reseed. Sampler, logits processors, stop-sequence position, parser state and remaining max_tokens are kept.
• Image prompts. A text suffix after a cached last image is prefilled in the chunked path, so it can be paused too.
• Resume policy. A paused request resumes when estimated free memory covers its need, with an anti-thrash rule, or alone (solo) when nothing else is admitted. Paused requests resume FIFO, ahead of newer ones.
• Fallback. When nothing can be paused, one victim is aborted instead of all. The emergency brake is unchanged.
• Never fits. A request whose weights + KV(prompt + max_tokens) + transient exceed capacity gets HTTP 400 context_exceeds_capacity, with seven numeric fields and retryable: false. A request that fits alone waits instead of failing. GET /v1/omlx/capacity and an X-oMLX-Max-Context-Tokens header expose the numbers, so an agent harness can compact before it hits the limit.

Defaults. Nothing changes unless the operator opts in:

• pause, resume, the single victim and the media route each have their own default-off switch;
• the never-fits 400, the fits-later wait and a progress-based admission-stall clock are on only with an explicit memory ceiling (memory_guard_tier: custom) or memory.memory_scheduling=on. On a dynamic tier, the ceiling follows other apps' memory, so promising what fits alone does not hold there.

Results

On current main (c0b1056), my fork rebased on it against main, one venv with main's pins for both, one server at a time, the same ceiling for both trees, one run per scenario:

Scenario	Branch	main (c0b1056)
Text decode, 44 GB	3/3 done, 0 aborts, 2 prefill parks resumed	0/3, all aborted
Image decode, 42 GB	3/3 done, 0 aborts, 2 decode pauses resumed	0/3, two abort-alls and one prefill-guard rejection
Image prefill, 44 GB	3/3 done, 0 aborts	3/3 done
Never fits, 42 GB	400 in 0.11 s with 7 fields; capacity endpoint; header	200, no endpoint, no header

main now completes the image-prefill case, so the difference there is gone; decode is where requests still get aborted. Against main c0b1056 the premises show no regression: 2 x 130k concurrent at a 51 GB peak against 116 GB, concurrency at parity, and single-request decode within noise.

Earlier, on main b7c7929 (both builds through the v0.7.0 app runtime; one run per pair):

Scenario	Branch	main (b7c7929)
Text decode, 48 GB	3/3 done, 0 aborts, 6 decode pauses and resumes	0/3, all aborted
Image decode, 48 GB	3/3 done, 0 aborts, 5 pauses and resumes	0/3, all aborted
Image prefill (branch 46 GB, main 48 GB)	3/3 done, 0 aborts	0/3, all aborted
Never fits, 42 GB	400 in 0.1 s with 7 fields; capacity endpoint; header	200, no endpoint, no header

Identity. With MTP off, 12 forced decode pauses (B=1) give identical tokens and 0.0 logprob drift for 3/3 requests of 2,000 tokens. With MTP on, outputs diverge at bf16 near-ties 42-79 tokens after a pause. The same prompts are identical with MTP off, which points to MTP verify regrouping (a hypothesis: it does not by itself rule out a pause-state effect that only shows with MTP on).

Size, and what I'd like to agree on

The cycle is about 13k changed lines, about two thirds of them tests (real scheduler and enforcer rigs on a tiny hybrid model and the real paged SSD cache). It also needs a prefill park and a prefill progress commit that are not on main. I propose a stack of PRs, each with its own tests, opened one at a time:

1. prefill progress commit;
2. prefill park instead of abort;
3. single victim and exact-chain cache primitives;
4. decode pause and resume;
5. memory-aware resume;
6. cached-media text suffix;
7. the memory_scheduling setting (resolved by the memory enforcer; nothing reads it yet);
8. never-fits gate, fits-later wait and capacity API;
9. resume under the hard watermark;
10. enforcer changes;
11. admission progress clock and solo admission;
12. two follow-ups: the hot cache on the never-fits line, and the stall clock across park holds.

I've built all of them on main 013bc3f: each applies on the previous one and passes the default suite. The never-fits slice is adapted to #3482 and #4252. The memory_scheduling setting comes before the slices it gates, and slices 8, 10 and 11 consult it from their first commit, so main's default admission stays the same after each merge unless an operator sets an explicit memory ceiling or memory_scheduling=on.

Questions:

• Is this direction welcome?
• Would you prefer the never-fits/capacity part (useful without pausing) first?
• Should the env-only switches become settings?

3. Slow suffix prefill after a restored TurboQuant prefix (PR ready)

A paged-SSD hit restores TurboQuant layers as quantized caches (#1842), and the scheduler keeps them quantized while it prefills the uncached suffix, so every chunk attends to the whole quantized prefix through the tiled quantized-attention path. On main c0b1056, a 36k-token suffix after a 74k-token restored prefix took 197.9 s against 82.2 s cold for the same positions (2.41x); with TurboQuant off the two match. The fix copies the restored layers into dense KV once before a suffix longer than one prefill chunk (0.43 s for 15 layers x 74k tokens) and re-uses the prefix codes unchanged at the end of prefill: 87.3 s against 94.7 s cold (0.92x). Output differs from cold only at near-ties, exactly as without the fix, and the prefill reservation now plans the dense layers (+2.23 GB). 2 commits, 16 tests: #4442.

4. Related findings (details in the write-up)

• Prefix-cache breaks from mid-conversation system updates. After a tool result, user_note_safe falls back to strict, which rewrites the prompt head. An opt-in policy that renders updates in place keeps the prefix (related: Prefix-cache match capped at ~25k tokens on verified linear (append-only) conversations — possibly cache_control breakpoint sensitivity #3608, Bug: Qwen3.6 false-positive mid-system preservation with leading developer block on /v1/responses #1966).
• Thinking stripped mid-loop by synthetic client messages. An agent client's synthetic "attached image" user message moves the template's last_query_index, so every image turn breaks the prefix. An opt-in thinking-retention guard gives 97.4-99.7% prefix hits on image turns (related: TurboQuant KV breaks preserve thinking #3378, preserve_thinking fails with Qwen3.6_35B on 0.3.7rc2 #900, Main Agent Session KV Prefix Cache Always Misses After Upgrading to 0.3.8.x (Working on 0.3.6) #1003).
• Decode fairness does not cover multimodal prompts. A request with an image anywhere in its history prefills in the synchronous external loop, so concurrent decode freezes for 25-46 s, and up to 155 s on cold 43k prefills (related: Decode fairness retains only 13–23% of solo generation throughput during concurrent long prefill on oMLX 0.6.3 #3209).

I'm happy to adapt the order or scope to your preferences, and to keep each PR small. Thanks for oMLX.

Victor Perez (victor.perez@e-evolution.com), lead maintainer of ADempiere (https://github.com/adempiere/adempiere/tree/hotfix/3.9.4.002)

[e-Evolution]

The pause-and-resume stack is now open as 12 PRs, rebased on current main (f3ad34d). Merge order:

1. feat(scheduler): commit completed prefill blocks to the prefix cache during prefill #4478 P-a: prefill progress commit (ready for review)
2. feat(scheduler): park a chunked prefill under memory pressure instead of aborting it #4479 P-b: prefill park instead of abort
3. feat(memory): abort one victim instead of all when nothing is parkable #4480 S1: single victim and exact-chain cache primitives
4. feat(scheduler): pause and resume a decoding row through an exact SSD prefix #4481 S2: decode pause and resume
5. feat(scheduler): resume a parked request when its memory need fits #4482 S3: memory-aware resume priority
6. feat(scheduler): prefill a cached-media prompt's text suffix in the chunked path #4483 S5: cached-media text suffix through the chunked path
7. feat(memory): add the memory_scheduling setting #4484 G: memory_scheduling setting
8. feat(memory): reject never-fits requests with the numbers and let fits-later requests wait #4485 S4: never-fits gate, fits-later wait, capacity API
9. fix(scheduler): resume parked requests under the hard watermark without thrashing #4486 S6: resume under the hard watermark
10. fix(memory): spare external-loop prefills and pause before the emergency abort-all #4487 S7: enforcer: external-loop prefills, no live victim, emergency pause-first
11. fix(scheduler): time the admission stall by progress and admit a lone head that fits #4488 S8: admission progress clock and solo admission
12. fix(memory): raise the never-fits line by the hot cache and keep the stall clock across park holds #4489 S9: hot cache in the never-fits line, stall clock across park holds

P-a is ready for review. The rest are drafts: each one lists the commits to review, and I rebase the next one after each merge. Every new behaviour is behind a setting that is off by default. Following #4396, the env-only switches became settings (memory.abort_single_victim, scheduler.media_text_suffix_chunked) and the debug-only forced-pause switch is gone. CI's command passes at every slice tip with only the 3 test_gemma4_rendering.py failures that main also has in my environment.

Thanks for merging #4442 and for #4453.

[jundot]

The direction here is right. Keeping in-flight work instead of aborting everything under memory pressure is worth having, and the analysis in #4440 made the problem much easier to see.

I can't take a 12-PR stack of about 17k lines for it, though. That's more new state in the scheduler and memory enforcer than I can maintain. I went back to the root cause and implemented a smaller version in #4509. It stores the live requests to the paged SSD cache at the hard watermark, drops the batch in one step, and resumes them from the stored prefix on the same stream. It also keeps the computed prefix when a client cancels, so a resend picks up where it stopped. The PR credits your idea.

Once #4509 is merged I'll close this stack. If you have feedback on the approach in the meantime, please leave it on #4509.

[jundot]

@e-Evolution I merged #4509, including the external-prefill cancellation and TurboQuant prefill fixes. Your analysis helped shape the approach, and your testing caught the TurboQuant gap before merge.

I'm closing #4478 through #4489 in favor of the merged implementation. The cached-media text-suffix route and never-fits/capacity API are not included, so I'd like to keep this discussion open for those proposals and agree on their scope before any new PRs.