[jvmenen]

Why

A conversation that comes back from the SSD cache, after it was evicted from memory or after a server restart, is processed much more slowly than one restored from memory. Its prefill runs at about a fifth of the speed and its decode is slower too, for the rest of that conversation.

How it works, in plain terms

MTPLX can lay out a conversation's KV cache in different ways, and each request asks for the layout that suits it: for a long chat, a dense cache. A restore from memory builds the cache with that request's layout. The restore from SSD did not: it built the runtime's default layout, which for a long context is the paged KV cache, much slower for this use. This change passes the request's layout to the SSD restore as well, so both restore paths build the same cache. The stored state is read into it unchanged; without a requested layout nothing changes.

What changes

• mtplx/session_bank.py: restore() passes the request's cache_factory and mtp_cache_factory to _restore_cold, which builds the restored cache with them, as the RAM restore does. Without a factory it keeps runtime.make_cache().
• tests/test_ssd_restore_cache_layout.py: new.
• CHANGELOG.md: entry.

Results

Qwen3.8-27B, server restarted with a warm SSD cache, each conversation continued for one turn, two runs per variant. Without this change every SSD restore built the paged cache, with it every restore built the dense cache.

• Time to first token: 24 to 36 s down to 7 to 9 s at 50K tokens, 67 s down to 7 to 13 s at 80K.
• Decode: 14 to 16 up to 23 to 27 tokens per second, close to a restore from memory (25 to 29).
• Peak memory at 50K: 3 to 4 GB lower.

Tests

tests/test_ssd_restore_cache_layout.py checks that an SSD restore builds the cache from the request's cache_factory (and the MTP history from mtp_cache_factory), and keeps the runtime default without one; the session bank and SSD spill tests pass unchanged.

🤖 Generated with Claude Code