Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163#!/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()