[jvmenen]

Why

When two clients take turns on one server, each with a long conversation, the idle conversation's SSD copy is often never written. Every commit queues that write on the idle lane, which runs only in a quiet window, and with one client's request always waiting that window does not come. When a request then does not fit in memory, the admission releases the idle conversation and cancels its queued write, so a conversation that existed only in RAM is lost. Its next turn re-reads the whole conversation, 80K to 100K tokens and several minutes, and the other client waits behind that prefill.

The same loss happens inside one conversation. A client's side call (a title or summary request) can arrive under the conversation's session id while the conversation's newest state waits in RAM with its write queued. When the side call's prefill trips the memory guard, the shed before an abort gives back the request's own entries it does not restore from, which is the conversation itself. In a field log on a 64 GB Mac, an 8K side request released its 23K conversation (9 GB, nothing on disk) and the next turn waited 69 s for its first token.

How it works, in plain terms

Before the memory guard lets a conversation go, it checks whether that conversation's newest state is still waiting for its SSD write. If so, it writes it out right away, then frees the memory as before. The write streams one tensor at a time in units of at most 32 MiB, so it needs almost no extra memory: the arrays it reads are the ones about to be freed. It has a time limit; if the disk cannot finish in time, or the SSD cache refuses the write, the conversation is dropped as it was before. The next turn of that conversation restores from disk instead of re-reading it.

What changes

• SessionBank.release_sessions(write_out_before_release=...); EngineSessionManager.release_idle_sessions passes it through. The memory guard sets it wherever it releases on the model-owner thread: the admission's idle-conversation and own-sibling steps, and the shed before a prefill abort (idle conversations and the request's own entries). During that shed the prefill waits between chunks, so the engine allocates nothing while the write runs.
• SessionBankColdTier.spill_entry(give_up_at_s=...): a deadline replaces the foreground signal for a caller that is itself the foreground work. Advertised by SUPPORTS_SPILL_DEADLINE.
• RELEASE_WRITE_OUT_MAX_S = 30.0 per release.
• Receipts gain written_out_entries and write_out_s (the prefill shed also own_session_written_out_entries); the eviction log gains release_write_out (outcome, time).
• Unchanged: the pressure trim, the allocation-failure shed, and everything when the SSD cache is off.

Results

Checked for correctness on a 27B with segmented KV. Two long conversations (about 60K tokens each) took turns while two small clients kept the server busy, so the idle lane never ran, and the memory limit was set so that each turn had to release the other conversation. All four releases wrote the conversation first (release_write_out outcome written, on_ssd_entries 1, dropped_entries 0). The next turn of each released conversation restored from SSD (ssd_clone, 97 to 98 % of the prompt cached) instead of re-reading it. While each write ran, sampled every 20 ms, MLX active memory rose 143 to 262 MiB in three writes and 1.3 GiB in one, and the system's available memory did not drop by more than 0.2 GiB. Without the change, the same release cancels the queued write and the conversation's next turn re-reads it.

The second commit (the prefill shed and the own-session release) is covered by the tests below; it has not been measured in a run like the one above yet.

Not changed here: the pressure trim can still drop an unwritten conversation (seen once in a run with a higher limit). It runs on another thread and is left for a separate change.

Tests

tests/test_release_writes_out_unwritten.py (13 tests, real SessionBank and SessionBankColdTier, the engine busy): the deadline against the foreground signal; nothing restorable left after a deadline; an unwritten entry is written before release and its queued job is gone; an entry pushed out of RAM with its write queued is written too; out of time it is dropped as before; without the flag the release still cancels; an entry already on disk is not rewritten; no SSD cache changes nothing; a busy session is skipped; the admission scenario, where the released conversation's newest entry is on SSD; and three side-request cases: the own-session release, the prefill shed on a death signature (an idle and the own conversation), and the admission's own-sibling step with the conversation and the request under one session id. The 8 release and admission tests fail on main, the three side-request tests fail on the first commit, and all pass with the change. Full suite passes.