From 373475f8bf76684b3777018387fb8343506c7d70 Mon Sep 17 00:00:00 2001 From: Paul Schulze Date: Sat, 4 Jul 2026 20:39:45 +0200 Subject: [PATCH] Add opt-in mlock of non-routed weights under DS4_MLOCK_NONROUTED 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. --- ds4.c | 64 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/ds4.c b/ds4.c index 640511eb0..d1afe57ef 100644 --- a/ds4.c +++ b/ds4.c @@ -1938,6 +1938,69 @@ static void parse_tensors(ds4_model *m, ds4_cursor *c) { } } +/* Env-gated: wire the mmap ranges of every non-routed-expert tensor so macOS + * cannot page them out under memory pressure. Routed experts still stream + * from the SSD-backed cache as usual. Requires DS4_MLOCK_NONROUTED=1. Large + * lock ranges typically need root (mlock user limits are tiny by default). */ +static bool tensor_name_is_routed_expert(ds4_str name) { + static const char *needles[] = { + ".ffn_gate_exps.", ".ffn_up_exps.", ".ffn_down_exps.", + }; + for (size_t i = 0; i < sizeof(needles) / sizeof(needles[0]); i++) { + const char *needle = needles[i]; + size_t nlen = strlen(needle); + if (name.len < nlen) continue; + for (uint64_t j = 0; j + nlen <= name.len; j++) { + if (memcmp(name.ptr + j, needle, nlen) == 0) return true; + } + } + return false; +} + +static void model_mlock_nonrouted(const ds4_model *m) { + const char *env = getenv("DS4_MLOCK_NONROUTED"); + if (!env || env[0] == '\0' || env[0] == '0') return; + + long page_size = sysconf(_SC_PAGESIZE); + if (page_size <= 0) page_size = 16384; + + uint64_t requested = 0, locked = 0, failed = 0, skipped = 0; + int last_errno = 0; + + for (uint64_t i = 0; i < m->n_tensors; i++) { + const ds4_tensor *t = &m->tensors[i]; + if (tensor_name_is_routed_expert(t->name)) { skipped += t->bytes; continue; } + if (t->bytes == 0) continue; + if (t->abs_offset > m->size || t->bytes > m->size - t->abs_offset) continue; + + uint64_t start = t->abs_offset; + uint64_t end = t->abs_offset + t->bytes; + uint64_t pstart = start & ~((uint64_t)page_size - 1); + uint64_t pend = (end + (uint64_t)page_size - 1) & ~((uint64_t)page_size - 1); + if (pend > m->size) pend = m->size; + uint64_t len = pend - pstart; + requested += len; + if (mlock((const uint8_t *)m->map + pstart, (size_t)len) == 0) { + locked += len; + } else { + failed += len; + last_errno = errno; + } + } + + const double gib = 1024.0 * 1024.0 * 1024.0; + fprintf(stderr, + "ds4: mlock non-routed weights: locked %.2f GiB, failed %.2f GiB, " + "skipped routed %.2f GiB\n", + (double)locked / gib, (double)failed / gib, (double)skipped / gib); + if (failed && last_errno) { + fprintf(stderr, "ds4: mlock error (last): %s (try running with sudo, " + "or raise the wired-memory limit)\n", + strerror(last_errno)); + } + (void)requested; +} + /* Open and map the GGUF once. Metal needs a shared mapping for no-copy * MTLBuffers; CPU uses a private read-only mapping to avoid Darwin VM stress. * Tokenizer-only callers pass prefetch_cpu=false so inspecting tokens never @@ -1988,6 +2051,7 @@ static void model_open(ds4_model *m, const char *path, bool metal_mapping, parse_tensors(m, &c); if (!metal_mapping && prefetch_cpu) model_prefetch_cpu_mapping(m); + if (metal_mapping) model_mlock_nonrouted(m); } static void print_size(uint64_t bytes) {