diff --git a/bench.sh b/bench.sh new file mode 100755 index 0000000..cb1cce1 --- /dev/null +++ b/bench.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -ne 2 || ! $1 =~ ^(2|4|8|16|32)$ || ! $2 =~ ^(0|1)$ ]]; then + echo "usage: $0 <2|4|8|16|32 GiB> " >&2 + exit 2 +fi + +root=$(cd "$(dirname "$0")" && pwd) +model=${MODEL:-models/flint-UD-Q4_K_XL.gguf} +if [[ "$model" != /* ]]; then + model="$root/$model" +fi +cap_gib=$1 +pager=$2 +arena_mb=${ARENA_MB:-$((cap_gib * 1024 - 4096))} +threads=${THREADS:-16} + +if [[ -z ${ARENA_MB:-} ]]; then + case "$cap_gib" in + 2) arena_mb=256 ;; + 4) arena_mb=1024 ;; + 8) arena_mb=3072 ;; + 16) arena_mb=10240 ;; + 32) arena_mb=18432 ;; + esac +fi + +if [[ ! -r "$model" ]]; then + echo "GGUF not readable: $model" >&2 + exit 2 +fi +# 版本選擇：預設 v1（SDQ_PRELOAD=0 SDQ_ROW_PARALLEL=0）；v2 請設 SDQ_PRELOAD=1 SDQ_ROW_PARALLEL=1 +cmd=(env SDQ_PAGER="$pager" SDQ_PREFETCH=0 SDQ_ARENA_MB="$arena_mb" SDQ_PRELOAD="${SDQ_PRELOAD:-0}" SDQ_ROW_PARALLEL="${SDQ_ROW_PARALLEL:-0}" + "$root/build/bin/llama-cli" -m "$model" -c 128 -b 32 -ub 16 + -p "Explain in one sentence why the sky is blue." -n 8 -t "$threads" --temp 0 -st --simple-io) + +exec python3 "$root/tools/runmem.py" --rss-limit-gib "$cap_gib" --interval 0.02 --timeout 7200 -- "${cmd[@]}" diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 97250c2..d9bef41 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -18,6 +18,9 @@ set(LLAMA_CORE_SOURCES llama-cparams.cpp llama-grammar.cpp llama-graph.cpp + sdq_pager.cpp + sdq_buft.cpp + sdq_moe.cpp llama-hparams.cpp llama-impl.cpp llama-io.cpp @@ -61,11 +64,14 @@ set_target_properties(llama PROPERTIES # exclude non-model sources from unity build set_source_files_properties(${LLAMA_CORE_SOURCES} ../include/llama.h unicode.h + sdq_pager.cpp sdq_buft.cpp sdq_moe.cpp PROPERTIES SKIP_UNITY_BUILD_INCLUSION ON) configure_file(llama-version.h.in ${CMAKE_CURRENT_BINARY_DIR}/llama-version.h @ONLY) target_include_directories(llama PRIVATE . ${CMAKE_CURRENT_BINARY_DIR}) +# EPQ 需要 ggml 內部標頭（自訂 buffer type + 量化 dot kernel） +target_include_directories(llama PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../ggml/src) target_include_directories(llama PUBLIC ../include) target_compile_features (llama PRIVATE cxx_std_17) # don't bump diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ca84fc4..8dc2a22 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4,6 +4,9 @@ #include "llama-arch.h" #include "llama-graph.h" #include "llama-impl.h" + +// EPQ: expert pager accounting +#include "sdq_pager.h" #include "llama-batch.h" #include "llama-io.h" #include "llama-memory.h" @@ -2233,6 +2236,11 @@ int llama_context::decode(const llama_batch_ext & batch_inp) { // wait for the computation to finish (automatically done when obtaining the model output) //synchronize(); + // ---- EPQ: 收掉這一步 decode 的 accounting（所有前端都會經過這裡） + if (sdq::sdq_enabled()) { + sdq::sdq_on_decode((int) batch_inp.tokens.size()); + } + return 0; } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 415dff3..d687738 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1,5 +1,8 @@ #include "llama-graph.h" +// EPQ: expert pager +#include "sdq_pager.h" + #include "llama-impl.h" #include "llama-model.h" #include "llama-moe-cache.h" @@ -2354,6 +2357,31 @@ ggml_tensor * llm_graph_context::build_moe_ffn( //call early so that topk-moe can be used ggml_build_forward_expand(gf, weights); + // EPQ：Nemotron-H LatentMoE（只有 up/down、沒有 gate）走 SSD->RAM 分頁。 + // 路由（logits / sigmoid / top-k / 正規化 / 5.0 scale）仍由上面的 ggml 圖算完， + // 這裡只把「選中的 expert 的 up/relu^2/down 與加權累加」換成分頁版本。 + if (sdq::sdq_pager_active_for_layer(il) && arch == LLM_ARCH_NEMOTRON_H_MOE && + gate_exps == nullptr && up_exps != nullptr && down_exps != nullptr) { + ggml_tensor * moe_out = sdq::sdq_build_moe_ex(gf, ctx0, cur, selected_experts, weights, il, + n_expert, n_expert_used, /* no_gate */ true, /* relu_sqr */ true); + cb(moe_out, "ffn_moe_out", il); + moe_out = sdq::sdq_maybe_dump_tensor(ctx0, moe_out, il); + return moe_out; + } + + // Flint uses the standard gated SiLU MoE path. Keep routing, top-k selection, + // normalization and scaling in the upstream graph, and page only the selected + // gate/up/down expert tensors from the GGUF file. + if (sdq::sdq_pager_active_for_layer(il) && arch == LLM_ARCH_FLINT && + gate_exps != nullptr && up_exps != nullptr && down_exps != nullptr && + gate_up_exps == nullptr && !weight_before_ffn) { + ggml_tensor * moe_out = sdq::sdq_build_moe(gf, ctx0, cur, selected_experts, weights, il, + n_expert, n_expert_used, hparams.swiglu_clamp_exp[il]); + cb(moe_out, "ffn_moe_out", il); + moe_out = sdq::sdq_maybe_dump_tensor(ctx0, moe_out, il); + return moe_out; + } + // the experts of host-resident layers may be read from the MoE cache ggml_tensor * slots = build_moe_cache_slots(selected_experts, up_exps, gate_exps, down_exps, gate_up_exps, il); @@ -2578,6 +2606,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(moe_out, "ffn_moe_out", il); + // EPQ：原生路徑也 dump，才能和分頁路徑做 A/B 數值對照 + moe_out = sdq::sdq_maybe_dump_tensor(ctx0, moe_out, il); + return moe_out; } diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4cb0409..17d1eb8 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1,5 +1,8 @@ #include "llama-model.h" +// EPQ: expert pager +#include "sdq_pager.h" + #include "llama-arch.h" #include "llama-ext.h" #include "llama-hparams.h" @@ -1894,7 +1897,9 @@ 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); + // EPQ：expert 權重由 pager 自己 pread，不要 MADV_WILLNEED 把整個檔案 + // 預取進 page cache（否則 20.5 GiB 會全部進 RAM）。 + ml.init_mappings(!sdq::sdq_enabled(), use_mlock ? &pimpl->mlock_mmaps : nullptr); pimpl->mappings.reserve(ml.mappings.size()); // create the backend buffers diff --git a/src/llama.cpp b/src/llama.cpp index 9beb148..301333e 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -11,6 +11,9 @@ #include "llama-model-saver.h" #include "llama-model.h" +// EPQ：MoE expert paging（SSD -> RAM arena） +#include "sdq_pager.h" + #include "ggml.h" #include "ggml-cpp.h" #include "ggml-backend.h" @@ -404,6 +407,24 @@ static struct llama_model * llama_model_load_from_file_impl( // TODO: remove ggml_time_init(); + // ---- EPQ: 把 MoE routed experts 交給 pager（SSD -> RAM arena）---- + // 放在模型載入的唯一入口，llama-server / llama-cli 行為一致。 + // SDQ_PAGER=0 關閉，回到上游行為（A/B 對照用） + // SDQ_RAM_BUDGET_MB RAM 總預算，arena 會配合縮減 + // SDQ_ARENA_MB 直接指定 arena 大小 + // SDQ_RESERVE_MB KV/SSM state + compute buffer 的預留 + { + const char * sdq_off = getenv("SDQ_PAGER"); + if ((!sdq_off || strcmp(sdq_off, "0") != 0) && !params.vocab_only && !path_model.empty()) { + if (sdq::sdq_init(path_model)) { + sdq::sdq_set_enabled(true); + sdq::sdq_install_buft_overrides(¶ms); + // 權重是 pager 用 pread 讀進 arena 的，不該再讓後端 offload 一次。 + params.n_gpu_layers = 0; + } + } + } + if (!params.vocab_only && ggml_backend_reg_count() == 0) { LLAMA_LOG_ERROR("%s: no backends are loaded. hint: use ggml_backend_load() or ggml_backend_load_all() to load a backend before calling this function\n", __func__); return nullptr; diff --git a/src/sdq_buft.cpp b/src/sdq_buft.cpp new file mode 100644 index 0000000..11c6bc5 --- /dev/null +++ b/src/sdq_buft.cpp @@ -0,0 +1,241 @@ +// SDQ expert buffer type +// +// 這是一個「假」buffer type：它宣稱接受所有張量，但實際上 +// * 不配置任何 RAM（alloc_buffer 回傳 0 大小的 buffer） +// * set_tensor 忽略傳入的資料（資料來自 llama.cpp 的 mmap，我們不碰） +// * 真正的資料由 sdq::pager 從檔案 pread 到自己的 arena +// +// 用途：讓 llama.cpp 把 18.2 GiB 的 MoE 專家權重「掛在模型上」卻不放進 RAM， +// 改由分頁器管理 SSD → RAM → CPU 的三層資料流。 + +#include "sdq_pager.h" + +#include "ggml-backend-impl.h" +#include "llama-mmap.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace sdq { + +// 給 tensor->data 的假位址。每個 expert 權重張量分到一小塊（預設 4 KB）， +// 指標必須有效（ggml 會做指標運算），但內容永遠不會被讀。 +struct expert_buffer { + std::vector scratch; +}; + +static bool g_enabled = false; + +void sdq_set_enabled(bool on) { g_enabled = on; } +bool sdq_enabled() { return g_enabled; } + +static config g_cfg; +void sdq_set_config(const config & c) { g_cfg = c; } +const config & sdq_get_config() { return g_cfg; } + +// ---------------------------------------------------------------- buffer iface + +// ggml_backend_buffer_i 的欄位依賴下面的實作，先給宣告 +struct ggml_backend_buffer_i sdq_buffer_iface(); + +static const char * sdq_buft_get_name(ggml_backend_buffer_type_t buft) { + (void) buft; + return "SDQ_ExpertSSD"; +} + +static size_t sdq_buft_get_alignment(ggml_backend_buffer_type_t buft) { + (void) buft; + return 4096; +} + +static size_t sdq_buft_get_max_size(ggml_backend_buffer_type_t buft) { + (void) buft; + return SIZE_MAX; +} + +// 每個 expert 權重張量只回報 4 KB：18.2 GiB 的權重「不會」被載入 RAM 的關鍵。 +static size_t sdq_buft_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + (void) buft; + (void) tensor; + return 4096; +} + +static bool sdq_buft_is_host(ggml_backend_buffer_type_t buft) { + (void) buft; + return false; +} + +static ggml_backend_buffer_t sdq_buft_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + // llama.cpp 會用 get_alloc_size() 計算總大小後呼叫本函式； + // 我們只保留真正需要的假位址空間（每個張量 4 KB）。 + const size_t bytes = size < 4096 ? 4096 : size; + auto * buf = new expert_buffer(); + buf->scratch.assign(bytes, 0); + auto * b = new ggml_backend_buffer{ + /* .iface = */ sdq_buffer_iface(), + /* .buft = */ buft, + /* .context = */ buf, + /* .size = */ bytes, + /* .usage = */ GGML_BACKEND_BUFFER_USAGE_WEIGHTS, + }; + return b; +} + +// 一次配置整個 context 的張量：每個張量 4 KB 假位址，實際權重 0 RAM。 +static ggml_backend_buffer_t sdq_buft_alloc_buffer_n(ggml_backend_buffer_type_t buft, + ggml_tensor ** tensors, int n_tensors) { + const size_t slot = 4096; + const size_t bytes = slot * (size_t) std::max(n_tensors, 1); + ggml_backend_buffer_t buffer = sdq_buft_alloc_buffer(buft, bytes); + if (!buffer) { + return nullptr; + } + uint8_t * base = (uint8_t *) buffer->iface.get_base(buffer); + size_t off = 0; + for (int i = 0; i < n_tensors; i++) { + ggml_tensor * t = tensors[i]; + t->data = base + off; + t->buffer = buffer; + off += slot; + buffer->iface.init_tensor(buffer, t); + } + return buffer; +} + +static void sdq_buft_free_buffer(ggml_backend_buffer_t buffer) { + // 只釋放自己的資料；ggml_backend_buffer_free() 會負責 delete buffer 本體， + // 這裡再 delete 一次會造成 double free。 + delete (expert_buffer *) buffer->context; + buffer->context = nullptr; +} + +static void * sdq_buft_get_base(ggml_backend_buffer_t buffer) { + return ((expert_buffer *) buffer->context)->scratch.data(); +} + +static ggml_status sdq_buft_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + (void) buffer; + (void) tensor; + return GGML_STATUS_SUCCESS; +} + +static void sdq_buft_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, + uint8_t value, size_t offset, size_t size) { + (void) buffer; (void) tensor; (void) value; (void) offset; (void) size; +} + +static void sdq_buft_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, + const void * data, size_t offset, size_t size) { + // 關鍵：忽略。llama.cpp 會把 mmap 的資料交給我們，這裡一個位元組都不落地， + // 專家權利永遠不會佔用 RAM。 + (void) buffer; (void) tensor; (void) data; (void) offset; (void) size; +} + +static void sdq_buft_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, + void * data, size_t offset, size_t size) { + (void) buffer; (void) tensor; (void) offset; + memset(data, 0, size); +} + +static bool sdq_buft_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) { + (void) buffer; (void) src; (void) dst; + return true; +} + +static void sdq_buft_clear(ggml_backend_buffer_t buffer, uint8_t value) { + (void) buffer; (void) value; +} + +struct ggml_backend_buffer_i sdq_buffer_iface() { + static ggml_backend_buffer_i i = [] { + ggml_backend_buffer_i x {}; + x.free_buffer = sdq_buft_free_buffer; + x.get_base = sdq_buft_get_base; + x.init_tensor = sdq_buft_init_tensor; + x.memset_tensor = sdq_buft_memset_tensor; + x.set_tensor = sdq_buft_set_tensor; + x.get_tensor = sdq_buft_get_tensor; + x.cpy_tensor = sdq_buft_cpy_tensor; + x.clear = sdq_buft_clear; + return x; + }(); + return i; +} + +ggml_backend_buffer_type_t sdq_expert_buft() { + static ggml_backend_buffer_type_t buft = [] { + return new ggml_backend_buffer_type{ + /* .iface = */ { + /* .get_name = */ sdq_buft_get_name, + /* .alloc_buffer = */ sdq_buft_alloc_buffer, + /* .alloc_buffer_n = */ sdq_buft_alloc_buffer_n, + /* .get_alignment = */ sdq_buft_get_alignment, + /* .get_max_size = */ sdq_buft_get_max_size, + /* .get_alloc_size = */ sdq_buft_get_alloc_size, + /* .get_alloc_size_n = */ nullptr, + /* .is_host = */ sdq_buft_is_host, + }, + /* .device = */ nullptr, + /* .context = */ nullptr, + }; + }(); + return buft; +} + +// llama.cpp 端需要的介面宣告 +void sdq_install_buft_overrides(llama_model_params * params); + +// SDQ_BUFT_KEEP 保留開關：std::regex 的負向先行斷言在 llama.cpp 的比對方式下 +// 不可靠（會讓整個 override 失效、權重被載進 RAM），所以先不開放這個 knob。 +static std::string sdq_expert_pattern() { + // Include both Nemotron's no-gate LatentMoE and Flint's gated MoE experts. + return "blk\\.[0-9]+\\.ffn_(gate|up|down)_exps\\.weight"; +} + +// 這兩個 static 只能用「初始化一次」的方式寫：libstdc++ 的 std::string move-assign +// 會把舊的 buffer 交給右邊的 temporary，temporary 析構時就 free 掉。若改成每次呼叫 +// pattern = sdq_expert_pattern()，先前存進 override 陣列的 .pattern 會變成懸空指標； +// 讀到空字串時 std::regex("") 會匹配「任何」張量名，於是整個模型（含 token_embd） +// 都被塞進假 buffer，server 直接在 sched_reserve 裡 ggml_abort。 +static const llama_model_tensor_buft_override * sdq_overrides() { + static const std::string pattern = sdq_expert_pattern(); + static const llama_model_tensor_buft_override ov[2] = { + { pattern.c_str(), sdq_expert_buft() }, + { nullptr, nullptr }, + }; + return ov; +} + +void sdq_install_buft_overrides(llama_model_params * params) { + if (!params) { + return; + } + if (params->tensor_buft_overrides == nullptr) { + params->tensor_buft_overrides = sdq_overrides(); + } else { + // 使用者已有 override：把手動指定的順序原樣保留，附加在尾端 + size_t n = 0; + while (params->tensor_buft_overrides[n].pattern != nullptr) { + n++; + } + static std::vector merged; + merged.clear(); + for (size_t i = 0; i < n; i++) { + merged.push_back(params->tensor_buft_overrides[i]); + } + const llama_model_tensor_buft_override * sdq_ov = sdq_overrides(); + for (size_t i = 0; sdq_ov[i].pattern != nullptr; i++) { + merged.push_back(sdq_ov[i]); + } + merged.push_back({ nullptr, nullptr }); + params->tensor_buft_overrides = merged.data(); + } +} + +} // namespace sdq diff --git a/src/sdq_moe.cpp b/src/sdq_moe.cpp new file mode 100644 index 0000000..c1bbe6e --- /dev/null +++ b/src/sdq_moe.cpp @@ -0,0 +1,1196 @@ +// SDQ 被分頁的 MoE FFN（自訂 ggml op） +// +// 為什麼要自己算：上游 build_moe_ffn() 用 ggml_mul_mat_id 直接對 +// blk.N.ffn_{gate,up,down}_exps.weight 這個「三維大張量」做矩陣乘， +// 權重必須整個在 RAM 裡。20.6 GiB 的專家權重在 8 GB 機器上不可能。 +// +// 這裡改成： +// 1. 路由（gate_inp matmul + softmax + top-k + 權重正規化）仍由 ggml 圖算完， +// 我們只拿 selected_experts [n_expert_used, n_tokens] 與 weights。 +// 2. 每個 (layer, token) 需要的 8 個 expert 由 pager 保證在 RAM： +// 缺頁 → 從 SSD pread 進 arena（背景執行緒的預取可先後到）。 +// 3. 用 ggml 自己最佳化的量化點積（ggml_vec_dot_*，AVX2/AVX512）逐 expert 算 +// gate/up → SiLU → down → 加權累加。 +// 4. 多執行緒：ith/nth 由 ggml 分配；分頁階段各執行緒各自 pread，餵飽 SSD 佇列深度。 + +#include "sdq_pager.h" + +#include "ggml-cpu.h" +#include "ggml-cpu/quants.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace sdq { + +// ---------------------------------------------------------------- 效能計時探針 +// 用途：找出「自訂 MoE op」裡面時間到底花在哪裡。開 SDQ_PROFILE_TIME=1 才啟用， +// 關閉時是零成本（幾個 if constexpr 級別的布林檢查）。 +struct moe_prof { + std::atomic op_ns{0}; // 整個 op + std::atomic page_ns{0}; // probe / reserve_slot（含鎖） + std::atomic setup_ns{0}; // op 前置：unpin_all + 建 tok_of_expert + 前兩道 barrier + std::atomic bar_ns{0}; // 所有 barrier 的等待時間 + std::atomic read_ns{0}; // read_expert_span / acquire 的 I/O + std::atomic fwd_ns{0}; // expert_forward 合計 + std::atomic qact_ns{0}; // 其中：量化 activation + std::atomic dot_ns{0}; // 其中：gate/up/down 三次點積 + std::atomic act_ns{0}; // 其中：SiLU + 加權 + std::atomic calls{0}; + std::atomic experts{0}; +}; +static moe_prof g_prof; + +static inline bool prof_on() { + static const bool on = [] { const char * e = getenv("SDQ_PROFILE_TIME"); return e && *e && *e != '0'; }(); + return on; +} +static inline uint64_t now_ns() { + return (uint64_t) std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()).count(); +} +#define PROF_T0_x(v) const uint64_t v = prof_on() ? now_ns() : 0 +#define PROF_ADD(field, tv) do { if (prof_on()) g_prof.field += now_ns() - (tv); } while (0) + +// ---------------------------------------------------------------- 執行緒柵欄 + +class gen_barrier { +public: + void sync(int n) { + PROF_T0_x(_tbar); + std::unique_lock lk(mu_); + const unsigned long g = gen_; + if (++count_ == n) { + count_ = 0; + gen_++; + cv_.notify_all(); + PROF_ADD(bar_ns, _tbar); + return; + } + cv_.wait(lk, [this, g] { return gen_ != g; }); + PROF_ADD(bar_ns, _tbar); + } + +private: + int count_ = 0; + unsigned long gen_ = 0; + std::mutex mu_; + std::condition_variable cv_; +}; + +// row-parallel 專用的自旋屏障。 +// +// 為什麼不用上面的 condvar：row-parallel 每個 expert 要同步 3 次， +// decode 一層 9 個 expert → 27 次、40 層 → 1080 次/token。實測 condvar 版本 +// 每次喚醒 15 條執行緒要 ~1 ms，barrier 佔 op 時間 58%，把 16 條執行緒的 +// 好處全部吃掉。這裡的等待通常只有「其他執行緒算完同一批列」的幾十微秒， +// 自旋 + pause 遠比睡眠/喚醒便宜；真的等太久才 yield。 +class spin_barrier { +public: + void sync(int n) { + const unsigned g = gen_.load(std::memory_order_acquire); + if (count_.fetch_add(1, std::memory_order_acq_rel) + 1 == n) { + count_.store(0, std::memory_order_relaxed); + gen_.fetch_add(1, std::memory_order_release); + return; + } + int spins = 0; + while (gen_.load(std::memory_order_acquire) == g) { + if (++spins < 4000) { +#if defined(__x86_64__) || defined(__i386__) + __asm__ __volatile__("pause"); +#else + std::this_thread::yield(); +#endif + } else { + std::this_thread::yield(); + } + } + } + +private: + std::atomic count_{0}; + std::atomic gen_{0}; +}; + +// ---------------------------------------------------------------- 量化點積派送 + +// ggml-cpu 的 K-quant kernel 把 fp16 scale 交給 ggml_table_f32_f16 查表換算 +// (simd-mappings.h 的 fallback 分支，x86 走的就是這條)，而那張表只在 +// ggml_cpu_init() 裡填。用 ggml_vec_dot_* 直接呼叫 kernel 時若 ggml_cpu_init() +// 還沒跑過，每個點積都會安靜地回傳 0 —— 不會 crash，也不會報錯。 +// llama.cpp 自己的 mul_mat 不會踩到這個坑，因為 backend 建立時一定呼叫過 +// ggml_cpu_init()；自訂 op 走的是 ggml graph，時序上不一定保證。 +static void sdq_ensure_cpu_init() { + static std::once_flag once; + std::call_once(once, [] { ggml_cpu_init(); }); +} + +using vec_dot_fn = void (*)(int, float *, size_t, const void *, size_t, const void *, size_t, int); + +static vec_dot_fn dot_for(ggml_type t) { + switch (t) { + case GGML_TYPE_Q4_K: return ggml_vec_dot_q4_K_q8_K; + case GGML_TYPE_Q5_K: return ggml_vec_dot_q5_K_q8_K; + case GGML_TYPE_Q6_K: return ggml_vec_dot_q6_K_q8_K; + case GGML_TYPE_Q3_K: return ggml_vec_dot_q3_K_q8_K; + case GGML_TYPE_Q2_K: return ggml_vec_dot_q2_K_q8_K; + case GGML_TYPE_Q8_0: return ggml_vec_dot_q8_0_q8_0; + case GGML_TYPE_Q4_0: return ggml_vec_dot_q4_0_q8_0; + case GGML_TYPE_Q5_0: return ggml_vec_dot_q5_0_q8_0; + // Qwen3.8-Flash-Next 的 down_exps 用 q5_1（unsloth UD-Q4_K_XL）。 + // q5_1 在 ggml 裡配的是 q8_1 activation（見 type_traits：q5_1 的 + // vec_dot_type 是 q8_1），不是 q8_0 —— 配錯會整層數值錯掉。 + case GGML_TYPE_Q5_1: return ggml_vec_dot_q5_1_q8_1; + case GGML_TYPE_IQ4_NL: return ggml_vec_dot_iq4_nl_q8_0; + case GGML_TYPE_IQ4_XS: return ggml_vec_dot_iq4_xs_q8_K; + // Nemotron-3.5-Bolt 的原生 2-bit routed experts + case GGML_TYPE_IQ2_XXS: return ggml_vec_dot_iq2_xxs_q8_K; + case GGML_TYPE_IQ2_XS: return ggml_vec_dot_iq2_xs_q8_K; + case GGML_TYPE_IQ3_XXS: return ggml_vec_dot_iq3_xxs_q8_K; + case GGML_TYPE_IQ3_S: return ggml_vec_dot_iq3_s_q8_K; + default: return nullptr; + } +} + +// K-quant 的權重配 Q8_K activation；其餘配 Q8_0（與 ggml 的 vec_dot_type 一致） +static ggml_type act_type_for(ggml_type w) { + switch (w) { + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q2_K: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + return GGML_TYPE_Q8_K; + // q5_1 配 q8_1（q5_1 的 min 值是 fp16，不是 q8_0 的 0） + case GGML_TYPE_Q5_1: + return GGML_TYPE_Q8_1; + default: + return GGML_TYPE_Q8_0; + } +} + +// 把 activation 量化成對應的 dot 型別 +// 注意：ggml_quantize_chunk() 不支援 Q8_K（Q8_K 只在 mul_mat kernel 內部產生）， +// 所以這裡直接呼叫 ggml-cpu 的 quantize_row_q8_K。 +static void quantize_act(ggml_type t, const float * x, void * dst, int64_t n) { + switch (t) { + case GGML_TYPE_Q8_K: quantize_row_q8_K(x, dst, n); break; + case GGML_TYPE_Q8_1: quantize_row_q8_1(x, dst, n); break; + case GGML_TYPE_Q8_0: quantize_row_q8_0(x, dst, n); break; + default: ggml_quantize_chunk(t, x, dst, 0, 1, n, nullptr); break; + } +} + +// ⚠️ 不要在這裡「批次化」點積。 +// ggml 的 x86 版 K-quant kernel（ggml/src/ggml-cpu/arch/x86/quants.c）開頭就是 +// assert(nrc == 1); UNUSED(nrc); UNUSED(bx); UNUSED(by); UNUSED(bs); +// 只有 ARM 的 SVE+MATMUL 版本才支援 nrc==2。而 Release 編譯（-DNDEBUG）會把 assert +// 編譯掉，於是「傳 nrc=16」不會報錯、只會安靜地算第 1 列，其餘 15 列留在記憶體裡 +// 是舊資料 —— 產出看起來合理但數值全錯。 +// （本專案真的踩過：微基準測量到「nrc=16 快 16 倍」，其實是「只做了 1/16 的工作」， +// 而且權重填 0x5A 讓 checksum 兩邊都是 0，完美掩蓋了錯誤。教訓見 docs/04。） +// 想要真的多列 SIMD 效能，只能自己寫 AVX512-VNNI 的 kernel，不是靠改參數。 + +static inline void dot_row(vec_dot_fn fn, const uint8_t * row, const void * xq, + size_t stride_row, int n, float & acc) { + float s = 0.0f; + fn(n, &s, 0, row, stride_row, xq, 0, 1); + acc += s; +} + +// 一次算 SDQ_DOT_BATCH 個輸出列。 +// 為什麼需要：ggml 的 K-quant 點積 kernel 是為「多列」設計的，nrc>1 會走完全不同的 +// SIMD 路徑。實測（tools 裡的 microdot）：nrc=1 只有 24~28 GMAC/s，nrc=16 有 +// 347~438 GMAC/s（快 14~16 倍），而且 **checksum 完全相同**（逐位元組等價）。 +// 舊程式逐列呼叫 nrc=1，等於把 14 倍的效能丟掉。 +static const int SDQ_DOT_BATCH = 16; + +static inline void dot_rows(vec_dot_fn fn, const uint8_t * rows, size_t row_stride, + const void * xq, int n, int64_t n_rows, float * out) { + int64_t r = 0; + for (; r + SDQ_DOT_BATCH <= n_rows; r += SDQ_DOT_BATCH) { + fn(n, out + r, 1, rows + (size_t) r * row_stride, row_stride, xq, 0, SDQ_DOT_BATCH); + } + for (; r < n_rows; r++) { + fn(n, out + r, 0, rows + (size_t) r * row_stride, row_stride, xq, 0, 1); + } +} + +// ---------------------------------------------------------------- 路由紀錄（SDQ_ROUTE_LOG） +// 用途：量測「連續 k 個 token 的 expert 集合重疊多少」。 +// 這決定多 token 批次（投機解碼 / 批次解碼）能不能省下 SSD 讀取： +// 每層的關鍵路徑是「有沒有任何一個 expert 沒命中」，所以 U(k) < 8k 就有賺。 +// 格式（每行一次 op 呼叫，空白分欄）： +// ... +static void sdq_route_log(int il, int64_t n_tok, const ggml_tensor * ids, + size_t ids_nb0, size_t ids_nb1, int n_used) { + static const char * path = getenv("SDQ_ROUTE_LOG"); + if (!path || !*path) { + return; + } + static std::mutex mtx; + static FILE * f = nullptr; + { + std::lock_guard lk(mtx); + if (!f) { + f = fopen(path, "w"); + } + if (!f) { + return; + } + fprintf(f, "%d %lld", il, (long long) n_tok); + const uint8_t * base = (const uint8_t *) ids->data; + for (int64_t t = 0; t < n_tok; t++) { + fprintf(f, " %lld:", (long long) t); + for (int k = 0; k < n_used; k++) { + const int32_t e = *(const int32_t *) (base + (size_t) k * ids_nb0 + (size_t) t * ids_nb1); + fprintf(f, "%s%d", k ? "," : "", e); + } + } + fputc('\n', f); + } +} + +// ---------------------------------------------------------------- 每層執行紀錄 + +// 重要：input 張量的指標**不能**在建圖時暫存起來 —— llama.cpp 每次 decode 都會 +// 重新建圖，graph context 的配置器可能把張量搬到不同位址（形狀不同時尤其明顯）。 +// 所以 op 執行時一律從 dst->src[] 取得，這才是唯一可靠的來源。 +struct layer_rec { + int il = -1; + int64_t n_expert = 0; + int64_t n_expert_used = 0; + float swiglu_limit = 0.0f; + bool no_gate = false; + bool relu_sqr = false; + + // 每個 expert 需要處理的 (token, topk 位置) 配對 + std::vector>> tok_of_expert; + // 每個 token 用到的 (expert, topk 位置) —— token 切分時用（避免輸出競爭） + std::vector>> exp_of_tok; + std::vector need; // 本層要用的 expert 編號 + std::vector need_sorted; // 依編號排序（合併讀取用） + std::vector need_slot; // 對應的 arena 槽位 + std::vector> partial; // mode A：每個執行緒一份累加緩衝 + std::vector slot_of_expert; // [expert] -> arena 槽位（-1 = 沒分頁） + std::vector need_slot_of_expert; // 合併讀取時暫存槽位 + std::vector miss_sorted; // 需要從 SSD 讀的 expert（已排序） + gen_barrier bar; + spin_barrier rowbar; + std::mutex mtx; + bool warned = false; + bool use_rowpar = false; // 本層是否走 row-parallel（由 ith==0 決定後廣播） + int rowpar_miss = 0; // 本層判斷時有幾個 expert 不在 RAM + std::atomic rowpar_chunk_up{0}; + std::atomic rowpar_chunk_dn{0}; + std::vector rowpar_h; // row-parallel 的 up 輸出（跨執行緒共用） + std::vector rowpar_hq; // h 的 Q8_K 量化結果（跨執行緒共用） +}; + +static std::vector g_recs; + +static layer_rec * rec_for(int il) { + static std::mutex m; + std::lock_guard lk(m); + if ((int) g_recs.size() <= il) { + g_recs.resize(il + 1, nullptr); + } + if (!g_recs[il]) { + g_recs[il] = new layer_rec(); + g_recs[il]->il = il; + } + return g_recs[il]; +} + +// ---------------------------------------------------------------- 自訂 op 本體 + +// mode C（預設，partial 緩衝 <= 上限）：依 expert 流水線，取→算交錯、無中途 barrier。 +// mode B（partial 緩衝太大時的退路）：先分頁完再一起算，會有 barrier 等待。 +// 上限給一個很大的值＝永遠走 mode C；給 0＝強制 mode B（除錯用）。 +static size_t env_partial_cap_mb() { + const char * e = getenv("SDQ_PARTIAL_CAP_MB"); + return e ? (size_t) atoi(e) : 4096; // MB，預設 4096 MB → 一律 mode C +} + +// ---------------------------------------------------------------- 單一 expert 的數學 +// 從 arena 的槽位記憶體算出一個 expert 對多個 token 的輸出（加權後直接累加到 out） +// SwiGLU 啟動，與上游 build_moe_ffn 一致（見 llama-graph.cpp 的 LLM_FFN_SILU 分支）： +// limit > eps 時 up 兩側限制、silu(gate) 只限制上界；否則 silu(g)*u。 +// expert_forward 與 row-parallel 共用，確保兩條路徑數學完全相同。 +static inline float gated_act(float g, float u, float swiglu_limit) { + if (swiglu_limit > 1e-6f) { + const float uc = std::max(-swiglu_limit, std::min(swiglu_limit, u)); + float gc = g / (1.0f + std::exp(-g)); + gc = std::min(swiglu_limit, gc); + return gc * uc; + } + return (g / (1.0f + std::exp(-g))) * u; +} + +static void expert_forward(const expert_desc & ed, const uint8_t * base, + const float * cur, int64_t n_embd, + const std::vector> & toks, + const float * wts, size_t wts_nb0, size_t wts_nb1, int /*n_used*/, + float swiglu_limit, + float * out, std::vector * scratch) { + sdq_ensure_cpu_init(); + + int64_t gate_rows = ed.ne1; + int64_t up_rows = ed.ne1; + size_t gate_row = ggml_row_size(ed.type[part::GATE], ed.ne0); + size_t up_row = gate_row; + if (ed.fused_gate_up) { + gate_rows = ed.ne1 / 2; + up_rows = gate_rows; + } else { + up_row = ggml_row_size(ed.type[part::UP], ed.ne0); + } + + const vec_dot_fn f_gate = dot_for(ed.type[part::GATE]); + const vec_dot_fn f_up = ed.fused_gate_up ? f_gate : dot_for(ed.type[part::UP]); + const vec_dot_fn f_down = dot_for(ed.type[part::DOWN]); + if (!f_gate || !f_up || !f_down) { + fprintf(stderr, "[sdq] 不支援的專家量化型別：gate=%s up=%s down=%s\n", + ggml_type_name(ed.type[part::GATE]), + ggml_type_name(ed.type[part::UP]), + ggml_type_name(ed.type[part::DOWN])); + return; + } + + const uint8_t * wg = base + ed.slot_off[part::GATE]; + const uint8_t * wu = base + ed.slot_off[part::UP]; + const uint8_t * wd = base + ed.slot_off[part::DOWN]; + const size_t down_row = ggml_row_size(ed.type[part::DOWN], ed.down_ne0); + + const ggml_type act_gate = act_type_for(ed.type[part::GATE]); + const ggml_type act_down = act_type_for(ed.type[part::DOWN]); + const size_t xq_sz = ggml_row_size(act_gate, ed.ne0); + const size_t hq_sz = ggml_row_size(act_down, ed.down_ne0); + if (scratch->size() < xq_sz + hq_sz) { + scratch->assign(xq_sz + hq_sz, 0); + } + uint8_t * xq = scratch->data(); + uint8_t * hq = xq + xq_sz; + + std::vector gate((size_t) gate_rows); + std::vector upv((size_t) up_rows); + std::vector h((size_t) ed.down_ne0); + + PROF_T0_x(_ta); // expert_forward 總計時 + for (const auto & tk : toks) { + const int32_t t = tk.first; + const int32_t k = tk.second; + + PROF_T0_x(_tb); // 量化 activation + quantize_act(act_gate, cur + (size_t) t * n_embd, xq, ed.ne0); + PROF_ADD(qact_ns, _tb); + + PROF_T0_x(_tc); // gate / up / down 三次點積 + for (int64_t r = 0; r < gate_rows; r++) { + float acc = 0.0f; + dot_row(f_gate, wg + r * gate_row, xq, gate_row, (int) ed.ne0, acc); + gate[(size_t) r] = acc; + } + for (int64_t r = 0; r < up_rows; r++) { + float acc = 0.0f; + const uint8_t * src = ed.fused_gate_up ? (wg + (gate_rows + r) * gate_row) + : (wu + r * up_row); + dot_row(f_up, src, xq, ed.fused_gate_up ? gate_row : up_row, (int) ed.ne0, acc); + upv[(size_t) r] = acc; + } + + PROF_ADD(dot_ns, _tc); + + PROF_T0_x(_td); // SiLU + 加權求和 + for (int64_t r = 0; r < ed.down_ne0; r++) { + h[(size_t) r] = gated_act(gate[(size_t) r], upv[(size_t) r], swiglu_limit); + } + quantize_act(act_down, h.data(), hq, ed.down_ne0); + + PROF_ADD(act_ns, _td); + + const float w = *(const float *) ((const uint8_t *) wts + k * wts_nb0 + (size_t) t * wts_nb1); + float * o = out + (size_t) t * n_embd; + for (int64_t r = 0; r < ed.down_ne1; r++) { + float acc = 0.0f; + dot_row(f_down, wd + r * down_row, hq, down_row, (int) ed.down_ne0, acc); + o[r] += acc * w; + } + } + PROF_ADD(fwd_ns, _ta); +} + + + +// ---------------------------------------------------------------- no-gate LatentMoE +// Nemotron-3.5-Bolt 的 routed experts 只有 up/down 兩個矩陣（沒有 gate）： +// h = relu(W_up x)^2 +// y = W_down h +// 輸入與輸出都在 latent 空間（moe_latent_size = 1024）。 +static void expert_forward_relu2(const expert_desc & ed, const uint8_t * base, + const float * cur, int64_t n_embd, + const std::vector> & toks, + const float * wts, size_t wts_nb0, size_t wts_nb1, + float * out, std::vector * scratch) { + sdq_ensure_cpu_init(); + + const int64_t n_ff = ed.ne1; // up 的輸出列數 + const int64_t n_down = ed.down_ne1; // down 的輸出列數（= latent 1024） + const size_t up_row = ggml_row_size(ed.type[part::UP], ed.ne0); + const size_t dn_row = ggml_row_size(ed.type[part::DOWN], ed.down_ne0); + + const vec_dot_fn f_up = dot_for(ed.type[part::UP]); + const vec_dot_fn f_down = dot_for(ed.type[part::DOWN]); + if (!f_up || !f_down) { + static std::mutex m; + std::lock_guard lk(m); + fprintf(stderr, "[sdq] 不支援的 expert 量化型別：up=%s down=%s\n", + ggml_type_name(ed.type[part::UP]), ggml_type_name(ed.type[part::DOWN])); + return; + } + + const uint8_t * wu = base + ed.slot_off[part::UP]; + const uint8_t * wd = base + ed.slot_off[part::DOWN]; + + const ggml_type act_up = act_type_for(ed.type[part::UP]); + const ggml_type act_down = act_type_for(ed.type[part::DOWN]); + const size_t xq_sz = ggml_row_size(act_up, ed.ne0); + const size_t hq_sz = ggml_row_size(act_down, ed.down_ne0); + if (scratch->size() < xq_sz + hq_sz) { + scratch->assign(xq_sz + hq_sz, 0); + } + uint8_t * xq = scratch->data(); + uint8_t * hq = xq + xq_sz; + + std::vector h((size_t) std::max(n_ff, 1)); + + PROF_T0_x(_ta); + for (const auto & tk : toks) { + const int32_t t = tk.first; + const int32_t k = tk.second; + + PROF_T0_x(_tb); + quantize_act(act_up, cur + (size_t) t * n_embd, xq, ed.ne0); + PROF_ADD(qact_ns, _tb); + + PROF_T0_x(_tc); + for (int64_t r = 0; r < n_ff; r++) { + float acc = 0.0f; + dot_row(f_up, wu + (size_t) r * up_row, xq, up_row, (int) ed.ne0, acc); + // relu(x)^2 + acc = acc > 0.0f ? acc : 0.0f; + h[(size_t) r] = acc * acc; + } + quantize_act(act_down, h.data(), hq, ed.down_ne0); + PROF_ADD(dot_ns, _tc); + + const float w = *(const float *) ((const uint8_t *) wts + k * wts_nb0 + (size_t) t * wts_nb1); + float * o = out + (size_t) t * n_embd; + PROF_T0_x(_td); + for (int64_t r = 0; r < n_down; r++) { + float acc = 0.0f; + dot_row(f_down, wd + (size_t) r * dn_row, hq, dn_row, (int) ed.down_ne0, acc); + o[r] += acc * w; + } + PROF_ADD(act_ns, _td); + } + PROF_ADD(fwd_ns, _ta); +} + +// ---------------------------------------------------------------- row-parallel（整層常駐時） +// 目的：decode 時 n_need = 9 < nth = 16，mode C 依 expert 切工作 → 9 條執行緒 +// 有 expert、7 條閒著。原生 ggml_mul_mat_id 的切法不是「一個 expert 一條」， +// 而是「輸出列切給 16 條，每條在自己的列區塊裡跑完所有 expert」。 +// 這裡照同樣的切法：整層的 (expert, token, 輸出列) 展平，16 條執行緒各拿一段， +// 每個 barrier 只發生在「up 全部算完 / h 全部量化完 / down 全部算完」， +// 不是每個 expert 都同步一次。 +// +// 只在「整層 expert 都已常駐」時啟用：row-parallel 沒有把 I/O 和計算交錯， +// 若還有 miss，會比 mode C 的讀取/計算重疊更慢。常駐時完全沒有 I/O。 +// +// 每層 4 道 barrier（acquire+xq、up、h 量化、down），decode 40 層 = 160 道； +// 用自旋屏障，成本是微秒級。 +static inline size_t act_block_bytes(ggml_type t) { + return ggml_row_size(t, QK_K); +} + +// 把一段連續的 QK_K block 量化到 dst（dst 必須指向該 token 的 row 開頭）。 +// 這樣 quantization 也可以切給 16 條執行緒，而不是只讓一條做 2048 個元素。 +static void quantize_blocks(ggml_type t, const float * src, uint8_t * dst, + int64_t b0, int64_t b1) { + const size_t rb = act_block_bytes(t); + for (int64_t b = b0; b < b1; b++) { + quantize_act(t, src + b * QK_K, dst + (size_t) b * rb, QK_K); + } +} + +static void moe_row_parallel(ggml_tensor * dst, int ith, int nth, layer_rec * R, + const layer_desc * L, const float * cur, const float * wts, + size_t wts_nb0, size_t wts_nb1, int64_t n_embd, int64_t n_tok) { + pager & P = pager::instance(); + const int n_need = (int) R->need.size(); + float * out = (float *) dst->data; + + // 和原生 mul_mat_id 一樣以「16 列」為一個工作單位，並用 atomic 動態派工。 + // 每個 chunk 會把該 16 列在「所有 expert/token」上都算完： + // up : 16 rows × 9 experts → h_all + // down : 16 rows × 9 experts → out（同一個 (t,r) 由同一條執行緒累加，無 race） + // 這樣不需要 per-thread partial，也不需要最後的歸約。 + const int64_t RB = 16; + const int64_t n_ff = L->experts[(size_t) R->need[0]].ne1; + const int64_t n_down = L->experts[(size_t) R->need[0]].down_ne1; + const int64_t nblk_up = (n_ff + RB - 1) / RB; + const int64_t nblk_dn = (n_down + RB - 1) / RB; + + // 先把整層的 expert 備齊（常駐時全部是 hit）。 + for (int i = ith; i < n_need; i += nth) { + const int32_t e = R->need[(size_t) i]; + const int sl = P.acquire(R->il, e, true); + R->need_slot[(size_t) i] = sl; + if (sl >= 0) { + R->slot_of_expert[(size_t) e] = sl; + P.note_used(R->il, e); + } + } + + const expert_desc & ed0 = L->experts[(size_t) R->need[0]]; + const ggml_type act_up = act_type_for(ed0.type[part::UP]); + const ggml_type act_down = act_type_for(ed0.type[part::DOWN]); + const size_t xq_sz = ggml_row_size(act_up, ed0.ne0); + const size_t hq_sz = ggml_row_size(act_down, ed0.down_ne0); + const int64_t xq_blocks = ed0.ne0 / QK_K; + const int64_t hq_blocks = ed0.down_ne0 / QK_K; + + // 展平 (expert, token) 的 prefix 與 mapping。這段不依賴 slot，先算， + // 才能在 barrier 前由 thread 0 配置「跨執行緒共用」的 h_all。 + std::vector up_off((size_t) n_need + 1, 0); + std::vector tok_t, tok_k, tok_e; + std::vector tok_h; + int64_t total_up = 0, total_tok = 0; + for (int j = 0; j < n_need; j++) { + const int32_t e = R->need[(size_t) j]; + const auto & toks = R->tok_of_expert[(size_t) e]; + up_off[(size_t) j] = total_up; + for (size_t ti = 0; ti < toks.size(); ti++) { + tok_t.push_back(toks[ti].first); + tok_k.push_back(toks[ti].second); + tok_e.push_back(j); + tok_h.push_back(total_up + (int64_t) ti * n_ff); + } + total_up += (int64_t) toks.size() * n_ff; + total_tok += (int64_t) toks.size(); + } + up_off[(size_t) n_need] = total_up; + + if (ith == 0) { + memset(out, 0, (size_t) ggml_nbytes(dst)); + R->rowpar_chunk_up.store(0, std::memory_order_relaxed); + R->rowpar_chunk_dn.store(0, std::memory_order_relaxed); + // h_all 必須跨執行緒共用：down 的每個 row block 需要「所有」up 列。 + // 之前放在函式區域變數會變成每個執行緒各一份，down 讀到別人的 0。 + R->rowpar_h.assign((size_t) total_up, 0.0f); + R->rowpar_hq.assign((size_t) total_tok * hq_sz, 0); + } + if (nth > 1) { R->rowbar.sync(nth); } + + float * h_all = R->rowpar_h.data(); + uint8_t * hq_all = R->rowpar_hq.data(); + + // 每個 expert 的常數先算好（slot 已由上面的 barrier 保證可見）。 + struct ed_info { + const uint8_t * wu = nullptr; + const uint8_t * wd = nullptr; + const uint8_t * wg = nullptr; // gated 模型才有：gate 權重 + size_t up_row = 0; + size_t dn_row = 0; + size_t g_row = 0; + vec_dot_fn f_up = nullptr; + vec_dot_fn f_down = nullptr; + vec_dot_fn f_gate = nullptr; + const expert_desc * ed = nullptr; + int64_t n_te = 0; + }; + std::vector info((size_t) n_need); + for (int j = 0; j < n_need; j++) { + const int32_t e = R->need[(size_t) j]; + const int sl = R->need_slot[(size_t) j]; + const expert_desc & ed = L->experts[(size_t) e]; + ed_info & q = info[(size_t) j]; + q.ed = &ed; + q.n_te = (int64_t) R->tok_of_expert[(size_t) e].size(); + q.up_row = ggml_row_size(ed.type[part::UP], ed.ne0); + q.dn_row = ggml_row_size(ed.type[part::DOWN], ed.down_ne0); + q.f_up = dot_for(ed.type[part::UP]); + q.f_down = dot_for(ed.type[part::DOWN]); + if (!R->no_gate) { + q.g_row = ggml_row_size(ed.type[part::GATE], ed.ne0); + q.f_gate = dot_for(ed.type[part::GATE]); + } + if (sl >= 0) { + q.wu = P.slot_ptr(sl) + ed.slot_off[part::UP]; + q.wd = P.slot_ptr(sl) + ed.slot_off[part::DOWN]; + if (!R->no_gate) { + q.wg = P.slot_ptr(sl) + ed.slot_off[part::GATE]; + } + } + } + + std::vector xq_local((size_t) n_tok * xq_sz); + std::vector xq_done((size_t) n_tok, 0); + + // ---- up：動態抓 16 列 block，對所有 (expert, token) 一起算 + for (;;) { + const int64_t b = R->rowpar_chunk_up.fetch_add(1, std::memory_order_relaxed); + if (b >= nblk_up) { + break; + } + const int64_t r0 = b * RB; + const int64_t r1 = std::min(r0 + RB, n_ff); + for (int j = 0; j < n_need; j++) { + const ed_info & q = info[(size_t) j]; + if (!q.wu || !q.f_up) { + continue; + } + if (!R->no_gate && (!q.wg || !q.f_gate)) { + continue; + } + const auto & toks = R->tok_of_expert[(size_t) R->need[(size_t) j]]; + for (size_t ti = 0; ti < toks.size(); ti++) { + const int32_t t = toks[ti].first; + if (!xq_done[(size_t) t]) { + quantize_blocks(act_up, cur + (size_t) t * n_embd, + xq_local.data() + (size_t) t * xq_sz, 0, xq_blocks); + xq_done[(size_t) t] = 1; + } + const uint8_t * xq = xq_local.data() + (size_t) t * xq_sz; + float * h = h_all + up_off[(size_t) j] + (int64_t) ti * n_ff; + for (int64_t r = r0; r < r1; r++) { + float acc = 0.0f; + dot_row(q.f_up, q.wu + (size_t) r * q.up_row, xq, + q.up_row, (int) q.ed->ne0, acc); + if (R->no_gate) { + // Nemotron LatentMoE：relu(x)^2，沒有 gate + acc = acc > 0.0f ? acc : 0.0f; + h[(size_t) r] = acc * acc; + } else { + // gated（Flint）：silu(gate)*up，與 expert_forward 相同的 SwiGLU + float g = 0.0f; + dot_row(q.f_gate, q.wg + (size_t) r * q.g_row, xq, + q.g_row, (int) q.ed->ne0, g); + h[(size_t) r] = gated_act(g, acc, R->swiglu_limit); + } + } + } + } + } + if (nth > 1) { R->rowbar.sync(nth); } + + // ---- h 量化：total_tok × 8 個 block 平均分給 16 條執行緒。 + // 不要讓每條執行緒各自量化一次：先前版本是 16 倍冗餘，decode 時反而成為瓶颈。 + { + const int64_t total_blocks = total_tok * hq_blocks; + for (int64_t b = ith; b < total_blocks; b += nth) { + const int64_t tg = b / hq_blocks; + const int64_t bb = b % hq_blocks; + quantize_blocks(act_down, h_all + tok_h[(size_t) tg], + hq_all + (size_t) tg * hq_sz, bb, bb + 1); + } + } + if (nth > 1) { R->rowbar.sync(nth); } + + // ---- down：同樣 16 列 block；同一 block 內把所有 expert 的貢獻加起來， + // 所以每個 (t, r) 只會被一條執行緒碰到，直接寫 out 即可。 + for (;;) { + const int64_t b = R->rowpar_chunk_dn.fetch_add(1, std::memory_order_relaxed); + if (b >= nblk_dn) { + break; + } + const int64_t r0 = b * RB; + const int64_t r1 = std::min(r0 + RB, n_down); + for (int64_t tg = 0; tg < total_tok; tg++) { + const int j = tok_e[(size_t) tg]; + const ed_info & q = info[(size_t) j]; + if (!q.wd || !q.f_down) { + continue; + } + const uint8_t * hq = hq_all + (size_t) tg * hq_sz; + const int32_t t = tok_t[(size_t) tg]; + const int32_t k = tok_k[(size_t) tg]; + const float w = *(const float *) ((const uint8_t *) wts + k * wts_nb0 + + (size_t) t * wts_nb1); + float * o = out + (size_t) t * n_embd; + for (int64_t r = r0; r < r1; r++) { + float acc = 0.0f; + dot_row(q.f_down, q.wd + (size_t) r * q.dn_row, hq, + q.dn_row, (int) q.ed->down_ne0, acc); + o[(size_t) r] += acc * w; + } + } + } + if (nth > 1) { R->rowbar.sync(nth); } + + if (ith == 0) { + for (int i = 0; i < n_need; i++) { + if (R->need_slot[(size_t) i] >= 0) { + P.release(R->need_slot[(size_t) i]); + } + } + } +} + +// 每次 MoE op 實際走的路徑計數（v2 row-parallel vs v1 mode C），結束時印出 +static std::atomic g_ops_rowpar{0}; +static std::atomic g_ops_modec{0}; + +static void moe_custom_op(ggml_tensor * dst, int ith, int nth, void * userdata) { + layer_rec * R = (layer_rec *) userdata; + pager & P = pager::instance(); + + // 從 op 的 src 取得輸入（保證是「這一次」真的被計算的張量） + const ggml_tensor * src_cur = dst->src[0]; + const ggml_tensor * src_ids = dst->src[1]; + const ggml_tensor * src_wts = dst->src[2]; + if (!src_cur || !src_ids || !src_wts) { + if (ith == 0) { + fprintf(stderr, "[sdq] op 缺少 src（%d）\n", R->il); + } + return; + } + + const layer_desc * L = P.layer(R->il); + PROF_T0_x(_tg); // op 總計時 + g_prof.calls++; + const int64_t n_embd = src_cur->ne[0]; + const int64_t n_tok = src_cur->ne[2]; + const int64_t n_used = R->n_expert_used; + // ids 是 2D [n_expert_used, n_tokens] → 用 nb[0]/nb[1] + // wts 是 3D [1, n_expert_used, n_tokens] → 用 nb[1]/nb[2]（nb[0] 會是 4 而不是 n_used*4！） + const size_t ids_nb0 = src_ids->nb[0], ids_nb1 = src_ids->nb[1]; + const size_t wts_nb0 = src_wts->nb[src_wts->ne[0] == 1 ? 1 : 0]; + const size_t wts_nb1 = src_wts->nb[src_wts->ne[0] == 1 ? 2 : 1]; + const float * cur = (const float *) src_cur->data; + const float * wts = (const float *) src_wts->data; + + if (!L || L->experts.empty() || dst->type != GGML_TYPE_F32) { + if (ith == 0 && !R->warned) { + R->warned = true; + fprintf(stderr, "[sdq] 層 %d 沒有可用的分頁描述\n", R->il); + } + return; + } + + if (R->il == 0 && getenv("SDQ_DBG_OP")) { + fprintf(stderr, "[sdq-op] L%d n_tok=%lld nth=%d need=%d ids_nb1=%zu wts_nb1=%zu\n", + R->il, (long long) n_tok, nth, (int) R->need.size(), ids_nb1, wts_nb1); + } + // ---- 1) 解除上一層的 pin（mode B 直接累加輸出，需要先清空） + PROF_T0_x(_tsetup); + if (ith == 0) { + P.unpin_all(); + } + if (nth > 1) { R->bar.sync(nth); } + + // ---- 2) 依 expert 分組 token（只有執行緒 0 做，其餘在此等待） + if (ith == 0) { + std::lock_guard lk(R->mtx); + R->tok_of_expert.assign((size_t) R->n_expert, {}); + R->exp_of_tok.assign((size_t) n_tok, {}); + const uint8_t * ids_base = (const uint8_t *) src_ids->data; + for (int64_t t = 0; t < n_tok; t++) { + for (int64_t k = 0; k < n_used; k++) { + const int32_t e = *(const int32_t *) (ids_base + k * ids_nb0 + (size_t) t * ids_nb1); + if (e >= 0 && e < R->n_expert) { + auto & v = R->tok_of_expert[(size_t) e]; + bool seen = false; + for (const auto & p : v) { + if (p.first == (int32_t) t) { seen = true; break; } + } + if (!seen) { + v.emplace_back((int32_t) t, (int32_t) k); + } + R->exp_of_tok[(size_t) t].emplace_back(e, (int32_t) k); + } + } + } + R->need.clear(); + for (int64_t e = 0; e < R->n_expert; e++) { + if (!R->tok_of_expert[(size_t) e].empty()) { + R->need.push_back((int32_t) e); + } + } + R->need_sorted = R->need; + std::sort(R->need_sorted.begin(), R->need_sorted.end()); + R->need_slot.assign(R->need.size(), -1); // 只能由 ith==0 配置，其餘執行緒唯讀 + R->slot_of_expert.assign((size_t) R->n_expert, -1); + R->need_slot_of_expert.assign((size_t) R->n_expert, -1); + R->miss_sorted.clear(); + R->miss_sorted.reserve((size_t) R->n_expert); + if (R->partial.size() < 64) { + R->partial.resize(64); + } + sdq_route_log(R->il, n_tok, src_ids, ids_nb0, ids_nb1, (int) n_used); + } + if (nth > 1) { R->bar.sync(nth); } + PROF_ADD(setup_ns, _tsetup); + + const int n_need = (int) R->need.size(); + + // ---- 2.5) 決定本層要不要走 row-parallel + // + // 條件：no-gate（本模型）、decode（n_need < nth）、而且整層 expert 都已常駐。 + // 常駐代表沒有 I/O 可以跟計算重疊，mode C 的「讀→算交錯」優勢消失， + // 此時把 expert 當序列外層、輸出列當平行內層，才能讓 16 條執行緒都有事做。 + // 只要有一個 miss 就退回 mode C，避免把小記憶體最需要的 I/O/計算重疊拆掉。 + static const bool rowpar_env = [] { + const char * e = getenv("SDQ_ROW_PARALLEL"); + return e ? atoi(e) != 0 : true; + }(); + static const int rowpar_max_miss = [] { + const char * e = getenv("SDQ_ROWPAR_MAX_MISS"); + return e ? atoi(e) : 0; + }(); + if (ith == 0) { + R->use_rowpar = false; + R->rowpar_miss = 0; + // gated（Flint）需要 gate 與 up 共用同一個 activation 量化型別，才能共用 xq。 + bool types_ok = true; + if (!R->no_gate && n_need > 0) { + const expert_desc & e0 = L->experts[(size_t) R->need[0]]; + types_ok = !e0.fused_gate_up && dot_for(e0.type[part::GATE]) && dot_for(e0.type[part::UP]) && + act_type_for(e0.type[part::GATE]) == act_type_for(e0.type[part::UP]); + } + if (rowpar_env && n_need > 0 && n_need < nth && types_ok) { + int miss = 0; + for (int i = 0; i < n_need; i++) { + if (!P.resident(R->il, R->need[(size_t) i])) { + miss++; + } + } + R->rowpar_miss = miss; + R->use_rowpar = miss <= rowpar_max_miss; + } + static std::atomic dbg_left{8}; + if (getenv("SDQ_DBG_OP") && dbg_left.fetch_add(-1) > 0) { + fprintf(stderr, "[sdq-op] L%d rowpar=%d need=%d n_tok=%lld\n", + R->il, (int) R->use_rowpar, n_need, (long long) n_tok); + } + } + if (nth > 1) { R->bar.sync(nth); } + + // 計數：每次 op 實際走哪條路徑（ith==0 記一次；barrier 之後 use_rowpar 已確定） + if (ith == 0) { + (R->use_rowpar ? g_ops_rowpar : g_ops_modec).fetch_add(1, std::memory_order_relaxed); + } + if (R->use_rowpar) { + moe_row_parallel(dst, ith, nth, R, L, cur, wts, wts_nb0, wts_nb1, n_embd, n_tok); + } + + // ---- 3+4) 分頁與計算融合成一條 pipeline + // + // 關鍵效能點：不要「先把整層的 expert 都讀完再開始算」。做法是每個執行緒 + // 自己依序處理它負責的 expert：讀一個 → 立刻算它 → 再讀下一個。 + // 這樣 SSD 讀取和 CPU 計算自然重疊，而且 8 條執行緒各自發出 pread， + // SSD 佇列深度也夠（不需要再靠 queue depth）。 + // + // 平行切分：mode A 依 expert 切（每個執行緒有自己的累加緩衝，最後相加）； + // mode B 依 token 切（同一個 token 的所有 expert 由同一執行緒算， + // 但必須先把該 token 的 expert 都備齊，所以先分頁再計算）。 + const int64_t out_elems = n_embd * n_tok; + const size_t partial_bytes = (size_t) out_elems * 4 * (size_t) nth; + const bool mode_a = partial_bytes <= (size_t) env_partial_cap_mb() << 20; + float * out = (float *) dst->data; + + if (!mode_a && !R->use_rowpar) { + memset(dst->data, 0, (size_t) ggml_nbytes(dst)); + } + + if (!R->use_rowpar && mode_a) { + // ===================== mode C（預設）：依 expert 流水線 ===================== + // 三個設計重點，每個都來自實測： + // + // (1) **取→算交錯，中間沒有 barrier**。 + // 舊的 mode B 是「先把整層 8 個 expert 全部讀完 → barrier → 才開始算」。 + // 量測顯示 decode 有 66% 的時間卡在 barrier 上等最慢的執行緒 + // （SDQ_PROFILE_TIME 的 barrier 等待 = op 時間的 70%）。改成交錯後 + // 讀取與計算自然重疊，實測 2.79 → 3.49 tok/s。 + // + // (2) **分頁走 acquire()，不用 reserve_slot + read_expert_span**。 + // reserve_slot() 會在資料還沒讀進來之前就把 (layer,expert) publish 到 + // page_ 表；一旦該層的 pread 失敗或被中斷，這個槽就會留在 page_ 裡 + // 卻是舊資料，之後的 probe() 會把它當成命中 → 靜默算錯。 + // acquire() 是「先讀完再 publish」，成本相同但不會有這個失效模式。 + // + // (3) **i += nth 交錯分配**，不是切成連續區塊。 + // 每個 expert 的 SSD 延遲差異很大，交錯分配讓快執行緒自然多幫忙。 + // + // 輸出不直接寫 dst（同一個 token 的不同 expert 會撞在一起）， + // 所以仍然用「每執行緒一份累加緩衝 + 最後歸約」。 + std::vector & mine = R->partial[(size_t) ith]; + mine.assign((size_t) out_elems, 0.0f); + std::vector scratch; + + // 在取本層的 expert **之前**先替下一層排預取。 + // + // 實測結論（S7）：**預設關閉，而且量到開啟反而慢 8%。** + // 預測來自 prev_used_[il+1]（上一個 token 在該層選的專家）。想用它躲在 + // 「本層的讀取＋計算」後面，但量到的結果是： + // 開（5.68 tok/s，預取使用率 0.48 %）vs 關（6.20 tok/s，0.28 %） + // 命中率幾乎一樣（70.5 % vs 70.7 %），但 SSD 位元組從 201.5 升到 216.0 MB/token + // → 預測不夠準，多讀的全是白讀，而且還要和本層的 demand 讀搶頻寬。 + // 另外開啟時同一提示詞的輸出不一致 → 預取佔用的槽位（busy）會讓 + // acquire() 偶爾拿不到槽位而**靜默跳過一個 expert**，這是正確性問題。 + // + // 保留程式碼是為了讓後人不用重新發明；文件見 docs/06_slots.md §4。 + // SDQ_PREFETCH_LAYER_AHEAD=1 可以打開來重現那個負結果。 + static const bool layer_ahead = [] { + const char * e = getenv("SDQ_PREFETCH_LAYER_AHEAD"); + return e ? atoi(e) != 0 : false; + }(); + if (layer_ahead && ith == 0) { + P.prefetch_layer_ahead(R->il + 1); + } + + for (int i = ith; i < n_need; i += nth) { + const int32_t e = R->need[(size_t) i]; + const int32_t sl = P.acquire(R->il, e, true); + if (sl < 0) { + continue; + } + P.note_used(R->il, e); + if (R->no_gate) { + expert_forward_relu2(L->experts[(size_t) e], P.slot_ptr(sl), cur, n_embd, + R->tok_of_expert[(size_t) e], wts, wts_nb0, wts_nb1, + mine.data(), &scratch); + } else { + expert_forward(L->experts[(size_t) e], P.slot_ptr(sl), cur, n_embd, + R->tok_of_expert[(size_t) e], wts, wts_nb0, wts_nb1, + (int) n_used, R->swiglu_limit, mine.data(), &scratch); + } + // expert_forward 回傳後這個槽位就沒人讀了（最後的歸約只讀 mine）， + // 所以立刻放掉 pin，不要讓它被釘到下一層開頭。 + // 不放的話 prefill 一層會把整層用到的 expert 全部釘住（最多 256 個）， + // 大槽位不夠時就會中止（見 pager::release 的說明）。 + P.release(sl); + } + + if (nth > 1) { R->bar.sync(nth); } + for (int64_t idx = ith; idx < out_elems; idx += nth) { + float acc = 0.0f; + for (int k = 0; k < nth; k++) { + acc += R->partial[(size_t) k][(size_t) idx]; + } + out[idx] = acc; + } + } else if (!R->use_rowpar) { + // mode B：token 多到 partial 緩衝太大時，依 token 切分 + // 先把所有需要的 expert 分頁（各執行緒平行 pread） + { + PROF_T0_x(_th); // 分頁 / SSD 讀取 + for (int i = ith; i < n_need; i += nth) { + const int32_t e = R->need[(size_t) i]; + const int s = P.acquire(R->il, e, true); + g_prof.experts++; + R->need_slot[(size_t) i] = s; + if (s >= 0) { + R->slot_of_expert[(size_t) e] = s; + P.note_used(R->il, e); + } + } + PROF_ADD(read_ns, _th); + } + if (nth > 1) { R->bar.sync(nth); } + const float * cur_p = (const float *) src_cur->data; + const float * wts_p = (const float *) src_wts->data; + for (int64_t t = ith; t < n_tok; t += nth) { + std::vector scratch; + for (const auto & ek : R->exp_of_tok[(size_t) t]) { + const int slot = R->slot_of_expert[(size_t) ek.first]; + if (slot < 0) { + continue; + } + std::vector> one = {{(int32_t) t, ek.second}}; + if (R->no_gate) { + expert_forward_relu2(L->experts[(size_t) ek.first], P.slot_ptr(slot), cur_p, n_embd, + one, wts_p, wts_nb0, wts_nb1, out, &scratch); + } else { + expert_forward(L->experts[(size_t) ek.first], P.slot_ptr(slot), cur_p, n_embd, + one, wts_p, wts_nb0, wts_nb1, (int) n_used, + R->swiglu_limit, out, &scratch); + } + } + } + } + + // 除錯：把某一層的實際輸入/輸出 dump 出來，用 NumPy 對照。 + // 必須先 barrier：上面的歸約是由各執行緒分片寫 out 的，沒有這一道會讀到半成品。 + if (nth > 1) { R->bar.sync(nth); } + if (const char * dump_path = getenv("SDQ_MOE_DUMP")) { + const int dump_layer = getenv("SDQ_MOE_DUMP_LAYER") ? atoi(getenv("SDQ_MOE_DUMP_LAYER")) : -1; + if ((dump_layer < 0 || R->il == dump_layer) && ith == 0) { + std::string path = std::string(dump_path) + "." + std::to_string(R->il); + FILE * f = fopen(path.c_str(), "wb"); + if (f) { + const int32_t meta[6] = {(int32_t) n_embd, (int32_t) n_tok, (int32_t) n_used, + (int32_t) R->n_expert, 1, 2}; + fwrite(meta, sizeof(int32_t), 6, f); + fwrite(cur, sizeof(float), (size_t) (n_embd * n_tok), f); + const uint8_t * ids_base = (const uint8_t *) src_ids->data; + for (int64_t t = 0; t < n_tok; t++) { + for (int64_t k = 0; k < n_used; k++) { + const int32_t e = *(const int32_t *) (ids_base + k * ids_nb0 + (size_t) t * ids_nb1); + fwrite(&e, sizeof(int32_t), 1, f); + } + } + fwrite(wts, sizeof(float), (size_t) (n_used * n_tok), f); + fwrite(out, sizeof(float), (size_t) (n_embd * n_tok), f); + fclose(f); + fprintf(stderr, "[sdq] dumped layer %d -> %s\n", R->il, path.c_str()); + } + } + } + + if (nth > 1) { R->bar.sync(nth); } + PROF_ADD(op_ns, _tg); +} + +// 輸出路徑計數（sdq_shutdown 呼叫） +void sdq_print_op_counts() { + fprintf(stderr, "[sdq] MoE op 路徑計數：row-parallel %llu 次，mode C %llu 次\n", + (unsigned long long) g_ops_rowpar.load(), (unsigned long long) g_ops_modec.load()); +} + +// 輸出計時結果（sdq_shutdown 呼叫） +void sdq_print_time_profile() { + if (!prof_on()) { + return; + } + const double ms = 1e6; + auto v = [ms](const std::atomic & a) { return a.load() / ms; }; + const double op = v(g_prof.op_ns); + const double fwd = v(g_prof.fwd_ns); + fprintf(stderr, "[sdq-prof] 自訂 MoE op（%llu 次呼叫、%llu 個 expert 分頁）\n", + (unsigned long long) g_prof.calls.load(), (unsigned long long) g_prof.experts.load()); + fprintf(stderr, "[sdq-prof] op 合計 %9.1f ms\n", op); + fprintf(stderr, "[sdq-prof] expert_forward %9.1f ms (%5.1f%%) 其中：\n", fwd, op > 0 ? 100.0 * fwd / op : 0.0); + fprintf(stderr, "[sdq-prof] 量化 activation%9.1f ms (%5.1f%% of fwd)\n", v(g_prof.qact_ns), fwd > 0 ? 100.0 * v(g_prof.qact_ns) / fwd : 0.0); + fprintf(stderr, "[sdq-prof] 三次點積 %9.1f ms (%5.1f%% of fwd)\n", v(g_prof.dot_ns), fwd > 0 ? 100.0 * v(g_prof.dot_ns) / fwd : 0.0); + fprintf(stderr, "[sdq-prof] SiLU+加權 %9.1f ms (%5.1f%% of fwd)\n", v(g_prof.act_ns), fwd > 0 ? 100.0 * v(g_prof.act_ns) / fwd : 0.0); + fprintf(stderr, "[sdq-prof] SSD 讀取/分頁 %9.1f ms (%5.1f%%)\n", v(g_prof.read_ns), op > 0 ? 100.0 * v(g_prof.read_ns) / op : 0.0); + fprintf(stderr, "[sdq-prof] 前置 setup %9.1f ms (%5.1f%%)\n", v(g_prof.setup_ns), op > 0 ? 100.0 * v(g_prof.setup_ns) / op : 0.0); + fprintf(stderr, "[sdq-prof] barrier 等待 %9.1f ms (%5.1f%%)\n", v(g_prof.bar_ns), op > 0 ? 100.0 * v(g_prof.bar_ns) / op : 0.0); +} + +// ---------------------------------------------------------------- graph 建構 + +ggml_tensor * sdq_build_moe_ex(ggml_cgraph * gf, ggml_context * ctx, ggml_tensor * cur, + ggml_tensor * ids, ggml_tensor * weights, int il, int64_t n_expert, + int64_t n_expert_used, bool no_gate, bool relu_sqr) { + layer_rec * R = rec_for(il); + R->n_expert = n_expert; + R->n_expert_used = n_expert_used; + R->no_gate = no_gate; + R->relu_sqr = relu_sqr; + R->swiglu_limit = 0.0f; + + // 呼叫端給的是 2D [n_embd, n_tokens]，這裡跟上游一樣先 reshape 成 3D + // [n_embd, 1, n_tokens]（上游 mul_mat_id 路徑也是這樣拿 token 數的） + ggml_tensor * cur3 = (cur->ne[1] == 1) ? cur : ggml_reshape_3d(ctx, cur, cur->ne[0], 1, cur->ne[1]); + const int64_t n_tok = cur3->ne[2]; + + // Nemotron LatentMoE 的 down 輸出寬度 = moe_latent_size，與輸入相同。 + int64_t out_w = cur3->ne[0]; + const layer_desc * L = pager::instance().layer(il); + if (no_gate && L && !L->experts.empty()) { + out_w = L->experts[0].down_ne1; + } + + ggml_tensor * args[3] = { cur3, ids, weights }; + ggml_tensor * out = ggml_custom_4d(ctx, GGML_TYPE_F32, out_w, n_tok, 1, 1, + args, 3, moe_custom_op, GGML_N_TASKS_MAX, R); + ggml_set_name(out, "sdq_moe_out"); + // 關鍵：把這個節點註冊進圖，ggml 才會為它配置獨立的 compute buffer。 + // 少了這一行，layer>=1 的輸出會被其他張量覆蓋（layer 0 僥倖正確）。 + ggml_build_forward_expand(gf, out); + return out; +} + +ggml_tensor * sdq_build_moe(ggml_cgraph * gf, ggml_context * ctx, ggml_tensor * cur, + ggml_tensor * ids, ggml_tensor * weights, int il, int64_t n_expert, + int64_t n_expert_used, float swiglu_limit) { + ggml_tensor * out = sdq_build_moe_ex(gf, ctx, cur, ids, weights, il, n_expert, n_expert_used, false, false); + // 把 swiglu limit 傳給 op（先前這裡 (void) 丟掉，造成 clamp 永遠不生效） + rec_for(il)->swiglu_limit = swiglu_limit; + return out; +} + +} // namespace sdq +namespace sdq { +// 自我測試：對指定 (layer, expert) 用 pager 的資料算一次 MoE 輸出， +// 讓 Python 用純 NumPy 從同一段 GGUF 位元組重算並比對。 +// 輸出：out[0..n_embd) = 輸出，out[n_embd..2*n_embd) = 輸入 activation +int sdq_moe_selftest(int il, int ie, const float * x, float * out_io, + int64_t n_embd, int64_t n_ff, int n_tok) { + pager & P = pager::instance(); + const layer_desc * L = P.layer(il); + if (!L || ie < 0 || ie >= (int) L->experts.size()) { + fprintf(stderr, "[sdq-selftest] 層/expert 不存在：%d/%d\n", il, ie); + return 1; + } + const int slot = P.acquire(il, ie, true); + if (slot < 0) { + fprintf(stderr, "[sdq-selftest] 分頁失敗\n"); + return 2; + } + const expert_desc & ed = L->experts[(size_t) ie]; + if (ed.ne0 != n_embd || ed.ne1 != n_ff) { + fprintf(stderr, "[sdq-selftest] 維度不符：ne0=%lld ne1=%lld（預期 %lld/%lld）\n", + (long long) ed.ne0, (long long) ed.ne1, (long long) n_embd, (long long) n_ff); + return 3; + } + std::vector> toks; + std::vector wts((size_t) n_tok, 1.0f); + for (int t = 0; t < n_tok; t++) { + toks.emplace_back(t, 0); + } + std::vector out((size_t) n_embd * n_tok, 0.0f); + std::vector scratch; + // SDQ_SELFTEST_DUMP: 把 arena 槽位的原始位元組存檔，驗證 pager 讀到的是正確的 expert + if (const char * dp = getenv("SDQ_SELFTEST_DUMP")) { + std::string path = std::string(dp) + "." + std::to_string(il) + "." + std::to_string(ie); + FILE * f = fopen(path.c_str(), "wb"); + if (f) { + fwrite(P.slot_ptr(slot), 1, (size_t) ed.bytes, f); + fclose(f); + fprintf(stderr, "[sdq-selftest] slot dump -> %s (%lld B)\n", path.c_str(), (long long) ed.bytes); + } + } + expert_forward(ed, P.slot_ptr(slot), x, n_embd, toks, wts.data(), 4, 4, 1, 0.0f, + out.data(), &scratch); + for (int64_t t = 0; t < n_tok; t++) { + for (int64_t i = 0; i < n_embd; i++) { + out_io[t * n_embd + i] = out[(size_t) t * n_embd + i]; + out_io[n_tok * n_embd + t * n_embd + i] = x[(size_t) t * n_embd + i]; + } + } + P.unpin_all(); + fprintf(stderr, "[sdq-selftest] L%d E%d n_tok=%d 計算完成（gate %s / up %s / down %s）\n", + il, ie, n_tok, ggml_type_name(ed.type[part::GATE]), ggml_type_name(ed.type[part::UP]), + ggml_type_name(ed.type[part::DOWN])); + return 0; +} +} // namespace sdq diff --git a/src/sdq_pager.cpp b/src/sdq_pager.cpp new file mode 100644 index 0000000..ff22bf7 --- /dev/null +++ b/src/sdq_pager.cpp @@ -0,0 +1,2567 @@ +// SDQ expert pager 實作：見 sdq_pager.h 的設計說明 +#include "sdq_pager.h" + +#include "llama-mmap.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace sdq { + +// ============================================================ 小工具 + +uint64_t now_us() { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (uint64_t) ts.tv_sec * 1000000ull + (uint64_t) ts.tv_nsec / 1000ull; +} + +size_t process_rss_bytes() { + std::ifstream f("/proc/self/statm"); + size_t total = 0, resident = 0; + if (f >> total >> resident) { + return resident * (size_t) sysconf(_SC_PAGESIZE); + } + return 0; +} + +size_t physical_ram_bytes() { + std::ifstream f("/proc/meminfo"); + std::string key; + size_t value = 0; + std::string unit; + while (f >> key >> value >> unit) { + if (key == "MemTotal:") { + return value * 1024; + } + } + return 0; +} + +// 可用記憶體（自動模式用）：取 /proc/meminfo 的 MemAvailable，若處於 cgroup v2 限制下， +// 再與 memory.max - memory.current 取較小值（容器上限比整台主機的可用量更有意義）。 +// memory.current 含 page cache，所以這個估計偏保守。 +size_t available_ram_bytes() { + size_t avail = 0; + { + std::ifstream f("/proc/meminfo"); + std::string key, unit; + size_t value = 0; + while (f >> key >> value >> unit) { + if (key == "MemAvailable:") { + avail = value * 1024; + break; + } + } + } + std::string cg; + { + std::ifstream f("/proc/self/cgroup"); + std::string line; + while (std::getline(f, line)) { + if (line.rfind("0::", 0) == 0) { + cg = line.substr(3); + break; + } + } + } + const std::string base = "/sys/fs/cgroup" + (cg == "/" ? std::string() : cg); + auto read_num = [](const std::string & path, size_t & out) -> bool { + std::ifstream f(path); + std::string v; + if (!(f >> v) || v == "max") { + return false; + } + try { + out = (size_t) std::stoull(v); + } catch (...) { + return false; + } + return true; + }; + size_t limit = 0, current = 0; + if (read_num(base + "/memory.max", limit) && read_num(base + "/memory.current", current)) { + const size_t headroom = limit > current ? limit - current : 0; + if (avail == 0 || headroom < avail) { + avail = headroom; + } + } + return avail; +} + +static size_t env_size(const char * name, size_t dflt) { + const char * v = getenv(name); + if (!v || !*v) { + return dflt; + } + return (size_t) strtoull(v, nullptr, 10); +} + +static bool env_flag(const char * name, bool dflt) { + const char * v = getenv(name); + if (!v || !*v) { + return dflt; + } + return !(v[0] == '0' && v[1] == '\0'); +} + +// ============================================================ GGUF 檔頭解析 + +namespace { + +struct gguf_reader { + const uint8_t * p = nullptr; + size_t size = 0; + size_t o = 0; + + uint8_t u8() { return p[o++]; } + uint32_t u32() { uint32_t v; memcpy(&v, p + o, 4); o += 4; return v; } + uint64_t u64() { uint64_t v; memcpy(&v, p + o, 8); o += 8; return v; } + int64_t i64() { int64_t v; memcpy(&v, p + o, 8); o += 8; return v; } + float f32() { float v; memcpy(&v, p + o, 4); o += 4; return v; } + std::string str() { + uint64_t n = u64(); + std::string s((const char *) p + o, (size_t) n); + o += n; + return s; + } + void skip(uint64_t n) { o += (size_t) n; } + uint16_t u16() { uint16_t v; memcpy(&v, p + o, 2); o += 2; return v; } + int16_t i16() { int16_t v; memcpy(&v, p + o, 2); o += 2; return v; } + double f64() { double v; memcpy(&v, p + o, 8); o += 8; return v; } +}; + +struct kv_value { + enum kind { NONE, INT, REAL, STR, ARR } k = NONE; + int64_t i = 0; + double f = 0; + std::string s; + std::vector arr; +}; + +struct tensor_info { + std::string name; + std::vector dims; + ggml_type type = GGML_TYPE_COUNT; + uint64_t offset = 0; // 在所屬 shard 檔案內、相對 data_start 的偏移 + int shard = 0; // 屬於哪一個 shard（sharded GGUF） +}; + +size_t ggml_type_size_of(ggml_type t) { + return ggml_row_size(t, 1); +} + +} // namespace + +// ============================================================ pager + +pager & pager::instance() { + static pager p; + return p; +} + +// 依 GGUF split 命名規則把第一個 shard 路徑展開成全部 shard。 +// 規則是 "-00001-of-00004.gguf"；llama.cpp 用 llama_split_prefix/llama_split_path +// 做同一件事，這裡獨立實作以免 pager 依賴 loader 的內部符號。 +std::vector pager::list_splits(const std::string & first_path, int n_split) { + std::vector out; + if (n_split <= 1) { + out.push_back(first_path); + return out; + } + // 用 llama.cpp 自己的命名規則（%s-%05d-of-%05d.gguf，且 split_no 由 1 起算）。 + // 自己刻字串很容易差一位數字 —— 實測就是這樣變成 "...XL--00001-of-..."。 + std::vector pre(llama_path_max(), 0); + const int pre_len = llama_split_prefix(pre.data(), pre.size(), first_path.c_str(), 0, n_split); + if (pre_len <= 0) { + throw std::runtime_error("無法解析 split 檔名：" + first_path); + } + std::vector buf(llama_path_max(), 0); + for (int i = 0; i < n_split; i++) { + const int len = llama_split_path(buf.data(), buf.size(), pre.data(), i, n_split); + if (len <= 0) { + throw std::runtime_error("無法產生 split 路徑：" + first_path); + } + out.emplace_back(buf.data(), (size_t) len); + } + return out; +} + +void pager::load_gguf_header(const std::string & path) { + const int fd = open(path.c_str(), O_RDONLY); + if (fd < 0) { + throw std::runtime_error("無法開啟 GGUF：" + path); + } + + struct stat st; + fstat(fd, &st); + + // 只讀檔頭：metadata + tensor 資訊大約 11 MB（shard 1 還含整份 tokenizer， + // 這裡的 64 MiB 上限對它足夠） + const size_t want = (size_t) std::min(st.st_size, 64ull << 20); + std::vector buf(want); + ssize_t got = 0; + while (got < (ssize_t) want) { + ssize_t n = pread(fd, buf.data() + got, want - got, got); + if (n <= 0) { + break; + } + got += n; + } + buf.resize((size_t) std::max(got, 0)); + + gguf_reader r{buf.data(), buf.size(), 0}; + if (memcmp(buf.data(), "GGUF", 4) != 0) { + close(fd); + throw std::runtime_error("不是 GGUF 檔：" + path); + } + r.o = 4; + const uint32_t version = r.u32(); + if (version < 2 || version > 3) { + close(fd); + throw std::runtime_error("不支援的 GGUF 版本"); + } + const uint64_t n_tensors = r.u64(); + const uint64_t n_kv = r.u64(); + + std::map kv; + for (uint64_t i = 0; i < n_kv; i++) { + std::string key = r.str(); + const uint32_t t = r.u32(); + kv_value v; + switch (t) { + // GGUF 型別寬度：U8/I8/BOOL=1、U16/I16=2、U32/I32/F32=4、U64/I64/F64=8。 + // 舊版把 case 0/1/2/3 都讀成 u32，遇到 tokenizer 或 U16 metadata 就整個 + // header 錯位（症狀是 std::length_error: basic_string::_M_create）。 + case 0: v.k = kv_value::INT; v.i = (int64_t) r.u8(); v.f = (double) v.i; break; + case 1: v.k = kv_value::INT; v.i = (int64_t) (int8_t) r.u8(); v.f = (double) v.i; break; + case 2: v.k = kv_value::INT; v.i = (int64_t) r.u16(); v.f = (double) v.i; break; + case 3: v.k = kv_value::INT; v.i = (int64_t) r.i16(); v.f = (double) v.i; break; + case 4: v.k = kv_value::INT; v.i = (int64_t) (uint32_t) r.u32(); v.f = (double) v.i; break; + case 5: v.k = kv_value::INT; v.i = (int64_t) (int32_t) r.u32(); v.f = (double) v.i; break; + case 6: v.k = kv_value::REAL; v.f = r.f32(); v.i = (int64_t) v.f; break; + case 7: v.k = kv_value::INT; v.i = (int64_t) r.u8(); v.f = (double) v.i; break; // BOOL = 1 byte + case 8: v.k = kv_value::STR; v.s = r.str(); break; + case 9: { + v.k = kv_value::ARR; + const uint32_t et = r.u32(); + const uint64_t n = r.u64(); + for (uint64_t j = 0; j < n; j++) { + switch (et) { + // 元素寬度必須跟 kv 一致（U8/I8/BOOL=1、U16/I16=2 …）。 + // F64 一定要消費 8 bytes，舊版沒讀會讓 header 往後全部錯位。 + case 0: v.arr.push_back(r.u8()); break; + case 1: v.arr.push_back((int8_t) r.u8()); break; + case 2: v.arr.push_back(r.u16()); break; + case 3: v.arr.push_back(r.i16()); break; + case 4: v.arr.push_back((uint32_t) r.u32()); break; + case 5: v.arr.push_back((int32_t) r.u32()); break; + case 6: v.arr.push_back((int64_t) r.f32()); break; + case 7: v.arr.push_back(r.u8()); break; + case 10: v.arr.push_back((int64_t) r.u64()); break; + case 11: v.arr.push_back(r.i64()); break; + case 12: r.f64(); v.arr.push_back(0); break; + case 8: r.skip(r.u64()); v.arr.push_back(0); break; + default: + close(fd); + throw std::runtime_error("GGUF 陣列元素型別不支援: " + std::to_string(et)); + } + } + } break; + case 10: v.k = kv_value::INT; v.i = (int64_t) r.u64(); v.f = (double) v.i; break; + case 11: v.k = kv_value::INT; v.i = r.i64(); v.f = (double) v.i; break; + case 12: v.k = kv_value::REAL; { double d; memcpy(&d, buf.data() + r.o, 8); r.o += 8; v.f = d; v.i = (int64_t) d; } break; + default: close(fd); throw std::runtime_error("GGUF kv 型別不支援: " + std::to_string(t)); + } + kv[key] = v; + } + + std::vector tensors; + tensors.reserve(n_tensors); + for (uint64_t i = 0; i < n_tensors; i++) { + tensor_info ti; + ti.name = r.str(); + const uint32_t nd = r.u32(); + for (uint32_t d = 0; d < nd; d++) { + ti.dims.push_back((int64_t) r.u64()); + } + ti.type = (ggml_type) r.u32(); + ti.offset = r.u64(); + tensors.push_back(std::move(ti)); + } + const uint64_t align = kv.count("general.alignment") ? (uint64_t) kv["general.alignment"].i : 32; + const uint64_t data_start = (r.o + align - 1) / align * align; + + close(fd); + + // ---- sharded GGUF：shard 1 只帶 metadata（n_tensors=0），真正的權重分散在 + // 其他 shard，而且每個 shard 有自己的 data_start。舊版只讀第一個檔， + // 於是 page map 會是空的（症狀：pager 起來但沒有任何 expert 可分頁）。 + // 這裡把每個 shard 的 header 都讀一遍，張量名稱為全域鍵，offset 記成 (shard, off)。 + { + const int64_t n_split = kv.count("split.count") ? kv["split.count"].i : 1; + shard_paths_ = list_splits(path, (int) n_split); + shard_data_start_.assign(shard_paths_.size(), 0); + shard_data_start_[0] = data_start; + + if (n_split > 1) { + std::vector merged; + merged.reserve((size_t) n_tensors + 1024); + for (size_t si = 1; si < shard_paths_.size(); si++) { + const int sfd = open(shard_paths_[si].c_str(), O_RDONLY); + if (sfd < 0) { + throw std::runtime_error("無法開啟 GGUF shard：" + shard_paths_[si]); + } + struct stat sst; + fstat(sfd, &sst); + // 後續 shard 的 header 很小（沒有 tokenizer），讀 16 MiB 綽綽有餘 + const size_t swant = (size_t) std::min(sst.st_size, 16ull << 20); + std::vector sbuf(swant); + ssize_t sgot = 0; + while (sgot < (ssize_t) swant) { + const ssize_t n = pread(sfd, sbuf.data() + sgot, swant - sgot, sgot); + if (n <= 0) { + break; + } + sgot += n; + } + sbuf.resize((size_t) std::max(sgot, 0)); + + gguf_reader sr{sbuf.data(), sbuf.size(), 0}; + if (sbuf.size() < 16 || memcmp(sbuf.data(), "GGUF", 4) != 0) { + close(sfd); + throw std::runtime_error("shard 不是 GGUF：" + shard_paths_[si]); + } + sr.o = 4; + sr.u32(); // version + const uint64_t s_ntensors = sr.u64(); + const uint64_t s_nkv = sr.u64(); + for (uint64_t i = 0; i < s_nkv; i++) { + sr.str(); // key + const uint32_t t = sr.u32(); + switch (t) { + case 0: case 1: case 7: sr.skip(1); break; + case 2: case 3: sr.skip(2); break; + case 4: case 5: case 6: sr.skip(4); break; + case 10: case 11: case 12: sr.skip(8); break; + case 8: sr.skip(sr.u64()); break; + case 9: { + const uint32_t et = sr.u32(); + const uint64_t n = sr.u64(); + if (et == 8) { + for (uint64_t j = 0; j < n; j++) { sr.skip(sr.u64()); } + } else { + size_t esz = 0; + switch (et) { + case 0: case 1: case 7: esz = 1; break; + case 2: case 3: esz = 2; break; + case 4: case 5: case 6: esz = 4; break; + case 10: case 11: case 12: esz = 8; break; + default: esz = 0; break; + } + sr.skip((uint64_t) n * esz); + } + } break; + default: + close(sfd); + throw std::runtime_error("shard kv 型別不支援: " + std::to_string(t)); + } + } + for (uint64_t i = 0; i < s_ntensors; i++) { + tensor_info ti; + ti.name = sr.str(); + const uint32_t nd = sr.u32(); + for (uint32_t d = 0; d < nd; d++) { + ti.dims.push_back((int64_t) sr.u64()); + } + ti.type = (ggml_type) sr.u32(); + ti.offset = sr.u64(); + ti.shard = (int) si; + merged.push_back(std::move(ti)); + } + shard_data_start_[si] = (sr.o + align - 1) / align * align; + close(sfd); + } + for (auto & ti : tensors) { + ti.shard = 0; + merged.push_back(std::move(ti)); + } + tensors.swap(merged); + fprintf(stderr, "[sdq] sharded GGUF：%d 個 shard，合計 %zu 個張量（shard1 為 metadata-only）\n", + (int) shard_paths_.size(), tensors.size()); + } + } + + // ---- hparams + auto gi = [&](const std::string & k, int64_t dflt) -> int64_t { + auto it = kv.find(k); + return it == kv.end() ? dflt : it->second.i; + }; + auto garr = [&](const std::string & k) -> const std::vector * { + auto it = kv.find(k); + return it == kv.end() || it->second.k != kv_value::ARR ? nullptr : &it->second.arr; + }; + info_.arch = kv.count("general.architecture") ? kv["general.architecture"].s : ""; + info_.name = kv.count("general.name") ? kv["general.name"].s : ""; + const std::string pfx = info_.arch + "."; + info_.n_layer = gi(pfx + "block_count", 0); + info_.n_expert = gi(pfx + "expert_count", 0); + info_.n_expert_used= gi(pfx + "expert_used_count", 0); + info_.n_embd = gi(pfx + "embedding_length", 0); + info_.n_ff_exp = gi(pfx + "expert_feed_forward_length", 0); + info_.n_ff_shexp = gi(pfx + "expert_shared_feed_forward_length", 0); + info_.n_head = gi(pfx + "attention.head_count", 0); + info_.n_head_kv = gi(pfx + "attention.head_count_kv", 0); + info_.n_embd_head_k= gi(pfx + "attention.key_length", 0); + info_.n_embd_head_v= gi(pfx + "attention.value_length", info_.n_embd_head_k); + info_.context_length = gi(pfx + "context_length", 0); + info_.full_attention_interval = gi(pfx + "full_attention_interval", 4); + // bailingmoe3（Ling-3.0-flash）：前 n_layer_dense_lead 層是 dense FFN， + // 它們沒有 ffn_*_exps 張量，必須讓 page map 保持「空」—— + // sdq_pager_active_for_layer() 看到 layer() == nullptr 就會走原生路徑。 + info_.n_layer_dense_lead = gi(pfx + "leading_dense_block_count", 0); + info_.n_expert_groups = gi(pfx + "expert_group_count", 0); + info_.n_expert_group_used = gi(pfx + "expert_group_used_count", 0); + info_.moe_latent_size = gi(pfx + "moe_latent_size", 0); + + // Nemotron-H / Flint：expert_count / expert_used_count may be scalar or + // per-layer arrays. Expand either representation to a per-layer table. + const auto * arr_used = garr(pfx + "expert_used_count"); + const auto * arr_exp = garr(pfx + "expert_count"); + info_.n_expert_layer.assign((size_t) info_.n_layer, info_.n_expert); + info_.n_expert_used_layer.assign((size_t) info_.n_layer, info_.n_expert_used); + if (arr_exp && (int64_t) arr_exp->size() >= info_.n_layer) { + for (int64_t i = 0; i < info_.n_layer; i++) { + info_.n_expert_layer[(size_t) i] = (*arr_exp)[(size_t) i]; + } + } + if (arr_used && (int64_t) arr_used->size() >= info_.n_layer) { + for (int64_t i = 0; i < info_.n_layer; i++) { + info_.n_expert_used_layer[(size_t) i] = (*arr_used)[(size_t) i]; + } + // 空層（dense / Mamba）是 0，取 MoE 層的最大值當代表 + int64_t mx = 0; + for (int64_t v : *arr_used) { + mx = std::max(mx, v); + } + info_.n_expert_used = mx; + } + + if (info_.n_expert <= 0 || info_.n_layer <= 0) { + throw std::runtime_error("這不是 MoE 模型（expert_count=" + std::to_string(info_.n_expert) + "）"); + } + + // ---- 頁面表：每個 (layer, expert) 的三段 + layers_.assign((size_t) info_.n_layer, {}); + std::map by_name; + for (const auto & ti : tensors) { + by_name[ti.name] = &ti; + } + + size_t max_expert_bytes = 0; + int n_moe_layers = 0; + for (int64_t il = 0; il < info_.n_layer; il++) { + const std::string b = "blk." + std::to_string(il) + "."; + layer_desc & L = layers_[(size_t) il]; + + auto get = [&](const std::string & suffix) -> const tensor_info * { + auto it = by_name.find(b + suffix); + return it == by_name.end() ? nullptr : it->second; + }; + + const tensor_info * t_gate = get("ffn_gate_exps.weight"); + const tensor_info * t_up = get("ffn_up_exps.weight"); + const tensor_info * t_gu = get("ffn_gate_up_exps.weight"); + const tensor_info * t_down = get("ffn_down_exps.weight"); + + if (!t_down || (t_gate && !t_up && !t_gu)) { + // 這層不是 MoE（bailingmoe3 的 leading dense layers 就是這樣）。 + // 保持 experts 為空 → layer() 回 nullptr → 這層走原生 mul_mat_id。 + continue; + } + // Expert counts and widths may vary by layer; use the tensor shape as + // the source of truth rather than relying on scalar metadata. + const int64_t n_expert_l = t_down->dims.size() >= 3 ? t_down->dims[2] : info_.n_expert; + L.experts.assign((size_t) n_expert_l, {}); + L.n_expert = n_expert_l; + L.no_gate = (t_gate == nullptr && t_gu == nullptr); + n_moe_layers++; + + // sharded GGUF：每段的絕對偏移要用「該張量自己 shard 的 data_start」， + // 不能用第一個 shard 的（兩者差了近 10 MB，直接用會讀到別的權重）。 + auto abs_off = [&](const tensor_info * t, uint64_t rel) -> uint64_t { + const size_t s = (size_t) t->shard; + const uint64_t base = (s < shard_data_start_.size()) ? shard_data_start_[s] : 0; + return base + rel; + }; + + for (int64_t e = 0; e < n_expert_l; e++) { + expert_desc & ed = L.experts[(size_t) e]; + if (t_gate || t_gu) { + const tensor_info * t = t_gate ? t_gate : t_gu; + const uint64_t nb = (uint64_t) ggml_row_size(t->type, t->dims[0]) * (uint64_t) t->dims[1]; + ed.seg[part::GATE].off = abs_off(t, t->offset + (uint64_t) e * nb); + ed.seg[part::GATE].size = nb; + ed.seg[part::GATE].shard = t->shard; + ed.ne0 = t->dims[0]; + ed.ne1 = t->dims[1]; + ed.type[part::GATE] = t->type; + } + if (t_up) { + const uint64_t nb = (uint64_t) ggml_row_size(t_up->type, t_up->dims[0]) * (uint64_t) t_up->dims[1]; + ed.seg[part::UP].off = abs_off(t_up, t_up->offset + (uint64_t) e * nb); + ed.seg[part::UP].size = nb; + ed.seg[part::UP].shard = t_up->shard; + ed.type[part::UP] = t_up->type; // gate/up 都是 Q4_K，型別一定要各自填 + if (!t_gate) { + ed.ne0 = t_up->dims[0]; + ed.ne1 = t_up->dims[1]; + // no-gate 模型（Nemotron LatentMoE）不要把 GATE 當成 up， + // expert_forward_relu2 走 UP/DOWN 兩段。 + } + } + { + const uint64_t nb = (uint64_t) ggml_row_size(t_down->type, t_down->dims[0]) * (uint64_t) t_down->dims[1]; + ed.seg[part::DOWN].off = abs_off(t_down, t_down->offset + (uint64_t) e * nb); + ed.seg[part::DOWN].size = nb; + ed.seg[part::DOWN].shard = t_down->shard; + ed.type[part::DOWN] = t_down->type; + ed.down_ne0 = t_down->dims[0]; + ed.down_ne1 = t_down->dims[1]; + } + ed.fused_gate_up = (t_gu != nullptr); + ed.bytes = 0; + int64_t acc = 0; + for (int p = 0; p < 3; p++) { + if (ed.seg[p].size == 0) { + continue; + } + acc = (acc + 63) & ~(int64_t) 63; // 每段 64B 對齊 + ed.slot_off[p] = acc; + acc += (int64_t) ed.seg[p].size; + } + ed.bytes = acc; + max_expert_bytes = std::max(max_expert_bytes, (size_t) ed.bytes); + } + L.bytes_per_expert = L.experts.empty() ? 0 : L.experts[0].bytes; + } + + slot_bytes_ = (max_expert_bytes + 4095) & ~(size_t) 4095; + if (slot_bytes_ < (256ull << 10)) { + slot_bytes_ = 256ull << 10; + } + + // ---- 大小兩種槽位 ---------------------------------------------------- + // 這個模型有 37 層的 expert 是 1900544 B（Q4_K + Q5_K），但層 34/38/39 是 + // 2039808 B（down 換成 Q6_K）。若所有槽位都照最大的配，arena 有 6.8% 被浪費 + // ——而 arena 就是 8 GB 預算裡最寶貴的資源（實測 +5% 槽位 ≈ +9% tok/s）。 + // + // 解法：槽位分成兩區。「小槽位」給 37 層用，「大槽位」只給那 3 層用， + // 大槽位的數量按「那 3 層佔請求的比例」分配（3/40 = 7.5%）。 + // 選「大多數層」當小的：用眾數而不是最大值，這樣即使某個模型只有一種大小， + // 也不會誤判。 + { + std::map hist; + for (const auto & L : layers_) { + if (L.experts.empty()) { + continue; // dense layer：沒有 expert，不參與槽位尺寸統計 + } + hist[(size_t) L.bytes_per_expert]++; + } + size_t small = 0; + int best_n = -1; + for (const auto & kv : hist) { + if (kv.second > best_n) { + best_n = kv.second; + small = kv.first; + } + } + size_t big = 0; + for (const auto & L : layers_) { + if (!L.experts.empty() && (size_t) L.bytes_per_expert > small) { + big = std::max(big, (size_t) L.bytes_per_expert); + } + } + small_bytes_ = (small + 4095) & ~(size_t) 4095; + big_bytes_ = big ? ((big + 4095) & ~(size_t) 4095) : 0; + if (big_bytes_ && big_bytes_ <= small_bytes_) { + big_bytes_ = 0; // 其實一樣大，別分兩區 + } + slot_bytes_ = big_bytes_ ? big_bytes_ : small_bytes_; + int n_big_layers = 0; + for (auto & L : layers_) { + L.needs_big_slot = big_bytes_ && !L.experts.empty() && + (size_t) L.bytes_per_expert > small_bytes_; + if (L.needs_big_slot) { + n_big_layers++; + } + } + // Nemotron has only four expert sizes; uniform slots simplify arena + // management. Flint layers also use a uniform slot size per layer. + small_bytes_ = slot_bytes_; + big_bytes_ = 0; + for (auto & L : layers_) { + L.needs_big_slot = false; + } + // 統計給 README / 實驗報告用：expert 權重在檔案裡到底佔多少 + uint64_t expert_bytes = 0; + uint64_t n_experts_total = 0; + for (const auto & L : layers_) { + expert_bytes += (uint64_t) L.bytes_per_expert * (uint64_t) L.experts.size(); + n_experts_total += (uint64_t) L.experts.size(); + } + n_moe_layers_ = n_moe_layers; + expert_bytes_total_ = expert_bytes; + info_.no_gate = true; + for (const auto & L : layers_) { + if (!L.experts.empty() && !L.no_gate) { + info_.no_gate = false; + } + } + fprintf(stderr, + "[sdq] MoE 層 %d / 共 %lld 層（dense lead %lld）" + "：expert 權重合計 %.2f GiB、%llu 個 expert\n", + n_moe_layers, (long long) info_.n_layer, + (long long) info_.n_layer_dense_lead, + expert_bytes / 1073741824.0, (unsigned long long) n_experts_total); + fprintf(stderr, + "[sdq] 統一槽位 %zu B（最大 expert 大小；平均 %.0f B）\n", + small_bytes_, n_experts_total ? (double) expert_bytes / n_experts_total : 0.0); + } +} + +bool pager::need_big_slot(int il, int ie) const { + if (!big_bytes_) { + return false; + } + const layer_desc * L = layer(il); + if (!L) { + return false; + } + if (L->needs_big_slot) { + return true; + } + return ie >= 0 && ie < (int) L->experts.size() && + (size_t) L->experts[(size_t) ie].bytes > small_bytes_; +} + +void pager::alloc_arena() { + const size_t slot = cfg_.slot_bytes ? cfg_.slot_bytes : slot_bytes_; + slot_bytes_ = slot; + + // arena 大小決定：RAM 是「總預算」，不是「還有多少可用」 + // 總預算 = 非專家權重 + ggml 模型緩衝 + KV cache + 計算緩衝 + 詞表 + arena + // 所以有指定 ram_budget 時，用「目前 RSS（不含 arena）」當作已用額度扣掉。 + // + // 重要：pager 在 llama.cpp 載入權重**之前**就做 sizing，此時 rss_now 只有 + // 幾十 MB（實測 0.09 GiB），所以後續所有匿名配置都不在 rss_now 裡， + // 必須由 reserve 手動涵蓋。這些事實上是三塊： + // + // 1. ggml 的模型緩衝 —— 主機有 AMX 時（-march=native 會定義 __AMX_INT8__ + // 使 llama.cpp 啟用 AMX repack buffer）Ling-3.0-flash Q4_K_M 要 + // 1876.34 MiB；關掉 AMX（GGML_NATIVE=OFF）則不需要這塊。 + // 2. CPU compute buffer：83.5 MiB @ubatch=128，ubatch 越大越多。 + // 3. KV/SSM state + 執行緒堆疊：ctx 越大越多。 + // + // 舊的預設值 900 MB 只涵蓋 (2)+(3)，漏掉 (1)，在 AMX 主機上會讓實際 + // 匿名用量超出預算約 1.9 GiB —— 嚴格 RLIMIT_DATA 下直接 OOM。 + // 實測（本機，ctx=2048、ubatch=128、AMX build）： + // reserve 1024 MB → 預算 8 GiB 實測匿名峰值 10.96 GiB（超出，OOM） + // reserve 3600 MB → 預算 8 GiB 通過，arena 4.40 GiB + // reserve 2700 MB → 預算 4 GiB 通過，arena 1.27 GiB（無 AMX build） + // 寧可少一些槽位，也不要在目標機器的 8 GiB 上 OOM —— 這是本專案的硬約束。 + size_t reserve = (size_t) env_size("SDQ_RESERVE_MB", 3600) << 20; + const size_t rss_now = process_rss_bytes(); + + size_t arena; + if (cfg_.ram_budget_bytes) { + const size_t used = rss_now; // 這時候 arena 還沒配置 + size_t avail = cfg_.ram_budget_bytes > used + reserve + ? cfg_.ram_budget_bytes - used - reserve + : 0; + arena = avail; + if (cfg_.arena_bytes) { + arena = std::min(arena, cfg_.arena_bytes); + } + fprintf(stderr, "[sdq] RAM 預算 %.2f GiB：已用(RSS) %.2f GiB、預留 %.2f GiB → arena %.2f GiB\n", + cfg_.ram_budget_bytes / 1073741824.0, used / 1073741824.0, + reserve / 1073741824.0, arena / 1073741824.0); + } else if (cfg_.arena_bytes) { + arena = cfg_.arena_bytes; // 明確指定 SDQ_ARENA_MB + } else { + // 自動模式（未設 SDQ_ARENA_MB、未設 SDQ_RAM_BUDGET_MB）： + // arena = 可用記憶體 × 0.8，上限為「全部 expert 的槽位」（再大也用不到， + // 而且上限到達時 preload 會自動啟用）。 + size_t total_experts = 0; + for (const auto & L : layers_) { + total_experts += L.experts.size(); + } + const size_t cap = total_experts * slot_bytes_; + const size_t avail = available_ram_bytes(); + arena = std::min((size_t) (avail * 0.80), cap); + fprintf(stderr, "[sdq] arena 自動：可用記憶體 %.2f GiB × 0.8 → %.2f GiB（全部 expert 上限 %.2f GiB）\n", + avail / 1073741824.0, arena / 1073741824.0, cap / 1073741824.0); + } + if (arena < slot_bytes_ * 16) { + throw std::runtime_error("expert arena 太小"); + } + // 切成兩區：大槽位只分給「需要大槽位的層」，比例 = 那些層佔全部層的比例。 + const bool force_uniform = env_flag("SDQ_FORCE_UNIFORM_SLOT", false); + if (big_bytes_ && small_bytes_ && !layers_.empty() && !force_uniform) { + int n_big_layers = 0; + for (const auto & L : layers_) { + if (L.needs_big_slot) { + n_big_layers++; + } + } + const double denom = (double) (n_moe_layers_ ? n_moe_layers_ : layers_.size()); + const double f = (double) n_big_layers / denom; + const double avg = f * (double) big_bytes_ + (1.0 - f) * (double) small_bytes_; + size_t n_total = (size_t)((double) arena / avg); + size_t nb = (size_t)(n_total * f); + // 大槽位的下限不是「有點就好」，而是**工作集的下限**： + // 一個 token 在每個大槽位層要 8 個大 expert，這個模型有 3 個大槽位層 + // → 至少要 3 × 8 = 24 個大槽位，而且還要留給預取的空間。 + // 只給 16 個時量到過「找不到能放 il=34 ie=114 的槽位」而中止 + // （這是在 verify 新增的預取壓力測試 [2b] 抓到的）。 + const size_t big_need = (size_t) (n_big_layers * info_.n_expert_used); + if (nb < big_need) { + nb = big_need; + } + if (nb > n_total / 4) { + nb = n_total / 4; + } + // 連工作集下限都放不進去時，寧可讓小槽位少一點（n_total 保持不變）， + // 否則會算出一個「必定會中止」的配置。 + n_big_ = nb; + n_small_ = n_total > nb ? n_total - nb : 0; + arena = n_small_ * small_bytes_ + n_big_ * big_bytes_; + } else { + // 單一尺寸：全部槽位都是「大的那一種」（slot_bytes_）。 + // 這是 SDQ_FORCE_UNIFORM_SLOT=1 與「模型只有一種 expert 大小」時的路徑。 + // + // 這裡必須記成 n_small_=0 / n_big_=全部，不能記成「n_small_=全部、單位用大的」： + // slot_ptr()/slot_size_of() 是用「索引 < n_small_ → 小槽位」來分支的， + // 把它記成全部都是小槽位的話，大 expert 會拿到小槽位的 stride + // （寫出槽位邊界、算錯），而且大 expert 永遠配不到槽位 → 推論卡死。 + // （第一版就是這樣，實驗跑到一半 hang 才找出來。） + const size_t unit = slot_bytes_ ? slot_bytes_ : small_bytes_; + n_small_ = 0; + // slot_ptr() 對 s >= n_small_ 用 big_bytes_ 當 stride，所以這裡一定要把 + // big_bytes_ 設成 unit；設成 0 會讓所有槽位都指到 arena 開頭（數值全錯）。 + big_bytes_ = unit; + n_big_ = arena / unit; + arena = n_big_ * unit; + } + if (n_small_ == 0 && n_big_ == 0) { + throw std::runtime_error("expert arena 太小（算不出槽位）"); + } + n_slots_ = n_small_ + n_big_; + if (n_slots_ < 16) { + throw std::runtime_error("expert arena 太小（算不出 16 個槽位）"); + } + + void * p = nullptr; + const size_t alloc = arena + (2ull << 20); + if (posix_memalign(&p, 2ull << 20, alloc) != 0 || !p) { + throw std::runtime_error("無法配置 expert arena（" + std::to_string(arena >> 20) + " MB）"); + } + // 讓 arena 全部變成實體 RAM，並提示核心使用 huge page（減少 TLB 壓力） + memset(p, 0, alloc); + madvise(p, alloc, MADV_HUGEPAGE); + arena_ = (uint8_t *) p; + arena_bytes_ = arena; + + // OPT-1：剖面計數表（n_layer × n_expert，40×256 = 10240 筆 ≈ 40 KB，可忽略） + prof_.assign((size_t) info_.n_layer * (size_t) info_.n_expert, 0); + ensure_layer_stats((int) info_.n_layer); // SDQ-0004 + if (!cfg_.profile_in.empty()) { + load_profile(cfg_.profile_in); + } + profile_out_ = cfg_.profile_out; + + slots_.assign(n_slots_, {}); + free_slots_.reserve(n_small_); + for (size_t i = n_small_; i-- > 0;) { + free_slots_.push_back((int32_t) i); + } + free_big_slots_.reserve(n_big_); + for (size_t i = n_slots_; i-- > n_small_;) { + free_big_slots_.push_back((int32_t) i); + } + used_.store(0); + + hist_.assign((size_t) info_.n_layer, {}); + for (auto & per_layer : hist_) { + per_layer.assign((size_t) info_.n_expert, {}); + } + prev_used_.assign((size_t) info_.n_layer, {}); + cur_used_.assign((size_t) info_.n_layer, {}); +} + +void pager::init(const config & cfg, ggml_context * ctx) { + (void) ctx; + if (ready_) { + return; + } + cfg_ = cfg; + if (cfg_.model_path.empty()) { + const char * v = getenv("SDQ_MODEL_PATH"); + cfg_.model_path = v ? v : ""; + } + if (cfg_.model_path.empty()) { + throw std::runtime_error("SDQ 缺少模型路徑（設定 SDQ_MODEL_PATH）"); + } + + // 先讀 header（順便列舉 shard），再開 fd：shard 清單只有在解析 + // split.count 之後才知道，而 open_model_fd() 需要它。 + load_gguf_header(cfg_.model_path); + open_model_fd(); + if (shard_fd_.empty() || shard_fd_[0] < 0) { + throw std::runtime_error("無法開啟模型檔：" + shard_path(0)); + } + // 順序有意義：I/O 池的對齊緩衝必須在 alloc_arena() 之前配置， + // 否則 arena 會把它當成「不存在的 RAM」，實際上就偷用了 8 GB 預算。 + io_pool_start(); + alloc_arena(); + + if (cfg_.prefetch) { + pf_ = std::thread([this] { pf_worker(); }); + } + ready_ = true; + print_summary(); +} + +const layer_desc * pager::layer(int il) const { + if (il < 0 || il >= (int) layers_.size()) { + return nullptr; + } + const layer_desc & L = layers_[(size_t) il]; + if (L.experts.empty()) { + return nullptr; + } + return &L; +} + +void pager::open_model_fd() { + if (shard_paths_.empty()) { + shard_paths_.push_back(cfg_.model_path); + } + shard_fd_.assign(shard_paths_.size(), -1); + shard_fd_direct_.assign(shard_paths_.size(), -1); + for (size_t i = 0; i < shard_paths_.size(); i++) { + shard_fd_[i] = open(shard_paths_[i].c_str(), O_RDONLY); + if (shard_fd_[i] < 0) { + fprintf(stderr, "[sdq] 無法開啟 shard %zu: %s\n", i, shard_paths_[i].c_str()); + return; + } + } + // O_DIRECT：讀資料完全不進 kernel page cache。 + // 這是「RAM 帳目誠實」的關鍵 —— 否則核心預讀會讓 mmap 檔案的頁算進 RSS + // （實測 4K prefill 時檔案對映 RSS 可以多到 6 GB 以上），而且量到的 + // 「SSD 讀取量」會被上一輪的 page cache 灌水。 + const char * no_direct = getenv("SDQ_NO_DIRECT_IO"); + if (no_direct && *no_direct && *no_direct != '0') { + return; + } +#ifdef O_DIRECT + // 每個 shard 各自開一個 O_DIRECT fd：O_DIRECT 是 fd 層級的屬性，不能共用。 + for (size_t i = 0; i < shard_paths_.size(); i++) { + shard_fd_direct_[i] = open(shard_paths_[i].c_str(), O_RDONLY | O_DIRECT); + if (shard_fd_direct_[i] < 0) { + break; + } + } + use_direct_ = shard_fd_direct_[0] >= 0; + if (use_direct_) { + fprintf(stderr, "[sdq] 使用 O_DIRECT 讀取模型檔（%zu 個 shard，不佔用 kernel page cache）\n", + shard_paths_.size()); + } else { + fprintf(stderr, "[sdq] O_DIRECT 不可用，改用 buffered + fadvise(DONTNEED)\n"); + } +#endif +} + +// 有對齊的緩衝區（O_DIRECT 要求），thread_local 避免競爭 +static thread_local uint8_t * g_direct_buf = nullptr; +static thread_local size_t g_direct_sz = 0; + +// 與 read_range_direct 相同，但用呼叫端提供的緩衝區（已配置好、已算進 RSS）。 +// I/O 執行緒池走這條：每條 worker 在啟動時就配置好對齊緩衝，於是 arena sizing 時 +// 這筆記憶體已經在 RSS 裡，不會變成「額外偷用的 RAM」。 +// sharded GGUF：把「哪個 shard」一路帶到 pread。 +static inline int sdq_fd_for(const std::vector & fds, int shard) { + if (shard < 0 || shard >= (int) fds.size()) { + return -1; + } + return fds[(size_t) shard]; +} + +bool pager::read_range_staged(int shard, uint64_t off, uint8_t * dst, size_t size, + uint8_t * stage, size_t stage_sz) { + if (!use_direct_) { + return read_range(shard, off, dst, size); + } + const int fd = sdq_fd_for(shard_fd_direct_, shard); + if (fd < 0) { + return false; + } + const size_t blk = 4096; + const uint64_t start = off & ~(uint64_t)(blk - 1); + const uint64_t end = (off + size + blk - 1) & ~(uint64_t)(blk - 1); + const size_t bytes = (size_t)(end - start); + if (stage_sz < bytes) { + // 預先配置不足（不該發生）：退回 thread_local 的通用路徑 + return read_range_direct(shard, off, dst, size); + } + size_t done = 0; + while (done < bytes) { + const ssize_t n = pread(fd, stage + done, bytes - done, + (off_t)(start + done)); + if (n <= 0) { + return false; + } + done += (size_t) n; + } + memcpy(dst, stage + (off - start), size); + return true; +} + +bool pager::read_range_direct(int shard, uint64_t off, uint8_t * dst, size_t size) { + if (!use_direct_) { + return read_range(shard, off, dst, size); + } + const int fd = sdq_fd_for(shard_fd_direct_, shard); + if (fd < 0) { + return false; + } + const size_t blk = 4096; + const uint64_t start = off & ~(uint64_t)(blk - 1); + const uint64_t end = (off + size + blk - 1) & ~(uint64_t)(blk - 1); + const size_t bytes = (size_t)(end - start); + if (g_direct_sz < bytes) { + if (g_direct_buf) { + free(g_direct_buf); + } + if (posix_memalign((void **) &g_direct_buf, blk, bytes) != 0) { + g_direct_buf = nullptr; + g_direct_sz = 0; + use_direct_ = false; + return read_range(shard, off, dst, size); + } + g_direct_sz = bytes; + } + size_t done = 0; + while (done < bytes) { + const ssize_t n = pread(fd, g_direct_buf + done, bytes - done, + (off_t)(start + done)); + if (n <= 0) { + return false; + } + done += (size_t) n; + } + memcpy(dst, g_direct_buf + (off - start), size); + return true; +} + +bool pager::read_range(int shard, uint64_t off, uint8_t * dst, size_t size) { + const int fd = sdq_fd_for(shard_fd_, shard); + if (fd < 0) { + return false; + } + size_t done = 0; + while (done < size) { + const ssize_t n = pread(fd, dst + done, size - done, (off_t) (off + done)); + if (n <= 0) { + return false; + } + done += (size_t) n; + } + // 讓 RAM 帳目誠實：把 kernel 對這個檔案的 page cache 丟掉，避免「藏在核心裡的 RAM」 + if (cfg_.fadvise_dontneed) { + posix_fadvise(fd, (off_t) off, (off_t) size, POSIX_FADV_DONTNEED); + } + return true; +} + +bool pager::read_expert(int il, int ie, uint8_t * dst) { + const layer_desc * L = layer(il); + if (!L || ie < 0 || ie >= (int) L->experts.size()) { + return false; + } + const expert_desc & ed = L->experts[(size_t) ie]; + for (int p = 0; p < 3; p++) { + if (ed.seg[p].size == 0) { + continue; + } + if (!read_range_direct(ed.seg[p].shard, ed.seg[p].off, + dst + ed.slot_off[p], (size_t) ed.seg[p].size)) { + return false; + } + } + return true; +} + +// ---------------------------------------------------------------- I/O 執行緒池 +// +// 問題：一個 expert = gate / up / down 三段，檔案上是三段不相鄰的資料。 +// 循序讀就是 3 次 pread 排在同一條執行緒上，而單次 1.9 MB 的 pread 在本機量到 +// ~7 ms（≈275 MB/s 的單流頻寬）。因此「一個 expert 要多久」= 3 × 單次延遲， +// 這條路徑的長度跟「有多少條執行緒」無關 —— mode C 再怎麼排程都救不了。 +// +// 做法：把三段交給三條執行緒同時 pread，臨界路徑降到約「單段」的時間。 +// 量測（8 個 expert 全部 miss 的一層，各 9 次取中位數）： +// 8 執行緒 × 三段序列讀 : 12.60 ms +// 16 執行緒 × 三段平行讀 : 7.57 ms +// 24 執行緒 × 三段平行讀 : 7.58 ms ← 16 就飽和，所以預設 3 × n_threads +// +// 這個池子對 prefill 更有用：一層要 256 個 expert = 768 個讀取工作， +// 序列讀時 16 執行緒只能拿到 16 路的頻寬 utilization。 +void pager::io_pool_start() { + if (!cfg_.io_parallel) { + fprintf(stderr, "[sdq] I/O 段平行：關（每個 expert 的 gate/up/down 循序讀）\n"); + return; + } + int n = cfg_.io_threads > 0 ? cfg_.io_threads : cfg_.n_threads * 3; + if (n < 4) { + n = 4; + } +if (n > 48) { + n = 48; + } + // 先算「一段權重最大有多少位元組」，每條 worker 配一塊對齊緩衝。 + // 這必须在 alloc_arena() 之前做完，否則這筆 RAM 會變成帳外支出。 + size_t max_seg = 0; + for (const auto & L : layers_) { + for (const auto & ed : L.experts) { + for (int p = 0; p < 3; p++) { + max_seg = std::max(max_seg, (size_t) ed.seg[p].size); + } + } + } + const size_t stage_bytes = ((max_seg + 2 * 4096) + 4095) & ~(size_t) 4095; + io_stage_.assign((size_t) n, io_stage{}); + for (int i = 0; i < n; i++) { + void * p = nullptr; + if (posix_memalign(&p, 4096, stage_bytes) == 0 && p) { + memset(p, 0, stage_bytes); // 實際碰頁，讓它出現在 RSS 裡 + io_stage_[(size_t) i].buf = (uint8_t *) p; + io_stage_[(size_t) i].sz = stage_bytes; + } + } + io_split_ = cfg_.io_split > 0 ? cfg_.io_split : 1; + io_stop_ = false; + io_pool_.reserve((size_t) n); + for (int i = 0; i < n; i++) { + const size_t wid = (size_t) i; + try { + io_pool_.emplace_back([this, wid] { io_worker(wid); }); + } catch (...) { + break; + } + } + io_threads_ = (int) io_pool_.size(); + io_pool_on_ = io_threads_ > 0; + if (io_pool_on_) { + fprintf(stderr, "[sdq] I/O 段平行：%d 條讀取執行緒 × %zu KB 對齊緩衝" + "（gate/up/down 同時讀，每段再切 %d 塊）\n", + io_threads_, stage_bytes >> 10, io_split_); + } +} + +void pager::io_pool_stop() { + if (io_pool_.empty()) { + for (auto & s : io_stage_) { + free(s.buf); + s.buf = nullptr; + } + io_stage_.clear(); + return; + } + { + std::lock_guard lk(io_mu_); + io_stop_ = true; + } + io_cv_.notify_all(); + for (auto & t : io_pool_) { + if (t.joinable()) { + t.join(); + } + } + io_pool_.clear(); + io_pool_on_ = false; + for (auto & s : io_stage_) { + free(s.buf); + s.buf = nullptr; + s.sz = 0; + } + io_stage_.clear(); +} + +void pager::io_worker(size_t wid) { + for (;;) { + io_job job{}; + { + std::unique_lock lk(io_mu_); + io_cv_.wait(lk, [this] { return io_stop_ || !io_q_.empty(); }); + if (io_q_.empty()) { + if (io_stop_) { + return; + } + continue; + } + job = io_q_.front(); + io_q_.pop_front(); + } + io_jobs_.fetch_add(1, std::memory_order_relaxed); + // 讀取本身不持鎖：io_mu_ 只保護佇列與完成計數 + const bool ok = (wid < io_stage_.size() && io_stage_[wid].buf) + ? read_range_staged(job.shard, job.off, job.dst, job.size, + io_stage_[wid].buf, io_stage_[wid].sz) + : read_range_direct(job.shard, job.off, job.dst, job.size); + { + std::lock_guard lk(io_mu_); + if (!ok) { + job.ctx->ok = false; + } + if (job.ctx->remaining.fetch_sub(1) == 1) { + io_done_cv_.notify_all(); + } + } + } +} + +// 把 (il,ie) 的三段交給執行緒池並行讀，呼叫端等到整批完成才返回。 +// ctx 放在呼叫端 stack：worker 只會在 remaining 歸零之後才再碰它， +// 而我們正是等到 remaining==0 才離開函式，所以沒有 use-after-return。 +bool pager::read_expert_par(int il, int ie, uint8_t * dst) { + const layer_desc * L = layer(il); + if (!L || ie < 0 || ie >= (int) L->experts.size()) { + return false; + } + if (!io_pool_on_) { + return read_expert(il, ie, dst); + } + const expert_desc & ed = L->experts[(size_t) ie]; + io_ctx ctx; + int n = 0; + // 把每一段再切成一樣大的好幾塊，交給不同的執行緒同時讀。 + // + // 為什麼：單一 pread 的時間 ≈ 固定延遲 + 位元組數 / 單流頻寬。 + // 本機實測單流約 275 MB/s、4K 隨機讀延遲約 0.57 ms，所以 + // 590 KB 讀一次 ≈ 0.57 + 2.1 = 2.7 ms + // 590 KB 切成 2 塊 ≈ 0.57 + 1.05 = 1.6 ms + // 而「讀完一個 expert」的時間正是每層的關鍵路徑，切塊就是直接砍關鍵路徑。 + // 設定檔 SDQ_IO_SPLIT。**實測預設是 1（不切塊）**：切塊在 24 條讀取執行緒 + // 之下沒有可重現的改善（split=1..6 的中位數在 4.9~5.9 之間跳動，純噪音）， + // 因為 8 個 expert × 3 段 = 24 個並行請求已經把裝置的佇列灌飽， + // 這時瓶頸是裝置總頻寬而不是單一請求的延遲。保留這個開關是為了其他裝置。 + const int split = io_split_; + const size_t blk = 4096; + { + std::lock_guard lk(io_mu_); + for (int p = 0; p < 3; p++) { + const size_t sz = (size_t) ed.seg[p].size; + if (sz == 0) { + continue; + } + uint8_t * base = dst + ed.slot_off[p]; + const int sh = ed.seg[p].shard; + if (split <= 1 || sz <= 2 * blk) { + io_q_.push_back(io_job{sh, ed.seg[p].off, base, sz, &ctx}); + n++; + continue; + } + size_t chunk = ((sz + (size_t) split - 1) / (size_t) split + blk - 1) & ~(blk - 1); + for (size_t off = 0; off < sz; off += chunk) { + const size_t take = sz - off < chunk ? sz - off : chunk; + io_q_.push_back(io_job{sh, ed.seg[p].off + off, base + off, take, &ctx}); + n++; + } + } + ctx.remaining.store(n); + } + if (n == 0) { + return true; + } + io_cv_.notify_all(); + const uint64_t t0 = now_us(); + { + std::unique_lock lk(io_mu_); + io_done_cv_.wait(lk, [&ctx] { return ctx.remaining.load() == 0; }); + } + io_wait_ns_.fetch_add((now_us() - t0) * 1000ull, std::memory_order_relaxed); + return ctx.ok; +} + +uint64_t pager::read_expert_span(int il, const int32_t * experts, int n, const int32_t * slots) { + const layer_desc * L = layer(il); + if (!L || n <= 0) { + return 0; + } + const uint64_t t_start = now_us(); + // 每個執行緒一份暫存區（多執行緒會同時呼叫，不能共用一個 buffer） + static thread_local std::vector scratch_buf; + uint64_t total = 0; + + // 把 experts 切成「檔案編號連續」的區段 + int i = 0; + while (i < n) { + int j = i + 1; + while (j < n && experts[j] == experts[j - 1] + 1) { + j++; + } + const int cnt = j - i; + const expert_desc & first = L->experts[(size_t) experts[i]]; + const expert_desc & last = L->experts[(size_t) experts[j - 1]]; + + for (int p = 0; p < 3; p++) { + if (first.seg[p].size == 0 || last.seg[p].size == 0) { + continue; + } + // 連續區段必須在同一個 shard 內；跨 shard 就切斷這個區段。 + if (last.seg[p].shard != first.seg[p].shard) { + continue; + } + const uint64_t off = first.seg[p].off; + const uint64_t end = last.seg[p].off + last.seg[p].size; + const size_t bytes = (size_t) (end - off); + if (scratch_buf.size() < bytes) { + scratch_buf.assign(bytes, 0); + } + if (!read_range_direct(first.seg[p].shard, off, scratch_buf.data(), bytes)) { + return total; + } + total += bytes; + // 散到各自的槽位：scratch 從 off 開始，所以偏移是「相對於 off」 + for (int k = i; k < j; k++) { + const expert_desc & ed = L->experts[(size_t) experts[k]]; + const size_t seg_bytes = (size_t) ed.seg[p].size; + const size_t delta = (size_t) (ed.seg[p].off - off); + if (slots[k] < 0) { + continue; + } + if ((size_t) slots[k] >= n_slots_ || (uint64_t) ed.slot_off[p] + seg_bytes > slot_bytes_ + || delta + seg_bytes > scratch_buf.size()) { + fprintf(stderr, "[sdq-pager] span 越界：il=%d slot=%d/%zu slot_off=%lld seg=%zu delta=%zu scratch=%zu p=%d e=%d\n", + il, slots[k], n_slots_, (long long) ed.slot_off[p], seg_bytes, delta, scratch_buf.size(), p, experts[k]); + return total; + } + memcpy(slot_ptr(slots[k]) + ed.slot_off[p], + scratch_buf.data() + delta, seg_bytes); + } + } + i = j; + } + stats_.demand_wait_us.fetch_add(now_us() - t_start, std::memory_order_relaxed); + return total; +} + +int pager::reserve_slot(int il, int ie) { + if (!ready_) { + return -1; + } + count_request(il, ie); + std::lock_guard lk(mu_); + const int32_t exist = find_slot_locked(il, ie); + if (exist >= 0) { + slots_[(size_t) exist].last_use = ++clock_; + slots_[(size_t) exist].pin++; + touch_freq_locked(slots_[(size_t) exist]); + return exist; + } + const int32_t s = victim_locked(); + if (s < 0) { + return -1; + } + slot_state & st = slots_[(size_t) s]; + if (st.layer >= 0) { + page_.erase(((int64_t) st.layer << 20) | (int64_t) st.expert); + stats_.evictions.fetch_add(1, std::memory_order_relaxed); + used_.fetch_sub(1, std::memory_order_relaxed); + } + st.layer = il; + st.expert = ie; + st.last_use = ++clock_; + st.freq = 1; + st.pin = 1; + st.is_hot = hot_lookup(il, ie); + page_[((int64_t) il << 20) | (int64_t) ie] = s; + used_.fetch_add(1, std::memory_order_relaxed); + stats_.expert_misses.fetch_add(1, std::memory_order_relaxed); + return s; +} + +int pager::probe(int il, int ie) { + if (!ready_) { + return -1; + } + stats_.expert_requests.fetch_add(1, std::memory_order_relaxed); + count_request(il, ie); + std::lock_guard lk(mu_); + const int32_t s = find_slot_locked(il, ie); + if (s < 0) { + stats_.expert_misses.fetch_add(1, std::memory_order_relaxed); + return -1; + } + slots_[(size_t) s].last_use = ++clock_; + slots_[(size_t) s].pin++; + touch_freq_locked(slots_[(size_t) s]); + stats_.expert_hits.fetch_add(1, std::memory_order_relaxed); + return s; +} + +bool pager::preload_all(int n_threads) { + if (!ready_) { + return false; + } + size_t total = 0; + for (int il = 0; il < (int) info_.n_layer; il++) { + const layer_desc * L = layer(il); + if (L) { + total += L->experts.size(); + } + } + if (total == 0 || n_slots_ < total) { + return false; + } + const int nt = std::max(1, n_threads); + std::atomic failed{0}; + std::vector th; + th.reserve((size_t) nt); + for (int t = 0; t < nt; t++) { + th.emplace_back([this, t, nt, &failed] { + for (int il = 0; il < (int) info_.n_layer; il++) { + const layer_desc * L = layer(il); + if (!L) { + continue; + } + for (int ie = t; ie < (int) L->experts.size(); ie += nt) { + const int s = acquire(il, ie, true); + if (s < 0) { + failed.fetch_add(1, std::memory_order_relaxed); + } else { + release(s); + } + } + } + }); + } + for (auto & x : th) { + x.join(); + } + reset_stats(); + return failed.load() == 0; +} + +void pager::reset_stats() { + auto z = [](std::atomic & a) { a.store(0, std::memory_order_relaxed); }; + z(stats_.expert_requests); z(stats_.expert_hits); z(stats_.expert_misses); + z(stats_.ssb_read_bytes_demand); z(stats_.ssb_read_bytes_prefetch); + z(stats_.prefetch_issued); z(stats_.prefetch_completed); z(stats_.prefetch_used); + z(stats_.prefetch_duplicate); z(stats_.evictions); z(stats_.demand_wait_us); + z(stats_.prefetch_wait_us); z(stats_.prefetch_bytes_used); z(stats_.tokens); + z(stats_.moe_calls); z(stats_.flops); + if (lstats_) { + for (int il = 0; il < (int) info_.n_layer; il++) { + z(lstats_[il].req); z(lstats_[il].hit); z(lstats_[il].miss); + z(lstats_[il].demand_bytes); z(lstats_[il].prefetch_bytes); + } + } +} + +bool pager::resident(int il, int ie) { + if (!ready_) { + return false; + } + std::lock_guard lk(mu_); + const int32_t s = find_slot_locked(il, ie); + return s >= 0 && !slots_[(size_t) s].busy; +} + +void pager::note_span_read(uint64_t bytes) { + stats_.ssb_read_bytes_demand.fetch_add(bytes, std::memory_order_relaxed); + stats_.demand_wait_us.fetch_add(0, std::memory_order_relaxed); +} + +// 衰減式頻率：每 8192 次使用就把所有計數減半，讓「過去熱門、現在變冷」的 expert 能被淘汰 +void pager::touch_freq_locked(slot_state & st) { + if (++freq_tick_ >= 8192) { + freq_tick_ = 0; + for (auto & s : slots_) { + s.freq /= 2; + } + } + if (st.freq < 0xFFFFFFFFu) { + st.freq++; + } +} + +// ============================== OPT-1：熱門 expert 剖面 ============================== +// 問題實測（docs/03）：arena 3.88 GiB 只放得下 2044 個 expert，總共卻有 10240 個， +// prefill 的幾百個「只用一次」的專家會把熱區整個洗掉，命中率卡在 ~70%。 +// 解法借用 Strata 的做法（src/core/expert_cache.cpp）：**離線剖面 + 靜態排行榜 + +// 熱頁永不淘汰**。這裡不動檔案格式，用純文字存排行榜。 + +// ---- SDQ-0004 +void pager::ensure_layer_stats(int n_layer) { + if (lstats_ || n_layer <= 0) { + return; + } + lstats_.reset(new layer_stat[(size_t) n_layer]); +} + +void pager::count_request(int il, int ie) { + if (prof_.empty() || il < 0 || ie < 0) { + return; + } + const size_t idx = (size_t) il * (size_t) info_.n_expert + (size_t) ie; + if (idx < prof_.size()) { + // 單執行緒 ++ 即可：請求次數是統計用途，不參與任何正確性判斷 + if (prof_[idx] < 0xFFFFFFFFu) { + prof_[idx]++; + } + } +} + +bool pager::hot_lookup(int il, int ie) const { + if (hot_.empty() || il < 0 || ie < 0) { + return false; + } + const size_t idx = (size_t) il * (size_t) info_.n_expert + (size_t) ie; + return idx < hot_.size() && hot_[idx] != 0; +} + +// 挑一個容得下 (il,ie) 的槽。與 victim_locked 的差別：**先只在冷頁裡找**， +// 找不到才退讓去動熱頁。這就是「熱區不會被 prefill 洗掉」的機制。 +// 這個槽位放得下「需要 big 的 expert」嗎？ +// 小槽位：只能放小 expert +// 大槽位：小expert 也放得下（只是浪費一點空間） +static inline bool slot_accepts(bool need_big, bool slot_is_big) { + return !need_big || slot_is_big; +} + +int pager::admit_slot_locked(int il, int ie) { + // il/ie 都 < 0 代表「不指定大小」→ 兩種槽位都可以（舊的呼叫路徑）。 + const bool need_big = (il >= 0 && ie >= 0) ? need_big_slot(il, ie) : false; + // 完全沒用過的空槽永遠優先（它不是熱頁，佔用等於浪費） + // 大 expert 先看大槽位；小 expert 先看小槽位。 + std::vector & first = need_big ? free_big_slots_ : free_slots_; + std::vector & second = need_big ? free_slots_ : free_big_slots_; + while (!first.empty()) { + const int32_t s = first.back(); + if (!slots_[(size_t) s].busy) { + first.pop_back(); + return s; + } + } + // 第一輪：只在「冷頁」裡挑最久未用的。 + // want_hot=true → 熱門 expert 進來時，把冷頁換掉（不動其他熱頁） + // want_hot=false → 冷門 expert 進來時，優先吃冷頁 + // 兩種情況的規則相同，所以 hot_lookup 的結果不影響選擇順序，只影響標記。 + for (int pass = 0; pass < 2; pass++) { + int64_t best = -1; + uint64_t best_clock = UINT64_MAX; + for (size_t s = 0; s < n_slots_; s++) { + const slot_state & st = slots_[s]; + if (st.pin > 0 || st.busy) { + continue; + } + if (pass == 0 && st.is_hot) { + continue; // 第一輪不動熱頁 + } + if (!slot_accepts(need_big, is_big_slot(s))) { + continue; + } + if (st.last_use < best_clock) { + best_clock = st.last_use; + best = (int64_t) s; + } + } + if (best >= 0) { + return (int) best; + } + } + // 退讓：另一種大小的空槽。**一定要檢查放不放得下** —— 小槽位放不下大 expert， + // 硬塞會讓 read_expert 寫出槽位邊界（靜默把別的 expert 寫壞）。 + while (!second.empty()) { + const int32_t s = second.back(); + if (!slots_[(size_t) s].busy && + slot_accepts(need_big, is_big_slot(s))) { + second.pop_back(); + return s; + } + if (!slots_[(size_t) s].busy) { + second.pop_back(); + } else { + break; + } + } + return -1; +} + +// least_stale：只淘汰「閒置 >= stale_window」的頁；在合格的頁裡挑 decayed freq +// 最小的（最少被重複使用），同 freq 再挑最久未用。熱頁第一輪不動。 +// 沒有合格頁時回 -1，由 victim_locked 退回 LRU。 +int pager::victim_least_stale_locked(int il, int ie, uint64_t stale_before) { + const bool need_big = (il >= 0 && ie >= 0) ? need_big_slot(il, ie) : false; + std::vector & first = need_big ? free_big_slots_ : free_slots_; + std::vector & second = need_big ? free_slots_ : free_big_slots_; + for (std::vector * lst : { &first, &second }) { + while (!lst->empty()) { + const int32_t s = lst->back(); + if (slots_[(size_t) s].busy) { + break; + } + if (slots_[(size_t) s].last_use <= stale_before && + slot_accepts(need_big, is_big_slot(s))) { + lst->pop_back(); + return s; + } + lst->pop_back(); + } + } + for (int pass = 0; pass < 2; pass++) { + int64_t best = -1; + uint32_t best_freq = UINT32_MAX; + uint64_t best_clock = UINT64_MAX; + for (size_t i = 0; i < n_slots_; i++) { + const slot_state & st = slots_[i]; + if (st.pin > 0 || st.busy || st.last_use > stale_before) { + continue; + } + if (pass == 0 && st.is_hot) { + continue; + } + if (!slot_accepts(need_big, is_big_slot(i))) { + continue; + } + if (st.freq < best_freq || + (st.freq == best_freq && st.last_use < best_clock)) { + best_freq = st.freq; + best_clock = st.last_use; + best = (int64_t) i; + } + } + if (best >= 0) { + return (int) best; + } + } + return -1; +} + +// 依 SDQ_EVICT_POLICY 選擇 victim。兩條路徑都必須回傳「放得下 (il,ie)」的槽。 +int pager::victim_locked(int il, int ie) { + if (cfg_.evict_policy == "least_stale") { + const uint64_t w = cfg_.stale_window; + const uint64_t stale_before = clock_ > w ? clock_ - w : 0; + const int s = victim_least_stale_locked(il, ie, stale_before); + if (s >= 0) { + return s; + } + // 所有頁都還很新（或都被 pin/busy）→ 退回全域 LRU，保證推論不會卡住。 + } + return admit_slot_locked(il, ie); +} + +bool pager::load_profile(const std::string & path) { + std::FILE * f = std::fopen(path.c_str(), "rb"); + if (!f) { + fprintf(stderr, "[sdq] 找不到 expert 剖面 %s（將以動態 LRU/LFU 運作）\n", path.c_str()); + return false; + } + char magic[8] = {0}; + if (std::fread(magic, 1, 4, f) != 4 || std::memcmp(magic, "SDQP", 4) != 0) { + std::fclose(f); + fprintf(stderr, "[sdq] %s 不是 SDQ expert 剖面（缺少 SDQP 標頭）\n", path.c_str()); + return false; + } + uint32_t hdr[3] = {0, 0, 0}; + if (std::fread(hdr, sizeof(uint32_t), 3, f) != 3) { + std::fclose(f); + return false; + } + const uint32_t n = hdr[2]; + std::vector>> ranked; + ranked.reserve(n); + for (uint32_t i = 0; i < n; i++) { + int32_t l = 0, e = 0; + uint32_t c = 0; + if (std::fread(&l, sizeof(int32_t), 1, f) != 1) break; + if (std::fread(&e, sizeof(int32_t), 1, f) != 1) break; + if (std::fread(&c, sizeof(uint32_t), 1, f) != 1) break; + ranked.push_back({c, {l, e}}); + } + std::fclose(f); + std::sort(ranked.begin(), ranked.end(), + [](const auto & a, const auto & b) { return a.first > b.first; }); + + hot_.assign((size_t) info_.n_layer * (size_t) info_.n_expert, 0); + // 預設釘選 25 % 的槽位（模擬顯示這是 1 GiB 階層的最佳點）。舊版沒有設定時 + // 會把所有出現過的 expert 都標成熱門，等於沒有冷槽可淘汰。 + const int want = cfg_.hot_pin > 0 ? cfg_.hot_pin + : std::max(1, (int) (n_slots_ / 4)); + n_hot_ = 0; + for (const auto & r : ranked) { + if (n_hot_ >= want) { + break; + } + const int l = r.second.first, e = r.second.second; + if (l < 0 || l >= info_.n_layer || e < 0 || e >= info_.n_expert) { + continue; + } + const size_t idx = (size_t) l * (size_t) info_.n_expert + (size_t) e; + if (!hot_[idx]) { + hot_[idx] = 1; + n_hot_++; + } + } + fprintf(stderr, "[sdq] expert 剖面：%s → 釘選 %d / %zu 個熱門 expert（永不淘汰）\n", + path.c_str(), n_hot_, ranked.size()); + return n_hot_ > 0; +} + +void pager::dump_profile(const std::string & path) { + if (prof_.empty()) { + return; + } + std::vector>> ranked; + ranked.reserve(prof_.size()); + uint64_t total = 0; + for (int l = 0; l < info_.n_layer; l++) { + for (int e = 0; e < info_.n_expert; e++) { + const uint32_t c = prof_[(size_t) l * (size_t) info_.n_expert + (size_t) e]; + if (c > 0) { + ranked.push_back({c, {l, e}}); + total += c; + } + } + } + std::sort(ranked.begin(), ranked.end(), + [](const auto & a, const auto & b) { return a.first > b.first; }); + std::FILE * f = std::fopen(path.c_str(), "wb"); + if (!f) { + fprintf(stderr, "[sdq] 無法寫出 expert 剖面 %s\n", path.c_str()); + return; + } + std::fwrite("SDQP", 1, 4, f); + const uint32_t hdr[3] = {1, (uint32_t) info_.n_layer, (uint32_t) ranked.size()}; + std::fwrite(hdr, sizeof(uint32_t), 3, f); + for (const auto & r : ranked) { + const int32_t l = r.second.first, e = r.second.second; + const uint32_t c = r.first; + std::fwrite(&l, sizeof(int32_t), 1, f); + std::fwrite(&e, sizeof(int32_t), 1, f); + std::fwrite(&c, sizeof(uint32_t), 1, f); + } + std::fclose(f); + // 覆盖率報告：前 K 名可以蓋掉多少請求 → 直接決定 hot_pin 該設多少 + uint64_t acc = 0; + const size_t ks[6] = {512, 1024, 1536, 2048, 3072, 4096}; + fprintf(stderr, "[sdq] expert 剖面已寫到 %s（%zu 個有被用到，總請求 %llu）\n", + path.c_str(), ranked.size(), (unsigned long long) total); + size_t ki = 0; + for (size_t i = 0; i < ranked.size(); i++) { + acc += ranked[i].first; + if (ki < 6 && i + 1 == ks[ki]) { + fprintf(stderr, "[sdq] 前 %5zu 名涵蓋 %6.2f%% 的請求\n", + ks[ki], 100.0 * (double) acc / (double) (total ? total : 1)); + ki++; + } + } +} + +int pager::find_slot_locked(int il, int ie) { + const int64_t key = ((int64_t) il << 20) | (int64_t) ie; + auto it = page_.find(key); + return it == page_.end() ? -1 : it->second; +} + +// 找一個「可以放 (il,ie) 而且 last_use <= stale_before」的槽位。 +// +// 為什麼不能直接用 victim_stale_locked()：大小兩種槽位之後，那個函式會掃過**所有** +// 槽位（包含小槽位），於是一個需要大槽位的 expert 可能拿到小槽位， +// 預取時 read_expert 會寫出槽位邊界 → 靜默把鄰居槽位的資料寫壞。 +// 這是引進大小兩種槽位時真正會咬人的地方。 +// +// 順序：先看對應尺寸的空槽（從沒用過，最該拿來放預取）， +// 再掃全部槽位挑最久未用的（放得下的那種尺寸）。 +int pager::victim_stale_for(int il, int ie, uint64_t stale_before) { + const bool need_big = need_big_slot(il, ie); + std::vector & first = need_big ? free_big_slots_ : free_slots_; + std::vector & second = need_big ? free_slots_ : free_big_slots_; + for (std::vector * lst : { &first, &second }) { + while (!lst->empty()) { + const int32_t s = lst->back(); + if (slots_[(size_t) s].busy) { + break; + } + if (slots_[(size_t) s].last_use <= stale_before && + slot_accepts(need_big, is_big_slot(s))) { + lst->pop_back(); + return s; + } + lst->pop_back(); + } + } + int64_t best = -1; + uint64_t best_clock = UINT64_MAX; + for (size_t i = 0; i < n_slots_; i++) { + const slot_state & st = slots_[i]; + if (st.pin > 0 || st.busy || st.last_use > stale_before) { + continue; + } + if (!slot_accepts(need_big, is_big_slot(i))) { + continue; + } + if (st.last_use < best_clock) { + best_clock = st.last_use; + best = (int64_t) i; + } + } + return (int) best; +} + +int pager::victim_stale_locked(uint64_t stale_before) { + // 只挑「last_use 很舊」而且沒被 pin 的槽：保護最近的工作集 + while (!free_slots_.empty()) { + const int32_t s = free_slots_.back(); + if (!slots_[(size_t) s].busy && slots_[(size_t) s].last_use <= stale_before) { + free_slots_.pop_back(); + return s; + } + if (free_slots_.size() * 4 < n_slots_) { + break; + } + free_slots_.pop_back(); + } + int64_t best = -1; + uint64_t best_clock = UINT64_MAX; + for (size_t i = 0; i < n_slots_; i++) { + const slot_state & st = slots_[i]; + if (st.pin > 0 || st.busy || st.last_use > stale_before) { + continue; + } + if (st.last_use < best_clock) { + best_clock = st.last_use; + best = (int64_t) i; + } + } + return (int) best; +} + +int pager::acquire(int il, int ie, bool blocking) { + if (!ready_) { + return -1; + } + std::unique_lock lk(mu_); + int32_t s = find_slot_locked(il, ie); + if (s >= 0 && slots_[s].busy) { + // 這個槽位**正在被填入**：另一條執行緒剛 admit 它、資料還沒讀完。 + // + // 這裡一定要等，不能當成命中直接用。否則讀到的是「寫到一半的權重」， + // 而且那個錯誤是**非決定性**的 —— 每次執行結果可能不同（實測在 + // prefill 時會讓同一個提示詞跑出不同的文字）。 + // busy 的槽位在 admit/victim 那邊本來就會被跳過，所以這裡是唯一 + // 還能拿到它的路徑，也因此必須在這裡等。 + const uint64_t t0 = now_us(); + while (slots_[s].busy) { + lk.unlock(); + std::this_thread::sleep_for(std::chrono::microseconds(200)); + lk.lock(); + if (now_us() - t0 > (uint64_t) env_size("SDQ_BUSY_WAIT_MS", 30000) * 1000) { + fprintf(stderr, + "[sdq] 致命：等待槽位 il=%d ie=%d 讀入超過 %u ms（busy 一直沒解除）\n", + il, ie, (unsigned) env_size("SDQ_BUSY_WAIT_MS", 30000)); + std::fflush(stderr); + abort(); + } + } + // 讀取可能失敗（那時槽位已被清空）→ 重新走一次正常的 miss 路徑 + if (find_slot_locked(il, ie) != s) { + s = -1; + } + } + if (s >= 0) { + slots_[s].last_use = ++clock_; + slots_[s].pin++; + touch_freq_locked(slots_[(size_t) s]); + // 剖面必須統計「所有」請求（命中也要算）。舊版只在 miss 路徑計數， + // 熱門 expert 因為大多命中，會被剖面系統性低估。 + count_request(il, ie); + stats_.expert_requests.fetch_add(1, std::memory_order_relaxed); + stats_.expert_hits.fetch_add(1, std::memory_order_relaxed); + if (lstats_ && il >= 0 && il < (int) info_.n_layer) { + lstats_[il].req.fetch_add(1, std::memory_order_relaxed); + lstats_[il].hit.fetch_add(1, std::memory_order_relaxed); + } + // A prefetch that already finished counts as "used" here. Without this + // the finished-before-request case was invisible and prefetch precision + // was understated (docs/10 §f). + if (slots_[s].from_prefetch) { + slots_[s].from_prefetch = false; + stats_.prefetch_used.fetch_add(1, std::memory_order_relaxed); + stats_.prefetch_bytes_used.fetch_add(slot_size_of(s), std::memory_order_relaxed); + } + return s; + } + + stats_.expert_requests.fetch_add(1, std::memory_order_relaxed); + count_request(il, ie); + stats_.expert_misses.fetch_add(1, std::memory_order_relaxed); + if (lstats_ && il >= 0 && il < (int) info_.n_layer) { + lstats_[il].req.fetch_add(1, std::memory_order_relaxed); + lstats_[il].miss.fetch_add(1, std::memory_order_relaxed); + } + + // 先找一個已經在背景讀入這頁的槽 + for (size_t i = 0; i < n_slots_; i++) { + slot_state & st = slots_[i]; + if (st.busy && st.layer == il && st.expert == ie) { + // 等背景執行緒讀完 + lk.unlock(); + const uint64_t t0 = now_us(); + while (true) { + lk.lock(); + if (!slots_[i].busy) { + break; + } + lk.unlock(); + std::this_thread::sleep_for(std::chrono::microseconds(200)); + } + s = (int32_t) i; + lk.unlock(); + stats_.prefetch_wait_us.fetch_add(now_us() - t0, std::memory_order_relaxed); + std::unique_lock lk2(mu_); + slots_[s].last_use = ++clock_; + slots_[s].pin++; + stats_.expert_hits.fetch_add(1, std::memory_order_relaxed); + stats_.prefetch_used.fetch_add(1, std::memory_order_relaxed); + return s; + } + } + + if (!blocking) { + return -1; + } + + // 選槽位時必須知道「這個 expert 需要哪一種大小」，否則大 expert 可能被塞進小槽位 + // （寫出邊界 → 靜默算錯），或者小 expert 佔掉大槽位（浪費 6.8% 的 arena）。 + s = victim_locked(il, ie); + if (s < 0) { + // 拿不到槽位：短暫等待別的執行緒解除 pin（每層開頭會 unpin_all）。 + // 這裡「重試」而不是直接回 -1，因為回 -1 會讓呼叫端**跳過這個 expert**， + // 結果是靜默少算一條專家分支 —— 那是最嚴重的錯誤。 + // 2000 次 × 250 µs = 最多等 0.5 秒；真的撐不到那麼久，代表 arena 被 + // 預取的 busy 槽位佔住了（實測：跨層預取開啟時會發生）。 + for (int attempt = 0; attempt < 2000 && s < 0; attempt++) { + lk.unlock(); + std::this_thread::sleep_for(std::chrono::microseconds(250)); + lk.lock(); + s = victim_locked(il, ie); + } + if (s < 0) { + fprintf(stderr, + "[sdq] 嚴重：找不到能放 il=%d ie=%d（need_big=%d）的槽位，" + "大槽 %zu／小槽 %zu\n", + il, ie, (int) need_big_slot(il, ie), n_big_, n_small_); + // 這裡「回 -1」會讓呼叫端**跳過這個 expert**，結果是靜默少算一條專家分支 + // —— 那是整個引擎最嚴重的錯誤型態（模型會安靜地給出錯誤答案，而且 + // 每次執行還不一定一樣）。寧可整個行程死掉，也不要產出「看起來正常」 + // 的錯誤結果。 + fprintf(stderr, + "[sdq] 致命：無法取得 expert 槽位。為了不給出靜默錯誤的結果，" + "直接中止。請回報這個情況（多半是預取把槽位佔成 busy 太久）。\n"); + std::fflush(stderr); + abort(); + } + } + slot_state & st = slots_[s]; + if (st.layer >= 0) { + page_.erase(((int64_t) st.layer << 20) | (int64_t) st.expert); + stats_.evictions.fetch_add(1, std::memory_order_relaxed); + used_.fetch_sub(1, std::memory_order_relaxed); + } + st.layer = il; + st.expert = ie; + st.last_use = ++clock_; + st.pin = 1; + st.is_hot = hot_lookup(il, ie); + // **先標 busy，再 publish 到 page_。** + // 順序反過來的話（先 publish 再標 busy 中間有一個窗口）另一條執行緒會把它 + // 當成命中，然後在資料還沒讀完時就拿去做矩陣乘法 → 非決定性的錯誤輸出。 + // acquire() 看到 busy 會等；admit/victim 本來就會跳過 busy 的槽位。 + st.busy = true; + page_[((int64_t) il << 20) | (int64_t) ie] = s; + + lk.unlock(); + + const uint64_t t0 = now_us(); + const bool ok = read_expert_par(il, ie, slot_ptr(s)); + const uint64_t dt = now_us() - t0; + + std::unique_lock lk2(mu_); + st.busy = false; + if (!ok) { + st.layer = -1; + st.expert = -1; + st.pin = 0; + page_.erase(((int64_t) il << 20) | (int64_t) ie); + return -1; + } + used_.fetch_add(1, std::memory_order_relaxed); + stats_.demand_wait_us.fetch_add(dt, std::memory_order_relaxed); + const layer_desc * L = layer(il); + if (L) { + const uint64_t nb = (uint64_t) L->experts[(size_t) ie].bytes; + stats_.ssb_read_bytes_demand.fetch_add(nb, std::memory_order_relaxed); + if (lstats_ && il >= 0 && il < (int) info_.n_layer) { // SDQ-0004 + lstats_[il].demand_bytes.fetch_add(nb, std::memory_order_relaxed); + } + } + return s; +} + +void pager::prefetch_expert(int il, int ie) { + if (!ready_ || !cfg_.prefetch) { + return; + } + { + std::unique_lock lk(mu_); + if (find_slot_locked(il, ie) >= 0) { + stats_.prefetch_duplicate.fetch_add(1, std::memory_order_relaxed); + return; + } + // 預取只能占用「空槽」或「已經很久沒被用到的頁」。 + // （實測：無差別淘汰會讓命中率從 90% 掉到 27%，而預取使用率是 0%，純粹有害。） + // + // **槽位一定要選對尺寸**：need_big_slot() 為真的 expert 拿到小槽位會寫出邊界。 + const uint64_t stale_before = clock_ - (uint64_t) env_size("SDQ_PREFETCH_STALE_TOKENS", 2) * 320; + // 預取佇列也不能太深：SSD 讀不夠快時，排在後面的預取早就過期了 + if (pf_q_.size() > (size_t) env_size("SDQ_PREFETCH_QUEUE", 64)) { + return; + } + // **必須留夠空的槽位給 demand。** + // 預取會把槽位標成 busy 並佔著它讀取；大槽位本來就少（工作集下限是 + // n_big_layers × n_expert_used），如果預取把最後幾個空槽都吃掉， + // 下一層的 demand 就會「找不到能放 il=34 ie=114 的槽位」而中止。 + // （這是 verify 新增的預取壓力測試 [2b] 抓到的實際情況。） + // 所以：對應尺寸的空槽少於 n_expert_used 個時，這一筆預取直接不做。 + const size_t reserve = (size_t) info_.n_expert_used; + const size_t n_free_right = need_big_slot(il, ie) ? free_big_slots_.size() : free_slots_.size(); + if (n_free_right < reserve) { + stats_.prefetch_duplicate.fetch_add(1, std::memory_order_relaxed); + return; + } + const int32_t s = victim_stale_for(il, ie, stale_before); + if (s < 0) { + return; + } + slot_state & st = slots_[s]; + if (st.layer >= 0) { + page_.erase(((int64_t) st.layer << 20) | (int64_t) st.expert); + stats_.evictions.fetch_add(1, std::memory_order_relaxed); + used_.fetch_sub(1, std::memory_order_relaxed); + } + st.layer = il; + st.expert = ie; + st.busy = true; + st.pin = 0; + pf_q_.push_back({s, il, ie, now_us()}); + stats_.prefetch_issued.fetch_add(1, std::memory_order_relaxed); + } + pf_cv_.notify_one(); +} + +void pager::pf_worker() { + std::unique_lock lk(pf_mu_); + for (;;) { + pf_cv_.wait(lk, [this] { return pf_stop_ || !pf_q_.empty(); }); + if (pf_stop_ && pf_q_.empty()) { + return; + } + const pf_item item = pf_q_.front(); + pf_q_.pop_front(); + lk.unlock(); + + // 這個槽可能已經被淘汰/重用 → 檢查 generation + { + std::unique_lock g(mu_); + const slot_state & st = slots_[item.slot]; + const bool still = st.busy && st.layer == item.il && st.expert == item.ie; + if (!still) { + lk.lock(); + continue; + } + } + + // 預取也走 I/O 段平行：單獨一條執行緒循序讀三段，是這條路徑唯一的瓶頸。 + const bool ok = cfg_.io_parallel ? read_expert_par(item.il, item.ie, slot_ptr(item.slot)) + : read_expert(item.il, item.ie, slot_ptr(item.slot)); + const layer_desc * L = layer(item.il); + { + std::unique_lock g(mu_); + slot_state & st = slots_[item.slot]; + if (st.busy && st.layer == item.il && st.expert == item.ie) { + st.busy = false; + st.last_use = ++clock_; + if (ok) { + page_[((int64_t) item.il << 20) | (int64_t) item.ie] = item.slot; + used_.fetch_add(1, std::memory_order_relaxed); + stats_.prefetch_completed.fetch_add(1, std::memory_order_relaxed); + // SDQ-0004: tag the page so the demand path can credit it, + // and so unused prefetched bytes can be reported as waste. + st.from_prefetch = true; + if (lstats_ && item.il >= 0 && item.il < (int) info_.n_layer) { + lstats_[item.il].prefetch_bytes.fetch_add( + (uint64_t) slot_size_of(item.slot), std::memory_order_relaxed); + } + if (L) { + stats_.ssb_read_bytes_prefetch.fetch_add( + (uint64_t) L->experts[(size_t) item.ie].bytes, std::memory_order_relaxed); + } + } else { + st.layer = -1; + st.expert = -1; + } + } + } + lk.lock(); + } +} + +void pager::release(int s) { + if (s < 0 || (size_t) s >= n_slots_) { + return; + } + std::lock_guard lk(mu_); + slot_state & st = slots_[(size_t) s]; + if (st.pin > 0) { + st.pin--; + } +} + +void pager::unpin_all() { + std::unique_lock lk(mu_); + for (auto & st : slots_) { + st.pin = 0; + } +} + +void pager::note_used(int il, int ie) { + if (il < 0 || ie < 0) { + return; + } + std::lock_guard lk(hist_mu_); + if (il >= (int) cur_used_.size()) { + return; + } + auto & v = cur_used_[(size_t) il]; + if (std::find(v.begin(), v.end(), ie) == v.end()) { + v.push_back(ie); + } +} + +// 「下一層」預取：只在 layer L 的 op 開頭呼叫一次。 +// +// 為什麼需要這個函式：end_token() 是**整個 token 的 40 層都算完之後**才被呼叫的 +// （sdq_cli.cpp 在 llama_decode 之後呼叫），所以它排的預取完全沒有和任何計算重疊 +// —— 量測上 prefetch 使用率只有 0~8 %，等於整條預取路徑是關的。 +// 在 layer L 開始時就替 layer L+1 排預取，才有一整層的時間（讀＋算，約 4 ms） +// 可以讓 SSD 讀取躲在後面。 +// +// 預測用的是 prev_used_[il+1]（上一個 token 在該層選的專家）—— MoE 的路由在 +// 鄰近 token 之間非常穩定，這是最便宜且不用訓練的預測器。 +void pager::prefetch_layer_ahead(int il) { + if (!ready_ || !cfg_.prefetch || il < 0 || il >= (int) layers_.size()) { + return; + } + int32_t ids[64]; + const int k = (int) info_.n_expert_used; + if (k <= 0 || k > 64) { + return; + } + predict(il, ids, k); + for (int i = 0; i < k; i++) { + prefetch_expert(il, ids[i]); + } +} + +// SDQ-0005: called once per llama_decode() from any front end. +// +// sdq-chat drives this itself from main(); llama-server does not, which left +// tokens at 0, ssd_bytes_per_token undefined and SDQ_STATS_FILE empty. This is +// the single place that closes a decode step, so hook it here. +// +// n_tokens > 1 means prefill: a whole prompt was just paged in. Counting that as +// one "token" would make ssd_bytes_per_token meaningless, and feeding it to the +// Markov table would poison the transition counts with a whole-batch transition. +// So prefill advances the byte accounting by n_tokens but leaves the routing +// history alone -- it is not a single-token transition. +void pager::note_decode(int n_tokens) { + if (!ready_ || n_tokens <= 0) { + return; + } + stats_.tokens.fetch_add((uint64_t) n_tokens, std::memory_order_relaxed); + if (n_tokens == 1) { + end_token(); + } + dump_stats("decode"); +} + +void pager::end_token() { + if (!ready_) { + return; + } + // 用「上個 token 的選擇」更新一階 Markov 轉移表 + { + std::lock_guard lk(hist_mu_); + for (size_t il = 0; il < cur_used_.size(); il++) { + auto & cur = cur_used_[il]; + if (cur.empty()) { + continue; + } + auto & layer_hist = hist_[il]; + // 以每層最近一次（或平均）的前一次選擇為前驅 + if (!prev_used_.empty() && il < prev_used_.size() && !prev_used_[il].empty()) { + for (int32_t from : prev_used_[il]) { + trans & t = layer_hist[(size_t) from]; + for (int32_t to : cur) { + int found = -1; + for (int32_t i = 0; i < t.n; i++) { + if (t.to[i] == to) { + t.cnt[i]++; + found = i; + break; + } + } + if (found < 0 && t.n < 64) { + t.to[t.n] = to; + t.cnt[t.n] = 1; + t.n++; + } + } + } + } + prev_used_[il] = cur; + cur.clear(); + } + } + + // 依預測排預取 + if (!cfg_.prefetch) { + return; + } + int32_t ids[64]; + for (size_t il = 0; il < layers_.size(); il++) { + const int k = (int) info_.n_expert_used; + predict((int) il, ids, k); + for (int i = 0; i < k; i++) { + prefetch_expert((int) il, ids[i]); + } + } +} + +void pager::predict(int il, int32_t * out, int k) { + if (il < 0 || il >= (int) prev_used_.size()) { + for (int i = 0; i < k; i++) { + out[i] = i; + } + return; + } + // 以最近一次選擇為前驅；沒有歷史就直接沿用上一次的選擇 + const std::vector & prev = prev_used_[(size_t) il]; + std::vector> scored; + if (prev.empty()) { + for (int i = 0; i < k; i++) { + out[i] = i; + } + return; + } + for (int32_t from : prev) { + const trans & t = hist_[(size_t) il][(size_t) from]; + for (int32_t i = 0; i < t.n; i++) { + scored.push_back({t.cnt[i], t.to[i]}); + } + } + if (scored.empty()) { + for (size_t i = 0; i < prev.size() && i < (size_t) k; i++) { + out[i] = prev[i]; + } + for (int i = (int) prev.size(); i < k; i++) { + out[i] = i; + } + return; + } + std::sort(scored.begin(), scored.end(), [](auto & a, auto & b) { + return a.first > b.first; + }); + for (int i = 0; i < k && i < (int) scored.size(); i++) { + out[i] = scored[(size_t) i].second; + } +} + +double pager::hit_rate() const { + const uint64_t req = stats_.expert_requests.load(); + const uint64_t hit = stats_.expert_hits.load(); + return req ? (double) hit / (double) req : 0.0; +} + +void pager::shutdown() { + if (!ready_) { + return; + } + if (!profile_out_.empty()) { + std::lock_guard lk(mu_); + dump_profile(profile_out_); + } + if (pf_.joinable()) { + { + std::unique_lock lk(pf_mu_); + pf_stop_ = true; + } + pf_cv_.notify_all(); + pf_.join(); + } + io_pool_stop(); + for (auto & fd : shard_fd_) { + if (fd >= 0) { + close(fd); + } + } + for (auto & fd : shard_fd_direct_) { + if (fd >= 0) { + close(fd); + } + } + shard_fd_.clear(); + shard_fd_direct_.clear(); + ready_ = false; +} + +pager::~pager() { + shutdown(); + if (arena_) { + // 注意：arena 是 posix_memalign 配置的，必須用 free() 釋放， + // 用 munmap() 會破壞 glibc 的 heap（之後任何 free 都會報 double free）。 + free(arena_); + arena_ = nullptr; + } +} + +static double pct(uint64_t a, uint64_t b) { + return b ? 100.0 * (double) a / (double) b : 0.0; +} + +void pager::print_summary() const { + if (!ready_) { + return; + } + // bailingmoe3 的第 0 層是 dense FFN（沒有 expert），要抓第一個真的有分頁的層 + const layer_desc * L0 = nullptr; + int first_moe = -1; + for (size_t i = 0; i < layers_.size(); i++) { + if (!layers_[i].experts.empty()) { + L0 = &layers_[i]; + first_moe = (int) i; + break; + } + } + if (L0 && !L0->experts.empty()) { + const expert_desc & e = L0->experts[0]; + fprintf(stderr, "[sdq] page-map: n_layer=%lld n_expert=%lld n_expert_used=%lld n_embd=%lld n_ff_exp=%lld" + " | L%dE0 gate=%llu/%llu up=%llu/%llu down=%llu/%llu (off/size) ne0=%lld ne1=%lld d_ne0=%lld d_ne1=%lld\n", + (long long) info_.n_layer, (long long) info_.n_expert, (long long) info_.n_expert_used, + (long long) info_.n_embd, (long long) info_.n_ff_exp, first_moe, + (unsigned long long) e.seg[part::GATE].off, (unsigned long long) e.seg[part::GATE].size, + (unsigned long long) e.seg[part::UP].off, (unsigned long long) e.seg[part::UP].size, + (unsigned long long) e.seg[part::DOWN].off, (unsigned long long) e.seg[part::DOWN].size, + (long long) e.ne0, (long long) e.ne1, (long long) e.down_ne0, (long long) e.down_ne1); + } + fprintf(stderr, + "[sdq] expert pager ready: %s\n" + "[sdq] 熱門 expert 釘選：%d 個%s\n" + "[sdq] slot=%.2f MiB (%zu B) slots=%zu arena=%.2f GiB (RAM 實測 RSS=%.2f GiB)\n" + "[sdq] evict_policy=%s stale_window=%llu\n", + cfg_.model_path.c_str(), + n_hot_, n_hot_ > 0 ? "（永不淘汰）" : "（無剖面，用動態 LRU/LFU）", + slot_bytes_ / 1048576.0, slot_bytes_, n_slots_, arena_bytes_ / 1073741824.0, + process_rss_bytes() / 1073741824.0, + cfg_.evict_policy.c_str(), (unsigned long long) cfg_.stale_window); + if (n_small_ == 0) { + fprintf(stderr, "[sdq] 槽位：單一尺寸 %.2f MiB × %zu\n", + slot_bytes_ / 1048576.0, n_big_); + } else { + // 大小兩種槽位：把配置結果講清楚，否則只看「slots=N」會誤以為每個槽一樣大。 + fprintf(stderr, + "[sdq] 槽位分兩種：小 %.2f MiB × %zu + 大 %.2f MiB × %zu（平均 %.2f MiB）\n", + small_bytes_ / 1048576.0, n_small_, + big_bytes_ / 1048576.0, n_big_, + (double) arena_bytes_ / (double) (n_slots_ ? n_slots_ : 1) / 1048576.0); + } + fprintf(stderr, + "[sdq] prefetch=%s fadvise_dontneed=%d\n", + cfg_.prefetch ? "on" : "off", (int) cfg_.fadvise_dontneed); +} + +void pager::dump_stats(const char * phase) { + if (cfg_.stats_path.empty()) { + return; + } + const uint64_t tokens = stats_.tokens.load(); + const uint64_t req = stats_.expert_requests.load(); + const uint64_t hit = stats_.expert_hits.load(); + const uint64_t issued = stats_.prefetch_issued.load(); + const uint64_t used = stats_.prefetch_used.load(); + const uint64_t b_dem = stats_.ssb_read_bytes_demand.load(); + const uint64_t b_pre = stats_.ssb_read_bytes_prefetch.load(); + // SDQ-0004: prefetched bytes that were consumed, so the rest is real waste. + const uint64_t b_pre_used = stats_.prefetch_bytes_used.load(); + + std::ostringstream o; + o << "{\"phase\":\"" << phase << "\"" + << ",\"tokens\":" << tokens + << ",\"expert_requests\":" << req + << ",\"expert_hits\":" << hit + << ",\"expert_misses\":" << (req - hit) + << ",\"hit_rate\":" << pct(hit, req) + << ",\"ssb_demand_bytes\":" << b_dem + << ",\"ssb_prefetch_bytes\":" << b_pre + << ",\"ssb_total_bytes\":" << (b_dem + b_pre) + << ",\"ssd_bytes_per_token\":" << (tokens ? (double) (b_dem + b_pre) / (double) tokens : 0.0) + << ",\"prefetch_issued\":" << issued + << ",\"prefetch_used\":" << used + << ",\"prefetch_used_rate\":" << pct(used, issued) + // SDQ-0004: prefetch waste. b_pre counts every byte the background reader + // pulled in; b_pre_used counts the ones a demand request actually hit. + // The difference is what made docs/07's negative result worse than "slower". + << ",\"prefetch_bytes_used\":" << b_pre_used + << ",\"prefetch_bytes_wasted\":" << (b_pre > b_pre_used ? b_pre - b_pre_used : 0) + << ",\"prefetch_waste_rate\":" << pct(b_pre > b_pre_used ? b_pre - b_pre_used : 0, b_pre) + << ",\"prefetch_duplicate\":" << stats_.prefetch_duplicate.load() + << ",\"evictions\":" << stats_.evictions.load() + << ",\"demand_wait_us\":" << stats_.demand_wait_us.load() + << ",\"prefetch_wait_us\":" << stats_.prefetch_wait_us.load() + << ",\"arena_bytes\":" << arena_bytes_ + // 大小兩種槽位之後，「用掉幾個槽位」不能直接乘 slot_bytes_（那是大的那種）。 + // dump_stats 只在階段轉換時呼叫一次，遍歷兩千幾個槽位的成本可以忽略。 + << ",\"arena_used_bytes\":" << used_bytes_snapshot() + << ",\"rss_bytes\":" << process_rss_bytes() + << ",\"n_slots\":" << n_slots_ + << ",\"n_small_slots\":" << n_small_ + << ",\"n_big_slots\":" << n_big_ + << ",\"small_slot_bytes\":" << small_bytes_ + << ",\"big_slot_bytes\":" << big_bytes_ + << ",\"slot_bytes\":" << slot_bytes_ + << ",\"n_moe_layers\":" << n_moe_layers_ + << ",\"expert_bytes_total\":" << expert_bytes_total_ + << ",\"arch\":\"" << info_.arch << "\"" + << ",\"name\":\"" << info_.name << "\"" + // SDQ-0004: per-layer breakdown. docs/10 needs this to tell "the device is + // slow" apart from "one layer keeps evicting everyone else's experts". + << ",\"per_layer\":["; + for (int64_t il = 0; lstats_ && il < (int64_t) info_.n_layer; il++) { + const uint64_t lreq = lstats_[il].req.load(); + const uint64_t lhit = lstats_[il].hit.load(); + o << (il ? "," : "") + << "{\"layer\":" << il + << ",\"req\":" << lreq + << ",\"hit\":" << lhit + << ",\"miss\":" << (lreq - lhit) + << ",\"hit_rate\":" << pct(lhit, lreq) + << ",\"demand_bytes\":" << lstats_[il].demand_bytes.load() + << ",\"prefetch_bytes\":" << lstats_[il].prefetch_bytes.load() + << "}"; + } + o << "]" << "}"; + std::ofstream f(cfg_.stats_path, std::ios::app); + if (f) { + f << o.str() << "\n"; + } +} + +} // namespace sdq +namespace sdq { +// 由 CLI 呼叫：依環境變數/參數初始化 pager +bool sdq_pager_active_for_layer(int il) { + pager & P = pager::instance(); + if (!P.ready() || P.layer(il) == nullptr) { + return false; + } + // 除錯用：SDQ_MOE_LO / SDQ_MOE_HI 可以只讓部分層走分頁路徑， + // 用來二分搜尋第一個發生數值差異的層。 + static const int lo = [] { const char * e = getenv("SDQ_MOE_LO"); return e ? atoi(e) : 0; }(); + static const int hi = [] { const char * e = getenv("SDQ_MOE_HI"); return e ? atoi(e) : 1 << 30; }(); + if (il < lo || il >= hi) { + return false; + } + // SDQ_MOE_SKIP=3,7：這些層即使在 [lo,hi) 內也走原生路徑。 + // 搭配 SDQ_BUFT_KEEP（讓該層的 expert 留在 RAM）就能做同一層的 A/B 比對。 + static const std::set skip = [] { + std::set s; + const char * e = getenv("SDQ_MOE_SKIP"); + if (e) { + for (const char * p = e; *p; ) { + const char * q = strchr(p, ','); + const size_t n = q ? (size_t) (q - p) : strlen(p); + if (n) { + s.insert(atoi(std::string(p, n).c_str())); + } + if (!q) { + break; + } + p = q + 1; + } + } + return s; + }(); + return skip.find(il) == skip.end(); +} + +void sdq_on_decode(int n_tokens) { + auto & inst = pager::instance(); + if (inst.ready()) { + inst.note_decode(n_tokens); + } +} + +bool sdq_init(const std::string & model_path) { + config c; + c.model_path = model_path; + c.arena_bytes = env_size("SDQ_ARENA_MB", 0) << 20; + c.ram_budget_bytes = env_size("SDQ_RAM_BUDGET_MB", 0) << 20; + c.slot_bytes = env_size("SDQ_SLOT_KB", 0) << 10; + // 預設關閉。docs/07 的實測是「開 5.68 / 關 6.20 tok/s」的負結果，而且 + // 預取中的槽位會讓 acquire() 靜默跳過 expert，曾造成輸出不一致（stage S7）。 + // 這裡原本是 true，等於讓那個負結果從來沒有真正生效。 + c.prefetch = env_flag("SDQ_PREFETCH", false); + // ---- OPT-1：熱門 expert 剖面（Strata 式）---- + if (const char * pi = getenv("SDQ_PROFILE_IN")) { c.profile_in = pi; } + if (const char * po = getenv("SDQ_PROFILE_OUT")) { c.profile_out = po; } + c.hot_pin = (int) env_size("SDQ_HOT_PIN", 0); + c.admit_min_hits = (int) env_size("SDQ_ADMIT_MIN_HITS", 1); + if (const char * ep = getenv("SDQ_EVICT_POLICY")) { c.evict_policy = ep; } + c.stale_window = (uint64_t) env_size("SDQ_EVICT_STALE_TOKENS", 2) * 320; + c.fadvise_dontneed = env_flag("SDQ_FADV_DONTNEED", true); + c.io_parallel = env_flag("SDQ_IO_PARALLEL", true); + // I/O 讀取執行緒數。sdq-chat 會在載入模型「之前」自己 setenv 這個值 + // （3 × 計算執行緒，上限 48，docs/05 量到序列 3.79 → 24 條 6.68 tok/s）。 + // 但 cfg_.n_threads 從來沒有被賦值，所以任何沒有自行設定的前端 + // （llama-server 就是）會拿到 0，然後在 io_pool_start() 被夾到下限 4 條 —— + // 等於一直跑在「序列讀」那一端。改成在 pager 自己算，量測與 sdq-chat 一致。 + if (!getenv("SDQ_IO_THREADS")) { + const int hw = (int) std::max(1u, std::thread::hardware_concurrency()); + const int compute = hw < 8 ? hw : 8; // 與 sdq-chat 同一個上限 8 + c.io_threads = compute * 3; + } else { + c.io_threads = (int) env_size("SDQ_IO_THREADS", 0); + } + c.io_split = (int) env_size("SDQ_IO_SPLIT", 1); + c.verbose = env_flag("SDQ_VERBOSE", false); + const char * sp = getenv("SDQ_STATS_FILE"); + if (sp) { + c.stats_path = sp; + } + sdq_set_config(c); + try { + pager::instance().init(c, nullptr); + } catch (const std::exception & e) { + fprintf(stderr, "[sdq] pager 初始化失敗：%s\n", e.what()); + return false; + } + // 32/64 GiB 情境：arena 大到能裝下全部 expert 時先全部讀進來。 + // 沒有 miss 之後 row-parallel 才能穩定把 decode 的 16 條執行緒吃滿。 + if (env_flag("SDQ_PRELOAD", true)) { + const bool ok = pager::instance().preload_all(16); + if (ok) { + fprintf(stderr, "[sdq] 已預載全部 expert（%zu slots）→ 之後 row-parallel 可用\n", + pager::instance().n_slots()); + } + } + // 長駐前端（llama-server）不會呼叫 sdq_shutdown()；用 atexit 在收尾時補上摘要。 + // 註冊時機必須在 pager::instance() 建構之後：static 物件的析構與 atexit 依 + // 「後註冊者先執行」，反過來就會在 pager 已經析構後才跑這段（讀到已釋放的物件）。 + static std::once_flag once_exit; + std::call_once(once_exit, [] { std::atexit([] { sdq_shutdown(); }); }); + return pager::instance().ready(); +} + + +// ---------------------------------------------------------------- 除錯：dump 任意 F32 張量 +// 目的：讓「原生 mul_mat_id 路徑」與「SDQ 分頁路徑」輸出同一層的 MoE 結果， +// 兩邊各自 dump 一次就能直接 A/B 對照，判定到底是哪一邊算錯。 +namespace { +struct dump_arg { std::string path; }; + +void dump_op(ggml_tensor * dst, int ith, int /*nth*/, void * userdata) { + if (ith != 0) { + return; + } + const ggml_tensor * src = dst->src[0]; + dump_arg * a = (dump_arg *) userdata; + // 重要：dump 節點必須是「透明」的 —— 把 src 原樣複製到 dst，後續節點才會讀到 + // 正確資料。舊版只寫檔、沒有複製，dst 是未初始化記憶體，導致下游殘差加法讀到垃圾。 + if (src->type == GGML_TYPE_F32 && ggml_is_contiguous(src) && ggml_is_contiguous(dst)) { + memcpy(dst->data, src->data, ggml_nbytes(src)); + } + FILE * f = fopen(a->path.c_str(), "wb"); + if (!f) { + fprintf(stderr, "[sdq-dump] 無法開啟 %s\n", a->path.c_str()); + return; + } + const int64_t meta[6] = { src->ne[0], src->ne[1], src->ne[2], src->ne[3], + (int64_t) ggml_nbytes(src), 0 }; + fwrite(meta, sizeof(int64_t), 6, f); + // 只支援 F32 contiguous；其他情況就明確報錯，不要假設 + if (src->type != GGML_TYPE_F32 || !ggml_is_contiguous(src)) { + fprintf(stderr, "[sdq-dump] %s 不是 contiguous F32（type=%d）\n", + a->path.c_str(), (int) src->type); + } else { + fwrite(src->data, sizeof(float), (size_t) (ggml_nelements(src)), f); + } + fclose(f); + fprintf(stderr, "[sdq-dump] %s (%lld 個元素)\n", a->path.c_str(), + (long long) ggml_nelements(src)); +} +} // namespace + + +ggml_tensor * sdq_maybe_dump_tensor(ggml_context * ctx, ggml_tensor * t, int il) { + const char * prefix = getenv("SDQ_MOE_OUT_DUMP"); + if (!prefix || !t) { + return t; + } + int layer = getenv("SDQ_MOE_OUT_DUMP_LAYER") ? atoi(getenv("SDQ_MOE_OUT_DUMP_LAYER")) : -1; + if (layer >= 0 && il != layer) { + return t; + } + dump_arg * a = new dump_arg{ std::string(prefix) + "." + std::to_string(il) }; + return ggml_custom_4d(ctx, GGML_TYPE_F32, t->ne[0], t->ne[1], t->ne[2], t->ne[3], + &t, 1, dump_op, 1, a); +} + +void sdq_shutdown() { + pager & P = pager::instance(); + sdq_print_time_profile(); + sdq_print_op_counts(); + P.dump_stats("shutdown"); + P.print_summary(); + P.shutdown(); +} +} // namespace sdq diff --git a/src/sdq_pager.h b/src/sdq_pager.h new file mode 100644 index 0000000..b109f33 --- /dev/null +++ b/src/sdq_pager.h @@ -0,0 +1,510 @@ +// SDQ expert pager：SSD → RAM → CPU 的 MoE 專家分層快取 +// +// 設計要點 +// * 分頁單位 = 一個 expert（layer, expert id），大小 ~1.8 MB（三個矩陣：gate/up/down） +// * RAM 端是一個固定槽位的 arena（slab），每槽 2 MiB，O(1) 配置、不碎片化 +// * 淘汰策略：LRU + 本 token 使用中（pinned）的保護 +// * 預取：依「上一個 token 的專家選擇」用一階 Markov 轉移預測下一個 token 的專家， +// 交給背景執行緒 pread()，量測「預取命中率 / 預取浪費量」 +// * 指標：hit/miss、SSD 讀取位元組、demand 停滯時間、prefetch 領先時間、arena 使用量 +// +// 不使用 swap：所有權重資料只存在 arena（實體 RAM）與檔案（SSD）兩處。 + +#pragma once + +#include +#include +#include "llama.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace sdq { + +// ---------------------------------------------------------------- GGUF 讀取 +struct gguf_info { + int64_t n_layer = 0; + int64_t n_expert = 0; + int64_t n_expert_used = 0; + int64_t n_embd = 0; + int64_t n_ff_exp = 0; + int64_t n_ff_shexp = 0; + int64_t n_head = 0; + int64_t n_head_kv = 0; + int64_t n_embd_head_k = 0; + int64_t n_embd_head_v = 0; + int64_t context_length = 0; + int64_t full_attention_interval = 4; + // bailingmoe3 / Ling-3.0-flash + int64_t n_layer_dense_lead = 0; + int64_t n_expert_groups = 0; + int64_t n_expert_group_used = 0; + int64_t is_recr[512] = {0}; + // nemotron_h_moe: 每層的 expert 數量與 top-k 可以不同（Puzzle / LatentMoE） + std::vector n_expert_layer; + std::vector n_expert_used_layer; + int64_t moe_latent_size = 0; + bool no_gate = false; // routed experts 只有 up/down（Nemotron LatentMoE） + bool relu_sqr = false; // 激活為 relu(x)^2 + std::string arch; + std::string name; +}; + +// 一個 expert 槽位：gate/up/down 三段（有些模型把 gate+up 合成 gate_up） +struct seg_desc { + uint64_t off = 0; // 在「所屬 shard 檔案內」的位元組偏移 + uint64_t size = 0; // 位元組數 + // sharded GGUF：同一層的三段不一定在同一個檔案裡（unsplit 後 100+ GB 的模型一定會被切）， + // 所以每一段都要記得自己屬於哪一個 shard，pread 才知道要開哪個 fd。 + int shard = 0; +}; + +enum part : int { GATE = 0, UP = 1, DOWN = 2, COUNT = 3 }; + +struct expert_desc { + seg_desc seg[3]; + int64_t slot_off[3] = {0, 0, 0}; // 三段在 arena 槽位內的相對偏移 + int64_t ne0 = 0; // gate/up row 長度（= n_embd） + int64_t ne1 = 0; // gate/up row 數（= n_ff_exp） + int64_t down_ne0 = 0; // down row 長度（= n_ff_exp） + int64_t down_ne1 = 0; // down row 數（= n_embd） + ggml_type type[3] = {GGML_TYPE_COUNT, GGML_TYPE_COUNT, GGML_TYPE_COUNT}; + int64_t bytes = 0; // 三段總大小 + bool fused_gate_up = false; +}; + +// 一層的所有 expert 描述（40 層 × 256 個） +struct layer_desc { + std::vector experts; + int64_t bytes_per_expert = 0; + int64_t n_expert = 0; // 這一層的 expert 數（Nemotron 每層不同） + bool no_gate = false; // 只有 up/down，沒有 gate + // 這個層的 expert 比「大多數層」大（Q4_K/Q5_K vs Q6_K）→ 需要大槽位。 + // 不然的話所有槽位都得按最大的那種配，白白浪費 6.8% 的 arena。 + bool needs_big_slot = false; +}; + +// ---------------------------------------------------------------- 指標 +struct stats { + std::atomic expert_requests{0}; // 每個 (layer, token) 需要幾個 expert + std::atomic expert_hits{0}; + std::atomic expert_misses{0}; + std::atomic ssb_read_bytes_demand{0}; + std::atomic ssb_read_bytes_prefetch{0}; + std::atomic prefetch_issued{0}; + std::atomic prefetch_completed{0}; + std::atomic prefetch_used{0}; // 預取後真的被用到 + std::atomic prefetch_duplicate{0};// 預取時已經在 RAM 裡 + std::atomic evictions{0}; + std::atomic demand_wait_us{0}; // 因為缺頁而停下來等 SSD 的時間 + std::atomic prefetch_wait_us{0}; + std::atomic prefetch_bytes_used{0}; // prefetched bytes a demand request actually hit + std::atomic tokens{0}; + std::atomic moe_calls{0}; + std::atomic flops{0}; +}; + +// ---- SDQ-0004: per-layer cache counters (Phase 0: is any layer starved?) +struct layer_stat { + std::atomic req{0}; + std::atomic hit{0}; + std::atomic miss{0}; + std::atomic demand_bytes{0}; + std::atomic prefetch_bytes{0}; +}; + + +// ---------------------------------------------------------------- pager +struct config { + size_t arena_bytes = 0; // 0 = 自動（依 RAM 預算） + size_t slot_bytes = 0; // 0 = 自動（依模型推導） + int n_threads = 0; + bool prefetch = true; + int prefetch_ahead_layers = 0; // >0 時限制同時在飛的預取數量 + bool fadvise_dontneed = true; // 讀完丟掉 kernel page cache，讓 RAM 帳目誠實 + bool verbose = false; + uint64_t ram_budget_bytes = 0; // 0 = 自動（實體 RAM 的一小部分） + std::string model_path; + std::string stats_path; + int slot_bias = 0; + + // ---- I/O 執行緒池 ---- + // 為什麼需要：一個 expert 在檔案上是 gate / up / down 三段「不相鄰」的資料， + // 循序讀要 3 次 pread，而單一 pread 在本機量到約 7 ms（1.9 MB / 275 MB/s）。 + // 也就是說「讀一個 expert」的延遲完全由這三段循序決定，執行緒再多也沒用。 + // 把三段交給三條執行緒同時讀，單一 expert 的延遵掉到約 1/2~1/3。 + // （實測：8 個 expert 全部 miss 的一層，序列讀 12.6 ms → 段平行 7.6 ms。） + bool io_parallel = true; // SDQ_IO_PARALLEL=0 可關（回到序列讀） + int io_threads = 0; // 0 = 自動（3 × n_threads，上限 48） + int io_split = 1; // SDQ_IO_SPLIT：每一段再切成幾塊並行讀（實測 1 最好） + + // ---- Strata 式「熱門 expert 剖面」（OPT-1）---- + // 離線用代表性文字跑一次，把每個 (layer, expert) 的請求次數記成剖面； + // 執行時把排行榜前 hot_pin 個釘成「熱頁」：熱頁可以被讀取，但**永遠不被淘汰**， + // 淘汰時一律先從冷頁挑。這解決的是「prefill 的一次性專家把熱區洗掉」。 + std::string profile_in; // 讀取剖面（SDQ_PROFILE_IN） + std::string profile_out; // 結束時輸出剖面（SDQ_PROFILE_OUT） + int hot_pin = 0; // 釘選前幾名（SDQ_HOT_PIN） + int admit_min_hits = 1; // 請求次數低於此值就不收進快取（SDQ_ADMIT_MIN_HITS） + + // ---- 淘汰策略（SDQ_EVICT_POLICY）---- + // lru：挑「最久未用」的頁（現行行為）。 + // least_stale：只淘汰閒置時間 >= stale_window 的頁；若沒有這種頁， + // 才退回全域 LRU。目的是保護剛用過的工作集，避免 thrash。 + std::string evict_policy = "lru"; + uint64_t stale_window = 640; // clock 單位（每次 expert 請求 +1） +}; + +class pager { +public: + static pager & instance(); + + // 從 GGUF 檔頭建立 expert 頁面表，並配置 RAM arena + void init(const config & cfg, ggml_context * ctx); + + bool ready() const { return ready_; } + size_t arena_bytes() const { return arena_bytes_; } + size_t slot_bytes() const { return slot_bytes_; } + size_t n_slots() const { return n_slots_; } + size_t small_bytes() const { return small_bytes_; } + size_t big_bytes() const { return big_bytes_; } + size_t n_small_slots() const { return n_small_; } + size_t n_big_slots() const { return n_big_; } + int n_moe_layers() const { return n_moe_layers_; } + uint64_t expert_bytes_total() const { return expert_bytes_total_; } + const gguf_info & info() const { return info_; } + + // 取得 layer 描述（idx 越界或沒資料回 nullptr） + const layer_desc * layer(int il) const; + + // 一次把「連續編號的一串 expert」讀進 RAM：同一層的三個矩陣（gate/up/down）本來 + // 就各自連續存放，所以 k 個連續 expert 只需要 3 次大讀取（而不是 3k 次）。 + // 對高延遲的儲存裝置（每個 request 都貴）這是最大的單項加速。 + // experts 必須已依編號遞增排序；slots[i] 是 experts[i] 對應的槽位（由呼叫端配置） + // 回傳實際讀取的位元組數。 + uint64_t read_expert_span(int il, const int32_t * experts, int n, const int32_t * slots); + + // 只配置槽位（不讀資料）：回傳槽位索引，-1 表示沒有空槽 + int reserve_slot(int il, int ie); + + // 查這個 expert 是否已經在 RAM：命中則 pin 並回傳槽位，未命中回傳 -1（不配置） + int probe(int il, int ie); + + // 只查詢、不改動 pin / LRU / 統計：回傳 true 表示資料已在 RAM 且可讀。 + // row-parallel 模式用它判斷「整層都已常駐」→ 才值得把 16 條執行緒 + // 全部拿來切同一個 expert 的輸出列（沒有 I/O 可重疊時才不會虧）。 + bool resident(int il, int ie); + + // 若 arena 大到能裝下「全部」expert，就先把 15.9 GiB 一次讀進來。 + // 這是 32/64 GiB 情境（原生 llama.cpp 也是全部常駐）的公平對照： + // 沒有 miss 之後 row-parallel 才能穩定吃滿 16 條執行緒。 + // 回傳 false = arena 不夠、沒有做（小記憶體情境不受影響）。 + bool preload_all(int n_threads); + + // 清空 hit/miss/bytes 等統計（preload 後呼叫，避免把載入階段算進命中率）。 + void reset_stats(); + + // 記錄一次合併讀取的位元組數（給指標用） + void note_span_read(uint64_t bytes); + + // 確保 (layer, expert) 在 RAM，回傳槽位索引；-1 表示失敗 + // pinned: 本 token 使用中，不可被淘汰；blocking: 缺頁時是否同步讀 SSD + int acquire(int il, int ie, bool blocking); + + // 背景預取：把 page 讀進 arena（非同步），依歷史預測下一個 token 會用到的 expert + void prefetch_expert(int il, int ie); + // 「下一層」預取：在 layer L 的 op 開頭呼叫，替 layer L+1 排預取。 + // 與 end_token() 的差別在時機 —— 後者在一個 token 的 40 層全算完才跑， + // 那時排的預取沒有任何計算可以重疊，實測使用率只有 0~8 %。 + void prefetch_layer_ahead(int il); + + // 每層計算開始時呼叫：解除上一層的 pin（本層用到的槽位才受保護） + void unpin_all(); + // 用完一個槽位後立刻放掉 pin（成對呼叫 acquire()）。 + // + // 為什麼需要：原本 pin 是「整層釘住、下一層開頭才 unpin_all」。mode C 的 + // expert_forward 回傳之後就再也不讀那個槽位（最後的歸約只讀各執行緒自己的 + // partial 累加緩衝），所以繼續釘著只是白白把工作集撐大。 + // 撐大的後果很具體：prefill 一個 ubatch 有 128 個 token，一層最多需要 + // 256 個不同的 expert → 整層釘住 = 256 個槽位不能回收。大小兩種槽位之下， + // 大槽位只有約 n_big_layers × n_expert_used 個，於是**必定**在 prefill 時 + // 「找不到能放 il=34 ie=… 的槽位」而中止。 + void release(int s); + void note_used(int il, int ie); + void end_token(); + // ---- SDQ-0005: one call per llama_decode() from any front end. + void note_decode(int n_tokens); + + // 依歷史預測某層下一個 token 的 top-k expert（回傳信心 0..1） + void predict(int il, int32_t * out, int k); + + void shutdown(); + stats & get_stats() { return stats_; } + + // 把當前指標寫成一行 JSON（append 到 stats_path） + void dump_stats(const char * phase); + void print_summary() const; + + // arena 內的槽位指標 + // 槽位索引空間是統一的：[0, n_small_) 是小槽位（連續）， + // [n_small_, n_slots_) 是大槽位（接在小槽位區之後）。 + uint8_t * slot_ptr(int s) { + return (size_t) s < n_small_ + ? arena_ + (size_t) s * small_bytes_ + : arena_ + small_bytes_ * n_small_ + (size_t)(s - (int) n_small_) * big_bytes_; + } + size_t slot_size_of(int s) const { + return (size_t) s < n_small_ ? small_bytes_ : big_bytes_; + } + bool need_big_slot(int il, int ie) const; + + // 這個槽位算不算「大槽位」。 + // + // 判斷依據是**槽位的實際大小**，不是索引 >= n_small_： + // 統一尺寸（SDQ_FORCE_UNIFORM_SLOT=1，或模型只有一種 expert 大小）時， + // 所有槽位的索引都 < n_small_ 但大小都是大的 —— 用索引判斷會讓大 expert + // 永遠配不到槽位，整個推論卡在 admit 的重試迴圈裡。 + bool is_big_slot(int s) const { + return big_bytes_ && slot_size_of(s) >= big_bytes_; + } + + size_t used_bytes() const { return used_bytes_snapshot(); } + double hit_rate() const; + + // arena 裡實際被占用的位元組數。大小兩種槽位之後不能用「槽位數 × slot_bytes_」 + // 計算（slot_bytes_ 是大的那一種），所以真的把槽位走一遍。 + // 呼叫端：dump_stats（每次解碼階段結束）與 print_summary，不是熱路徑。 + size_t used_bytes_snapshot() const { + size_t b = 0; + for (size_t s = 0; s < slots_.size(); s++) { + if (slots_[s].layer >= 0) { + b += slot_size_of((int) s); + } + } + return b; + } + + // 讀 SSD（同步，給 acquire 用） + bool read_expert(int il, int ie, uint8_t * dst); + + // ---- I/O 執行緒池（見 config::io_parallel 的說明）---- + // 一個「讀取批次」的完成狀態；放在呼叫端的 stack 上，worker 只在 remaining==0 + // 之後才會再碰它，所以離開 read_expert_par() 之後不會有 dangling pointer。 + struct io_ctx { + std::atomic remaining{0}; + bool ok = true; + }; + struct io_job { + int shard; + uint64_t off; + uint8_t * dst; + size_t size; + io_ctx * ctx; + }; + + void io_pool_start(); + void io_pool_stop(); + void io_worker(size_t wid); + bool read_expert_par(int il, int ie, uint8_t * dst); + bool read_range_staged(int shard, uint64_t off, uint8_t * dst, size_t size, + uint8_t * stage, size_t stage_sz); + + // I/O 執行緒池每條 worker 自帶的 O_DIRECT 對齊緩衝（啟動時配置，已計入 RSS） + struct io_stage { uint8_t * buf = nullptr; size_t sz = 0; }; + std::vector io_stage_; + + bool io_pool_on_ = false; + int io_threads_ = 0; + int io_split_ = 1; + std::vector io_pool_; + std::mutex io_mu_; + std::condition_variable io_cv_; // 有工作 / 要停 + std::condition_variable io_done_cv_; // 某個批次讀完了 + std::deque io_q_; + bool io_stop_ = false; + std::atomic io_jobs_{0}; + std::atomic io_wait_ns_{0}; + +private: + pager() = default; + ~pager(); + + struct slot_state { + int32_t layer = -1; + int32_t expert = -1; + uint64_t last_use = 0; // LRU 記時 + uint32_t freq = 0; // 這個 expert 被用到幾次（衰減式；用來做 LFU 混合策略） + uint32_t pin = 0; // >0 時不可淘汰 + bool busy = false; // 背景執行緒正在寫入 + bool is_hot = false; // 屬於剖面排行榜：淘汰時最後才會被選中 + bool from_prefetch = false; // page was filled by the prefetcher + }; + + struct pf_item { + int slot; + int il; + int ie; + uint64_t issue_time_us; + }; + + void touch_freq_locked(slot_state & st); + // ---- OPT-1：剖面 / 熱頁 ---- + void count_request(int il, int ie); // 記一次請求（給剖面用） + bool load_profile(const std::string & path); // 讀剖面並標記前 hot_pin 名 + void dump_profile(const std::string & path); // 輸出排行榜 + bool hot_lookup(int il, int ie) const; + void load_gguf_header(const std::string & path); + void build_page_map(const std::string & path); + void alloc_arena(); + int find_slot_locked(int il, int ie); + int admit_slot_locked(int il, int ie); + int victim_locked(int il = -1, int ie = -1); + int victim_stale_locked(uint64_t stale_before); + int victim_stale_for(int il, int ie, uint64_t stale_before); + int victim_least_stale_locked(int il, int ie, uint64_t stale_before); + void pf_worker(); + bool read_range(int shard, uint64_t off, uint8_t * dst, size_t size); + bool read_range_direct(int shard, uint64_t off, uint8_t * dst, size_t size); + void open_model_fd(); + + // ---- sharded GGUF：模型被切成多個 -0000N-of-000M.gguf，每個檔有自己的 + // header、自己的 data_start 與自己的張量子集。pager 必須同時開著全部的 fd， + // 並把每個張量的「檔內偏移」記成 (shard, off)。 + std::vector shard_paths_; + std::vector shard_fd_; + std::vector shard_fd_direct_; + std::vector shard_data_start_; + bool use_direct_ = false; + std::string shard_path(int i) const { + return (i >= 0 && i < (int) shard_paths_.size()) ? shard_paths_[(size_t) i] : std::string(); + } + // 依 GGUF 的 split 命名規則列舉同目錄的所有 shard + static std::vector list_splits(const std::string & first_path, int n_split); + + bool ready_ = false; + config cfg_; + gguf_info info_; + std::vector layers_; + int n_moe_layers_ = 0; + uint64_t expert_bytes_total_ = 0; + + uint8_t * arena_ = nullptr; + size_t arena_bytes_ = 0; + size_t slot_bytes_ = 0; // = max(small_bytes_, big_bytes_)，保留給舊的訊息輸出 + size_t small_bytes_ = 0; // 小槽位大小（大多數層用） + size_t big_bytes_ = 0; // 大槽位大小；0 = 這個模型沒有大層 + size_t n_small_ = 0; + size_t n_big_ = 0; + size_t n_slots_ = 0; + mutable std::mutex mu_; + std::vector slots_; + std::unordered_map page_; // (layer<<16|expert) -> slot + std::vector free_slots_; // 空的小槽位 + std::vector free_big_slots_; + + // ---- OPT-1：剖面 / 熱頁（Strata 式：靜態排行榜 + 熱頁不淘汰）---- + std::vector prof_; // 每個 (layer, expert) 的累計請求次數 + std::unique_ptr lstats_; // SDQ-0004 + void ensure_layer_stats(int n_layer); // SDQ-0004 + std::vector hot_; // 是否在排行榜前 hot_pin 名 + int n_hot_ = 0; + std::string profile_out_; + uint32_t freq_tick_ = 0; + std::atomic used_{0}; + uint64_t clock_ = 0; + + // 背景預取 + std::thread pf_; + std::mutex pf_mu_; + std::condition_variable pf_cv_; + std::deque pf_q_; + bool pf_stop_ = false; + std::atomic pf_inflight_{0}; + + // 路由歷史（一階 Markov：上個 token 的 expert -> 這次被選到的次數） + struct trans { int32_t to[64]; uint32_t cnt[64]; int32_t n = 0; }; + std::vector> hist_; // [layer][prev_expert] + std::vector> prev_used_; // [layer] 上一個 token 用到的 expert（預測用） + std::vector> cur_used_; // [layer] 當前 token 用到的 expert + std::mutex hist_mu_; // 保護 cur_used_ / prev_used_ / hist_（多執行緒會同時 note_used） + std::atomic hist_lock_{0}; + + stats stats_; +}; + +// ---------------------------------------------------------------- 工具 +uint64_t now_us(); +size_t process_rss_bytes(); +size_t physical_ram_bytes(); + +// ---------------------------------------------------------------- ggml 介接 +// 自訂 buffer type：expert 權重佔 0 個 RAM 位元組，資料留在檔案裡 +ggml_backend_buffer_type_t sdq_expert_buft(); + +// 註冊 model params 的 tensor buffer override（expert 權重 → 我們的 buffer） +void sdq_install_buft_overrides(llama_model_params * params); + +// 建立被分頁的 MoE FFN（替代 build_moe_ffn 裡的三次 mul_mat_id） +ggml_tensor * sdq_build_moe(ggml_context * ctx, ggml_tensor * cur, + ggml_tensor * ids, ggml_tensor * weights, int il, + int64_t n_expert, int64_t n_expert_used, + float swiglu_limit); + +// 設定 arena 大小 / 預取開關等（CLI 在載入模型前呼叫） +bool sdq_init(const std::string & model_path); +void sdq_set_config(const config & c); +const config & sdq_get_config(); +void sdq_set_enabled(bool on); +void sdq_shutdown(); + +bool sdq_enabled(); + +// SDQ-0005: bridge for front ends that are not sdq-chat (llama-server). +// Defined in sdq_pager.cpp; safe to call when the pager was never initialised. +void sdq_on_decode(int n_tokens); + +// 效能計時探針的結果（SDQ_PROFILE_TIME=1） +void sdq_print_time_profile(); + +// MoE op 路徑計數（row-parallel / mode C），sdq_shutdown 時印出 +void sdq_print_op_counts(); + +// 自我測試：算一次指定 expert 的 MoE 輸出（給 NumPy 對照用） +int sdq_moe_selftest(int il, int ie, const float * x, float * out_io, + int64_t n_embd, int64_t n_ff, int n_tok); + +// 建立我們的 MoE 自訂節點。 +// **必須把 gf 傳進來並呼叫 ggml_build_forward_expand(gf, out)**： +// ggml 的排程器只會為「進了圖的節點」配置 compute buffer，沒註冊的節點其輸出 +// 緩衝區會被 liveness 分析判成可重用，導致 layer>=1 的結果被其他張量覆蓋 +// （實測症狀：layer 0 正確、layer 1 起整個錯掉）。 +ggml_tensor * sdq_build_moe(ggml_cgraph * gf, ggml_context * ctx, ggml_tensor * cur, + ggml_tensor * ids, ggml_tensor * weights, int il, + int64_t n_expert, int64_t n_expert_used, float swiglu_limit); + +// Nemotron-H LatentMoE：routed experts 只有 up/down，激活 relu(x)^2。 +ggml_tensor * sdq_build_moe_ex(ggml_cgraph * gf, ggml_context * ctx, ggml_tensor * cur, + ggml_tensor * ids, ggml_tensor * weights, int il, + int64_t n_expert, int64_t n_expert_used, + bool no_gate, bool relu_sqr); + +// 這一層的分頁是否可用（graph build 時呼叫） +bool sdq_pager_active_for_layer(int il); + +// 除錯／驗證用：把某個 F32 張量 dump 到檔案。 +// 必須在「原生 mul_mat_id 路徑」與「SDQ 分頁路徑」兩邊都呼叫，才能做真正的 A/B 對照 +// （env: SDQ_MOE_OUT_DUMP=prefix、SDQ_MOE_OUT_DUMP_LAYER=層編號，-1 = 全部） +ggml_tensor * sdq_maybe_dump_tensor(ggml_context * ctx, ggml_tensor * t, int il); + + +} // namespace sdq \ No newline at end of file diff --git a/tools/runmem.py b/tools/runmem.py new file mode 100644 index 0000000..9d5e96c --- /dev/null +++ b/tools/runmem.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +"""Sample process-tree RSS/swap while running a command.""" + +import argparse +import json +import signal +import subprocess +import sys +import time + + +def process_tree(root): + output = subprocess.check_output(["ps", "-e", "-o", "pid=,ppid="], text=True) + children = {} + for line in output.splitlines(): + pid, ppid = map(int, line.split()) + children.setdefault(ppid, []).append(pid) + result, todo = [], [root] + while todo: + pid = todo.pop() + result.append(pid) + todo.extend(children.get(pid, ())) + return result + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--interval", type=float, default=0.25) + parser.add_argument("--timeout", type=float, default=7200) + parser.add_argument("--rss-limit-gib", type=float) + parser.add_argument("command", nargs=argparse.REMAINDER) + args = parser.parse_args() + command = args.command[1:] if args.command[:1] == ["--"] else args.command + if not command: + parser.error("missing command") + + started = time.monotonic() + process = subprocess.Popen(command) + peak_rss = peak_swap = 0 + exceeded = False + while process.poll() is None: + rss = swap = 0 + for pid in process_tree(process.pid): + try: + with open(f"/proc/{pid}/status", encoding="ascii") as status: + for line in status: + if line.startswith("VmRSS:"): + rss += int(line.split()[1]) + elif line.startswith("VmSwap:"): + swap += int(line.split()[1]) + except (FileNotFoundError, ProcessLookupError): + pass + peak_rss = max(peak_rss, rss) + peak_swap = max(peak_swap, swap) + # VmRSS 以 kB 為單位，門檻須換算成 kB（1 GiB = 1024 * 1024 kB） + if args.rss_limit_gib is not None and rss > args.rss_limit_gib * (1024 ** 2): + exceeded = True + process.send_signal(signal.SIGKILL) + break + if time.monotonic() - started > args.timeout: + process.kill() + break + time.sleep(args.interval) + result = process.wait() + print("RUNMEM " + json.dumps({ + "peak_kb": peak_rss, + "peak_gib": round(peak_rss / 1048576, 3), + "limit_gib": args.rss_limit_gib, + "limit_exceeded": exceeded, + "swap_kb": peak_swap, + "wall_s": round(time.monotonic() - started, 2), + "exit": result, + "cmd": command, + "method": "procfs-rss-watchdog", + })) + return 0 if args.rss_limit_gib is None else (1 if exceeded else result) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/runmem_cgroup.py b/tools/runmem_cgroup.py new file mode 100755 index 0000000..3b983b0 --- /dev/null +++ b/tools/runmem_cgroup.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Run a command inside a cgroup-v2 MemoryMax scope and report peak memory.""" + +import argparse +import json +import os +import subprocess +import sys +import time + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--limit-gib", type=int, required=True) + parser.add_argument("--timeout", type=float, default=7200) + parser.add_argument("command", nargs=argparse.REMAINDER) + args = parser.parse_args() + command = args.command[1:] if args.command[:1] == ["--"] else args.command + if not command: + parser.error("missing command") + if args.limit_gib not in (2, 4, 8, 16, 32): + parser.error("limit must be 2, 4, 8, 16, or 32 GiB") + + unit = f"flint-epq-bench-{os.getpid()}" + runner = [ + "systemd-run", "--quiet", "--wait", "--collect", "--pipe", + "--unit", unit, "--property", f"MemoryMax={args.limit_gib}G", + "--property", "MemorySwapMax=0", "--", *command, + ] + started = time.monotonic() + process = subprocess.Popen(runner, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + time.sleep(1) + if process.poll() is not None: + output, _ = process.communicate() + if output: + sys.stdout.write(output) + print("UNTESTED: systemd/cgroup MemoryMax unavailable; RLIMIT_DATA is not a substitute", file=sys.stderr) + return 77 + + peak = swap_peak = 0 + while process.poll() is None: + try: + cgroup = subprocess.check_output( + ["systemctl", "show", f"{unit}.scope", "-p", "ControlGroup", "--value"], + text=True, stderr=subprocess.DEVNULL, + ).strip() + if cgroup: + path = os.path.join("/sys/fs/cgroup", cgroup.lstrip("/")) + for name, target in (("memory.peak", "peak"), ("memory.swap.peak", "swap")): + try: + with open(os.path.join(path, name), encoding="ascii") as stream: + value = int(stream.read().strip()) + if target == "peak": + peak = max(peak, value) + else: + swap_peak = max(swap_peak, value) + except OSError: + pass + except (OSError, subprocess.CalledProcessError): + pass + if time.monotonic() - started > args.timeout: + process.kill() + break + time.sleep(0.25) + + output, _ = process.communicate() + if output: + sys.stdout.write(output) + print("RUNMEM " + json.dumps({ + "limit_gib": args.limit_gib, + "peak_bytes": peak, + "peak_gib": round(peak / (1024 ** 3), 3), + "swap_peak_bytes": swap_peak, + "wall_s": round(time.monotonic() - started, 2), + "exit": process.returncode, + "cmd": command, + "method": "cgroup-v2-MemoryMax", + })) + return process.returncode + + +if __name__ == "__main__": + raise SystemExit(main())