Simply Patrick

tiny-llm-runner 深入解讀 (3):tensor.rs —— 用生命週期蓋一層薄殼

featured.svg

本文由 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_rowtensor.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 的隱含好處

TensorViewCopy,意味著:

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(因為 SelfTensorView<'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 出迴圈,但 matvecpar_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 深入解讀 (2):dequant.rs —— 量化、Block、與內積核心

featured.svg

本文由 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
}

數學上的等價性:

$$\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 —— 把 GGUF metadata 變成 Rust 強型別

featured.svg

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

tiny-llm-runner 介紹文裡我用一張表把九個檔案掃過去,但每個檔案值得講的東西其實都不止一段。所以我打算開一個小型的「深入解讀」系列,一個檔案一篇,把每個檔案的概念、演算法、Rust 用法、以及未來最佳化空間都拆開來講清楚。

第一篇從最樸素的 config.rs 開始——它只有 100 多行,乍看之下沒什麼好講的,但其實它隱含了好幾個值得展開的細節。


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

GGUF 檔頭裡的 metadata 區塊是一個 (key, value) 的扁平 map,key 是字串、value 是一個 sum type(int、float、string、array…)。但要跑 Llama 推論,我們需要的是一個有型別、有不變式、欄位齊全的 struct。config.rs 的工作就是把前者轉成後者:

pub struct LlamaConfig {
    pub n_ctx: usize,
    pub n_embd: usize,
    pub n_layer: usize,
    pub n_head: usize,
    pub n_head_kv: usize,
    pub n_ff: usize,
    pub vocab_size: usize,
    pub rms_eps: f32,
    pub rope_freq_base: f32,
    pub rope_dim_count: usize,
}

這 10 個欄位涵蓋了整個 forward pass 需要的所有形狀資訊。後面 runner.rs 不需要再去查 metadata,它只需要拿著一個 &LlamaConfig 就能算出所有 buffer 大小、KV cache 維度、GQA group 數。


概念一:什麼是 hyperparameter?為什麼要從 metadata 讀?

對沒做過 LLM 的人來說,「為什麼這些參數不能寫死?」是個合理的疑問。答案是:Llama 是一個架構家族,不是單一模型。同樣是 Llama 架構,你可以有:

變體 n_layer n_embd n_head n_head_kv
TinyLlama 1.1B 22 2048 32 4
Llama-2 7B 32 4096 32 32
Llama-2 13B 40 5120 40 40
Llama-3 8B 32 4096 32 8
Llama-3 70B 80 8192 64 8

如果把這些寫死在程式裡,你的 runner 就只能跑某一個特定模型。所以 GGUF 把這些「形狀相關的常數」存在檔頭,runner 啟動時動態讀取——這就是 metadata 的用途。


演算法核心:型別轉換的容錯邏輯

get_usize 是這個檔案唯一稱得上有「演算法」的地方:

fn get_usize(g: &GgufFile, key: &str) -> Result<usize> {
    let v = g.metadata.get(key)
        .with_context(|| format!("missing metadata {key}"))?;
    match v {
        Value::Uint8(x)  => Ok(*x as usize),
        Value::Uint16(x) => Ok(*x as usize),
        Value::Uint32(x) => Ok(*x as usize),
        Value::Uint64(x) => Ok(*x as usize),
        Value::Int8(x)   => Ok(*x as usize),
        // ... 其他整數型別
        _ => Err(anyhow!("metadata {key} is not an integer")),
    }
}

為什麼要列這麼多分支?因為 GGUF metadata 的型別是寫檔時決定的llama.cpp 在不同版本可能會把 n_layer 存成 u32u64;某個訓練框架可能會把它存成 i32。我們不能只認一種——但所有整數型別都能安全轉成 usize(因為這些值通常很小),所以一個大 match 就解決了。

這是 Rust enum 派發的典型用法:用 match 列舉所有可能,用 _ 兜底拒絕。如果哪天 GGUF 規格新增了一種整數型別,Rust 編譯器會在所有 match 上提醒你(因為 _ 通配符會吃下新型別),這時候你就要決定到底要不要支援它。


Rust 用法:with_context 與錯誤鏈

let v = g.metadata.get(key)
    .with_context(|| format!("missing metadata {key}"))?;

這一行用了 anyhow 套件的 with_context。它做的事是:如果這個 ResultErr,就把錯誤 wrap 在一個有上下文的新錯誤裡;如果是 Ok,那個 closure 根本不會執行(所以 format! 不會被付出代價)。

為什麼用 closure 而不是字串?lazy evaluation。如果直接寫:

.context(format!("missing metadata {key}"))   // 每次都會 format!

那麼即使在 happy path(找得到 metadata)下,format! 也會被執行——對 hot path 來說這是純粹的浪費。with_context 接 closure 的版本只在錯誤時才執行,這在效能敏感的程式碼裡是個重要的習慣。

? 運算子的小巧思

? 不只是「if Err return」,它還會自動呼叫 From::from。所以 anyhow::Error 可以從任何實作了 std::error::Error 的型別轉換過來——這就是為什麼 parse_gguf 回傳的錯誤可以無縫轉成我函式的 anyhow::Result


不變式檢查:為什麼要在 load 時就驗?

if n_embd % n_head != 0 {
    bail!("n_embd {n_embd} not divisible by n_head {n_head}");
}
if n_head % n_head_kv != 0 {
    bail!("n_head {n_head} not divisible by n_head_kv {n_head_kv}");
}

這兩個不變式是 Llama 架構的硬性要求:

  1. n_embd / n_head = head_dim——每個 attention head 平均分配 hidden dim。
  2. n_head / n_head_kv = gqa_groups——GQA 下,每組 query head 共用一個 K/V head。

如果這些不變式不成立,後面 runner.rs 算出來的 offset 全部會錯,可能會產生未定義行為(讀到別的張量的 bytes)。在 load 階段就 bail 掉,比在 forward pass 跑到一半 panic 友善多了。

這是 Rust 「fail fast at boundaries」哲學的具體實踐:在外部資料進入系統的邊界(這裡是 GGUF metadata)就驗證,往後的程式碼可以放心假設不變式成立。


衍生方法:head_dim()kv_dim()gqa_groups()

pub fn head_dim(&self) -> usize { self.n_embd / self.n_head }
pub fn kv_dim(&self)   -> usize { self.head_dim() * self.n_head_kv }
pub fn gqa_groups(&self) -> usize { self.n_head / self.n_head_kv }

這三個方法是 derived 量,理論上每次都能從 base 欄位算出來。為什麼不直接存 head_dim 而是每次算?因為:

  1. 避免不一致:如果同時存 n_headn_embdhead_dim,理論上有人可能改了其中一個忘了改另外兩個,三者就對不起來。只存 base 量是 single source of truth。
  2. divide 的成本可忽略:這些方法不在 hot path 上,runner 是在 new() 時呼叫一次來算 buffer 大小,不會在 inner loop 反覆呼叫。

這是 Rust struct design 的一個小哲學:只存 fundamental fields,derived 量寫成方法


從 hyperparameter 算出模型大小

在動手算之前,先用一張圖把這些欄位放到推論流程上看看,就會比較清楚每個參數實際在控制哪一段:

hyperparams-flow.svg

有了這 10 個欄位,你其實可以直接算出整個模型有幾個參數,連模型都還不用載入。這對於估記憶體、估推論速度,或者單純滿足好奇心,都很實用。

Llama 的權重主要藏在三個地方:embedding、每一層 transformer block、最後的 LM head。把它們一個一個拆開來看:

1. Token embedding 與 LM head

每個 token 對應一個長度 n_embd 的向量,總共 vocab_size 個 token:

  • Embedding:vocab_size × n_embd
  • LM head:vocab_size × n_embd(有些模型會跟 embedding 共享權重,Llama 家族通常不共享)

這兩個加起來通常是模型裡單一最大塊的權重,特別是 vocab 很大的時候(Llama-3 把 vocab 從 32K 拉到 128K,就是這裡膨脹得最兇)。

2. 每一層 transformer block

每層裡面有 attention 和 FFN 兩個子模組。先看 attention 的四個 projection:

  • Q projection:n_embd × n_embd
  • K projection:n_embd × kv_dim(GQA 下 K 的維度是縮小的)
  • V projection:n_embd × kv_dim
  • Output projection:n_embd × n_embd

加起來是 2 × n_embd² + 2 × n_embd × kv_dim

FFN 是 SwiGLU 結構,有三個矩陣:

  • Gate projection:n_embd × n_ff
  • Up projection:n_embd × n_ff
  • Down projection:n_ff × n_embd

加起來是 3 × n_embd × n_ff

另外每層還有兩個 RMSNorm 的 scale 向量(各長 n_embd),加上最後一層 final norm 也是 n_embd——這些跟矩陣相比小到可以忽略。

3. 把公式拼起來

total_params ≈ 2 × vocab_size × n_embd                # embedding + LM head
             + n_layer × (
                   2 × n_embd²                        # Q + O
                 + 2 × n_embd × kv_dim                # K + V
                 + 3 × n_embd × n_ff                  # SwiGLU
               )

套到 TinyLlama 1.1B 驗算一次

從前面的表格抓出 TinyLlama 的 hyperparameters:n_layer=22, n_embd=2048, n_head=32, n_head_kv=4, n_ff=5632, vocab_size=32000

先算衍生量:head_dim = 2048 / 32 = 64kv_dim = 64 × 4 = 256

部分 公式 數值
Embedding 32000 × 2048 ≈ 65.5M
LM head 32000 × 2048 ≈ 65.5M
每層 attention 2 × 2048² + 2 × 2048 × 256 ≈ 9.4M
每層 FFN 3 × 2048 × 5632 ≈ 34.6M
全部 22 層 22 × (9.4M + 34.6M) ≈ 968M
總和 ≈ 1.10B

剛好對得上 “1.1B”——這就是模型名字裡那個數字的來源。

從這個算式還能看出兩件事

  1. FFN 才是主導參數量的部分3 × n_embd × n_ff 通常比 4 × n_embd × head_dim × (n_head + n_head_kv) 大上 3~4 倍。所以模型放大的時候,最先膨脹的是 FFN,再來才是層數。
  2. 記憶體佔用可以反推。fp16 下每個參數 2 bytes,1.1B × 2 ≈ 2.2 GB;換成 4-bit 量化理論上只剩 1/4,約 550 MB——這就是 tiny-llm-runner 能直接 mmap 一個 q4_k 模型在普通筆電上跑起來的關鍵。

下次拿到一個陌生的 GGUF 檔,不用去翻 Hugging Face card,光看 metadata 就能算出大概多大、能不能塞進你的記憶體。


總結:什麼時候輪到 config.rs 上場

整個 tiny-llm-runner 的生命週期裡,config.rs 只被呼叫一次:在 main.rs 的開頭。它的角色是型別系統的入口閘——把外部世界的 schemaless metadata 轉成內部世界的 type-safe struct。一旦過了這道閘,後面所有檔案就可以放心吃 &LlamaConfig,不必再碰 GGUF 的細節。

下一篇我會講 dequant.rs,那是這個專案的「重武器庫」——所有量化解碼和點積核心都藏在那裡。


系列文章:


純手工 vs. 全功能框架:tiny-llm-runner 與 candle 的架構對照

featured.svg

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

上一篇我把 tiny-llm-runner 拆開來介紹完之後,自然會冒出一個問題:

Rust 圈不是已經有一個叫做 candle 的專案了嗎?Hugging Face 親兒子、CUDA / Metal 後端齊全、連 Llama / Mistral / Whisper / Stable Diffusion 都跑得起來——那我這個小小的 tiny-llm-runner 到底有什麼存在意義?

這篇文章就是要正面回答這個問題。比較的重點不是「誰比較好」——candle 在功能、效能、生態系上樣樣都把我的玩具碾壓——而是想看一件更有意思的事:當你給自己訂的目標不一樣時,最後寫出來的架構會差到什麼程度?

兩個專案都用 Rust 寫、都吃 GGUF、都能跑 Llama,但你只要把它們的 Cargo.tomlsrc/ 並排放,就會發現它們其實在解非常不同的問題。


先把比較的基準擺正

要公平比較,得先承認兩者的「設計願景」根本不在同一個量級:

維度 tiny-llm-runner candle
目標 教學版:每一行 forward pass 都要看得懂 通用 ML 框架:要當 PyTorch 在 Rust 的對應物
範圍 只跑 Llama 1/2/3 推論 跑各式架構(Llama、Mistral、Mixtral、Whisper、SD…)
後端 只有 CPU + rayon CPU(可選 MKL/Accelerate)、CUDA、Metal
量化格式 5 種(F32/F16/Q8_0/Q4_0/Q6_K) 完整 GGML 量化家族 + 自家 QTensor
訓練 沒有 有(autograd)
程式碼規模 < 2,000 行 數萬行(多個 crate)

tiny-llm-runner 拿來和 candle 比,差不多就像是拿一台「自製真空管收音機」去和一台「現代多媒體電視」比較——兩者都會發出聲音,但它們存在的理由完全不同。所以下面我不會去比效能或功能(那只會是一面倒),而是要看「為了達成各自的目標,他們在哪些設計選擇上做了不一樣的取捨」。


差異一:Tensor 抽象的厚度

最容易看出兩者哲學差異的地方,是「Tensor 這個型別到底是什麼」。

tiny-llm-runner:32-byte 的薄殼

我的 TensorView 全部就只有這幾個欄位:

pub struct TensorView<'a> {
    pub data: &'a [u8],   // 直接借用 mmap 的 bytes
    pub ggml_type: GgmlType,
    pub dims: [u64; 4],
    pub n_dims: usize,
}

不是一個 Tensor——它是「mmap 裡某段 byte 的解釋方式」。它沒有 device、沒有 dtype 抽象(只有 GgmlType 列舉)、沒有自動微分的計算圖、沒有 broadcasting 規則、沒有 lazy execution。它只暴露兩個方法:dot_rowdequant_row

candle:完整的 PyTorch 風格 Tensor

candle_core::Tensor 的概念負擔則完全是另一個世界。它要追蹤:

  • Shape:可變維度的張量形狀,支援 broadcasting。
  • DTypeF32F16BF16U8U32I64F64
  • DeviceCpuCuda(...)Metal(...),每個張量都「住」在某一個裝置上。
  • Storage:實際資料的 reference-counted handle,支援共享和切片。
  • Op tracking:autograd 要用——每個 a + b 都會記下「我是從 a 和 b 加出來的」這條邊,反向傳播時才能還原計算圖。

換句話說,candle::Tensor 是一個有身分證的對象,它知道自己住哪、長什麼樣、從哪裡來。TensorView 則是一個臨時護照,過了這個 scope 就什麼都不剩了。

這個差異會傳染到整個 codebase

當你的 Tensor 抽象很厚時,所有運算都得透過它走。candle 裡你寫 x.matmul(&w),回傳的是新的 Tensor,背後可能是 CPU 的 BLAS、CUDA 的 cuBLAS、或 Metal 的 MPS——你都不需要知道。

但當你的抽象很薄(像我)時,matvec 直接吃 &mut [f32]&TensorView,沒有任何派發層,整個函式 inline 起來幾乎沒有 overhead

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);   // 直接 match ggml_type 派發
    });
}

代價是:這個 matvec 永遠只會在 CPU 上、只能算 f32×量化、只能做 row-parallel。要支援 GPU?沒辦法。要支援 batch-2 matrix×matrix?得重寫。


差異二:權重從 GGUF 到記憶體的路徑

這條路徑兩邊走得很不一樣,也是兩種哲學最具體的展現。

tiny-llm-runner:mmap → 不解壓 → row-by-row dot

我的 LlamaModel::load 只做兩件事:

  1. 把 GGUF 的 tensor 列表掃一遍,建一個 name → TensorInfo 的 HashMap。
  2. 對每個我需要的張量名(token_embd.weightblk.0.attn_q.weight…),建一個 TensorView

整個過程沒有讀任何權重資料——TensorView 只記住「這個張量住在 mmap 的哪一段 bytes」。等到 forward() 真的呼叫 dot_row(i, x) 時,才會去讀那一個 row 的 bytes、解碼、做內積。

整個生命週期裡,Q4_0 的權重就一直是 Q4_0 在硬碟上躺著,從來沒有被解壓成 f32 矩陣放在 RAM 裡。一個 Q4_0 的 7B 模型大約是 4 GB,如果你解壓成 f32 就是 28 GB——這對筆電來說是天差地別。

candle:QTensor → QMatMul → 編譯期分派

candle 的處理方式精緻很多。它有專門的 candle_core::quantized 模組,裡面有:

  • QTensor:量化張量的內部表示。
  • QMatMul:把量化權重和 f32 啟動值相乘的高階介面。
  • gguf_file::Content:解析 GGUF 檔。

載入流程大致是:

  1. gguf_file::Content::read(...) 把整個 metadata 讀進來。
  2. 對每個張量,用 tensor.qmatmul(...)tensor.dequantize(...) 取得 QTensor
  3. QTensor 包進 QMatMul,當作一個可呼叫的層。

QMatMul 在 forward 時會根據量化型別、後端(CPU/CUDA/Metal)派發到對應的 SIMD/CUDA kernel。CPU 上是 AVX2/NEON 手工最佳化過的;CUDA 上是專門的 dequant+gemm fused kernel。

兩條路的代價

tiny-llm-runner 來說,「不解壓」是一個雙面刃:

  • ✅ 記憶體用量極低、不需要 warmup。
  • ✅ 大檔案啟動秒開(mmap 立即返回)。
  • ❌ 沒有任何 SIMD/向量化——我的 dot_q4_0 是純標量迴圈,比 candle 慢一個數量級以上。

candle 來說,多走一層 QTensor 抽象的代價也很實在:

  • ✅ 同一份程式碼可以在 CPU/CUDA/Metal 上跑。
  • ✅ SIMD 最佳化由底層 kernel 統一處理。
  • ❌ 程式碼複雜度暴增——要看懂 QMatMul::forward 在哪個分支會做什麼事,得跨好幾個 trait、好幾個 cfg flag。

差異三:模型架構是寫死的還是組裝的

這個維度大概是最直觀的「規模差距」。

tiny-llm-runner:把 Llama 寫死進去

我的 runner.rs 整個 forward pass 是一段大約 100 行的指令式程式碼,每一行都假設模型是 Llama

// attention norm → qkv → RoPE → GQA attention → output proj
rmsnorm(&mut self.xb, &self.x, &layer.attn_norm, cfg.rms_eps);
matvec(&mut self.q, &layer.wq, &self.xb);
matvec(krow, &layer.wk, &self.xb);
matvec(vrow, &layer.wv, &self.xb);
apply_rope(&mut self.q, pos, head_dim, cfg.rope_dim_count, ...);
// ... 以下省略

要支援 Mistral?得改 RoPE 的 sliding window mask。要支援 Mixtral?得加 MoE 路由邏輯。每加一個架構,就是改 runner.rs 一次

candle:用 nn::Module trait 組合

candle 走的是 PyTorch 風格:每一層是一個實作了 Module trait 的 struct,forward 是它的方法。Llama 的整個架構在 candle_transformers::models::llama 裡是這樣組起來的:

// 大約的形狀(為了清楚我簡化過)
pub struct Block {
    rms_1: RmsNorm,
    attn: CausalSelfAttention,
    rms_2: RmsNorm,
    mlp: Mlp,
}

impl Module for Block {
    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let h = (x + self.attn.forward(&self.rms_1.forward(x)?)?)?;
        let h = (&h + self.mlp.forward(&self.rms_2.forward(&h)?)?)?;
        Ok(h)
    }
}

每一層都是獨立的元件,換一個架構只是換組合方式。Mistral 用的是同一套 RmsNormMlp,只是把 CausalSelfAttention 換成支援 sliding window 的版本。Mixtral 把 Mlp 換成 SparseMoeBlock

VarBuilder:把名字和張量解耦的關鍵

candle 在「載入權重」這件事上有一個我覺得很漂亮的抽象:VarBuilder。它是一個 hierarchical 的權重來源,可以這樣用:

let vb = VarBuilder::from_gguf("model.gguf", &device)?;
let block_0 = Block::load(vb.pp("blk.0"), &cfg)?;
let block_1 = Block::load(vb.pp("blk.1"), &cfg)?;

vb.pp("blk.0") 會回傳一個「prefix 是 blk.0 的子 builder」,子 builder 再去要 attn_q.weight 時,實際上會去找 blk.0.attn_q.weight。這讓模型的層級結構和權重的命名空間自然對應——你不用在每一層的 load 函式裡手動拼字串。

我的 model.rs 則是另一個極端:

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)?,
        // ...
    });
}

直接把 layer index 用 format! 拼進去。簡單粗暴,但只能應付「我知道我要找的張量叫什麼名字」這種情境。換一個架構,這整段都得改寫。


差異四:autograd —— 有沒有反向傳播

這個是最大、也最一刀切的差異。

candle 是一個訓練框架——它的 Tensor 在做運算時會自動建立計算圖,呼叫 loss.backward() 就能拿到所有參數的梯度。這也是為什麼它的 Tensor 必須記住「我從哪些 Tensor 來」、必須有 reference counting、必須處理生命週期。

tiny-llm-runner 完全不知道反向傳播是什麼。所有運算都是純函式式的:吃 &[f32] 進去、寫 &mut [f32] 出來,過了就忘。沒有梯度、沒有計算圖、沒有 backward。

這個取捨的影響比想像中還深:

  • 沒有 autograd → 不需要 op 物件 → 不需要 trait object → 不需要 reference counting → 不需要 Arc<Tensor> → 不需要 thread-safe 設計。
  • 結果就是:我可以用 &mut [f32] 切來切去,編譯器幫我擋住所有別名問題,程式碼長得像極了 C,但完全沒有 unsafe

換句話說,「我只做推論」這個決定,讓整個 codebase 的型別系統可以瘦一個量級。


差異五:執行模型 —— eager 還是 graph?

tiny-llm-runner 是純 eager:每個 matvec 呼叫就是一個立即發生的計算。沒有圖、沒有 fusion、沒有調度器。

candle 大致也是 eager 的(這也是它和 PyTorch 對應的點),但它的 Tensor 抽象為未來的 graph 最佳化留了空間。已經有些情境下 candle 會做 op fusion(例如 softmax + attention mask 在某些 backend 上是同一個 kernel)。

這個差異的代價是:candle 在 hot path 上,每個運算都要過一層型別/裝置/dtype 派發。對小張量來說這個 overhead 不可忽略;對大張量(典型 LLM 推論)來說相對微不足道。我則是從定義上就不可能有這個 overhead——因為我根本沒有派發。


一張表總結兩者的設計選擇

flowchart LR subgraph TLR[tiny-llm-runner] T1[mmap bytes] T2[TensorView
32 bytes] T3[matvec/rmsnorm/
RoPE 純函式] T4[Runner.forward
寫死 Llama] T5[只 CPU + rayon] T1 --> T2 --> T3 --> T4 --> T5 end subgraph C[candle] C1[mmap bytes] C2[QTensor / Tensor
多 backend] C3[Module trait
Linear/RmsNorm/...] C4[VarBuilder 載入
架構組合] C5[CPU + CUDA + Metal
autograd] C1 --> C2 --> C3 --> C4 --> C5 end

兩個都是同一條路徑——bytes → tensor → ops → model → backend——但每一段抽象的厚度差距,乘起來就是「兩千行 vs. 數萬行」、「一個架構 vs. 一打架構」、「教學版 vs. 生產級」的差距。


那什麼時候該用哪一個?

我對自己的定位很清楚:

candle 當你在做以下事情:

  • 想要在 Rust 裡跑非 Llama 的模型(Whisper、SD、各種 vision encoder)。
  • 需要 GPU 加速(CUDA / Metal)。
  • 想要 fine-tune,或在 Rust 裡訓練。
  • 想要把同一份模型程式碼跑在多種 backend 上。

tiny-llm-runner 這種東西的價值:

  • 教學:把每一個量化格式、每一個 RoPE 變體、每一個 attention 細節都自己摸過一遍。
  • 可審計:當 forward pass 出 bug 時,你不需要去翻 30 個檔案、跨 5 個 trait——全部就在 runner.rs 那 100 行裡。
  • 超低依賴:整個 Cargo.toml 只有 anyhowclaphalfmemmap2rayon,再加我自己的 llm-gguf-parser。沒有 BLAS、沒有 CUDA toolkit、沒有 LLVM。
  • 特殊嵌入場景:如果你要把推論塞到一個極小的環境(IoT、WASM、安全沙盒),tiny-llm-runner 這種「沒有外部依賴的 forward pass」是個還算合理的起點。

我從這次比較中學到的事

寫了 tiny-llm-runner 再去看 candle 的程式碼,最大的收穫不是「我發現 candle 哪裡寫得不好」(事實上它寫得很好),而是讓我清楚看到「為了通用性需要付出多少抽象成本」

candle 的每一層抽象——TensorDeviceDTypeModuleVarBuilderQMatMul——都不是憑空長出來的,都是在解一個「如果不抽象,使用者就會痛」的問題。但反過來說,如果你的使用情境不需要那種通用性(例如「我只想跑 Llama Q4_0、只在 CPU 上、只做 inference」),那這些抽象對你而言就是純粹的閱讀負擔。

這也是為什麼我覺得 tiny-llm-runner 並不是 candle 的「失敗版」——它是在另一個座標點上的一個合理選擇。抽象是手段,不是目的。 當你的目標是「我要把 forward pass 看懂」時,零抽象的版本反而是最對的設計。

下次如果有空,我想拿 tiny-llm-runner 的 forward pass 對 candle 的 models::llama::Llama::forward 做一個逐行對照——應該會是另一篇有趣的文章。


原始碼:

相關閱讀:


用純 Rust 跑 Llama:從 GGUF 到 Token 的零 C 依賴推論引擎

featured.svg

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

之前的 llm-local-studio 系列我都是把 llama.cpp 當作黑盒子,透過 FFI 把 GGUF 檔案丟進去、Token 拿出來。雖然這條路最快、效能也最好,但每次看到「forward pass」這個詞從文件裡飄過去,心裡總是會冒出一個聲音:「裡面到底發生了什麼事?我能不能自己用 Rust 從頭寫一遍?」

於是有了這次的 tiny-llm-runner:一個純 Rust、零 C 依賴的 Llama 架構推論引擎,直接吃我先前寫的 llm-gguf-parser 解析出來的 GGUF 檔,跑出和 llama.cpptemperature=0逐位元組一致的輸出。

這篇文章不打算逐行講 Transformer 數學(網路上太多了),而是要回答一個更實際的問題:要把一個會跑 LLM 的專案寫出來,到底需要哪些零件?這些零件之間怎麼接? 我會用 tiny-llm-runner 的九個檔案當作骨架,一塊一塊看下去。


整體骨架:九個檔案、各司其職

整個專案不到兩千行 Rust,但拆得算清爽。先看一眼 src/ 底下長什麼樣:

檔案 角色
config.rs 從 GGUF metadata 撈 Llama 的超參數(層數、頭數、context length…)
dequant.rs 量化資料的解碼與點積核心(F32/F16/Q8_0/Q4_0/Q6_K)
tensor.rs 一個薄薄的 TensorView,蓋在 mmap 切片上
model.rs 把每個 Llama 張量(token_embdblk.N.*output*)找出來、放好
ops.rs RMSNorm、matvec、softmax、RoPE、SiLU、向量加法
runner.rs 真正的 forward pass + KV cache
tokenizer.rs SentencePiece-BPE encode/decode,含 byte-fallback
sampler.rs Greedy / temperature / top-k 採樣
main.rs CLI:載檔 → encode → prefill → decode → 印 tok/s

從上往下看就是一個資料流:檔案 → metadata 與張量視圖 → 數學運算 → 一步 forward → token → 字串。下面就照這個順序,一站一站走過去。


第一站:把 GGUF 當作 mmap 來看

main.rs 的開頭很單純,但其實已經把整個專案的「省記憶體」哲學立起來了:

let file = File::open(&args.model)?;
let mmap = unsafe { Mmap::map(&file)? };
let gguf = parse_gguf(&mmap)?;

這裡的 Mmap 來自 memmap2,它把整個 GGUF 檔(可能是好幾 GB)對應到處理程序的虛擬位址空間,但沒有真的把資料讀進 RAM。作業系統會在你真正讀取某個位址時才把對應的硬碟分頁載進來。這個性質非常重要,因為 LLM 模型檔案 99.9% 的空間都是張量權重,而我們的 forward pass 每次只會用到一小塊。

接下來的這一行是樞紐:

let blob = &mmap[gguf.data_offset as usize..];
let model = LlamaModel::load(config.clone(), &gguf.tensors, blob)?;

gguf.data_offset 是 GGUF 檔頭結束、張量資料開始的位置。從這個 offset 之後的 byte 切片就是「所有權重的原始 bytes」。整個專案後面的所有計算都是直接在這個 blob 上做切片,從來沒有把模型完整解壓進記憶體過。


第二站:config.rs —— 把 metadata 變成型別

LlamaConfig 是個樸素的 struct:

pub struct LlamaConfig {
    pub n_ctx: usize,       // context length
    pub n_embd: usize,      // hidden dim
    pub n_layer: usize,     // transformer block 數
    pub n_head: usize,      // query head 數
    pub n_head_kv: usize,   // K/V head 數(GQA 時 < n_head)
    pub n_ff: usize,        // FFN 中間維度
    pub vocab_size: usize,
    pub rms_eps: f32,
    pub rope_freq_base: f32,
    pub rope_dim_count: usize,
}

from_gguf 做的事就是把 llama.embedding_lengthllama.attention.head_count 這些 GGUF metadata key 一個一個撈出來,順便驗一些不變式(例如 n_embd % n_head == 0n_head % n_head_kv == 0)。

值得一提的是 n_head_kv:當它小於 n_head 時,模型用的是 Grouped-Query Attention(GQA)——多個 query head 共用同一組 K/V。這是 Llama 2/3 為了壓縮 KV cache 大小常用的招數,後面 runner.rs 在算 attention 時會直接用到 gqa_groups() 這個輔助方法。


第三站:dequant.rs —— 量化格式的搬運工

這是整個專案最「接地氣」的一塊,也是和 llama.cpp 對齊得最仔細的一塊。GGML 的量化格式長這樣:

格式 一個 block 的元素數 block 大小 怎麼存
F32 1 4 bytes 直接 little-endian
F16 1 2 bytes half-float
Q8_0 32 34 bytes 1 個 fp16 scale + 32 個 int8
Q4_0 32 18 bytes 1 個 fp16 scale + 16 個 packed int4
Q6_K 256 210 bytes ql(128B) + qh(64B) + 16 個 i8 sub-scale + 1 個 fp16 super-scale

每一種格式我都實作了兩個函式:

  • dequant_row_*(q, out):把整個 row 還原成 f32 陣列。
  • dot_*(q, x):直接算這個 row 跟一個 f32 向量的內積,不額外解壓。

第二種才是 hot path。一次推論要做幾百萬次內積,如果每次都先把整個 row 解壓出來再做 dot,記憶體頻寬會吃不消。dot_q4_0 的精髓是這一段:

for _ in 0..n_blocks {
    let d = read_f16(&q[qp..qp + 2]);   // 這個 block 的 scale
    qp += 2;
    let mut s = 0.0f32;
    for i in 0..QK4_0 / 2 {
        let byte = q[qp + i];
        let lo = (byte & 0x0F) as i32 - 8;   // 低 nibble
        let hi = (byte >> 4) as i32 - 8;     // 高 nibble
        s += lo as f32 * x[xp + i];
        s += hi as f32 * x[xp + i + QK4_0 / 2];
    }
    acc += d * s;   // 整個 block 共用一個 scale
}

注意那個 (nibble) - 8:Q4_0 的 4-bit 值代表的是 [-8, 7] 的有號整數,scale 統一乘在最後。這也是 Q4_0 之所以叫做「對稱量化」的原因——沒有 zero point。

至於為什麼還要支援 Q6_K?因為 llama.cpp 在量化模型時,output.weight(lm_head)幾乎都會保留成 Q6_K,即使其他層是 Q4_0 也一樣。這是因為最後那層直接決定 logits 分佈,敏感度最高,省那點空間反而會讓品質明顯下降。所以如果你不支援 Q6_K,TheBloke 上幾乎所有 Q4_0 GGUF 都跑不起來。


第四站:tensor.rs —— mmap 上的薄殼

TensorView 大概是整個專案最 Rust 的部分:

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

沒有擁有任何資料——'a 是 mmap 的生命週期。整個 LlamaModel 裡幾十個 TensorView 加起來也只是幾百 byte 的中繼資料,真正的權重 bytes 還在硬碟上躺著。

這個型別只暴露兩個有意思的方法:

pub fn dot_row(&self, i: usize, x: &[f32]) -> f32;
pub fn dequant_row(&self, i: usize, out: &mut [f32]);

第一個是 hot path(matvec 用),第二個是 cold path(embedding lookup 用)。內部就是一個 match,根據 ggml_type 派發到 dequant.rs 對應的核心去。換句話說,所有「這個張量是哪種量化」的條件分支只發生在 row 邊界,不會滲進內層迴圈。


第五站:model.rs —— 把 GGUF 張量名對應回 Llama 結構

GGUF 把張量存成扁平的清單,每個張量有個名字字串,像是 blk.0.attn_q.weightblk.0.ffn_gate.weighttoken_embd.weightoutput_norm.weight…。model.rs 的工作就是把這些字串對應回有結構的 LlamaModel

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>>,
}

注意一個小細節:兩種 norm 權重(attn_normffn_normoutput_norm)我直接 dequant 成 Vec<f32>,但 Q/K/V/O、gate/up/down 我留成 TensorView。為什麼?因為 norm 的權重維度很小(就是 n_embd 而已,幾千個 float),一次解出來放在記憶體裡完全沒負擔;但 Q/K/V 那些是 [n_embd, n_embd] 的大矩陣,留在原本量化格式裡才能享受 dot_row 的加速。

還有一個容錯:

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

有些模型會把 lm_head 和 token embedding 綁定(tied embeddings),GGUF 裡就不會單獨存 output.weight。這時候我們直接重用 token_embd 當作輸出投影即可。


第六站:ops.rs —— Transformer 的六種樂高

這個檔案大概只有 100 行,但它就是整個 Transformer 的所有原語:

  • rmsnorm(out, x, w, eps):根均方歸一化。注意 ssf64 累加,這是和 llama.cpp 一致的細節,不然在大維度下浮點誤差會被放大。

  • matvec(out, w, x):矩陣乘向量。這是整個專案的效能熱點,所以我用 rayon 對 row 平行化:

    out.par_iter_mut().enumerate().for_each(|(i, o)| {
        *o = w.dot_row(i, x);
    });

    每個執行緒拿一個 row 的 slice,獨立呼叫 dot_row——零鎖、零共享,scale 得很乾淨。

  • softmax(x):標準的 max-shift + exp + normalize。

  • apply_rope(...):Rotary Position Embedding。這個之後會獨立講。

  • silu(x)x * sigmoid(x),FFN 的活化函數。

  • add_inplace(a, b):殘差連接的左半邊。

RoPE 的兩種約定

apply_rope 有個容易踩坑的地方——它支援兩種完全不同的旋轉約定:

pub enum RopeStyle {
    Llama,   // (x[2i], x[2i+1]),舊版 convert.py 用的「permuted」格式
    Neox,    // (x[i], x[i+d/2]),新版 convert_hf_to_gguf.py 用的格式
}

兩者的數學本質一樣(都是把 head_dim 兩兩配對成複數來旋轉),但配對方式不同llama.cppconvert.py 時代會把 Q/K 矩陣的 row 重排(permute)成 adjacent-pair 格式,之後在執行階段就直接用 Llama 風格旋轉;而 HF 原生格式不做 permute,所以要用 Neox 風格旋轉。

如果你拿到的 GGUF 是 TheBloke 那批早期轉檔(包括我用來驗證的 TinyLlama Q4_0),就要用 --rope llama;如果是現代 convert_hf_to_gguf.py 轉的,就用 --rope neox搞錯了不會 crash,但生出來的內容會像被酒駕的 LLM 寫的——這個 bug 我親自踩過。


第七站:runner.rs —— Forward Pass 的指揮中心

Runner 是把所有零件串起來的那一層。它持有 LlamaModel 的借用、KV cache、以及一堆預先配好的 scratch buffer:

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>>,
    pos: usize,

    // 複用的 scratch:避免每次 forward 重新配記憶體
    x: Vec<f32>, xb: Vec<f32>, xb2: Vec<f32>,
    hb: Vec<f32>, hb2: Vec<f32>,
    q: Vec<f32>, att: Vec<f32>,
    logits: Vec<f32>,
}

每呼叫一次 forward(token),就執行一個 token 的完整推論。流程拆開來看是這樣的:

flowchart TD A[token id] --> B[model.embed → x] B --> C{每層 0..n_layer} C --> D[rmsnorm x → xb] D --> E[wq·xb → q
wk·xb → kcache row
wv·xb → vcache row] E --> F[apply_rope on q & k] F --> G[每個 head:
q·K 算 attention scores
softmax
加權 V → xb] G --> H[wo·xb → xb2
x += xb2 殘差] H --> I[rmsnorm x → xb
FFN: silu·gate * up
w_down → xb2
x += xb2] I --> C C --> J[output_norm → xb] J --> K[lm_head: output·xb → logits] K --> L[回傳 logits slice]

幾個值得特別講的設計:

  1. K/V 直接寫進 cachematvec(krow, &layer.wk, &self.xb) 裡的 krowkcache[l]pos 位置的 slice。換句話說,K 算完之後根本不需要另一個中間變數——它就是 cache 的那一行。RoPE 也是直接在 cache 上 in-place 做,整個 attention 過程不會把 K/V 複製過第二次。
  2. GQA 的對應:每個 query head h 對應到 kv_head = h / gqa,所以多個 query head 會去讀同一份 K/V。這在記憶體佈局上很自然——你只需要算對 offset 就好。
  3. 單序列、單 batch:這個 runner 不做 prompt batching,prefill 階段是一個 token 接一個 token 跑。對教學版來說這很值得——因為 batch 會把 forward 變成 matrix 乘 matrix,遮罩和記憶體佈局會複雜一個量級。

第八站:tokenizer.rs —— SentencePiece-BPE 的迴圈

Llama 的 tokenizer 是 SentencePiece 訓練出來的 BPE 變種。GGUF 把 token 字串和分數都存在 metadata 裡,所以 from_gguf 只是把它們撈出來建一個 HashMap<String, u32>

encode 的演算法本身很短,但很容易寫錯:

loop {
    // 在所有相鄰 pair 中,找出「合併後是 vocab 裡的 token」
    // 且分數最高的那一對。
    let (best_idx, best_id) = ...;
    match best_idx {
        Some(i) => { ids[i] = best_id; ids.remove(i + 1); }
        None => break,
    }
}

這就是「最高分鄰接合併」的標準 SentencePiece 編碼迴圈。容易踩的雷有兩個:

  • 空白要先換成 (U+2581):SentencePiece 用這個特殊符號代表 word boundary,整個輸入還要在最前面加一個
  • byte fallback:如果某個字元(例如中文、Emoji)不在 vocab 裡,就把它的 UTF-8 bytes 一個一個用 <0x00> ~ <0xFF> 這些 byte token 接起來。decode 時要反過來把這些 byte token 重新拼成 UTF-8——所以 decode() 我是先收集到 Vec<u8>,最後用 String::from_utf8_lossy 一次轉,避免單獨 decode 每個 byte 時切到 codepoint 中間造成亂碼。

第九站:sampler.rs —— 從 Logits 拿到下一個 Token

這個檔案最短,但它是「讓模型有溫度」的地方:

pub fn sample(&mut self, logits: &mut [f32]) -> u32 {
    if self.temperature <= 0.0 {
        return argmax(logits) as u32;   // greedy
    }
    // 1. 套用溫度
    for v in logits.iter_mut() { *v *= 1.0 / self.temperature; }
    // 2. top-k cutoff
    if self.top_k > 0 && self.top_k < logits.len() {
        // 用 select_nth_unstable_by 做 partial sort
        ...
    }
    // 3. softmax → 累積分佈 → 抽樣
    softmax(logits);
    let r = self.next_f32();
    let mut acc = 0.0f32;
    for (i, &p) in logits.iter().enumerate() {
        acc += p;
        if acc >= r { return i as u32; }
    }
}

PRNG 用的是極小的 xorshift64,目的是讓「同一個 seed → 同一個輸出」變得可預期,方便和 llama.cpp 對拍。我自己驗證的方式很粗暴:把 temperature 設 0(greedy),跑 64 個 token,然後 diff 我的輸出和 llama.cpp 的輸出。只要有一個 byte 不一樣,就一定有 bug 等著我抓——因為 greedy 是完全 deterministic 的。


把零件接起來:main.rs 的旅程

把上面九個零件串起來其實只有薄薄一層 CLI:

// 1. 載入:mmap → parse_gguf → LlamaConfig + LlamaModel + Tokenizer
let mmap = unsafe { Mmap::map(&file)? };
let gguf = parse_gguf(&mmap)?;
let config = LlamaConfig::from_gguf(&gguf)?;
let model = LlamaModel::load(config.clone(), &gguf.tensors, &mmap[gguf.data_offset as usize..])?;
let tokenizer = Tokenizer::from_gguf(&gguf)?;

// 2. Encode
let prompt_ids = tokenizer.encode(&args.prompt, !args.no_bos);

// 3. Prefill:把 prompt 一個一個餵進去,丟掉 logits(最後一個保留)
let mut runner = Runner::new(&model, rope_style);
let mut last_logits = None;
for &tok in &prompt_ids {
    last_logits = Some(runner.forward(tok).to_vec());
}

// 4. Decode:迴圈採樣 → 印出 → 餵回去
let mut logits = last_logits.unwrap();
for _ in 0..args.n_predict {
    let next = sampler.sample(&mut logits);
    if next == tokenizer.eos { break; }
    print!("{}", tokenizer.decode(&[next]));
    logits = runner.forward(next).to_vec();
}

整個 token 流程從 byte 進來到 byte 出去,每一段都是可審計的 Rust,沒有任何 C library 在底下偷偷做事。


我從這個專案學到的事

寫這個 runner 之前,我一直以為 LLM 推論的「神祕」程度應該很高。實際做完才發現:真正讓人卡住的不是數學,是各種約定(convention)

  • RoPE 有兩種配對方式,搞錯了模型會講胡話。
  • Q4_0 nibble 是先低後高、要減 8、scale 是 fp16,每一條都是會把人坑半天的細節。
  • output.weight 常常和其他層用不同的量化格式。
  • Tokenizer 的 替換、byte fallback、合併迴圈——任何一步寫錯都會生出和 llama.cpp 不一樣的 token 序列。

但反過來說,正因為這些約定都是 ggml/llama.cpp 既有檔案決定好的,只要你願意一條一條對齊,純 Rust 是真的可以跑。我在 temperature=0 下用 TheBloke 的 tinyllama-1.1b-chat-v1.0.Q4_0.gguf 比對過,輸出逐字逐句和 llama.cpp 一樣。對我這種喜歡「把黑盒子拆開」的工程師來說,這比效能再快一倍都還令人滿足。

當然,這個專案明確不做的事情也很多:K-quants(除了 Q6_K)、SIMD、GPU 後端、prompt batching、beam search、sliding window、MoE、其他架構……這些都是 llama.cpp 多年來辛苦累積的工程成果,要在純 Rust 裡複製一遍是另一個量級的工作。但作為一個「我自己讀得懂每一行 forward pass」的學習版,tiny-llm-runner 已經達到了它的目的。

下一步如果有空,我想試試看把 matvec 改成 SIMD 版本(用 std::simdwide),或者把 Q4_K_M 也加進來——後者才是現在 GGUF 世界的真主流。


原始碼:github.com/p47t/rust-52-projects/tree/master/tiny-llm-runner

相關閱讀: