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、合併迴圈——任何一步寫歪，生出來的 token 序列就跟 llama.cpp 對不上了。

不過反過來想，正因為這些約定都是 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 的反面實驗）