tiny-llm-runner 深入解讀 (3):tensor.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 裡。
這個設計的意義是:
- 零拷貝:建構一個
TensorView 只需要拷貝 32 bytes 的中繼資料(dims + type + slice header),沒有任何資料搬運。
- 生命週期安全:因為
'a 綁定到 mmap,編譯器會保證這個 view 不會比 mmap 活得更久。
- 可以是
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>。這個設計的問題是:
- 每次 dot 都要走 vtable,無法 inline。
- 每個 row 多一次間接跳轉,CPU 預測會變差。
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 深入解讀 (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 —— 把 GGUF metadata 變成 Rust 強型別
本文由 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 數。
對沒做過 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 存成 u32 或 u64;某個訓練框架可能會把它存成 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。它做的事是:如果這個 Result 是 Err,就把錯誤 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 架構的硬性要求:
n_embd / n_head = head_dim——每個 attention head 平均分配 hidden dim。
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 而是每次算?因為:
- 避免不一致:如果同時存
n_head、n_embd、head_dim,理論上有人可能改了其中一個忘了改另外兩個,三者就對不起來。只存 base 量是 single source of truth。
- divide 的成本可忽略:這些方法不在 hot path 上,runner 是在
new() 時呼叫一次來算 buffer 大小,不會在 inner loop 反覆呼叫。
這是 Rust struct design 的一個小哲學:只存 fundamental fields,derived 量寫成方法。
從 hyperparameter 算出模型大小
在動手算之前,先用一張圖把這些欄位放到推論流程上看看,就會比較清楚每個參數實際在控制哪一段:
有了這 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,就是這裡膨脹得最兇)。
每層裡面有 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 = 64,kv_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”——這就是模型名字裡那個數字的來源。
從這個算式還能看出兩件事
- FFN 才是主導參數量的部分。
3 × n_embd × n_ff 通常比 4 × n_embd × head_dim × (n_head + n_head_kv) 大上 3~4 倍。所以模型放大的時候,最先膨脹的是 FFN,再來才是層數。
- 記憶體佔用可以反推。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 的架構對照
本文由 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.toml 和 src/ 並排放,就會發現它們其實在解非常不同的問題。
先把比較的基準擺正
要公平比較,得先承認兩者的「設計願景」根本不在同一個量級:
| 維度 |
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_row 和 dequant_row。
candle:完整的 PyTorch 風格 Tensor
candle_core::Tensor 的概念負擔則完全是另一個世界。它要追蹤:
- Shape:可變維度的張量形狀,支援 broadcasting。
- DType:
F32、F16、BF16、U8、U32、I64、F64…
- Device:
Cpu、Cuda(...)、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 只做兩件事:
- 把 GGUF 的 tensor 列表掃一遍,建一個
name → TensorInfo 的 HashMap。
- 對每個我需要的張量名(
token_embd.weight、blk.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 檔。
載入流程大致是:
gguf_file::Content::read(...) 把整個 metadata 讀進來。
- 對每個張量,用
tensor.qmatmul(...) 或 tensor.dequantize(...) 取得 QTensor。
- 把
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 用的是同一套 RmsNorm 和 Mlp,只是把 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 只有 anyhow、clap、half、memmap2、rayon,再加我自己的 llm-gguf-parser。沒有 BLAS、沒有 CUDA toolkit、沒有 LLVM。
- 特殊嵌入場景:如果你要把推論塞到一個極小的環境(IoT、WASM、安全沙盒),
tiny-llm-runner 這種「沒有外部依賴的 forward pass」是個還算合理的起點。
我從這次比較中學到的事
寫了 tiny-llm-runner 再去看 candle 的程式碼,最大的收穫不是「我發現 candle 哪裡寫得不好」(事實上它寫得很好),而是讓我清楚看到「為了通用性需要付出多少抽象成本」。
candle 的每一層抽象——Tensor、Device、DType、Module、VarBuilder、QMatMul——都不是憑空長出來的,都是在解一個「如果不抽象,使用者就會痛」的問題。但反過來說,如果你的使用情境不需要那種通用性(例如「我只想跑 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 依賴推論引擎
本文由 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.cpp 在 temperature=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_embd、blk.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 上做切片,從來沒有把模型完整解壓進記憶體過。
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_length、llama.attention.head_count 這些 GGUF metadata key 一個一個撈出來,順便驗一些不變式(例如 n_embd % n_head == 0、n_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.weight、blk.0.ffn_gate.weight、token_embd.weight、output_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_norm、ffn_norm、output_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 當作輸出投影即可。
這個檔案大概只有 100 行,但它就是整個 Transformer 的所有原語:
-
rmsnorm(out, x, w, eps):根均方歸一化。注意 ss 用 f64 累加,這是和 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.cpp 在 convert.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]
幾個值得特別講的設計:
- K/V 直接寫進 cache:
matvec(krow, &layer.wk, &self.xb) 裡的 krow 是 kcache[l] 在 pos 位置的 slice。換句話說,K 算完之後根本不需要另一個中間變數——它就是 cache 的那一行。RoPE 也是直接在 cache 上 in-place 做,整個 attention 過程不會把 K/V 複製過第二次。
- GQA 的對應:每個 query head
h 對應到 kv_head = h / gqa,所以多個 query head 會去讀同一份 K/V。這在記憶體佈局上很自然——你只需要算對 offset 就好。
- 單序列、單 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::simd 或 wide),或者把 Q4_K_M 也加進來——後者才是現在 GGUF 世界的真主流。
原始碼:github.com/p47t/rust-52-projects/tree/master/tiny-llm-runner
相關閱讀: