Simply Patrick Hopes can always go up, tears can only come down.

本文由 AI Agent（Claude）代筆撰寫，文中的「我」指的是 AI Agent。Patrick 只有在文章最後做過潤飾調整。

上一篇看完了所有 Transformer 原語，這一篇要看的是把它們編排成一次完整 forward pass 的核心檔案：runner.rs。

整個專案最值得讀的就是這個檔案——大約 170 行的 Rust，把 Llama 的整個推論流程鋪展開來，每一步都看得清清楚楚。

概念一：autoregressive 推論的本質

LLM 推論不是「一次處理整段文字」，而是「一次處理一個 token」：

prompt: "The capital of France"
        → token ids [464, 3139, 286, 4881]

prefill (處理 prompt):
  forward(464)  → logits     (我們不用，但要建 KV cache)
  forward(3139) → logits
  forward(286)  → logits
  forward(4881) → logits     ← 用這個 logits 抽下一個 token

decode (生成):
  sampler(logits) → " is"
  forward(" is")  → logits
  sampler(logits) → " Paris"
  forward(" Paris") → logits
  ... 一直到 EOS

每次 forward(token) 是一個 step。step 之間透過 KV cache 累積上下文。

概念二：KV Cache —— 為什麼 attention 要 cache 過去？

Attention 的數學是這樣的：

在 autoregressive 推論時，第 t 個 token 的 Q 只需要和前 t 個 token 的 K/V 做運算（因為未來的 token 還沒出現）。如果每次 forward 都重新計算所有過去的 K/V，那 prefill 階段是 $O(n^2)$，decode 階段每次都是 $O(n)$ 的重做——太浪費了。

KV cache 的作法：每算過一次某個 position 的 K/V 就存起來，下次直接用。這樣每次 forward 只需要算當前這個 token 的 Q/K/V，K 和 V 寫進 cache、Q 拿來和整個 cache 點積。

概念三：Context 和 KV Cache 的關係

很多人聽到「context length 2048」會直覺想成「模型有個地方存著最近 2048 個 token」。這不正確——模型實際上只存了 2048 個 token 被各層 attention 算出來的 K 和 V。Context 的這個容量限制就是 KV cache 的容量限制。

關係可以這樣理解：

context window = 「我能看見多遠的過去」的能力上限
KV cache       = 這個能力實際的儲存形式
n_ctx          = KV cache 在 position 維度上的大小

當你看到 n_ctx = 2048，背後真正配出來的記憶體就是：

kcache: [n_layer][n_ctx × kv_dim] f32
vcache: [n_layer][n_ctx × kv_dim] f32

對 TinyLlama (n_layer=22, kv_dim=256) 來說，這是 22 × 2 × 2048 × 256 × 4 = 92 MB。對 Llama-2 7B (n_layer=32, kv_dim=4096) 來說，是 32 × 2 × 2048 × 4096 × 4 = 2 GB——KV cache 比模型本身的量化權重還大。

Token 一旦進入 cache，就再也找不回來了。Context window 滿了之後，要嘛截斷舊 token、要嘛丟掉新 token，沒有第三個選擇——因為原始 token id 早就被 attention 投影成 K/V 之後丟掉了。

概念四：為什麼不存 token 本身？為什麼不存 embedding？為什麼是 K/V？

這個問題的答案會徹底解釋 KV cache 的設計。讓我們從「最樸素」的方案開始往下思考：

方案 A：只存 token id

「我把過去的 token id 都存起來，下次推論時整段重跑就好。」

問題：這個方案下，每次 forward 都要從 token id 開始，重算 embedding、重算 22 層 transformer。對第 t 個 token 來說是 $O(t \times \text{layers} \times \text{matmul})$ 的工作，整個 prompt 的 prefill 是 $O(n^2)$，decode 累積下來也是 $O(n^2)$。完全沒有 cache 的意義。

方案 B：存 embedding（每個 token 對應一個 n_embd 向量）

「那我存 token embedding，這樣可以跳過 embedding lookup。」

問題：embedding 只是 forward pass 的輸入。從 embedding 到「attention 真正用的東西」，還要走完整層的 RMSNorm + Q/K/V 投影。也就是說每層的 attention 看到的是「自己這一層的輸入經過 W_k 和 W_v 投影後的 K/V」，不是 token embedding。如果只存 token embedding：

• 每次 forward 都要重新跑 22 層的 RMSNorm、重新算 K/V 投影。
• 但這些工作的結果和 token 位置 t 無關——每次重算都會得到同一個 K_t 和 V_t。
• 純粹的浪費。

更糟的是：第 1 層的輸入是 token embedding，但第 2 層的輸入是「第 1 層的輸出」。所以你不能用「token embedding」這個單一概念來代表「每一層 attention 需要的東西」。每層之間隔著 attention + FFN 的非線性轉換，「embedding」這個詞只在最開始那一層有意義。

方案 C：存每層的 hidden state（每層都有自己的 [n_layer][n_ctx][n_embd]）

「那我存每層的輸出 hidden state？」

問題：你存了 hidden state，但 attention 真正需要的是 K = W_k(hidden_state) 和 V = W_v(hidden_state)。每次 forward 還是要重算 K/V 投影（一個 [n_embd, kv_dim] 的 matvec）。白存了，因為投影沒省到。

而且這個方案佔的記憶體比 K/V cache 大：hidden state 是 n_embd 寬，K/V 在 GQA 下只有 kv_dim = head_dim × n_head_kv 寬（TinyLlama 是 256，比 n_embd=2048 小 8 倍）。

方案 D：直接存 K 和 V（最終答案）

K 和 V 才是 attention 真正消費的東西。$\text{attn}(Q, K, V)$ 的公式裡沒有 hidden state、沒有 embedding、沒有 token id——只有 Q、K、V。把過去的 K 和 V 存起來，attention 就完整了。

而且 K/V 有兩個迷人的性質：

1. 它們和 Q 是解耦的：K_t 和 V_t 一旦算出來就和「未來會用什麼 Q 來查它」無關。所以可以放心 cache。
2. 它們不需要被未來修改：因為 LLM 是 causal——位置 t 的 K/V 只被 position ≥ t 的 Q 查詢，但這些查詢不會回頭改 K_t/V_t。一旦寫進 cache 就是 immutable 的。

這就是為什麼 KV cache 是 transformer 推論的「最佳化終點」——它剛好存了 attention 需要的最少資訊，再少就會破壞語義，再多就是浪費。

從「資訊保留鏈」看這件事

換個角度：模型做 forward pass 是一條把「token id」逐步轉成「下一個 token 機率分佈」的資訊流。中途有很多中間表示：

token id → embedding → layer 1 hidden → layer 1 K/V → ...
                                     → layer 2 hidden → layer 2 K/V → ...
                                     ...
                                     → final hidden → logits

每一個中間表示都包含了關於這個 token 的某種資訊。對 attention 來說，這條鏈上唯一一個「跨 token 互動」的點就是它把 K/V 拿來算內積。其他每一步（RMSNorm、Q/K/V 投影、FFN…）都是 pure-position 的——只處理當前這個 token 的資料。

所以：

• 那些「pure-position」的步驟不需要 cache，因為每個 token 的計算和其他 token 無關。
• 那個「跨 token 互動」的步驟必須 cache，因為它要看到所有過去 token 的資訊。

而 attention 跨 token 看到的東西就是 K 和 V。所以這就是該被 cache 的東西。

一個簡單的比喻

把 LLM 比喻成一個圖書館：

• Token id 是書的索書號。
• Embedding 是書的封面（提供入口）。
• 每層 hidden state 是書頁的內容（給讀者看）。
• K 和 V 是給其他書「互相引用」用的索引卡。

Attention 是書與書之間的對話。書本身（hidden state）很厚，但對話只透過索引卡（K/V）進行。把對話用過的卡片留在桌上，下次新書進來就可以直接和它們對話——不需要把整本書重新搬出來看。

概念五：為什麼 KV cache 必須是 per-layer，而不只是 per-token？

注意 Runner 的 cache 是 Vec<Vec<f32>>，第一個維度是 layer，第二個才是 token position：

kcache: Vec<Vec<f32>>,   // [n_layer][n_ctx × kv_dim]
vcache: Vec<Vec<f32>>,

對 22 層的 TinyLlama 來說，這是 22 份獨立的 cache。為什麼不能共用同一份？為什麼一個 token 不只對應一組 K/V，而是對應 n_layer 組？

因為「同一個 token」在每一層的 K/V 是完全不同的東西

關鍵在 forward pass 的結構。回頭看看每一層 attention 在算什麼：

layer 1: x¹  = embedding(token)
         k¹  = W_k¹(rmsnorm(x¹))     ← 第 1 層的 K
         v¹  = W_v¹(rmsnorm(x¹))     ← 第 1 層的 V
         x²  = x¹ + attn(...) + ffn(...)

layer 2: k²  = W_k²(rmsnorm(x²))     ← 第 2 層的 K（input 不同、權重不同）
         v²  = W_v²(rmsnorm(x²))
         x³  = x² + attn(...) + ffn(...)

layer 3: k³  = W_k³(rmsnorm(x³))     ← 第 3 層的 K
         ...

每一層的 K/V 同時受兩件事決定：

1. 這一層的 input hidden state —— 上一層 attention + FFN 完成後傳下來的東西，每層都不同。
2. 這一層自己的權重 W_k^l / W_v^l —— 每層獨立訓練的不同矩陣。

也就是說 layer 1 的 K 是「token 在剛 embed 時的樣子被 W_k¹ 看出來的特徵」，layer 5 的 K 是「token 經過 4 層 attention + FFN 整合上下文後的樣子被 W_k⁵ 看出來的特徵」。這兩個 K 既不是同一個東西、也不能互相替代。

換個角度：每一層都有自己的 attention

Transformer 的設計本身就讓每層 attention 是獨立的計算單元。每層問「我這層的 query 和我這層 cache 裡的 key 像不像」，答案決定我這層怎麼把 V 加總起來。底下幾層通常學「文法、近距離 token 關係」，中層學「短語結構」，高層學「語意、長距離依賴」——這些功能只有在它們各自的 K/V 空間裡才有意義。

如果你硬要說「我只想存一份 K/V 給所有層用」，那等於是在說「所有層都做同一件事」——這就退化成一個非常淺的模型了。Transformer 的深度價值就在於每層做不同的轉換，而每層的 K/V 是這個轉換的 fingerprint。

從層之間的依賴鏈看為什麼不能省

更技術一點：layer l 的 K/V 是 layer l-1 的輸出的函式。整個 forward pass 是這樣的串聯：

        x¹              x²              x³
embed → ─→ attn¹+ffn¹ ─→ ─→ attn²+ffn² ─→ ─→ ...
        │                │                │
        └─→ k¹, v¹       └─→ k², v²       └─→ k³, v³

每一層的 K/V 都「只能在這一層用」——下一層的 attention 不會去看 k¹，因為它要看的是經過這一層 transform 過的世界。所以這 22 份 K/V 是 22 個獨立的記憶體，不是一份共用的。

一個 token 在 cache 裡到底佔多少？

把上面所有事情串起來，一個 token 進入 cache 後實際佔的記憶體是：

single token → n_layer × 2 × kv_dim × sizeof(f32)
              ↑          ↑   ↑
              層數       K + V

對 TinyLlama (n_layer=22, kv_dim=256, f32) 來說：

22 × 2 × 256 × 4 = 45 KB per token

對 Llama-2 7B (n_layer=32, kv_dim=4096) 來說：

32 × 2 × 4096 × 4 = 1 MB per token

這就是為什麼長 context 推論這麼吃記憶體——不只是 token 數量乘上一個小的數字，而是 token 數量乘上 n_layer × 2 × kv_dim。一個 8K context 的 7B 模型光 KV cache 就要 8 GB，比模型本身還大。

如果 KV cache 可以「per-token only」（共用所有層），同樣的 8K context 只需要 32 MB——但那樣的模型已經不是 Transformer 了。這個「per-layer」的代價，就是 Transformer 之所以是 Transformer 的代價。

「per-layer × per-token」 是 cache 的最小完備形式

可以這樣總結：要讓 attention 在每一層都能正確算出來，你需要的最少資訊是：

把這三個維度切片乘起來，就是 [n_layer][n_ctx][2][kv_dim] 的 4D 張量——這就是 KV cache 的最小完備形狀。再省任何一個維度都會破壞 attention 的語義；多存任何一個維度都是浪費（因為其他資訊都可以重新計算）。

回到我的 Rust 程式碼：

kcache: Vec<Vec<f32>>,   // [n_layer][n_ctx × kv_dim]
vcache: Vec<Vec<f32>>,   // [n_layer][n_ctx × kv_dim]

外層 Vec 是 layer 維度、內層的扁平陣列裡 n_ctx × kv_dim 是「token position × K/V 寬度」。這兩個 Vec<Vec<f32>> 加起來就是 attention 完整需要的最小狀態，一塊不多、一塊不少。

Runner struct 的記憶體佈局

pub struct Runner<'a, 'm> {
    model: &'m LlamaModel<'a>,
    rope_style: RopeStyle,
    kcache: Vec<Vec<f32>>,   // [n_layer][n_ctx * kv_dim]
    vcache: Vec<Vec<f32>>,
    pub pos: usize,

    // 預先配好的 scratch buffer
    x: Vec<f32>, xb: Vec<f32>, xb2: Vec<f32>,
    hb: Vec<f32>, hb2: Vec<f32>,
    q: Vec<f32>, att: Vec<f32>,
    logits: Vec<f32>,
}

幾個重要的設計決定：

兩個生命週期 'a 和 'm

• 'a 是 mmap 的生命週期。
• 'm 是 LlamaModel 的生命週期，且必須 'm: 'a（model 不能比 mmap 活得更久）。

Runner 借用 &'m LlamaModel<'a>，自己沒有持有任何模型資料，所以這個結構是 cheap to construct——所有大型 buffer 都是 scratch 用途。

KV cache 是 Vec<Vec<f32>> 而不是 Vec<f32>

每層一個獨立的 Vec，而不是把所有層拼在同一個 Vec 裡。這是為了：

1. 生命週期分離：不同層的 cache 可以獨立管理（雖然目前我沒這麼做）。
2. 記憶體大小靈活：理論上不同層可以有不同的 KV dim（例如某些 MoE 變體），雖然 Llama 沒有這個需求。

代價是多一層間接（access 要走 kcache[l][...]）。對 LLM 推論來說 negligible。

一堆 scratch buffer 一次配好

x: vec![0.0; n_embd],     // 主 hidden state
xb: vec![0.0; n_embd],    // 暫存 a (attention/FFN 內部)
xb2: vec![0.0; n_embd],   // 暫存 b (殘差用)
hb: vec![0.0; n_ff],      // FFN gate 的中間結果
hb2: vec![0.0; n_ff],     // FFN up 的中間結果
q: vec![0.0; n_embd],     // 當前 token 的 Q
att: vec![0.0; n_head * n_ctx],  // attention scores
logits: vec![0.0; vocab],  // 輸出 logits

這些 buffer 在 new() 一次配好，整個推論過程不再 allocate。每次 forward() 進來都直接覆寫這些 buffer。沒有 GC、沒有 alloc 壓力。

對比一下 PyTorch/Candle 的做法：每個 op 回傳新 Tensor，靠 reference counting 回收。對訓練很合理（autograd 需要保留中間值），但對只做 inference 的 runner 來說，預配 + 覆寫是更精簡、更可預測的做法。

演算法核心：forward 的整體結構

forward(token) 的結構是這樣的：

每個 block 的內部精細化看起來會更清楚：

// attention norm
rmsnorm(&mut self.xb, &self.x, &layer.attn_norm, cfg.rms_eps);

// qkv projections
matvec(&mut self.q, &layer.wq, &self.xb);
let krow = &mut kc[pos * kv_dim..(pos + 1) * kv_dim];
let vrow = &mut vc[pos * kv_dim..(pos + 1) * kv_dim];
matvec(krow, &layer.wk, &self.xb);    // K 直接寫進 cache
matvec(vrow, &layer.wv, &self.xb);    // V 也是

// RoPE
apply_rope(&mut self.q, pos, head_dim, ...);
apply_rope(krow, pos, head_dim, ...);

// attention
for h in 0..cfg.n_head {
    let kv_head = h / gqa;          // GQA 對應
    // 算 attention scores
    // softmax
    // 加權 V
}

// 輸出投影 + 殘差
matvec(&mut self.xb2, &layer.wo, &self.xb);
add_inplace(&mut self.x, &self.xb2);

// FFN: x = x + Wdown(silu(Wgate(norm)) * Wup(norm))
rmsnorm(&mut self.xb, &self.x, &layer.ffn_norm, cfg.rms_eps);
matvec(&mut self.hb, &layer.w_gate, &self.xb);
matvec(&mut self.hb2, &layer.w_up, &self.xb);
for i in 0..self.hb.len() {
    self.hb[i] = silu(self.hb[i]) * self.hb2[i];
}
matvec(&mut self.xb2, &layer.w_down, &self.hb);
add_inplace(&mut self.x, &self.xb2);

演算法核心一：K/V 直接寫進 cache

let krow = &mut kc[pos * kv_dim..(pos + 1) * kv_dim];
let vrow = &mut vc[pos * kv_dim..(pos + 1) * kv_dim];
matvec(krow, &layer.wk, &self.xb);
matvec(vrow, &layer.wv, &self.xb);

這四行有個非常重要的設計選擇：K/V 計算的輸出 buffer 直接是 cache 的某一行，而不是先算到 scratch buffer 再 copy 進 cache。這省下了一次 n_embd × 4 bytes 的記憶體拷貝。

matvec 的簽章 fn matvec(out: &mut [f32], ...) 接受任何 mutable slice，cache 的 slice 完全可以塞進去。Rust 的借用檢查器會驗證 kc 和 xb、q 等其他 buffer 不會 alias。

演算法核心二：RoPE 的兩階段套用

apply_rope(&mut self.q, pos, head_dim, ...);
apply_rope(krow, pos, head_dim, ...);

注意這裡 RoPE 是套在 Q 和當前的 K 上，不是套在 cache 裡所有 K 上。為什麼？

因為 RoPE 是「絕對位置」資訊（每個 K_t 套用 t 對應的 RoPE），而每個 K_t 在它被計算的那個 forward call 裡套用一次就夠了。等到後面要做 attention 時，cache 裡的 K_t 已經帶著 RoPE，不必再做。

V 不需要 RoPE，因為 attention 對 V 的處理是線性 weighted sum，position 資訊在 Q·K 那裡就已經注入。

演算法核心三：GQA 的 head 對應

for h in 0..cfg.n_head {
    let kv_head = h / gqa;          // GQA 對應
    let q_off = h * head_dim;
    let q = &self.q[q_off..q_off + head_dim];

    let att = &mut self.att[h * cfg.n_ctx..h * cfg.n_ctx + (pos + 1)];
    for (t, score) in att.iter_mut().enumerate() {
        let k_off = t * kv_dim + kv_head * head_dim;
        let k = &self.kcache[l][k_off..k_off + head_dim];
        let mut s = 0.0f32;
        for i in 0..head_dim {
            s += q[i] * k[i];
        }
        *score = s * scale;
    }
    softmax(att);
    // ... 加權 V
}

GQA 的 trick：在 vanilla MHA（multi-head attention）裡，每個 query head 對應一個獨立的 K/V head。但 GQA 把 query head 分組，每組共用同一個 K/V head：

n_head = 32, n_head_kv = 4 → gqa = 8
query head 0..7   → K/V head 0
query head 8..15  → K/V head 1
query head 16..23 → K/V head 2
query head 24..31 → K/V head 3

kv_head = h / gqa 就是這個對應。這個技巧把 KV cache 的大小從 n_head × head_dim × n_ctx 縮成 n_head_kv × head_dim × n_ctx，對長 context 推論很重要——KV cache 是長 context 推論的記憶體大頭。

演算法核心四：attention scores 的記憶體佈局

let att = &mut self.att[h * cfg.n_ctx..h * cfg.n_ctx + (pos + 1)];

att buffer 是 [n_head, n_ctx] 的展平。對 head h 的 attention scores 佔 att[h*n_ctx .. h*n_ctx + n_ctx]。但每次 forward 我們只用前 pos+1 個（因為只有過去到現在的 token 有 score）。

這個設計的好處是：buffer 是固定大小（為 worst case n_ctx 配好），不需要 dynamic resize。代價是有些空間沒用到——對 TinyLlama 來說 n_head * n_ctx * 4 = 32 * 2048 * 4 = 256 KB，可以接受。

Rust 用法：借用檢查器與切片

這個 forward pass 裡有大量 &mut 切片操作：

let kc = &mut self.kcache[l];
let krow = &mut kc[pos * kv_dim..(pos + 1) * kv_dim];
matvec(krow, &layer.wk, &self.xb);

注意我必須先取 &mut self.kcache[l] 存進局部變數 kc，再對 kc 切片。為什麼？因為 Rust 的借用檢查器需要看到「我們對 self.kcache 的 mutable 借用」是清晰的。

這是 Rust 寫多 mutable buffer 程式碼最常踩的坑之一。如果你寫：

matvec(&mut self.kcache[l][pos*kv_dim..(pos+1)*kv_dim], &layer.wk, &self.xb);

而其他地方還有 &self.x、&self.xb 之類的不可變借用，編譯器可能會抗議「不能同時 mutable 和 immutable 借用 self」。先把 mutable slice 提取出來，再做後續操作，是讓借用範圍縮小的標準技巧。

split_at_mut 的那種招數

更進階的場合可能要 split_at_mut——例如同時拿 K cache 和 V cache 的 mutable ref：

let (kc, vc) = (&mut self.kcache[l], &mut self.vcache[l]);
// 這個寫法借用檢查器會接受，因為 self.kcache 和 self.vcache 是不同的 field

但 kcache[l] 和 vcache[l] 是不同的 field，所以可以同時 mutable borrow。如果是同一個 Vec 的兩段切片，就需要 split_at_mut：

let (left, right) = my_vec.split_at_mut(mid);
// left 和 right 都是 &mut [T]，但編譯器知道它們不重疊

效能最佳化空間

runner.rs 是整個專案的效能熱點集中地，有大量改進空間：

1. Attention 的 head-parallel

我目前的 attention loop 是 sequential for h in 0..n_head：

for h in 0..cfg.n_head {
    // 算 attention scores、softmax、加權 V → 寫進 self.xb 的某段
}

每個 head 寫進 self.xb 的不同 slice，是 disjoint 的。可以用 rayon 的 chunks_exact_mut 平行：

self.xb.par_chunks_exact_mut(head_dim).enumerate().for_each(|(h, out)| {
    // 算這個 head 的 attention
});

但 att buffer 是共享的（self.att），平行版需要每個 thread 獨立的 attention scratch。head 數量通常很小（32-64），fork-join 的 overhead 可能吃掉收益——這個值得 benchmark 但不一定划算。

2. FlashAttention 風格的 fused attention

我目前的 attention 是「先算所有 score、再 softmax、再加權 V」三步。FlashAttention 把它們 fuse 成一個 streaming pass，記憶體占用從 O(n_ctx) 降到 O(head_dim)，且 cache 友善。但實作複雜得多——這是 candle/ggml 級別的優化。

3. KV cache 的量化

KV cache 是長 context 的記憶體大頭。對 7B 模型、2k context、F32 cache，每層 cache 是 2 * 2048 * 1024 * 4 = 16 MB，32 層 = 512 MB。如果把 cache 也量化成 Q8（1 byte/element），可以降到 128 MB。llama.cpp 已經支援這個（-ctk q8_0 -ctv q8_0）。

4. Prefill batching

目前 prefill 是 sequential——一個 token 一個 forward。但 prompt 的所有 token 都已知，可以拼成一個矩陣做 batched forward：

seq forward: O(L * (matvec for each layer))
batch forward: O(matmul of [L, n_embd] for each layer)

GEMM 的 cache locality 比 GEMV 好得多，prefill 速度可以提升 5-10×。但要重新設計 attention（causal mask 變成必須的）、要修改 KV cache 寫入方式（一次寫 L 個 token）。這是中量級的改動。

5. Speculative decoding

更進階的優化：用一個小的 “draft model” 一次猜 K 個 token，用主模型 verify。如果猜對 K 個就只花了一次 forward 的時間做 K 個 token。這是 inference latency 的革命性技術，但需要兩個模型協同。

6. Continuous batching

如果 runner 要 serve 多個請求，可以做 continuous batching——不同請求的 token 進入同一個 batch，動態加減。但這要求 KV cache 是「per-request」的，記憶體管理更複雜。vLLM 是這條路徑的代表。

一個有趣的權衡：為什麼不寫得更 functional？

我可以把 forward 寫得更 functional，例如把每個 layer 寫成一個閉包、用 fold 串起來。但目前這個 indicative、命令式的寫法反而更容易讀——你能一眼看出每一步在改哪個 buffer、用哪個權重。

LLM 推論的程式碼有個特殊性：90% 的時間在做矩陣運算，10% 在做 control flow。寫得太 functional 會把 control flow 的成本顯露在前景，反而蓋住了真正重要的計算。

總結：runner.rs 的角色

• 概念上：把 token id 映射到 logits 的單一函式，內部管理 KV cache 和層級遞進。
• 實作上：fixed-size scratch buffers + sequential layer loop + manual KV slice manipulation。
• 設計上：把所有複雜性集中在這個檔案裡，其他檔案（ops、tensor）保持簡單。

下一篇講 tokenizer.rs，把 byte 流變成 token id 的細節。

系列文章：

• tiny-llm-runner 介紹
• (1) config ・ (2) dequant ・ (3) tensor ・ (4) model ・ (5) ops
• (6) runner.rs（本篇）

本文由 AI Agent（Claude）代筆撰寫，文中的「我」指的是 AI Agent。Patrick 只有在文章最後做過潤飾調整。

上一篇看完了 LlamaModel 怎麼把張量組起來，這一篇要看的是 Transformer 真正的「動詞」們：ops.rs。

這個檔案只有 100 行，但它就是整個 Transformer 的所有原語。一個現代 LLM 的所有計算，最終都會被分解成這六種運算的組合。

樂高積木一：RMSNorm

pub fn rmsnorm(out: &mut [f32], x: &[f32], w: &[f32], eps: f32) {
    let mut ss = 0.0f64;
    for &v in x {
        ss += v as f64 * v as f64;
    }
    ss /= x.len() as f64;
    ss += eps as f64;
    let scale = (1.0 / ss.sqrt()) as f32;
    for i in 0..x.len() {
        out[i] = w[i] * (x[i] * scale);
    }
}

算法：

也就是「先把 x 除以它自己的 RMS，再用 w 逐元素 scale」。RMSNorm 是 LayerNorm 的簡化版（沒有 mean-shift、沒有 bias），Llama 用它取代 LayerNorm 是 2019-2020 年的設計趨勢。

為什麼用 f64 累加？

注意 ss 是 f64，不是 f32。這不是隨便寫的：累加大量 f32 值容易產生「精度吞噬」（catastrophic cancellation）。一個 4096 維的 hidden state，每個 f32 大約 1e-1 量級，平方後是 1e-2，4096 個累加起來大概 40-100。在這個量級下，每個新加進來的 1e-2 增量在 f32 上會有相對誤差 ~1e-7 × 100 = 1e-5，4096 次累積下來會放大。

llama.cpp 在這條路徑用了 f64 累加器，我必須跟它對齊——否則 RMSNorm 的輸出會和它差千分之一，雖然每層只差千分之一，但 22 層累積下來輸出 token 就會不一樣。這是「對齊既有實作」的另一個例子。

樂高積木二：matvec —— 用 rayon 平行化

pub fn matvec(out: &mut [f32], w: &TensorView<'_>, x: &[f32]) {
    out.par_iter_mut().enumerate().for_each(|(i, o)| {
        *o = w.dot_row(i, x);
    });
}

這是整個專案的效能瓶頸所在，也是最簡潔的一段程式碼。讓我們拆解這幾行做了什麼：

1. par_iter_mut() 來自 rayon，把 &mut [f32] 變成一個平行 iterator。
2. enumerate() 加上 row index。
3. for_each(|(i, o)| ...) 對每個 row 平行執行 closure。

每個 row 是獨立的（out[i] 只被當前 closure 寫入），所以不需要任何鎖。rayon 的 work-stealing 排程器會自動把 row 切成 chunk 分給可用的執行緒。

為什麼 row-parallel 而不是 element-parallel？

如果改成 element-parallel（每個輸出元素一個 task），task 太細會被 rayon 的 task overhead 吃掉。row-parallel 的 granularity 剛好——每個 task 是一個完整的 dot product（4096 個乘加），夠粗到值得啟動執行緒，又夠細到能平均負載。

Rust 的 par_iter_mut 設計上避免了一個很容易踩的雷：你不能讓兩個 thread 同時寫 out 的同一個位置。par_iter_mut 的型別簽章保證每個 closure 拿到的是 &mut f32（單一元素的 mutable ref），編譯期就排除了 race condition。

樂高積木三：softmax —— max-shift 防止 overflow

pub fn softmax(x: &mut [f32]) {
    let mut max = f32::NEG_INFINITY;
    for &v in x.iter() {
        if v > max { max = v; }
    }
    let mut sum = 0.0f32;
    for v in x.iter_mut() {
        *v = (*v - max).exp();
        sum += *v;
    }
    let inv = 1.0 / sum;
    for v in x.iter_mut() {
        *v *= inv;
    }
}

數學上，softmax 是：

但 e^{x} 對大 x 會 overflow（f32 的 exp 在 x > ~88 就溢位）。標準做法是先減 max：

數學上等價（分子分母同乘 $e^{-\max(x)}$），但所有 exp 的輸入都在 $(-\infty, 0]$，永遠不會 overflow。這是寫 softmax 的不可省的步驟——少了它，遇到極端 logits 就會 NaN。

樂高積木四：RoPE —— 相對位置的旋轉

這是最微妙的一個運算。RoPE（Rotary Position Embedding）是 Llama 用來把位置資訊塞進 attention 的方法，數學上是把每個 head 的 Q/K 向量看成「複數對」並旋轉一個和位置相關的角度。

pub fn apply_rope(
    vec: &mut [f32], pos: usize, head_dim: usize,
    rot_dim: usize, base: f32, style: RopeStyle,
) {
    let half = rot_dim / 2;
    let n_heads = vec.len() / head_dim;
    for h in 0..n_heads {
        let off = h * head_dim;
        for i in 0..half {
            let freq = base.powf(-((2 * i) as f32) / rot_dim as f32);
            let theta = pos as f32 * freq;
            let (sin, cos) = theta.sin_cos();
            let (i0, i1) = match style {
                RopeStyle::Llama => (2 * i, 2 * i + 1),
                RopeStyle::Neox  => (i, i + half),
            };
            let v0 = vec[off + i0];
            let v1 = vec[off + i1];
            vec[off + i0] = v0 * cos - v1 * sin;
            vec[off + i1] = v0 * sin + v1 * cos;
        }
    }
}

核心想法：每對 (v[i0], v[i1]) 是一個 2D 向量，被旋轉一個角度 $\theta = \text{pos} \cdot \text{freq}_i$。不同的 i 用不同的 freq——i 越小 freq 越大（位置感越強）、i 越大 freq 越小（接近不變）。

頻率公式：

base 通常是 10000，rot_dim 是 head_dim（或某個比例）。i = 0 時 freq = 1（每移動一個 position 旋轉 1 弧度），i = rot_dim/2 - 1 時 freq ≈ 1/10000（要 6000 個 position 才旋轉 1 弧度）。

這個多尺度的旋轉讓 attention 既能感知近距離的細微位置差異、也能感知遠距離的大致位置。

為什麼有兩種 style？

let (i0, i1) = match style {
    RopeStyle::Llama => (2 * i, 2 * i + 1),         // 鄰接成對
    RopeStyle::Neox  => (i, i + half),              // 上下半成對
};

這兩個是完全不同的「複數對」配對方式：

• Llama：把 v[0], v[1] 當一對、v[2], v[3] 當一對……（鄰接 pair）
• Neox：把 v[0], v[half] 當一對、v[1], v[half+1] 當一對……（上下 split）

兩者旋轉的數學是一樣的，但配對方式不同。為什麼兩者並存？

歷史上，HuggingFace 的 transformers 採用 NeoX 風格。但 llama.cpp 早期的 convert.py 在轉檔時會把 Q/K 矩陣的每對 row permute 成鄰接格式，這樣就可以用 Llama 風格做 RoPE，記憶體 access 更線性。

新版的 convert_hf_to_gguf.py 不做 permute，所以你必須用 NeoX 風格。搞錯了不會 crash 但模型會講胡話——這是我親自踩過的雷。

sin_cos() 是 f32::sin_cos()——同時算 sin 和 cos，硬體上比分別算快一倍。

樂高積木五：SiLU —— SwiGLU 的活化函數

#[inline]
pub fn silu(x: f32) -> f32 {
    x / (1.0 + (-x).exp())
}

數學定義：

$\sigma$ 是 sigmoid。SiLU 也叫 Swish。Llama 的 FFN 用的是 SwiGLU（Swish-Gated Linear Unit）：

⊙ 是逐元素相乘。也就是先把 x 過 gate 和 up 兩個投影，gate 過 SiLU 之後逐元素乘上 up，再過 down 投影。比起傳統 ReLU FFN，SwiGLU 多了一個 gate 投影，但實驗顯示效果好得多。

#[inline] 是給編譯器的提示，因為這個函式會在 inner loop 被呼叫成千上萬次，inline 進去能避免函式呼叫的 overhead。

樂高積木六：add_inplace —— 殘差連接

pub fn add_inplace(a: &mut [f32], b: &[f32]) {
    for (x, y) in a.iter_mut().zip(b.iter()) {
        *x += *y;
    }
}

最樸素的一個原語：把 b 加進 a。Transformer 的每個 sublayer 都會做：

x = x + Attention(LayerNorm(x))
x = x + FFN(LayerNorm(x))

那兩個加號就是 add_inplace。

注意我用 iter_mut().zip(iter()) 而不是用 index。這個寫法的好處是：

1. 沒有越界檢查——iterator 自動知道何時停止。
2. 編譯器更容易自動向量化——iterator 模式對 LLVM 友善。
3. 零成本抽象——這個寫法和手寫 raw loop 編譯出來幾乎一樣的組合語言。

Rust 用法：#[derive] 的小幫手

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RopeStyle {
    Llama,
    Neox,
}

Copy 讓我可以把 RopeStyle 當值傳遞，不需要 &。Eq 讓我可以在 match 比較。Debug 讓 eprintln!("{:?}", style) 可以印出來除錯。不到一秒的工作，省下 10 行 boilerplate——這是 Rust derive 機制的常規用法。

效能最佳化空間

ops.rs 是另一個重要的最佳化戰場。每個函式都有改進空間：

1. RMSNorm 的 SIMD

那個 for &v in x { ss += v * v; } 是個 reduction，可以用 SIMD 平行化：

use std::simd::f32x8;
let mut acc = f32x8::splat(0.0);
for chunk in x.chunks_exact(8) {
    let v = f32x8::from_slice(chunk);
    acc += v * v;
}
let ss = acc.reduce_sum() as f64;

f64 vs f32 累加器的精度問題在 SIMD 環境下變得有趣——AVX-512 有 f64 SIMD，AVX2 沒有。最務實的做法是維持 f32 SIMD reduction，因為實作簡單且誤差可接受（畢竟 4096 個 ~1e-2 量級的數，誤差不會大到讓 token 不一致）。

2. matvec 的 GEMV → GEMM 升級

我目前的 matvec 是 GEMV（matrix × vector），每次 forward pass 對每層做兩次（attention 的 Q + FFN 的 down）。但prefill 階段的 prompt 是序列 token，可以批次處理——把 N 個 token 的 hidden state 拼成 [seq_len, n_embd] 矩陣，做一次 GEMM。

GEMM 比 GEMV 快很多（典型 5-10×），因為它能 reuse 權重（每次 load 一個 row 的權重，可以對 seq_len 個輸入做計算）。但這需要重新設計 Runner——目前是單 token forward，要改成 batch forward。

3. softmax 的 SIMD exp

(*v - max).exp() 是逐元素的，可以 SIMD 化。但 SIMD exp 比較麻煩——需要用多項式逼近（典型用 Cephes 或 Padé approximant 的 SIMD 版本）。sleef crate 有現成的 SIMD math 函式，可以接上。

4. RoPE 的 sin/cos 預計算

每次 forward pass 都呼叫 theta.sin_cos()。但 freq_i 只和 i 有關，theta = pos * freq_i 也可以分解。如果預計算所有可能的 (pos, i) 的 sin/cos 並存成 table，runtime 只需要查表：

let cos_table: Vec<Vec<f32>> = (0..n_ctx).map(|pos| {
    (0..rot_dim/2).map(|i| {
        let freq = base.powf(-((2*i) as f32) / rot_dim as f32);
        (pos as f32 * freq).cos()
    }).collect()
}).collect();

對 TinyLlama 是 n_ctx × rot_dim/2 × 4 bytes = 2048 × 32 × 4 = 256 KB 的表，啟動時花一次性的時間算好，runtime 都是查表。L2 cache 就能裝下。

5. add_inplace 也可以 SIMD

雖然這是個簡單的逐元素加，但 LLVM 對 iter_mut().zip() 模式的 auto-vectorization 不一定總是會發生。明確的 f32x8::from_slice(...) + f32x8::from_slice(...) 能保證 SIMD 化。

6. 把 forward pass 內整層 fuse 起來

最大的最佳化潛力是 kernel fusion。例如 attention + softmax 可以 fuse（FlashAttention 就是這個思路）；rmsnorm + matvec 也可以 fuse 成一個 pass。但 fusion 會大幅複雜化程式碼，違背 tiny-llm-runner 的「易讀」目標——所以這是 candle/ggml 的世界，不是這裡。

一個 Rust 特有的優勢：debug 與 release 的 debug_assert

我每個函式裡都有 debug_assert_eq!(out.len(), x.len()) 之類的。Rust 的 debug_assert! 在 release 會被完全消除——既能在 dev 抓 bug，又不付出 runtime 成本。對 hot path 來說這是個重要的工具。

C/C++ 也有 assert + NDEBUG，但 Rust 把這個習慣做成了語言預設——你不需要記得 #define NDEBUG，cargo build --release 自動處理。

總結：ops.rs 的角色

• 概念上：Transformer 的所有原語（norm、matmul、softmax、RoPE、activation、add）。
• 實作上：純函式、輸出參數、明確簽章、零隱藏狀態。
• 設計上：把派發（量化型別 match）藏在 TensorView::dot_row 裡，ops.rs 自己只看到 &[f32]，乾乾淨淨。

下一篇講 runner.rs，那是把所有 ops 編排起來的指揮中心。

系列文章：

• tiny-llm-runner 介紹
• (1) config ・ (2) dequant ・ (3) tensor ・ (4) model
• (5) ops.rs（本篇）

本文由 AI Agent（Claude）代筆撰寫，文中的「我」指的是 AI Agent。Patrick 只有在文章最後做過潤飾調整。

上一篇看完了 TensorView 那層薄殼，這一篇要看的是 model.rs：把這些 view 組合成一個有結構、有層級的 LlamaModel。

GGUF 把所有張量存成一個扁平 list，每個張量靠字串名字識別。但 Llama 是一個層級結構：每一層有 9 個張量、整個模型有 N 層加上幾個全域張量。model.rs 的工作就是這個 flat list → tree 的對應。

概念一：Llama 的張量命名規則

打開任何一個 Llama 架構的 GGUF，你會看到這些張量名字：

token_embd.weight              # token embedding 表
output_norm.weight             # 最後的 RMSNorm
output.weight                  # lm_head（有時不存在 = tied embeddings）

blk.0.attn_norm.weight         # 第 0 層的 attention 前 RMSNorm
blk.0.attn_q.weight            # Q 投影
blk.0.attn_k.weight            # K 投影
blk.0.attn_v.weight            # V 投影
blk.0.attn_output.weight       # output 投影
blk.0.ffn_norm.weight          # FFN 前 RMSNorm
blk.0.ffn_gate.weight          # SwiGLU 的 gate
blk.0.ffn_up.weight            # SwiGLU 的 up
blk.0.ffn_down.weight          # SwiGLU 的 down
blk.1.attn_norm.weight
... (重複 N 次)

這個命名約定是 llama.cpp 訂的。每一層有 9 個張量，前綴是 blk.{layer_id}.，後綴對應 attention 或 FFN 的某個元件。model.rs 的工作就是按這個規則去 GGUF 裡找。

概念二：n_ff、n_layer 與 Transformer block 的關係

從上面的張量命名可以看到一個很整齊的結構：每一個 blk.{l} 是一個完整的 Transformer block，整個模型就是 n_layer 個這種 block 疊起來。但每個 block 內部到底是什麼？以及 n_ff 在這裡扮演什麼角色？

一個 Transformer block 的兩個子層

每個 block 內部有兩個 sublayer，依序執行：

1. Attention sublayer：attn_norm → Q/K/V 投影 → RoPE → attention → attn_output 投影 → 殘差。
2. FFN sublayer (SwiGLU)：ffn_norm → ffn_gate 和 ffn_up 兩個並列投影 → silu(gate) * up → ffn_down 投影 → 殘差。

attention 負責「跨 token 的資訊互換」，FFN 負責「在每個 token 內做非線性變換」。實證上這兩件事都很重要——你不能只有 attention（會少了 token 內部的資料豐富度），也不能只有 FFN（每個 token 各做各的、沒有上下文）。

n_ff 是什麼？

n_ff（也寫作 feed_forward_length 或 intermediate_size）是FFN 的中間維度。具體來說：

hidden state x        : 形狀 [n_embd]            = TinyLlama 的 2048

ffn_gate(x)  : matmul → 形狀 [n_ff]              = 5632   ↑
ffn_up(x)    : matmul → 形狀 [n_ff]              = 5632   │ FFN 內部「膨脹」的維度
silu(gate) * up       : 形狀 [n_ff]              = 5632   ↓

ffn_down(...) : matmul → 形狀 [n_embd]           = 2048   (回到 hidden 大小)

也就是說 FFN 把資料先膨脹再壓回去：

n_embd → n_ff → n_embd
 2048    5632    2048   (TinyLlama)
 4096    14336   4096   (Llama-3 8B)

膨脹後的 n_ff 維度才是 FFN 真正「思考」的空間。在這個更高維的空間裡套 SiLU 非線性，再投影回 n_embd。

為什麼 n_ff 比 n_embd 大？

理論上 FFN 沒有規定中間維度要多大。但實證上幾乎所有 Transformer 都選 n_ff ≈ 4 × n_embd（或對 SwiGLU 來說是 ≈ 8/3 × n_embd ≈ 2.7 ×，因為 SwiGLU 比傳統 FFN 多一個投影，作者要讓總參數量持平）。直覺是：

1. 更高維度提供更多「特徵組合」的空間。一層 FFN 等於是「在更寬的空間裡做一個 lookup table」，太窄就學不到複雜的模式。
2. 大部分模型容量集中在 FFN，不是 attention。Attention 的權重矩陣是 [n_embd × n_embd]（每個 Q/K/V/O 一個），FFN 是 [n_embd × n_ff]（gate/up/down 三個）。當 n_ff = 4 × n_embd 時，FFN 一個 block 大概佔了 12 × n_embd²，attention 佔 4 × n_embd²——FFN 是 attention 的 3 倍。

n_layer 和 n_ff 是兩個獨立的 scaling 維度

這兩個參數常常被一起調，但它們的角色完全不同：

兩者的權重總量都會直接影響模型大小：

總參數量 ≈ n_layer × (4·n_embd² + 3·n_embd·n_ff)
                    ↑                ↑
                 attention         FFN (SwiGLU 三個矩陣)

對 TinyLlama 1.1B：

22 × (4·2048² + 3·2048·5632)
= 22 × (16.8M + 34.6M)
= 22 × 51.4M
≈ 1.13B

換成 Llama-3 8B：

32 × (4·4096² + 3·4096·14336) ≈ 32 × (67M + 176M) ≈ 7.8B

數字對得上實際模型大小（差距是 token embedding 和 norm 那些佔的份量）。

在 GGUF 命名裡的對應

n_ff 直接決定了 FFN 三個權重矩陣的形狀：

blk.{l}.ffn_gate.weight  : [n_embd, n_ff]    → 投影到中間維度
blk.{l}.ffn_up.weight    : [n_embd, n_ff]    → 投影到中間維度
blk.{l}.ffn_down.weight  : [n_ff, n_embd]    → 投影回 hidden 維度

n_layer 則決定了這組張量會出現幾次——blk.0.* 到 blk.{n_layer-1}.*，每層獨立一份權重。所以這兩個數字的乘積（加上 attention 那部分）決定了整個模型 9 × n_layer 個 block 級權重的數量。

理解了這個關係，你再看 LlamaConfig 裡 n_layer = 22、n_ff = 5632，就能在腦子裡馬上算出「FFN 中間有 5632 維、總共有 22 個這樣的 block 串起來、每個 block 的 FFN 佔 ~34.6M 參數、整個 FFN 部分大概是 760M 參數」——這就是 Llama 模型 size 的來源。

演算法核心：建立 name → TensorInfo 的索引

第一步是把扁平 list 轉成可查的 HashMap：

let by_name: HashMap<&str, &TensorInfo> =
    tensors.iter().map(|t| (t.name.as_str(), t)).collect();

這個 one-liner 用了 Rust collection 的一個漂亮特性：collect() 會根據目標型別自動選擇收集方式。HashMap<K, V> 從 Iterator<Item = (K, V)> 收集就會建一個 hash table。

注意這裡的型別是 HashMap<&str, &TensorInfo>——key 和 value 都是借用，沒有任何 String/Vec 拷貝。整個索引大概只佔幾 KB（每個 entry 是兩個指標 + hash 雜湊），相對於 4 GB 的權重來說可以忽略。

演算法核心：tied embeddings 的容錯

let output = match by_name.get("output.weight") {
    Some(info) => TensorView::from_info(info, blob)?,
    None => token_embd,
};

這段在處理一個 LLM 的細節：有些模型會把 lm_head 和 token embedding 共用同一份權重（叫做 “tied embeddings”）。這樣可以省下 vocab_size × n_embd 的權重——TinyLlama 1.1B 來說大概是 32000 × 2048 × 4 bytes = 262 MB，省下來是非常實在的。

GGUF 在這種情況下會省略 output.weight，runner 必須自己知道要 fallback 到 token_embd.weight。注意我直接 = token_embd——因為 TensorView 是 Copy，這只是一個 32-byte 拷貝，沒有生命週期問題。如果是 Box<TensorView> 或 Rc<TensorView>，這行就會麻煩很多。

演算法核心：分批載入 norm 權重

fn load_f32_vec(by_name: &HashMap<&str, &TensorInfo>, name: &str, blob: &[u8])
    -> Result<Vec<f32>>
{
    let info = by_name.get(name)?;
    if info.tensor_type != GgmlType::F32 {
        bail!("expected {name} to be F32, got {:?}", info.tensor_type);
    }
    let elems: u64 = info.dimensions.iter().product();
    let mut out = vec![0.0f32; elems as usize];
    dequant::dequant_row_f32(&blob[start..end], &mut out);
    Ok(out)
}

注意這裡的策略：RMSNorm 的權重直接 dequant 成 Vec<f32>，但 Q/K/V/FFN 矩陣保持成 TensorView。

為什麼？因為它們的大小差距很大：

• RMSNorm 權重：n_embd 個 f32，TinyLlama 是 2048 × 4 bytes = 8 KB，整個模型 N 層 × 2 個 norm + 1 個 final = 約 360 KB。
• 注意力矩陣：[n_embd, n_embd] 個 Q4_0 元素，TinyLlama 是 2048 × 2048 × 0.5 bytes = 2 MB，整個模型大概 4 GB。

對 norm 來說，多花 360 KB 換來「forward pass 不必每次重新解碼」是極划算的。分清楚什麼東西該 eager dequant、什麼該 lazy dequant，是這個檔案的關鍵設計決定。

而且 norm 權重必須是 F32（這個 bail 不只是規範性檢查，而是事實——llama.cpp 從不量化 norm，因為 norm 的數值範圍小到不值得量化）。

Rust 用法：lifetime 在 struct 上的擴散

pub struct LayerWeights<'a> {
    pub attn_norm: Vec<f32>,
    pub wq: TensorView<'a>,
    pub wk: TensorView<'a>,
    pub wv: TensorView<'a>,
    pub wo: TensorView<'a>,
    pub ffn_norm: Vec<f32>,
    pub w_gate: TensorView<'a>,
    pub w_up: TensorView<'a>,
    pub w_down: TensorView<'a>,
}

pub struct LlamaModel<'a> {
    pub config: LlamaConfig,
    pub token_embd: TensorView<'a>,
    pub output_norm: Vec<f32>,
    pub output: TensorView<'a>,
    pub layers: Vec<LayerWeights<'a>>,
}

每個 struct 都帶 'a，因為它們都間接持有 TensorView<'a>。這個 'a 一路傳到 LlamaModel 的層級，意思是：整個 LlamaModel 不能比 mmap 活得久。

這在實務上會長這樣：

fn main() -> Result<()> {
    let mmap = unsafe { Mmap::map(&file)? };       // mmap: 'mmap
    let model = LlamaModel::load(..., &mmap[...])?; // model: LlamaModel<'mmap>
    let mut runner = Runner::new(&model, ...);     // runner 借用 &model

    // 現在 runner、model、mmap 全部都借用鏈到 mmap 上
    // 編譯器保證它們的生命週期是 mmap ⊃ model ⊃ runner
}

如果你不小心想把 model 從 main 回傳出去：

fn load_model() -> LlamaModel<'???> { ... }   // 編譯失敗

編譯器會直接拒絕——因為 'a 必須繫結到一個外部生命週期。這是 Rust 在「結構體借用資料」這個模式上的招牌。

Rust 用法：用 format! 拼字串 vs 用 const generics

for l in 0..config.n_layer {
    layers.push(LayerWeights {
        attn_norm: load_f32_vec(&by_name, &format!("blk.{l}.attn_norm.weight"), blob)?,
        wq:        view(&by_name, &format!("blk.{l}.attn_q.weight"), blob)?,
        ...
    });
}

這段每呼叫一次 format! 都會 allocate 一個 String。對 22 層的 TinyLlama 來說是 22 × 9 = 198 次 allocation，整個 load 過程大概幾百 KB——對 startup 來說微不足道。

但有個更省的寫法：用 write! macro 寫進一個複用的 buffer：

let mut buf = String::with_capacity(64);
for l in 0..config.n_layer {
    buf.clear();
    write!(buf, "blk.{l}.attn_norm.weight").unwrap();
    let attn_norm = load_f32_vec(&by_name, &buf, blob)?;
    ...
}

不過這個改動會犧牲可讀性換 0.1 ms 啟動時間，這就是過度最佳化的典型例子。我選擇可讀性。

演算法核心：embed 函式的 row-major 觀念

pub fn embed(&self, token: u32, out: &mut [f32]) {
    self.token_embd.dequant_row(token as usize, out);
}

這只有一行，但有個微妙細節：token embedding 矩陣的 row 是 n_embd、token 數量是 dim1。也就是說：

token_embd shape (GGUF dim 順序): [n_embd, vocab_size]
                                    ^^^^^^  ^^^^^^^^^^
                                    row 寬度  row 數量

我們對 token id t 的 lookup 就是「拿第 t 個 row」。所以 dequant_row(token, out) 直接把 token 的 embedding vector 寫進 out。

這個 embedding lookup 是整個 forward pass 中唯一會做 dequant 而不是 dot 的地方。原因很合理：embedding 是「直接拿值」而不是「拿值算內積」，所以 dequant 是必要的。

效能最佳化空間

model.rs 不在 hot path 上——它只在啟動時跑一次。但有幾個方向值得想：

1. 平行載入 norm 權重

目前 load_f32_vec 是序列的，22 層 × 2 個 norm + 1 個 = 45 次小 dequant。可以用 rayon 把它們平行掉：

use rayon::prelude::*;
let norms: Result<Vec<_>> = (0..n_layer).into_par_iter()
    .map(|l| load_f32_vec(&by_name, &format!("blk.{l}.attn_norm.weight"), blob))
    .collect();

但 norm 權重很小（每個 8 KB 左右），平行化的 overhead 可能比實際工作還大——不見得會更快。

2. 把 token_embd 也 eager dequant 成 f32

目前 embed 每次都會去 dequant 一個 row。對 7B 模型的 prefill 來說，假設 prompt 是 100 token、每 token 都 dequant 一次 row，這是 100 個 small dequant call。

但如果你 eager 把整個 token embedding 解壓成 f32，會佔用 vocab_size × n_embd × 4 bytes——TinyLlama 是 256 MB，但 Llama-3 8B 是 128k × 4096 × 4 = 2 GB。這個交換在小 vocab 模型划算，在大 vocab 模型不划算。一個折衷方案是 LRU cache：只 cache 最近用過的 K 個 token 的 embedding。

3. 自訂 hashmap 改善 locality

HashMap<&str, &TensorInfo> 的 default hasher 是 SipHash，安全但慢。對啟動時的 lookup 來說，可以換成 ahash 或 fxhash：

use std::collections::HashMap;
type FastMap<K, V> = HashMap<K, V, ahash::RandomState>;

不過這同樣是「啟動時間 0.5 → 0.3 秒」的微優化，相對於 mmap warmup 的時間（取決於檔案大小，可能是幾秒）算不上痛點。

4. 預先驗證所有張量都存在

目前如果某層的 attn_q.weight 缺失，會在 load 那一層時才 bail。可以在 LlamaModel::load 一開始就用一個 list 驗證所有預期張量都在：

let expected: Vec<String> = (0..n_layer).flat_map(|l| {
    [
        format!("blk.{l}.attn_norm.weight"),
        format!("blk.{l}.attn_q.weight"),
        ...
    ]
}).collect();
for name in &expected {
    if !by_name.contains_key(name.as_str()) {
        bail!("missing {name}");
    }
}

這對使用者體驗有改善——一次回報所有缺失，而不是一個一個踩雷。

一個被隱藏的設計選擇：為什麼 norm 是 Vec<f32> 而不是 &[f32]？

理論上 norm 也可以是個 view：

pub attn_norm: &'a [f32],   // 直接借用 mmap 的 bytes（如果 alignment 對齊）

這樣連 dequant 都不必做，只要把 mmap 的 bytes 重新解釋成 &[f32] 就行。但有兩個問題：

1. endianness：GGUF 規定 little-endian，現代電腦也都是 LE，但 Rust 的 align_to 不保證安全做這個轉換。
2. alignment：mmap 是 page-aligned 的，但 norm 張量的 offset 不一定 4-byte 對齊。

更安全的做法是 dequant_row_f32 走 f32::from_le_bytes，每個元素都明確做 byte 轉換。代價是一次性的解碼成本（一個 7B 模型大概 1 MB f32 norm，在啟動時花幾毫秒解掉）。

如果未來想做 zero-copy norm，可以驗 alignment 然後用 bytemuck::cast_slice，這是一個成熟的「safe transmute」套件。

總結：model.rs 的角色

• 概念上：扁平的 GGUF 張量 list 對應到結構化的 Llama 模型樹。
• 實作上：HashMap 索引 + 字串拼接 + 條件性 fallback（tied embeddings）。
• 設計上：分清楚什麼 eager dequant（norm，小且常用）、什麼 lazy dequant（大矩陣）。

下一篇講 ops.rs，那是 Transformer 的「樂高積木」。

系列文章：

• tiny-llm-runner 介紹
• (1) config.rs ・ (2) dequant.rs ・ (3) tensor.rs
• (4) model.rs（本篇）

本文由 AI Agent（Claude）代筆撰寫，文中的「我」指的是 AI Agent。Patrick 只有在文章最後做過潤飾調整。

上一篇講完 dequant.rs 的量化核心，這一篇要看的是 tensor.rs——一個只有 150 多行的小檔案，但它展示了 Rust 在系統程式設計上最有特色的能力之一：用生命週期把「資料在誰手上」這件事編譯期就講清楚。

這個檔案要解決什麼問題？

dequant.rs 提供的所有函式都是「裸的」：你給它一段 &[u8] 和一個 &[f32]，它做點積。但在 forward pass 裡，我們不想讓 runner.rs 直接看到 byte slice——那太底層、太容易出錯。我們想要一個更高階的抽象，例如：

「這是一個 [K, N] 的 Q4_0 矩陣，請對它的第 i row 和輸入向量 x 做點積。」

TensorView 就是這個抽象。

概念一：什麼是「view」？為什麼不擁有資料？

#[derive(Clone, Copy)]
pub struct TensorView<'a> {
    pub data: &'a [u8],
    pub ggml_type: GgmlType,
    pub dims: [u64; 4],
    pub n_dims: usize,
}

注意這個 'a 生命週期參數。TensorView 不是「擁有」這些 bytes，它只是「看著」這些 bytes。實際的 bytes 還在 mmap 裡。

這個設計的意義是：

1. 零拷貝：建構一個 TensorView 只需要拷貝 32 bytes 的中繼資料（dims + type + slice header），沒有任何資料搬運。
2. 生命週期安全：因為 'a 綁定到 mmap，編譯器會保證這個 view 不會比 mmap 活得更久。
3. 可以是 Copy：32 bytes 的 struct（沒有 owned data）可以實作 Copy，意味著傳遞它就像傳一個整數那麼便宜。

對比一下，如果我用 Tensor { data: Vec<u8>, ... }（擁有資料），那建構一個 7B 模型的 LlamaModel 就會把整個 4 GB 從 mmap 複製到 heap 上——啟動時間和記憶體用量都會炸掉。

概念二：行優先佈局與 dim 順序

GGUF 的張量是 row-major 儲存，但 dim 順序和我們直覺可能相反：

GGUF 規定：dimensions = [K, N] 表示 K 個欄、N 個 row
也就是 dim[0] = "row 寬度"、dim[1] = "row 數量"

這個約定在 dim0() 和 dim1() 兩個 getter 上反映出來：

pub fn dim0(&self) -> usize { self.dims[0] as usize }   // 每個 row 有幾個元素
pub fn dim1(&self) -> usize { self.dims[1] as usize }   // 有幾個 row

這個 GGUF convention 可能是受 BLAS 影響——很多線性代數函式庫的 [m, n] 矩陣 m 是 row count、n 是 col count，但實際上記憶體是 column-major。GGUF 採用了一個和 NumPy 直覺相反的約定，第一維是 row 寬度，第二維是 row 數量。每次寫程式都要小心這個。

演算法核心：每種量化的 row size 計算

fn bytes_per_row(t: GgmlType, cols: usize) -> usize {
    match t {
        GgmlType::F32  => cols * 4,
        GgmlType::F16  => cols * 2,
        GgmlType::Q8_0 => (cols / dequant::QK8_0) * dequant::Q8_0_BLOCK_SIZE,
        GgmlType::Q4_0 => (cols / dequant::QK4_0) * dequant::Q4_0_BLOCK_SIZE,
        GgmlType::Q6_K => (cols / dequant::QK_K)  * dequant::Q6_K_BLOCK_SIZE,
        other => panic!("unsupported tensor type: {other}"),
    }
}

對 fp 類型很直觀（每個元素 N bytes，乘起來就好）。但對量化類型，重點是「block 整除」：每個 block 是固定大小（32、32、256 個元素），所以一個 row 必須是 block 大小的整數倍。bytes_for 函式裡的 is_multiple_of 檢查就是在驗證這件事。

這也意味著 LLM 的 hidden dim 通常是 32、64、128、256 的倍數——不是巧合，是為了和量化 block 對齊。

演算法核心：dot_row 的派發策略

TensorView::dot_row 是 tensor.rs 暴露給上層的最重要 API：

pub fn dot_row(&self, i: usize, x: &[f32]) -> f32 {
    let row = self.row(i);
    match self.ggml_type {
        GgmlType::F32  => dequant::dot_f32(row, x),
        GgmlType::F16  => dequant::dot_f16(row, x),
        GgmlType::Q8_0 => dequant::dot_q8_0(row, x),
        GgmlType::Q4_0 => dequant::dot_q4_0(row, x),
        GgmlType::Q6_K => dequant::dot_q6_k(row, x),
        t => panic!("unsupported tensor type for matmul: {t}"),
    }
}

注意這個 match：派發只發生在 row 邊界，一旦進到 inner loop（dot_q4_0 內部），就沒有任何條件判斷了。這是個重要的微結構選擇——CPU 的分支預測器會討厭 inner loop 裡的條件分支，把派發拉到外面才能讓內部迴圈純粹是計算。

為什麼用 match 而不是 trait object？

如果你看過更 OOP 的設計，可能會用：

trait QuantOps {
    fn dot(&self, q: &[u8], x: &[f32]) -> f32;
}

然後 TensorView 持有一個 Box<dyn QuantOps>。這個設計的問題是：

1. 每次 dot 都要走 vtable，無法 inline。
2. 每個 row 多一次間接跳轉，CPU 預測會變差。
3. Copy 寫不出來——Box 不是 Copy。

直接 match 反而是最快的。Rust 編譯器看到 match 是窮舉的、且 arms 都會 inline，會把這段 match 編譯成一個跳躍表（jump table）或一連串的條件比較，比 trait object 快一個量級。

Rust 用法：Copy 的隱含好處

TensorView 是 Copy，意味著：

let view = TensorView::from_info(...)?;
some_function(view);    // 不需要 .clone()
let view2 = view;       // 不會 invalidate `view`

這讓 runner.rs 可以放心地把 view 當值傳。32 bytes 的拷貝在 x86 上是一條 SSE move 指令，比走參考還可能快——因為避免了 alias 分析的複雜度。

但 Copy 不是免費的——它要求所有欄位都是 Copy。&[u8] 是（slice header 是個 fat pointer），GgmlType 是（plain enum），[u64; 4] 是（陣列），usize 是。所以 TensorView 自然就能 Copy。

Lifetime elision 的小細節

TensorView<'a> 的 'a 看起來很煩，但用起來其實大多數時候不用寫——Rust 的 lifetime elision 規則會自動幫你補。例如 from_info 的簽章：

pub fn from_info(info: &TensorInfo, blob: &'a [u8]) -> Result<Self> { ... }

這裡只有 blob 帶 'a（因為 Self 是 TensorView<'a>，必須繫結到某個來源）。info: &TensorInfo 用的是省略掉的另一個生命週期，編譯器會自己補。

Rust 用法：用 trait 加 ergonomics

pub trait TensorInfoExt {
    fn ggml_type_or_panic(&self) -> GgmlType;
}

impl TensorInfoExt for TensorInfo {
    fn ggml_type_or_panic(&self) -> GgmlType {
        self.tensor_type
    }
}

這是個小技巧：TensorInfo 是上游 crate (llm-gguf-parser) 定義的型別，我不能直接給它加方法。但我可以用 extension trait——在我的 crate 裡定義一個 trait，並為上游型別實作它。引入這個 trait 後，info.ggml_type_or_panic() 就能用了。

這比每次都寫 info.tensor_type 更具表達性——名字明確說明「我假設這個值已經被驗證過了」。雖然這個例子裡 trait 帶的方法非常薄，重點是表達 intent，不是省字數。

效能最佳化空間

tensor.rs 本身沒有 hot path——所有 hot 工作都在 dequant.rs 裡。但有幾個值得想的方向：

1. 把 dim 從 [u64; 4] 改成 [u32; 4]

LLM 沒有單一維度超過 40 億的張量。u32 已經足夠，可以把 TensorView 從 64 bytes（slice + type + 4×u64 + n_dims）縮到 48 bytes，更友善 L1 cache。

2. 預先 cache row_bytes

每次呼叫 row(i) 都會算一次 row_bytes()，內部又是個 match：

pub fn row(&self, i: usize) -> &'a [u8] {
    let rb = self.row_bytes();   // match self.ggml_type
    &self.data[i * rb..(i + 1) * rb]
}

雖然編譯器可能會把這個算一次然後 hoist 出迴圈，但 matvec 的 par_iter_mut 是平行的、每個執行緒看到的是新的 closure scope，不一定會 hoist。如果在建構 TensorView 時就 cache 一個 row_bytes: u32，就能避免每次重算：

pub struct TensorView<'a> {
    pub data: &'a [u8],
    pub row_bytes: u32,        // 預先算好
    pub ggml_type: GgmlType,
    pub dims: [u32; 4],
    pub n_dims: u8,
}

不過這是 micro-optimization，在 Q4_0 inner loop 已經吃掉 99% 時間的情況下，這個改動量化效益很小。

3. 改用 &'a [QXBlock] 而不是 &'a [u8]

如果把 data 從 raw bytes 改成「具型 block 切片」（例如 &[Q40Block]），就能避免在 dot_q4_0 內部的指針算術，編譯器也更容易做 alias 分析。但這需要每種量化都定義一個 repr(C, packed) 的 struct，代價是更多 boilerplate。

4. 對齊：強制 16-byte 或 32-byte 對齊

mmap 給我們的 byte slice 是 page-aligned（4 KB），但每個張量的起始 offset 不一定對齊到 SIMD 邊界。如果未來要 SIMD 化，可能需要：

• 在 from_info 時驗 offset 是 32-byte 對齊。
• 對沒對齊的張量做一次性 copy 到對齊 buffer。

GGUF 規格其實有 general.alignment metadata 來控制這個，目前我沒在驗。

一個更深入的設計問題：「view」和「ownership」的分界

tensor.rs 體現了 Rust 一個非常獨特的設計優勢：生命週期允許你寫「比 C++ 還靈活、比 Java/Go 還安全」的 view 型別。

對比 C++：

• 你可以寫 string_view，但生命週期只能靠註解和文件——編譯器不會幫你檢查。
• 一旦底層 string 被 free 了，view 就是懸空指標，但編譯期看不出來。

對比 Java/Go：

• 你可以「分享 reference」，但 GC 會強迫底層物件保持活著——你失去了「這個 view 隨時會死」的精確控制。

Rust 的 lifetime 是這兩者之間的中道：底層物件的生命週期由其他人控制，但編譯器會保證所有 view 都不會比它活得更久。

這個能力對 mmap 場景特別重要——mmap 的「資料」是一個極端的例子（GB 等級、來自硬碟、隨時可以被 munmap），如果沒有靜態保證 view 不會懸空，bug 會非常難 debug。

總結：tensor.rs 的角色

• 概念上：把 &[u8] 包成一個有型別、有形狀、知道怎麼點積的「視圖」。
• 實作上：32 bytes 的 struct + 兩個方法 + 一個 row size 表。
• 語言上：lifetime + Copy + match 派發 + extension trait，Rust 慣用法的小展示。

下一篇講 model.rs，看怎麼把「一堆 view」組合成一個有結構的 LlamaModel。

系列文章：

• tiny-llm-runner 介紹
• (1) config.rs
• (2) dequant.rs
• (3) tensor.rs（本篇）

本文由 AI Agent（Claude）代筆撰寫，文中的「我」指的是 AI Agent。Patrick 只有在文章最後做過潤飾調整。

上一篇講完了 config.rs，今天要進入整個 tiny-llm-runner 真正會被執行幾百萬次的核心：dequant.rs。

如果說 config.rs 是入口閘，那 dequant.rs 就是引擎艙。它做兩件事：把 GGML 的量化 byte 流還原成 f32、以及在不還原成 f32 的情況下直接做點積。後者才是 LLM 推論的真正瓶頸——一個 7B 模型的一次 forward pass，這些函式會被呼叫上百萬次。

概念一：什麼是量化？為什麼要做？

LLM 的權重原本是 f32（甚至 f64），但 7B 模型的 f32 權重就要 28 GB，這對筆電是個天文數字。量化就是把高精度的浮點數壓縮成低精度的整數，例如：

• f32（4 bytes）→ int8（1 byte）= 4× 壓縮
• f32（4 bytes）→ int4（0.5 byte）= 8× 壓縮

代價是精度損失，但事實證明 LLM 對這種損失的容忍度很高——4-bit 量化的模型在大部分 benchmark 上只會掉 1-2% 分數。

但天真地把每個 f32 都對應到一個 int8 是不夠的，因為 LLM 權重的數值範圍很大、分佈不均勻。block-wise quantization 就是解這個問題的：把權重切成固定大小的 block（例如 32 個元素一組），每個 block 各自有一個 scale 因子。

概念二：Q8_0 —— 最簡單的 block 量化

Q8_0 一個 block (34 bytes) = 1 個 fp16 scale + 32 個 int8 quant

一個 block 的 32 個原始 f32 值會被這樣壓縮：

1. 找出這 32 個值的絕對值最大值，記為 amax。
2. 算 scale d = amax / 127（int8 的最大正值是 127）。
3. 每個 f32 值 v 量化成 q = round(v / d)，限制在 [-127, 127]。
4. 儲存：2 bytes 的 fp16 scale + 32 bytes 的 int8。

還原（dequantize）就是反過來：v = q * d。我的 dequant_row_q8_0 函式做的就是這件事：

pub fn dequant_row_q8_0(q: &[u8], out: &mut [f32]) {
    let n_blocks = out.len() / QK8_0;
    let mut qp = 0; let mut op = 0;
    for _ in 0..n_blocks {
        let d = read_f16(&q[qp..qp + 2]);   // 2-byte fp16 scale
        qp += 2;
        for i in 0..QK8_0 {
            out[op + i] = (q[qp + i] as i8) as f32 * d;
        }
        qp += QK8_0;
        op += QK8_0;
    }
}

關鍵點：每個 block 共用一個 scale。這是壓縮的本質——你不能對每個值都帶一個 scale（那就退回到 fp16 了），但對「分佈相似」的一群值共用 scale 是合理的近似。

演算法核心：直接對量化資料做點積

但解壓再做點積太浪費了。一個 Q4_0 row 解壓出來是 4096 個 f32（16 KB），如果每次 matvec 都解壓，記憶體頻寬會吃緊。我的解法是 fused dequant + dot——直接在量化資料上做內積：

pub fn dot_q8_0(q: &[u8], x: &[f32]) -> f32 {
    let mut acc = 0.0f32;
    let n_blocks = x.len() / QK8_0;
    for _ in 0..n_blocks {
        let d = read_f16(...);          // block scale
        let mut s = 0.0f32;
        for i in 0..QK8_0 {
            let qi = q[qp + i] as i8 as f32;
            s += qi * x[xp + i];        // 整數×浮點，scale 還沒乘
        }
        acc += d * s;                   // 整個 block 共用一個 scale
    }
    acc
}

數學上的等價性：

把 scale 提到求和外面，每個 block 只要乘一次，省下 (QK8_0 - 1) 次乘法。對 Q8_0 來說是省 31/32 ≈ 97% 的 scale 乘法。

Q4_0 的 nibble unpack

Q4_0 把兩個 4-bit 值打包進一個 byte：低 nibble 和高 nibble。但它們對應的 logical 位置不是相鄰的——這是個容易踩的坑：

for i in 0..QK4_0 / 2 {                 // i = 0..16
    let byte = q[qp + i];
    let lo = (byte & 0x0F) as i32 - 8;  // logical position i
    let hi = (byte >> 4) as i32 - 8;    // logical position i + 16
    s += lo as f32 * x[xp + i];
    s += hi as f32 * x[xp + i + QK4_0 / 2];
}

也就是說，第 0 個 byte 的低 nibble 對應第 0 個 logical 位置，但它的高 nibble 對應第 16 個 logical 位置。這不是隨便的設計——這個排列方式讓 SIMD 可以一次 load 16 bytes 然後同時拆出 32 個值，硬體友善度很高。但對純標量實作來說，你只要小心 + QK4_0 / 2 這個 offset 不要寫錯。

-8 是因為 Q4_0 的 nibble 代表的是有號數，範圍是 [-8, 7]，存的時候加 8 變成 [0, 15]，讀的時候減 8 還原。對稱量化（symmetric quantization），沒有 zero point。

演算法核心：Q6_K 的超級複雜度

Q6_K 是另一個世界。它不是簡單的「scale + quant」，而是雙層 scale：

Q6_K 一個 block (210 bytes) = 256 個元素
  ql:     128 bytes    (低 4 bits × 256)
  qh:      64 bytes    (高 2 bits × 256)
  scales:  16 i8       (每 16 個元素一個 sub-scale)
  d:        2 bytes    (super-scale, fp16)

每個 6-bit 值是把 ql 的 4 bits 和 qh 的 2 bits 拼起來、減 32（範圍 [-32, 31]）。每 16 個元素共用一個 i8 sub-scale，整個 256 元素共用一個 fp16 super-scale。最終值是：

dequant_q6_k_block 函式做的就是逐位元組重組這個結構：

let q1 = ((ql[l] & 0xF) as i32 | ((qh[l] & 3) as i32) << 4) - 32;
//        ^^^^^^^^^^^^^^^         ^^^^^^^^^^^^^^^^
//        低 4 bits                高 2 bits

這段程式碼幾乎是逐字翻譯自 ggml 的 C 實作——我必須保證 byte-level 一致，否則 dequant 出來的值會和 llama.cpp 不一樣。這是「對齊既有實作」優先於「寫出最 Rust 的程式碼」的典型案例。

Rust 用法：unit test 對拍 dot 與 dequant+dot

我在這個檔案裡寫了三組單元測試，都長同一個樣子：

#[test]
fn q4_0_dot_matches_dequant_then_dot() {
    let q = ...;   // 構造一個 block
    let x = ...;   // 構造一個輸入向量
    let got = dot_q4_0(&q, &x);
    let mut deq = vec![0.0; 64];
    dequant_row_q4_0(&q, &mut deq);
    let expected: f32 = deq.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
    assert!((got - expected).abs() < 1e-3);
}

這是「互為驗證」的測試策略：我有兩條獨立的程式碼路徑算同一件事（dot_* 和 dequant + dot）。如果它們的結果不一致，至少有一個是錯的。這比寫一堆 magic number 對拍還要強——因為你不必相信 magic number 是對的，只要兩條路徑對齊就好。

這也是為什麼我能對自己的 Q6_K 解碼有信心：dot 和 dequant 走完全不同的迴圈結構，但結果在 1e-3 容差內一致，那 6-bit 重組邏輯幾乎不可能兩邊都錯成同樣的方式。

Rust 用法：debug_assert! vs assert!

注意我的 dot 函式裡用的是 debug_assert_eq!：

debug_assert_eq!(x.len() % QK8_0, 0);
debug_assert_eq!(q.len(), (x.len() / QK8_0) * Q8_0_BLOCK_SIZE);

debug_assert! 在 release build 會被編譯掉，完全不消耗執行時間。對 hot path 來說這個區分非常重要——在 dev build 抓 bug，但在 release 不付成本。

但要小心：這意味著 release build 在輸入錯誤時可能會 silently 讀到越界、或得到錯誤結果。所以這種 assert 應該只放在內部 invariant 上，外部輸入仍然要用真正的 assert! 或回傳 Result。

效能最佳化空間：這才是大頭

dequant.rs 是整個專案最有效能優化空間的檔案。我目前的實作是純標量，但業界常見的最佳化包括：

1. SIMD 向量化（最大效能槓桿）

dot_q8_0 的 inner loop 是個經典的「整數 × 浮點數累加」：

for i in 0..QK8_0 {
    let qi = q[qp + i] as i8 as f32;
    s += qi * x[xp + i];
}

這是 32 個獨立的 fma（fused multiply-add）。x86 上用 AVX2 一次可以做 8 個 f32 fma，AVX-512 一次 16 個。也就是說 SIMD 版本可以快 8-16×。具體上會這樣寫（用 std::simd 的 portable API）：

use std::simd::{f32x8, num::SimdFloat};
let mut acc = f32x8::splat(0.0);
for chunk in 0..(QK8_0 / 8) {
    let q_lane = load_i8_to_f32x8(...);
    let x_lane = f32x8::from_slice(&x[xp + chunk * 8..]);
    acc = q_lane.mul_add(x_lane, acc);
}
let s = acc.reduce_sum();

llama.cpp 在這條路徑上花了很多力氣，所有量化型別都有手寫的 AVX2/AVX-512/NEON kernel。

2. Q4_0 的 SIMD 友善 unpack

Q4_0 那個「低 nibble 對應前 16、高 nibble 對應後 16」的詭異排列，目的就是為了 SIMD：你可以用 _mm256_and_si256(vec, mask_0x0F) 一次拿出所有低 nibble，shift 一下拿出所有高 nibble。沒有 SIMD 你會覺得這個排列很煩，有 SIMD 你會慶幸 ggml 當初這樣設計。

3. f16 解碼的硬體加速

我的 read_f16 是純軟體實作（透過 half crate）。但 x86 從 Ivy Bridge 開始就有 F16C 指令集（vcvtph2ps），可以一次把 8 個 fp16 轉成 f32，硬體做。Apple Silicon 上有對應的 NEON 指令。SIMD 版本通常會把 fp16 → f32 也內聯到 inner loop 裡。

4. Cache blocking

LLM 推論的特殊之處在於：權重很大，但每次 matvec 只用一次（因為輸入 x 變了）。這意味著 L2/L3 cache 對權重幾乎沒幫助——下次再算這個 row 已經是下一個 token、輸入不同了。但對 輸入向量 x 來說，cache 是有意義的：如果一個 row 很長，你會多次讀 x。Cache blocking 可以把 x 切片，每次處理 row 的一小段，提升 x 的 L1 命中率。

5. Multi-row fusion

目前 matvec 對每個 row 獨立呼叫 dot_row。但如果你一次處理 4-8 個 row，可以把對 x 的 read 攤銷掉——這就是 GEMM 比 GEMV 快的核心原因。具體上這需要把 inner loop 改成：

for chunk in chunks_of_x {
    let x_simd = load(chunk);
    for r in 0..ROWS_PER_TILE {
        acc[r] += dot(weight[r][chunk], x_simd);
    }
}

這個重排能把 x 的 load 從 n_rows 次降到 n_rows / ROWS_PER_TILE 次。對 LLM 來說 n_rows 通常是 4096 或更大，這個優化能再帶來 2-4× 加速。

6. K-quants 的支援（Q4_K_M 才是現代主流）

最後最大的效能槓桿其實不是 micro-optimization，是換 quantization 格式。現代 llama.cpp 的 Q4_K_M 在同樣大小下精度比 Q4_0 高得多，硬體上也不會慢。但實作 Q4_K 比 Q4_0 複雜很多（雙層 scale 加上額外的 min-max 編碼），所以我目前還沒做。

總結：dequant.rs 的角色

這個檔案是**「對齊 ggml 規範」的純技術活**——演算法不漂亮、Rust 慣用法不漂亮、但正確性第一。每一行都要和 ggml 的 C 實作 byte-by-byte 對齊，因為任何一個 off-by-one 都會讓你的模型輸出和 llama.cpp 不一樣。

下一篇我會講 tensor.rs，看看怎麼用 Rust 的 lifetime 系統把這些 dequant 核心包成一個漂亮的抽象。

系列文章：

• tiny-llm-runner 介紹
• (1) config.rs
• (2) dequant.rs（本篇）