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

本文由 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/ 底下長什麼樣：

從上往下看就是一個資料流：檔案 → 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_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 的量化格式長這樣：

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

• 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 當作輸出投影即可。

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

這個檔案大概只有 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 的完整推論。流程拆開來看是這樣的：

幾個值得特別講的設計：

1. K/V 直接寫進 cache：matvec(krow, &layer.wk, &self.xb) 裡的 krow 是 kcache[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::simd 或 wide），或者把 Q4_K_M 也加進來——後者才是現在 GGUF 世界的真主流。

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

相關閱讀：

• 用 Rust 手寫 GGUF 解析器（這個專案的上游）
• 用 Rust 打造本地 LLM 工作站（用 FFI 包 llama.cpp 的反面實驗）