#!/usr/bin/env python3
"""
Patches vendor/llama.cpp for macOS Unified Memory SSD-streamed MoE (--cpu-moe -nr -lzm on):
1. llama-model-loader.cpp (init_mappings): includes 3D host-overridden expert tensors (t->ne[2] > 1)
in file_lazy_ranges so llama_mmap skips startup POSIX_MADV_WILLNEED on the 17.0 GiB experts
and marks them POSIX_MADV_RANDOM instead.
2. llama-mmap.cpp: sets fcntl(fd, F_RDAHEAD, 0) on macOS when lazy_ranges is non-empty to prevent
APFS sequential cluster read-ahead amplification on expert page faults.
3. ggml-cpu.c (ggml_compute_forward_mul_mat_id): AFTER ggml_barrier (once ith==0 has populated
matrix_row_counts), all nth threads issue parallel posix_madvise(..., POSIX_MADV_WILLNEED)
460 KiB cluster reads (shared vm_map_lock_read in XNU) across all active experts for both
prefill and decode, and for decode (ids->ne[1] <= 2) compute disjoint active experts in parallel.
"""
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent / "vendor" / "llama.cpp"
def patch_model_loader() -> None:
p = ROOT / "src" / "llama-model-loader.cpp"
txt = p.read_text(encoding="utf-8")
if "file_lazy_ranges" in txt:
print("llama-model-loader.cpp already patched")
return
old = """ std::unique_ptr<llama_mmap> mapping = std::make_unique<llama_mmap>(file.get(), prefetch_size, is_numa,
lazy.for_file(idx));"""
new = """ llama_mmap::ranges file_lazy_ranges = lazy.for_file(idx);
if (tensor_buft_overrides) {
for (const auto & [ctx_key, ctx_ptr] : ctx_map) {
if (!ctx_ptr || !ggml_backend_buft_is_host(ctx_key.buft)) {
continue;
}
for (ggml_tensor * t = ggml_get_first_tensor(ctx_ptr.get()); t != nullptr; t = ggml_get_next_tensor(ctx_ptr.get(), t)) {
const auto * w = get_weight(ggml_get_name(t));
if (w && w->idx == idx && t->ne[2] > 1) {
file_lazy_ranges.emplace_back(w->offs, w->offs + ggml_nbytes(t));
}
}
}
}
std::unique_ptr<llama_mmap> mapping = std::make_unique<llama_mmap>(file.get(), prefetch_size, is_numa,
file_lazy_ranges);"""
assert old in txt, "Target block not found in llama-model-loader.cpp"
p.write_text(txt.replace(old, new, 1), encoding="utf-8")
print("Patched llama-model-loader.cpp")
def patch_mmap() -> None:
p = ROOT / "src" / "llama-mmap.cpp"
txt = p.read_text(encoding="utf-8")
if "F_RDAHEAD" in txt:
print("llama-mmap.cpp already patched")
return
old = """ int fd = file->file_id();
int flags = MAP_SHARED;
if (numa) { prefetch = 0; }"""
new = """ int fd = file->file_id();
if (numa) { prefetch = 0; }
int flags = MAP_SHARED;
#if defined(__APPLE__)
if (!lazy_ranges.empty()) {
fcntl(fd, F_RDAHEAD, 0);
}
#endif"""
assert old in txt, "Target block not found in llama-mmap.cpp"
p.write_text(txt.replace(old, new, 1), encoding="utf-8")
print("Patched llama-mmap.cpp")
def patch_ggml_cpu() -> None:
p = ROOT / "ggml" / "src" / "ggml-cpu" / "ggml-cpu.c"
txt = p.read_text(encoding="utf-8")
if "POSIX_MADV_WILLNEED" in txt:
print("ggml-cpu.c already patched")
return
old_inc = "#if defined(__gnu_linux__)\n#include <syscall.h>\n#endif"
new_inc = old_inc + "\n#if defined(__APPLE__)\n#include <sys/mman.h>\n#endif"
assert old_inc in txt, "Include target not found in ggml-cpu.c"
txt = txt.replace(old_inc, new_inc, 1)
old_barrier = """ ggml_barrier(params->threadpool);
for (int cur_a = 0; cur_a < n_as; ++cur_a) {"""
new_barrier = """ ggml_barrier(params->threadpool);
#if defined(__APPLE__)
if (nb02 >= 65536 && n_as <= 512) {
int active_experts[512];
int n_active = 0;
for (int cur_a = 0; cur_a < n_as; ++cur_a) {
if (matrix_row_counts[cur_a] > 0) {
active_experts[n_active++] = cur_a;
}
}
// Issue parallel 460 KiB cluster prefetch across all active experts (shared vm_map_lock_read)
for (int idx_a = ith; idx_a < n_active; idx_a += nth) {
const int cur_a = active_experts[idx_a];
posix_madvise((char *) src0->data + cur_a * nb02, nb02, POSIX_MADV_WILLNEED);
}
if (ids->ne[1] <= 2) {
const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata;
const size_t row_size = ggml_row_size(vec_dot_type, ne10);
for (int idx_a = ith; idx_a < n_active; idx_a += nth) {
const int cur_a = active_experts[idx_a];
const int64_t cne1 = matrix_row_counts[cur_a];
const char * src0_cur = (const char *) src0->data + cur_a * nb02;
ggml_compute_forward_mul_mat_id_one_chunk(
dst, src0, src1, ids, cur_a,
0, ne01, 0, cne1,
src0_cur, matrix_rows, row_size, src1_cont, wdata
);
}
return;
}
}
#endif
for (int cur_a = 0; cur_a < n_as; ++cur_a) {"""
assert old_barrier in txt, "Barrier target not found in ggml-cpu.c"
txt = txt.replace(old_barrier, new_barrier, 1)
p.write_text(txt, encoding="utf-8")
print("Patched ggml-cpu.c (parallel cluster prefetch for prefill+decode + per-expert decode)")
def patch_context() -> None:
p = ROOT / "src" / "llama-context.cpp"
txt = p.read_text(encoding="utf-8")
if "POSIX_MADV_WILLNEED" in txt:
print("llama-context.cpp already patched")
return
old_inc = '#include "llama-context.h"'
new_inc = old_inc + "\n#if defined(__APPLE__)\n#include <sys/mman.h>\n#endif"
assert old_inc in txt, "Include target not found in llama-context.cpp"
txt = txt.replace(old_inc, new_inc, 1)
old_loop = """ // group consecutive experts and copy them together
for (int64_t first = 0; first < n_expert; ) {"""
new_loop = """#if defined(__APPLE__)
for (int64_t first = 0; first < n_expert; ) {
if (!st.used[first]) { first++; continue; }
int64_t last = first;
while (last + 1 < n_expert && st.used[last + 1]) { last++; }
posix_madvise((char *) src->data + first * expert_size, (last + 1 - first) * expert_size, POSIX_MADV_WILLNEED);
first = last + 1;
}
#endif
// group consecutive experts and copy them together
for (int64_t first = 0; first < n_expert; ) {"""
assert old_loop in txt, "Loop target not found in llama-context.cpp"
txt = txt.replace(old_loop, new_loop, 1)
p.write_text(txt, encoding="utf-8")
print("Patched llama-context.cpp (prefill cluster prefetch in sched_copy_experts)")
if __name__ == "__main__":
patch_model_loader()
patch_mmap()
patch_ggml_cpu()
patch_context()