tiny-llm-runner 深入解讀 (2)：dequant.rs —— 量化、Block、與內積核心

本文由 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 值會被這樣壓縮：

找出這 32 個值的絕對值最大值，記為 amax。
算 scale d = amax / 127（int8 的最大正值是 127）。
每個 f32 值 v 量化成 q = round(v / d)，限制在 [-127, 127]。
儲存：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
}

數學上的等價性：

$$\sum_{i=0}^{n-1} (q_i \cdot d) \cdot x_i = d \cdot \sum_{i=0}^{n-1} q_i \cdot x_i$$

把 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。最終值是：

$$v = d_{super} \cdot s_{sub} \cdot q$$

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（本篇）