Blame for src/sdq_moe.cpp

// 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 <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <condition_variable>
#include <cstdio>
#include <cstring>
#include <mutex>
#include <utility>
#include <string>
#include <vector>

namespace sdq {

// ---------------------------------------------------------------- 效能計時探針
// 用途：找出「自訂 MoE op」裡面時間到底花在哪裡。開 SDQ_PROFILE_TIME=1 才啟用，
// 關閉時是零成本（幾個 if constexpr 級別的布林檢查）。
struct moe_prof {
    std::atomic<uint64_t> op_ns{0};        // 整個 op
    std::atomic<uint64_t> page_ns{0};       // probe / reserve_slot（含鎖）
    std::atomic<uint64_t> setup_ns{0};      // op 前置：unpin_all + 建 tok_of_expert + 前兩道 barrier
    std::atomic<uint64_t> bar_ns{0};        // 所有 barrier 的等待時間
    std::atomic<uint64_t> read_ns{0};       // read_expert_span / acquire 的 I/O
    std::atomic<uint64_t> fwd_ns{0};        // expert_forward 合計
    std::atomic<uint64_t> qact_ns{0};       // 其中：量化 activation
    std::atomic<uint64_t> dot_ns{0};        // 其中：gate/up/down 三次點積
    std::atomic<uint64_t> act_ns{0};        // 其中：SiLU + 加權
    std::atomic<uint64_t> calls{0};
    std::atomic<uint64_t> 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::nanoseconds>(
        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<std::mutex> 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<int> count_{0};
    std::atomic<unsigned> 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 呼叫，空白分欄）：
//   <layer> <n_tok> <t0:e0,e1,...> <t1:...> ...
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<std::mutex> 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<std::vector<std::pair<int32_t, int32_t>>> tok_of_expert;
    // 每個 token 用到的 (expert, topk 位置) —— token 切分時用（避免輸出競爭）
    std::vector<std::vector<std::pair<int32_t, int32_t>>> exp_of_tok;
    std::vector<int32_t> need;             // 本層要用的 expert 編號
    std::vector<int32_t> need_sorted;      // 依編號排序（合併讀取用）
    std::vector<int32_t> need_slot;        // 對應的 arena 槽位
    std::vector<std::vector<float>> partial; // mode A：每個執行緒一份累加緩衝
    std::vector<int32_t> slot_of_expert;      // [expert] -> arena 槽位（-1 = 沒分頁）
    std::vector<int32_t> need_slot_of_expert; // 合併讀取時暫存槽位
    std::vector<int32_t> 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<int> rowpar_chunk_up{0};
    std::atomic<int> rowpar_chunk_dn{0};
    std::vector<float> rowpar_h;     // row-parallel 的 up 輸出（跨執行緒共用）
    std::vector<uint8_t> rowpar_hq;  // h 的 Q8_K 量化結果（跨執行緒共用）
};

static std::vector<layer_rec *> g_recs;

static layer_rec * rec_for(int il) {
    static std::mutex m;
    std::lock_guard<std::mutex> 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<std::pair<int32_t, int32_t>> & toks,
                           const float * wts, size_t wts_nb0, size_t wts_nb1, int /*n_used*/,
                           float swiglu_limit,
                           float * out, std::vector<uint8_t> * 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<float> gate((size_t) gate_rows);
    std::vector<float> upv((size_t) up_rows);
    std::vector<float> 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<std::pair<int32_t, int32_t>> & toks,
                                 const float * wts, size_t wts_nb0, size_t wts_nb1,
                                 float * out, std::vector<uint8_t> * 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<std::mutex> 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<float> h((size_t) std::max<int64_t>(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<int64_t> up_off((size_t) n_need + 1, 0);
    std::vector<int32_t> tok_t, tok_k, tok_e;
    std::vector<int64_t> 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<ed_info> 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<uint8_t> xq_local((size_t) n_tok * xq_sz);
    std::vector<uint8_t> 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<int64_t>(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<int64_t>(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<uint64_t> g_ops_rowpar{0};
static std::atomic<uint64_t> 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<std::mutex> 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<int> 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<float> & mine = R->partial[(size_t) ith];
        mine.assign((size_t) out_elems, 0.0f);
        std::vector<uint8_t> 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<uint8_t> 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<std::pair<int32_t, int32_t>> 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<uint64_t> & 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<std::pair<int32_t, int32_t>> toks;
    std::vector<float> wts((size_t) n_tok, 1.0f);
    for (int t = 0; t < n_tok; t++) {
        toks.emplace_back(t, 0);
    }
    std::vector<float> out((size_t) n_embd * n_tok, 0.0f);
    std::vector<uint8_t> 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