From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 From: HelloSun Date: Thu, 7 Oct 2026 17:00:00 +0200 Subject: [PATCH] st: SmallThinker expert 級 SSD 分頁 熱的 expert 權重留在 RAM，冷的用 madvise + posix_fadvise 丟回 SSD。 計算路徑完全不動（仍是 llama.cpp 自己的 mul_mat_id），所以輸出與上游逐字相同。 --- diff --git a/.st-patched b/.st-patched new file mode 100644 index 000000000..e69de29bb diff --git a/.st-unpatched b/.st-unpatched new file mode 100644 index 000000000..e69de29bb diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index afdaddc79..31b2c2896 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -18,6 +18,7 @@ set(LLAMA_CORE_SOURCES llama-cparams.cpp llama-grammar.cpp llama-graph.cpp + st_pager.cpp llama-hparams.cpp llama-impl.cpp llama-io.cpp diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 1112ad885..e335a9d5f 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1,5 +1,7 @@ #include "llama-graph.h" +#include "st_pager.h" + #include "llama-impl.h" #include "llama-model.h" #include "llama-batch.h" @@ -2205,6 +2207,14 @@ ggml_tensor * llm_graph_context::build_moe_ffn( //call early so that topk-moe can be used ggml_build_forward_expand(gf, weights); + // ST expert pager：在真正算 expert 之前，先知道這一層要用哪些 expert， + // 然後把用不到的 expert 權重丟回 SSD。計算本身完全不動 —— + // 下面仍然是 llama.cpp 自己的 mul_mat_id，所以輸出與上游一致。 + if (ggml_tensor * st_note = st::build_routing_op(ctx0, selected_experts, il, + (int) n_expert, (int) n_expert_used)) { + ggml_build_forward_expand(gf, st_note); + } + cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens); if (weight_before_ffn) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 1f3b80b08..c7b77476e 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1,5 +1,7 @@ #include "llama-model.h" +#include "st_pager.h" + #include "llama-arch.h" #include "llama-ext.h" #include "llama-hparams.h" @@ -1510,6 +1512,8 @@ void llama_model_base::load_vocab(llama_model_loader & ml) { } bool llama_model_base::load_tensors(llama_model_loader & ml) { + // ST: 讀環境變數要在 init_mappings 之前（它決定要不要 MAP_POPULATE） + st::set_config_from_env(); const auto & split_mode = params.split_mode; const bool use_mlock = params.load_mode == LLAMA_LOAD_MODE_MLOCK || params.load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK; const auto & tensor_split = params.tensor_split; @@ -1821,7 +1825,13 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // per-tensor activation precision policy prec_policy.load(ml, *this); - ml.init_mappings(true, use_mlock ? &pimpl->mlock_mmaps : nullptr); + // ST expert pager：開著分頁時**不能** prefetch。 + // llama.cpp 預設 prefetch=true 會用 MAP_POPULATE 把整個 GGUF（2.45 GiB） + // 在「載入模型」那一瞬间全部 fault 進 page cache —— 那一刻的 RSS 峰值 + // 就已經超過任何分頁預算了，後面的 madvise 再怎麼丟都沒用（峰值已經發生）。 + // 這是量測方式／時機的坑，不是分頁機制本身有問題。 + // 註意 st::set_config_from_env() 必須在這之前呼叫（見 load_tensors 開頭）。 + ml.init_mappings(!st::will_page(arch_name().c_str()), use_mlock ? &pimpl->mlock_mmaps : nullptr); pimpl->mappings.reserve(ml.mappings.size()); // create the backend buffers @@ -1969,6 +1979,35 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } } + // ST expert pager：把 MoE expert 權重記成 (layer, expert) 的檔案/記憶體區間， + // 讓 build_moe_ffn 插入的 op 能在用不到時把它們丟回 SSD。 + // 註冊失敗不是致命錯誤 —— 那只是不分頁，模型照上游路徑跑。 + { + st::install_signal_handler(); + + void * map_base = nullptr; + size_t map_size = 0; + int map_fd = -1; + if (!pimpl->mappings.empty() && !ml.files.empty()) { + map_base = pimpl->mappings[0]->addr(); + map_size = pimpl->mappings[0]->size(); + map_fd = ml.files[0]->file_id(); + } + + std::map by_name; + for (const auto & [nm, t] : tensors_by_name) { + by_name.emplace(nm, t); + } + + if (getenv("ST_VERBOSE") && !map_base) { + LLAMA_LOG_INFO("%s: [st] 沒有 mapping（mappings=%zu files=%zu），無法分頁\n", + __func__, pimpl->mappings.size(), ml.files.size()); + } + if (map_base && !st::register_model(arch_name().c_str(), map_base, map_size, map_fd, by_name)) { + LLAMA_LOG_INFO("%s: [st] expert 分頁未啟用（arch=%s）\n", __func__, arch_name().c_str()); + } + } + return true; } diff --git a/src/st_pager.cpp b/src/st_pager.cpp new file mode 100644 index 000000000..12da88c88 --- /dev/null +++ b/src/st_pager.cpp @@ -0,0 +1,588 @@ +// ST expert pager —— 熱 expert 留在 RAM，冷 expert 丟回 SSD +// +// 設計取捨（為什麼不照抄 sddqwen35a3b_v01 的 pread arena）： +// sddqwen35a3b 做的是「把 expert 權重 pread 進自己的 arena，用自訂 ggml op +// 重寫量化點積」。那在 22 GB 權重 / 小 RAM 機器上是必要的，代價是要自己重寫 +// Q4_K 的點積 kernel，數值正確性得另外驗證。 +// SmallThinker 只有 2.45 GiB，權重可以用 file-backed mmap 留在 SSD， +// 所以本專案改成「**計算路徑完全不動，只把用不到的 expert 頁丟掉**」： +// - 計算仍是 llama.cpp 自己的 kernel → 輸出與上游一致，沒有重寫 kernel 的風險。 +// - 淘汰的最小單位是「一個 expert 的三段權重」，不是 page。 +// 代價：不能像 arena 那樣用 prefetch 提前重疊 I/O，SSD 讀取是 demand-paging +// （第一次碰到缺頁才讀）。這是本專案目前的取捨，量測見 docs。 +// +// 淘汰的兩步驟（順序不能換，這是實測踩出來的）： +// a. madvise(MADV_DONTNEED) → 拿掉本行程 PTEs，行程總 RSS 立刻下降。 +// b. posix_fadvise(DONTNEED) → 核心的 invalidate_mapping_pages() 會跳過 +// 「仍被某個 VMA 映射」的頁，所以 (a) 必須先做；(b) 做完這些頁才真的從 +// page cache 消失，下次碰到才是真的讀 SSD。 +// 只做 (a)：RSS 下降但 read_bytes 不變（頁還在 page cache）。 +// 只做 (b)：page cache 下降但行程 RSS 不變（VMA 還映射著）。 + +#include "st_pager.h" + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace st { + +// ------------------------------------------------------------------ 設定 + +static bool g_enabled = false; +static bool g_verbose = false; +static size_t g_budget = 0; // expert 權重可以常駐 RAM 的上限（bytes） +static int g_hot_tokens = 8; // 最近 N 個 tick 用過的 expert 視為熱 + +struct span { + uint8_t * addr = nullptr; + size_t bytes = 0; + size_t file_off = 0; +}; + +struct slot { + span parts[st::N_PART]; + uint64_t last_used = 0; + uint64_t uses = 0; + bool resident = false; +}; + +// 一個 expert 的三段位元組數**每層可能不同**（實測：down_exps 有些層是 Q6_K、 +// 有些是 Q4_K），所以逐 slot 算，不能用一個全模型常數。 +static inline size_t slot_bytes_of(const slot & sl) { + return sl.parts[0].bytes + sl.parts[1].bytes + sl.parts[2].bytes; +} + +static std::vector g_slots; // [layer * n_expert + expert] +static int g_n_layer = 0; +static int g_n_expert = 0; +static size_t g_slot_bytes = 0; // 三段加總（= 一個 expert 的權重） +static size_t g_part_bytes[st::N_PART] = { 0, 0, 0 }; +static uint8_t * g_map_base = nullptr; +static int g_fd = -1; +static size_t g_map_size = 0; +static uint64_t g_tick = 0; +static std::mutex g_mu; +static stats g_st; + +static size_t page_size() { + static const size_t ps = (size_t) sysconf(_SC_PAGESIZE); + return ps; +} + +static uint64_t env_u64(const char * name, uint64_t dflt) { + const char * e = getenv(name); + if (!e || !*e) return dflt; + char * end = nullptr; + const unsigned long long v = strtoull(e, &end, 10); + if (end == e) return dflt; + return v; +} + +void set_config_from_env() { + g_enabled = env_u64("ST_PAGER", 1) != 0; + g_verbose = env_u64("ST_VERBOSE", 0) != 0; + + const uint64_t budget_mb = env_u64("ST_RAM_BUDGET_MB", 0); + const uint64_t reserve_mb = env_u64("ST_RESERVE_MB", 0); + const uint64_t arena_mb = env_u64("ST_ARENA_MB", 0); + + uint64_t mb; + if (arena_mb > 0) { + mb = arena_mb; // llama_server.sh 已經算好 + } else if (budget_mb > reserve_mb) { + mb = budget_mb - reserve_mb; + } else { + mb = budget_mb; + } + g_budget = (size_t) mb * 1024 * 1024; + + g_hot_tokens = (int) env_u64("ST_HOT_TOKENS", 8); + if (g_hot_tokens < 1) g_hot_tokens = 1; +} + +bool enabled() { return g_enabled && !g_slots.empty(); } + +// 這個查詢會在 init_mappings() **之前**被呼叫（那時還沒載入權重， +// register_model 當然還沒跑過），所以不能只看 enabled()。 +// 它只能回答「這個模型**預計**會不會被分頁」。 +bool will_page(const char * arch) { + if (!g_enabled) return false; + return arch && strcmp(arch, "smallthinker") == 0; +} + +bool paging_active() { return enabled(); } + +// ------------------------------------------------------------------ 統計 + +stats get_stats() { + std::lock_guard lk(g_mu); + stats s = g_st; + s.budget_bytes = g_budget; + size_t rb = 0; + for (const auto & sl : g_slots) { + if (sl.resident) rb += slot_bytes_of(sl); + } + s.resident_bytes = rb; + return s; +} + +// 讀自己的 /proc/self/io。外部工具（例如 tools/ab_test.py）在子行程結束後 +// 再讀 /proc//io 已經來不及了 —— 行程一退出，/proc/ 就消失， +// 結果會是「io_delta: {}」，看起來像「完全沒有 I/O」。 +// 所以這個數字只能由**行程自己**在退出前寫下來。 +static bool self_io(uint64_t * read_bytes, uint64_t * rchar) { + *read_bytes = *rchar = 0; + FILE * f = fopen("/proc/self/io", "r"); + if (!f) return false; + char k[64]; + unsigned long long v; + // 注意開頭那個空格：沒有它的話 %[^:] 會把上一行的 '\n' 一起讀進 key， + // strcmp 永遠比不到 "read_bytes"，結果永遠是 0。這個 bug 真的踩過， + // 症狀是「量到 read_bytes = 0」，看起來像「完全沒讀 SSD」。 + while (fscanf(f, " %63[^:]: %llu", k, &v) == 2) { + if (!strcmp(k, "read_bytes")) *read_bytes = v; + else if (!strcmp(k, "rchar")) *rchar = v; + } + fclose(f); + return true; +} + +void dump_stats(const char * path) { + const stats s = get_stats(); + uint64_t rb = 0, rc = 0; + self_io(&rb, &rc); + FILE * f = fopen(path, "w"); + if (!f) return; + fprintf(f, + "{\n" + " \"routing_calls\": %llu,\n" + " \"expert_touches\": %llu,\n" + " \"hits\": %llu,\n" + " \"misses\": %llu,\n" + " \"evictions\": %llu,\n" + " \"evicted_bytes\": %llu,\n" + " \"sweeps\": %llu,\n" + " \"resident_bytes\": %llu,\n" + " \"budget_bytes\": %llu,\n" + " \"hit_rate\": %.4f,\n" + " \"proc_self_read_bytes\": %llu,\n" + " \"proc_self_rchar\": %llu\n" + "}\n", + (unsigned long long) s.routing_calls, + (unsigned long long) s.expert_touches, + (unsigned long long) s.hits, + (unsigned long long) s.misses, + (unsigned long long) s.evictions, + (unsigned long long) s.evicted_bytes, + (unsigned long long) s.sweeps, + (unsigned long long) s.resident_bytes, + (unsigned long long) s.budget_bytes, + s.expert_touches ? (double) s.hits / (double) s.expert_touches : 0.0, + (unsigned long long) rb, + (unsigned long long) rc); + fclose(f); +} + +// 行程退出前自動寫一次統計（atexit） +static void install_atexit() { + static const char * path = getenv("ST_STATS_FILE"); + if (!path || !*path) return; + static char buf[1024]; + snprintf(buf, sizeof(buf), "%s", path); + static std::string keep(buf); + std::atexit([] { dump_stats(keep.c_str()); }); +} + +// ------------------------------------------------------------------ 註冊 + +static bool expert_layout_ok(const ggml_tensor * t, const char * what) { + // 關鍵前提：expert 維度必須是最外層且整段連續，單一 expert 才是 + // 「一段連續位元組」。若上游改了 layout，這裡直接擋下來， + // 而不是算出錯誤的位址把別人的記憶體丟掉。 + // + // 注意：量化張量的 nb[0] 是「一個 block 的位元組數」，不是 1， + // 所以 nb[1] != ne[0]*nb[0]。唯一要檢查的是 nb[2] == ne[1]*nb[1] + // —— expert 之間沒有任何 padding/stride。 + if (t->ne[2] <= 1) return false; + if (t->nb[2] != (size_t) t->ne[1] * t->nb[1]) return false; + if (g_verbose) { + fprintf(stderr, "[st] %s: ne=[%lld,%lld,%lld] nb=[%zu,%zu,%zu] 每 expert %zu bytes\n", + what, (long long) t->ne[0], (long long) t->ne[1], (long long) t->ne[2], + t->nb[0], t->nb[1], t->nb[2], t->nb[2]); + } + return true; +} + +bool register_model(const char * arch, + void * mapping_base, size_t g_map_size_hint, int fd, + const std::map & by_name) { + if (!g_enabled) return false; + if (strcmp(arch, "smallthinker") != 0) { + if (g_verbose) fprintf(stderr, "[st] arch=%s 不支援 expert 分頁\n", arch); + return false; + } + + g_map_base = (uint8_t *) mapping_base; + g_map_size = g_map_size_hint; + g_fd = fd; + + auto get = [&](int il, int p) -> ggml_tensor * { + char nm[64]; + snprintf(nm, sizeof(nm), "blk.%d.ffn_%s_exps.weight", il, + p == st::GATE ? "gate" : (p == st::UP ? "up" : "down")); + auto it = by_name.find(nm); + return it == by_name.end() ? nullptr : it->second; + }; + + // 層數由第一層之後往後掃到空為止。 + std::vector> layers; + int n_expert = 0; + size_t slot_bytes = 0; // 三段加總 + + for (int il = 0; ; il++) { + std::array row{}; + bool ok = true; + for (int p = 0; p < st::N_PART; p++) { + row[p] = get(il, p); + if (!row[p]) { ok = false; break; } + } + if (!ok) break; + + for (int p = 0; p < st::N_PART; p++) { + char what[80]; + snprintf(what, sizeof(what), "blk.%d part %d", il, p); + if (!expert_layout_ok(row[p], what)) { + fprintf(stderr, "[st] %s 的 expert 維度不連續，放棄分頁（回到上游行為）\n", what); + return false; + } + const int64_t ne2 = row[p]->ne[2]; + const size_t nb2 = row[p]->nb[2]; + if (il == 0) { + n_expert = (int) ne2; + g_part_bytes[p] = nb2; + } + // **逐層**累加，不能用 layer 0 的大小代表全部： + // 實測這個模型有 16 層的 down_exps 是 Q6_K、16 層是 Q4_K， + // 用 layer 0 會把總量高估 148 MiB（2092 → 2241 MiB，差 7%）。 + slot_bytes += nb2; + + if ((int) ne2 != n_expert) { + // 只要求 expert **數量**一致。三段的量化型別可以逐層不同 + // （實測 down_exps 有些層是 Q6_K、有些是 Q4_K）， + // 所以位元組數逐 slot 存，不拿 layer 0 當基準。 + fprintf(stderr, "[st] layer %d part %d 的 expert 數量不一致（%lld vs %d）\n", + il, p, (long long) ne2, n_expert); + return false; + } + } + layers.push_back(row); + } + + if (layers.empty()) return false; + + g_n_expert = n_expert; + g_slot_bytes = slot_bytes; + g_n_layer = (int) layers.size(); + const size_t total_expert_bytes = slot_bytes * (size_t) g_n_expert; + + g_slots.assign((size_t) g_n_layer * g_n_expert, slot{}); + for (int il = 0; il < g_n_layer; il++) { + for (int e = 0; e < g_n_expert; e++) { + auto & sl = g_slots[(size_t) il * g_n_expert + e]; + for (int p = 0; p < st::N_PART; p++) { + const ggml_tensor * t = layers[il][p]; + uint8_t * a = (uint8_t *) t->data + (size_t) e * t->nb[2]; + sl.parts[p].addr = a; + sl.parts[p].bytes = t->nb[2]; + sl.parts[p].file_off = (size_t) (a - g_map_base); + + // 診斷：確認 t->data 真的在 mmap 區間內。 + // 若不在，madvise 打在匿名記憶體上 = MADV_DONTNEED 會把它歸零 + // → 輸出直接變成亂碼。這是本專案真的踩到的坑。 + { + const uintptr_t mb = (uintptr_t) g_map_base; + const uintptr_t msz = (uintptr_t) g_map_size; + const uintptr_t ua = (uintptr_t) a; + const int in_map = (ua >= mb && ua < mb + msz) ? 1 : 0; + if (!in_map) { + // 這是本專案踩過最嚴重的一個坑： + // 若 ggml 用 CPU_REPACK 把權重轉成 repack 格式放在**匿名**緩衝區， + // t->data 就不在 mmap 內。這時對它 madvise(MADV_DONTNEED) + // 不是「讀 SSD」，而是**把那塊記憶體歸零** → 輸出直接變亂碼， + // 而且看起來「有在跑」、速度也只慢一點，完全不像出錯。 + // 所以這裡必須硬性擋下，並且把 buffer 型別印出來。 + const ggml_tensor * tt = layers[il][p]; + fprintf(stderr, + "[st] 權重不在 mmap 內（layer %d part %d expert %d: %p，buffer=%s）" + " → 分頁停用（需要 -DGGML_CPU_REPACK=OFF，否則 madvise 會歸零權重）\n", + il, p, e, (void *) a, + (tt->buffer ? ggml_backend_buffer_name(tt->buffer) : "(null)")); + return false; + } + } + } + sl.resident = true; // llama.cpp 剛載入時全部在 RAM + } + } + + fprintf(stderr, + "[st] expert pager 就緒：%d 層 × %d experts = %d 個 expert，" + "expert 權重總量 %.1f MiB（平均每個 %.2f MiB；逐層量化型別不同，" + "以 layer 0 為例是 gate %zu + up %zu + down %zu = %.2f MiB），" + "預算 %.1f MiB，hot window %d tick\n", + g_n_layer, g_n_expert, g_n_layer * g_n_expert, + (double) total_expert_bytes / (1024.0 * 1024.0), + (double) total_expert_bytes / (double) (g_n_layer * g_n_expert) / (1024.0 * 1024.0), + g_part_bytes[0], g_part_bytes[1], g_part_bytes[2], + (double) (g_part_bytes[0] + g_part_bytes[1] + g_part_bytes[2]) / (1024.0 * 1024.0), + (double) total_expert_bytes / (1024.0 * 1024.0), + (double) g_budget / (1024.0 * 1024.0), g_hot_tokens); + return true; +} + +// ------------------------------------------------------------------ 淘汰 + +// 診斷用：行程自己的 RSS（總量 / 匿名 / 檔案對映）。 +// 這是「分頁到底有沒生效」唯一可信的數字 —— 外部量測取樣太粗， +// 而且 prefill 會讓 peak 失真。 +struct self_rss { size_t total = 0, anon = 0, file = 0; }; + +static self_rss read_self_rss() { + self_rss r; + FILE * f = fopen("/proc/self/status", "r"); + if (f) { + char line[256]; + while (fgets(line, sizeof line, f)) { + unsigned long long v; + if (sscanf(line, "VmRSS: %llu", &v) == 1) r.total = v * 1024; + } + fclose(f); + } + f = fopen("/proc/self/smaps_rollup", "r"); + if (f) { + char line[256], k[64]; + unsigned long long v; + while (fgets(line, sizeof line, f)) { + if (sscanf(line, "%63[^:]: %llu", k, &v) != 2) continue; + if (!strcmp(k, "Anonymous")) r.anon = v * 1024; + else if (!strcmp(k, "Private_Clean") || !strcmp(k, "Private_Dirty")) r.file += v * 1024; + } + fclose(f); + } + return r; +} + + +// ST_EVICT_MODE 診斷用（實測發現只做其中一步會導致輸出損毀，見 docs）： +// both（預設）| madvise | fadvise | none +static int evict_mode() { + static const int m = [] { + const char * e = getenv("ST_EVICT_MODE"); + if (!e || !*e) return 0; + if (!strcmp(e, "none")) return 3; + if (!strcmp(e, "madvise")) return 1; + if (!strcmp(e, "fadvise")) return 2; + return 0; + }(); + return m; +} + +static void evict_slot(const slot & sl) { + const size_t ps = page_size(); + const int mode = evict_mode(); + if (mode == 3) return; + + for (int p = 0; p < st::N_PART; p++) { + const span & s = sl.parts[p]; + if (!s.addr || !s.bytes) continue; + + if (mode == 0 || mode == 1) { + // (a) 拿掉本行程的 PTE + const uintptr_t a = (uintptr_t) s.addr; + const uintptr_t a0 = a & ~(uintptr_t) (ps - 1); + madvise((void *) a0, s.bytes + (size_t) (a - a0), MADV_DONTNEED); + } + + if ((mode == 0 || mode == 2) && g_fd >= 0) { + // (b) 把 page cache 裡的頁真的丟掉 + const size_t fo0 = s.file_off & ~(ps - 1); + posix_fadvise(g_fd, (off_t) fo0, + (off_t) (s.bytes + (s.file_off - fo0)), POSIX_FADV_DONTNEED); + } + } +} + +static void enforce_budget(uint64_t now) { + size_t resident = 0; + for (const auto & sl : g_slots) { + if (sl.resident) resident += slot_bytes_of(sl); + } + if (resident <= g_budget) return; + + g_st.sweeps++; + + // 候選 = 還在 RAM、且「不是這個 tick 用到的」。 + // last_used == 0 代表「從來沒被用到過」→ 一定是冷的。 + // 注意不能用 `last_used < now - hot_tokens`：現在 tick 才剛開始時 + // 那個門檻會是 0，而 last_used 也是 0，於是**永遠沒有候選**， + // 第一輪 sweep 必定什麼都不做（這個 bug 真的踩過）。 + std::vector cand; + cand.reserve(g_slots.size()); + const uint64_t hot_floor = now >= (uint64_t) g_hot_tokens ? now - (uint64_t) g_hot_tokens : 0; + for (size_t i = 0; i < g_slots.size(); i++) { + const auto & sl = g_slots[i]; + if (!sl.resident) continue; + if (sl.last_used == now) continue; // 這輪正在用，不能動 + if (sl.last_used != 0 && sl.last_used > hot_floor) continue; // hot window 內 + cand.push_back(i); + } + + // 淘汰順序：先用得少的（LFU）→ 比較久沒用的（LRU）。 + std::sort(cand.begin(), cand.end(), [&](size_t a, size_t b) { + const auto & sa = g_slots[a]; + const auto & sb = g_slots[b]; + if (sa.uses != sb.uses) return sa.uses < sb.uses; + return sa.last_used < sb.last_used; + }); + + std::vector to_evict; + size_t freed = 0; + for (size_t i : cand) { + if (resident - freed <= g_budget) break; + to_evict.push_back(i); + freed += slot_bytes_of(g_slots[i]); + } + for (size_t i : to_evict) { + g_slots[i].resident = false; + evict_slot(g_slots[i]); + g_st.evictions++; + g_st.evicted_bytes += slot_bytes_of(g_slots[i]); + } + + if (g_verbose && !to_evict.empty()) { + const self_rss rss = read_self_rss(); + fprintf(stderr, + "[st] sweep #%llu：丟 %zu 個 expert（%.1f MiB）→ 帳 resident %.1f MiB / budget %.1f MiB" + " ｜ 實際 RSS total %.1f MiB（anon %.1f + file %.1f）\n", + (unsigned long long) g_st.sweeps, to_evict.size(), + freed / (1024.0 * 1024.0), + (resident - freed) / (1024.0 * 1024.0), + g_budget / (1024.0 * 1024.0), + rss.total / (1024.0 * 1024.0), + rss.anon / (1024.0 * 1024.0), + rss.file / (1024.0 * 1024.0)); + } +} + +void install_signal_handler() { + static const char * path = getenv("ST_STATS_FILE"); + if (!path || !*path) return; + static char buf[1024]; + snprintf(buf, sizeof(buf), "%s", path); + + struct sigaction sa {}; + sa.sa_handler = [](int) { + // signal handler 裡只能呼叫 async-signal-safe 的東西； + // dump_stats 用 stdio（有鎖），實務上這裡只會在沒人在同一瞬間 + // 寫檔時觸發，風險可接受 —— 但不要從這裡再呼叫 exit()。 + st::dump_stats(buf); + }; + sigemptyset(&sa.sa_mask); + sigaction(SIGUSR2, &sa, nullptr); + + install_atexit(); +} + +// ------------------------------------------------------------------ routing op + +namespace { + +struct routing_ctx { + int layer = -1; + int n_expert = 0; + int n_expert_used = 0; +}; + +// ids 一定要在 op 執行時從 dst->src[0] 拿，不能在建圖時把指標存起來： +// llama.cpp 每次 decode 都會重新建圖，graph context 的配置器可能把張量 +// 搬到不同位址（形狀不同時尤其明顯）。這是本專案踩過的坑。 +void routing_op(ggml_tensor * dst, int /*ith*/, int /*nth*/, void * userdata) { + auto * rc = (routing_ctx *) userdata; + if (!dst || dst->op != GGML_OP_CUSTOM) return; + + ggml_tensor * ids = dst->src[0]; + if (!ids || !ids->data) return; + + const int n_used = (int) ids->ne[0]; + const int n_tok = (int) (ids->ne[1] * ids->ne[2] * ids->ne[3]); + const int32_t * p = (const int32_t *) ids->data; + + std::lock_guard lk(g_mu); + g_st.routing_calls++; + const uint64_t now = ++g_tick; + + // 這一層、這批 token 用到的所有 expert（去重） + bool touched[256] = { false }; + std::vector uniq; + for (int t = 0; t < n_tok; t++) { + for (int k = 0; k < n_used; k++) { + const int e = p[(size_t) k * ids->nb[0] / sizeof(int32_t) + + (size_t) t * ids->nb[1] / sizeof(int32_t)]; + if (e < 0 || e >= rc->n_expert) continue; + g_st.expert_touches++; + if (!touched[e]) { touched[e] = true; uniq.push_back(e); } + } + } + + for (int e : uniq) { + auto & sl = g_slots[(size_t) rc->layer * rc->n_expert + e]; + sl.last_used = now; + sl.uses++; + if (sl.resident) { + g_st.hits++; + } else { + // 這一段會被 matmul 讀到 → 核心會從 SSD 補回來。 + // 標記 resident 只是為了讓 hit/miss 統計誠實： + // 實際上 miss 之後這一層的 matmul 結束前不會再被丟掉。 + g_st.misses++; + sl.resident = true; + } + } + + enforce_budget(now); +} + +} // namespace + +ggml_tensor * build_routing_op(ggml_context * ctx, ggml_tensor * selected_experts, + int layer, int n_expert, int n_expert_used) { + if (!enabled()) return nullptr; + if (!selected_experts || selected_experts->ne[0] != n_expert_used) return nullptr; + + auto * rc = new routing_ctx{ layer, n_expert, n_expert_used }; + + // 輸出 1 個元素：我們只要副作用，不要資料。 + // 不要用 ne0 = 0 —— 有些 ggml 版本對空張量的排程處理會直接跳過這個 node。 + ggml_tensor * args[1] = { selected_experts }; + return ggml_custom_4d(ctx, GGML_TYPE_I32, 1, 1, 1, 1, + args, 1, routing_op, 1, rc); +} + +} // namespace st \ No newline at end of file diff --git a/src/st_pager.h b/src/st_pager.h new file mode 100644 index 000000000..336e37de3 --- /dev/null +++ b/src/st_pager.h @@ -0,0 +1,84 @@ +// ST expert pager —— 熱 expert 留在 RAM，冷 expert 丟回 SSD +// +// 對外介面。llama.cpp 端（llama-model.cpp / llama-graph.cpp）只include 這個檔案。 +// +// 設計取捨（為什麼不照抄 sddqwen35a3b_v01 的 pread arena）： +// sddqwen35a3b 做的是「把 expert 權重 pread 進自己的 arena，用自訂 ggml op +// 重寫量化點積」。那在 22 GB 權重 / 小 RAM 機器上是必要的，代價是要自己重寫 +// Q4_K 的點積 kernel，數值正確性得另外驗證。 +// SmallThinker 只有 2.45 GiB，權重可以用 file-backed mmap 留在 SSD， +// 所以本專案改成「**計算路徑完全不動，只把用不到的 expert 頁丟掉**」： +// - 計算仍是 llama.cpp 自己的 kernel → 輸出與上游一致，沒有重寫 kernel 的風險。 +// - 淘汰的最小單位是「一個 expert 的三段權重」（gate/up/down），不是 page。 +// 代價：不能像 arena 那樣用 prefetch 提前重疊 I/O，SSD 讀取是 demand-paging +// （第一次碰到缺頁才讀）。這是本專案目前的取捨。 +// +// 淘汰的兩步驟（順序不能換，這是實測踩出來的）： +// a. madvise(MADV_DONTNEED) → 拿掉本行程 PTEs，行程總 RSS 立刻下降。 +// b. posix_fadvise(DONTNEED) → 核心的 invalidate_mapping_pages() 會跳過 +// 「仍被某個 VMA 映射」的頁，所以 (a) 必須先做；(b) 做完這些頁才真的從 +// page cache 消失，下次碰到才是真的讀 SSD。 +// 只做 (a)：RSS 下降但 read_bytes 不變（頁還在 page cache）。 +// 只做 (b)：page cache 下降但行程 RSS 不變（VMA 還映射著）。 + +#pragma once + +#include +#include +#include +#include + +struct ggml_context; +struct ggml_tensor; + +namespace st { + +// 三段權重（對應 llama.cpp 的 smallthinker 命名） +enum part { GATE = 0, UP = 1, DOWN = 2, N_PART = 3 }; + +struct stats { + uint64_t routing_calls = 0; // 每層每個 ubatch 算一次 + uint64_t expert_touches = 0; // 「某個 expert 被某個 token 選中」的總次數 + uint64_t hits = 0; // 選中時它還在 RAM + uint64_t misses = 0; // 選中時它已被丟回 SSD → 這次要重讀 + uint64_t evictions = 0; + uint64_t evicted_bytes = 0; + uint64_t sweeps = 0; + size_t resident_bytes = 0; + size_t budget_bytes = 0; +}; + +// 讀環境變數。必須在 register_model() 之前呼叫一次。 +// ST_PAGER=0 完整關閉（回到上游行為，做 A/B 用） +// ST_RAM_BUDGET_MB RAM 預算（llama_server.sh 算好） +// ST_RESERVE_MB KV + compute 預留，從預算裡扣掉 +// ST_ARENA_MB 直接指定 expert 可常駐的量（有值就優先用） +// ST_HOT_TOKENS=8 最近 N 個 tick 用過的 expert 不淘汰 +// ST_VERBOSE=1 印淘汰紀錄 +// ST_STATS_FILE=path 收到 SIGUSR2 時把統計寫成 JSON +void set_config_from_env(); + +// 模型載入完之後呼叫。tensors_by_name 是 llama.cpp 的 tensors_by_name。 +// 回傳 false 表示這個模型沒有被分頁（呼叫端繼續走上游路徑）。 +bool register_model(const char * arch, + void * mapping_base, size_t mapping_size, int fd, + const std::map & tensors_by_name); + +// 在每層 top-k 選完 expert 之後插入的副作用 op。 +// 回傳 nullptr 表示不插入。呼叫端要把回傳值 ggml_build_forward_expand 進圖。 +ggml_tensor * build_routing_op(ggml_context * ctx, ggml_tensor * selected_experts, + int layer, int n_expert, int n_expert_used); + +bool enabled(); +// 分頁有註冊成功 → 必須關掉 MAP_POPULATE（否則載入時就會 fault 整個檔案） +bool paging_active(); +// 只看「有沒有打算分頁」，不看權重有沒有載入。init_mappings() 在載入之前， +// 那時 register_model 還沒跑，所以只能靠 arch 判斷。 +bool will_page(const char * arch); +stats get_stats(); +void dump_stats(const char * path); + +// 讓 SIGUSR2 可以觸發 dump（由 llama.cpp 呼叫端註冊） +void install_signal_handler(); + +} // namespace st \ No newline at end of file