Opt-in mlock of non-routed weights for SSD streaming (DS4_MLOCK_NONROUTED)#499
Opt-in mlock of non-routed weights for SSD streaming (DS4_MLOCK_NONROUTED)#499hi-pauls wants to merge 1 commit into
Conversation
When --ssd-streaming keeps routed MoE experts on disk, only the routed expert cache buffer is Metal-wired; the non-routed weights (attention, shared experts, output head, embeddings) are mapped from the GGUF and remain evictable. On memory-constrained machines the OS evicts those pages under pressure and every decode step re-reads them from SSD, collapsing throughput. This patch adds a small env-gated pass in model_open() that, when DS4_MLOCK_NONROUTED=1, iterates the parsed tensor table and mlock()s the page-aligned mmap range of every tensor whose name is not one of the routed-expert arrays (.ffn_gate_exps., .ffn_up_exps., .ffn_down_exps.). Routed experts continue to stream from SSD via the existing cache; only their mmap regions are skipped. A summary line reports locked/failed/skipped bytes, and mlock failures fall back gracefully so the process keeps running. Measured on a 64 GB M1 Max running the Q2 Flash imatrix quant with 25 GB expert cache and 256K context, decode at 2048-token frontier went from 1.27 tok/s (baseline) to 4.54 tok/s (+257%) once the non-routed portion (8.22 GiB) stayed resident. The gain scales with how aggressively the OS was previously evicting those weights.
|
@hi-pauls do you have any other Macs to test this on? I don't know about q2-q4-imatrix and q4-imatrix, and DeepSeek V4 Pro would easily run on something like an M3 Ultra 256GB |
|
Hi @hi-pauls , I think you are trying to optimize a suboptimal configuration to begin with... The plot down below has been taken on a M2 max 32GB of RAM, you can see that up to 15GB of expert cache budget it runs at about 6 t/s, as you keep increasing the cache size the kernel starts evicting the pages that contain the unrouted weights as you reported. If you mlock those weights, the system will start swapping other programs or event critical systems components at the limit it can cause OOM errors and crash your system. Therefore I do not think that including the option of mlocking is safe. The safest thing to do is just avoid extreme memory pressure, which will result in better performances overall without the risk of OOM. 
Best, |
3 participants
Summary
Adds an opt-in
DS4_MLOCK_NONROUTED=1environment flag thatmlock()s the mmap ranges of every non-routed-expert tensor afterparse_tensors(). Routed experts (.ffn_gate_exps.,.ffn_up_exps.,.ffn_down_exps.) are explicitly skipped so they keep streaming from SSD via the existing cache path. Everything else — attention projections, shared experts, output head, embeddings, norms — becomes non-evictable.Motivation
On memory-constrained machines (64 GB unified memory on Apple Silicon, tested), the current
--ssd-streamingpath only wires the routed expert cache buffer. Non-routed weights are file-backed mmap; the OS evicts them under pressure and every decode step re-reads them from SSD, tanking throughput. mlock'ing them keeps the "hot skeleton" resident while routed experts continue to stream normally.Benchmark
64 GB M1 Max, Q2 Flash imatrix quant, 25 GB expert cache, 256K ctx, gen 8192 (single 2048-token frontier):
DS4_MLOCK_NONROUTED=1+257% decode with 8.22 GiB additional wired. Broader cross-domain sweep (coding, hardware, rocketry, optics, RAG at 10/15/20/25/30 GB caches) showed consistent double-digit percentage gains, with 30 GB cache + mlock landing at ~7.7 tok/s average vs ~4.9 tok/s without.
Design notes
abs_offset..abs_offset+bytesto page boundaries so mlock accepts the requestTest plan
make)DS4_MLOCK_NONROUTED=0(default): no behavior change, no log lineDS4_MLOCK_NONROUTED=1: logslocked X GiB, failed 0 GiB, skipped routed Y GiB, verifiedX + Y == total tensor bytes(80.76 GiB for Q2 Flash) — no unclassified tensorsvmmap -summaryconfirms the non-routed portion moves from evictable to wiredds4-agent --non-interactive -p ...returns same greedy tokens)DeepSeek-V4-Flash-IQ2XXS-…-imatrix.gguf); classification is name-based so Q4 Flash and Pro variants should work identically