量化 llama 的一次前向

主线位置：推理主线的中段，ModelWeights::from_gguf 和 forward(input, pos)。生成一个 token 已经给读 GGUF、建模型、prefill、每个 decode 步分段计时，并讲了 Metal 上 decode 走 sdpa、冷缓存时编译 .metal 的代价。本页往里走一层：from_gguf 要哪些名字，forward 里每一步在做什么，一个 decode 步的时间按层里的哪一段分，以及 index_pos 和 KV 缓存对不上时会发生什么。代码在 candle-transformers/src/models/quantized_llama.rs（0.11.0 时 692 行，含测试）。

flowchart TB
  ids["token ids 1 x T"] --> emb["Embedding f32"]
  emb --> n1["RmsNorm"]
  n1 --> qkv["QMatMul q k v"]
  qkv --> rot["rope 用 index_pos"]
  rot --> kv["cat 进 KV 缓存"]
  kv -->|"Metal 且 T 为 1"| sdpa["ops sdpa"]
  kv -->|其他| mm["repeat_kv matmul mask softmax"]
  sdpa --> wo["QMatMul o 加残差"]
  mm --> wo
  wo --> n2["RmsNorm"]
  n2 --> mlp["QMatMul gate up down"]
  mlp -->|下一层| n1
  mlp --> out["最后一个位置 logits"]

本页的实验在 models-exp 的 qllama 和 attncore 里（release），模型是本地的 SmolLM2-135M-Instruct（Q8_0 GGUF：30 层、隐层 576、9 个查询头、3 个 KV 头、head_dim 64、词表 49152），平台 M2 Max、macOS 27、rustc 1.99.0。Metal 的数字都是热缓存（.metal 库已经在 macOS 的缓存里）。

from_gguf：键名写死了 llama.

ModelWeights::from_gguf 先用一个闭包 md_get 从元数据取超参数，找不到就报「cannot find {s} in metadata」：

键名都写死了 llama. 前缀。GGUF 按架构给键加前缀，quantized_qwen2.rs 读的就是 qwen2.attention.head_count；所以真正的 Qwen2 文件在这里第一个 ? 就失败，要用 quantized_qwen2。

然后按名字逐个 Content::tensor（读文件、在目标设备上建 QTensor，见 GGUF 文件）：

文件里的 QMatMul 是本文件自己定义的一个包装：一个 candle::quantized::QMatMul 加一个名为 qmatmul 的 tracing span，和 with_tracing.rs 里的同名类型重复。QMatMul::from_qtensor 按 dtype 决定是保持量化还是反量化成 f16 / f32，见 QTensor 与 QMatMul。

precomput_freqs_cis 一次算好 MAX_SEQ_LEN = 4096 个位置的 cos / sin 表（(4096, rope_dim / 2)），每层克隆一份 Tensor（Arc，不复制数据）。index_pos + seq_len 超过 4096 时，narrow 报错：

index_pos 4096: Err(narrow invalid args start + len > dim_len: [4096, 32], dim: 0, start: 4096, len:1)

这个上限和模型训练时的上下文长度无关，是文件里的一个常量。

RoPE：两种约定，按架构名选

from_gguf 还读 general.architecture，在一个列表里（qwen、qwen2、falcon、phi3、stablelm 等 20 个）就用 rotary_emb::rope（NEOX 约定：第 i 维和第 i + d/2 维配对），否则用 rope_i（第 2i 和 2i+1 维配对，Llama 的约定）。两个函数和它们的融合内核见常用层。

既然 llama.* 键名把真正的 Qwen2 文件挡在了门外，这个列表只在「架构名不是 llama、键却是 llama.*」的文件上起作用。约定选错是什么样子：把 SmolLM2 的 general.architecture 在内存里改成 qwen2（rope），贪心生成 20 个 token：

llama (rope_i):  "Rust is a modern, high-performance, and object-oriented programming language developed by Brian K"
as qwen2 (rope): "Rust\nRustin\nRustin\nRustin the world\nRust\n"

不报错，只是位置信息被打乱，模型很快开始重复。

forward：mask 什么时候建

ModelWeights::forward：

    pub fn forward(&mut self, x: &Tensor, index_pos: usize) -> Result<Tensor> {
        let (_b_sz, seq_len) = x.dims2()?;
        let mask = if seq_len == 1 {
            None
        } else {
            Some(self.mask(seq_len, index_pos, x.device())?)
        };
        let _enter = self.span.enter();
        let mut layer_in = self.tok_embeddings.forward(x)?;
        for layer in self.layers.iter_mut() {
            let x = layer_in;
            let residual = &x;
            let x = layer.attention_norm.forward(&x)?;
            let attn = layer.forward_attn(&x, mask.as_ref(), index_pos)?;
            let x = (attn + residual)?;

            // MLP
            let _enter = layer.span_mlp.enter();
            let residual = &x;
            let x = layer.ffn_norm.forward(&x)?;
            let x = layer.mlp_or_moe.forward(&x)?;
            let x = (x + residual)?;
            layer_in = x
        }
        let x = self.norm.forward(&layer_in)?;
        let x = x.i((.., seq_len - 1, ..))?;
        let _enter = self.span_output.enter();
        self.output.forward(&x)
    }

• 只有输入多于一个 token 时才要 mask。 decode 时（seq_len == 1）唯一的查询位置本来就可以看所有键，不需要遮。
• mask 按 (seq_len, kv_len) 缓存。 mask 先查 self.masks: HashMap<(usize, usize), Tensor>，没有才调 build_causal_mask：在主机上建一个 Vec<u8>（1 表示不能看），from_slice 传到设备上。kv_len = index_pos + seq_len，前 index_pos 列全 0（缓存里的前缀都能看），后 seq_len 列是因果三角。这个表只增不减，clear_kv_cache 不清它。
• mask 里存的是 kv_len，不是缓存的真实长度。 它由调用方传的 index_pos 算出来；缓存实际有多长，forward_attn 里 cat 完才知道。两者不一致时，见下文「index_pos 传错」。
• 只算最后一个位置的 logits。 x.i((.., seq_len - 1, ..)) 在输出层之前就把其余位置丢了，prefill 13 个 token 也只返回一行 (1, vocab)。要每个位置的 logits（例如算困惑度）得自己改。

forward_attn：自己拼 KV

forward_attn 的骨架在生成一个 token 和 KV 缓存与注意力里都引过，这里只补没讲的几点：

• q、k、v 三个投影各是一个 QMatMul，输入 (1, T, 576)，输出 reshape 成 (1, T, heads, 64) 再 transpose(1, 2)。只有 v 马上 contiguous()，注释说这是为了下面的快速内核，「actually a no-op except when processing the initial prompt」：T = 1 时转置后仍然连续。q、k 在 apply_rotary_emb 里 contiguous()，那里用 index_pos 从 cos / sin 表里 narrow 出 T 行。
• KV 缓存是每层一个 Option<(Tensor, Tensor)>，index_pos == 0 时丢掉旧的，否则在序列维 cat。
• 非 sdpa 分支先 repeat_kv 把 3 个 KV 头复制成 9 个（cat 再 reshape，源码注释说比广播快，见 #2043），再 matmul、masked_fill、softmax_last_dim、matmul。masked_fill 用 where_cond 把 mask 为 1 的位置换成预先建好的 -inf 标量。
• 这条路上每个张量都要按缓存长度复制：cat 复制旧缓存，repeat_kv 把它再复制 3 份。下一节量这一点。

一个 decode 步里时间去了哪

quantized_llama 在 ModelWeights 上有 model、output 两个 span，每层有 attn、attn-rot、attn-mlp 三个，每个 QMatMul 一个 qmatmul，每个 RmsNorm 一个 rms-norm（span 怎么创建、为什么要在建模型之前装 subscriber，见 BERT）。实验装一个自己写的 tracing_subscriber::Layer，在 on_enter / on_exit 记时间，按名字汇总；「自身」是减去在它进入期间进入的其他 span 后的时间。先完整生成两遍预热，再量 prefill 和第 20 个 decode 步（位置 32）：

墙钟时间：CPU 上 prefill 20.8 ms、decode 5.0 ms（CPU 上 forward 返回时就算完了）；Metal 上 prefill 9.9 ms，decode 的 forward 0.94 ms 就返回，synchronize 之后是 2.25 ms。这台机器同时在跑别的实验，CPU 的 decode 在 5 到 8 ms 之间波动，表里是负载最低的一次。

几点：

• prefill 时 attn-rot 很贵：60 次 4.6 ms，每次约 77 µs，而 decode 时每次 1 µs。CPU 上的 rope 按 par_chunks 交给 rayon，和下面 softmax 是同一个问题：RAYON_NUM_THREADS=1 时同样的 60 次合计 1.6 ms。
• span 的个数就是结构。 qmatmul 211 次 = 30 层 × 7 个投影 + 输出层；rms-norm 61 次 = 30 × 2 + 最后一个；attn-rot 60 次 = 30 × (q, k)。attn-mlp 的总时间里含着 ffn_norm 和三个 qmatmul。
• CPU 上 decode 约一半的时间不在矩阵乘里。 attn 的自身时间（除去 4 个投影和 rotary）和全部 211 个 qmatmul 相当，这一次还多出 40%（另一次负载较高的运行是 3.6 对 3.8 ms）。单 token 时每个投影只是一次量化矩阵乘向量，很快；剩下的时间花在注意力那十几个小 op 上。
• Metal 上 span 量到的是编码时间。 forward 返回时 GPU 还没算完，span 统计的是 CPU 把内核编进命令缓冲区的时间；synchronize 之后的墙钟时间才是 GPU 真正算完的时间。想在 Metal 上按层分 GPU 时间，span 做不到，要每层后面 synchronize（会改变提交节奏，见 Metal 后端）或者用 Xcode 的 GPU 工具。

把 CPU 上非 sdpa 分支的注意力拆开，逐个 op 计时（attncore：照着 forward_attn 的写法用随机张量重放一层，每个 op 后计时，300 次平均）：

• 上下文 32 时，一层注意力里最贵的是 softmax_last_dim：只有 9 行、每行 33 个数，却要几十微秒。它的 CPU 实现按行 par_chunks 交给 rayon 线程池，唤醒和调度的开销远大于这点计算；RAYON_NUM_THREADS=1 时降到 9 µs 左右。
• 上下文 1024 时，cat、repeat_kv 和两次 matmul 都随缓存长度线性增长。只是复制数据的 cat 加 repeat_kv 合计 129 µs，比两次真正计算的 matmul（56 µs）还多。
• 30 层乘起来，上下文 1024 时光这几个 op 就是每个 token 好几毫秒，这是下一节 CPU decode 随上下文变慢的主要原因。

线程数的影响，同一个程序只改 RAYON_NUM_THREADS 和 CANDLE_NUM_THREADS（两个池分别管 softmax 等 par_* 算子和量化矩阵乘，见 CPU 后端）：

对这个 135M 的模型，decode 时 4 个线程最快，1 个线程也比 8 个快：每个 op 都太小，分到 8 个线程上的收益抵不过调度。prefill 的 op 大一些，单线程就慢了一倍。最合适的线程数要按模型大小和机器自己量。

prefill 与 decode 的吞吐

throughput：prefill 送 T 个 token（index_pos = 0），量到 argmax 读回为止；decode 先 prefill 一段上下文，再连续 32 步，取平均。每种设置先整轮跑一遍预热，下面是第二轮：

• prefill 摊薄了固定开销。 Metal 上送 4 个和送 64 个 token 都是 9 到 10 ms，送 128 到 256 个时每秒七千七百多 token，512 个时回落（注意力的 T × T 部分开始显眼）。送 1 个 token 走的是 decode 的 sdpa 路和矩阵乘向量内核，所以只要 2.3 ms。CPU 上 prefill 的吞吐在每秒 420 到 550 token 封顶，和 decode 的每秒 140 个左右只差三四倍。
• decode 随上下文变慢。 CPU 上从上下文 16 到 1024，每个 token 的时间涨了一倍多，原因是上一节的 cat、repeat_kv 和注意力本身；Metal 上走 sdpa，只从 2.6 ms 涨到 3.4 ms。
• 每一步都 cat 整段缓存，生成 n 个 token 一共复制约 n²/2 行，这在 KV 缓存与注意力里和预分配的 KvCache 对比过。

index_pos 传错

index_pos 是调用方自己维护的计数器，forward 不检查它和缓存长度是否一致。正确的值是「缓存里已经有多少个位置」。故意传错（pos，CPU；Metal 上输出逐字相同）：

correct index_pos:      "Rust is a modern, high-performance, and object-oriented programming language developed by Brian K"
decode index_pos +100:  "RRRRRRRRRRRRRRRRRRRR"
decode index_pos -3:    "R\nRust is a modern, object-oriented programming language developed by the Rust initiative, which"
decode index_pos 0:     "R\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n"
chunk 2 with index_pos 5: Err(cannot broadcast [5, 10] to [1, 9, 5, 13])
chunk 2 with index_pos 8: ok, next token "R"

前四行 prefill 都正确（第一个 token「R」来自 prefill），只有之后每个 decode 步的 index_pos 被改了：

• decode 时传错不报错。 单 token 没有 mask，index_pos 只用来从 rotary 表里取哪一行，以及判断是不是 0。
• 偏大 100：每个新 token 的位置编码都比真实位置远 100，模型看到的「自己在哪」和缓存里的键对不上，退化成一直重复。
• 偏小 3：差得少，相对位置还大致对，输出看起来像样，但已经和正确的那一行不同。这种错最难发现。
• 传 0：每一步都走 index_pos == 0 那一支，把缓存整个丢掉，模型每步只看见刚生成的那一个 token。
• 分两段 prefill 时传错会报错。 先送 8 个 token，再送剩下 5 个时传 index_pos = 5：mask 按 kv_len = 5 + 5 = 10 建，缓存拼完却有 8 + 5 = 13 列。错误信息说的是广播，不提 index_pos。

换 prompt 要不要 clear_kv_cache

clear_kv_cache 把每层的缓存设成 None，是 PR #3536（「Add clear_kv_cache to quantized_llama and quantized_qwen2」，首次进入 0.11.0）加的，提交说明是「Lets callers free cached attention state between independent conversations without recreating the model」。

动脑筋答案

cache，CPU（Metal 上逐字相同）：

fresh model, prompt 2:           "Python is a high-level, interpreted programming language developed by Guido van Rossum in 1"
after prompt 1, no clear, pos 0: "Python is a high-level, interpreted programming language developed by Guido van Rossum in 1"
after prompt 1, continue at pos 33: "Python is a high-level, interpreted programming language developed by Guido van Rossum in 1"
Error: cannot broadcast [13, 46] to [1, 9, 13, 13]

• 不清缓存、从 0 开始：和新模型完全相同。 forward_attn 看到 index_pos == 0 就不拼旧缓存，第二个 prompt 根本看不到第一段对话。所以对 quantized_llama 来说，「换 prompt 忘了 clear_kv_cache」不会让输出变乱；clear_kv_cache 的作用是在两次对话之间提前释放 30 层的缓存张量。
• 接着 33 往后送也正常（第三行）：这时缓存里真的有前一段对话，mask 是 13 × 46 的长方形，模型在前文之后回答了新问题，这里恰好和单独回答一样。
• 清了缓存却不把 index_pos 归零：报错（第四行）。mask 按 kv_len = 33 + 13 = 46 建，缓存却只剩这次的 13 个位置。如果第二段只送一个 token（没有 mask），连这个错误都没有，而且输出和从 0 开始完全一样：缓存清空后这个 token 只和自己做注意力，q、k 被 RoPE 转过同一个角度，点积不变。真正悄悄算错的是同一段对话里计数器和缓存对不上，见上一节和常见坑「清了缓存却沿用旧的 index_pos」。

真正要守住的不变式是：index_pos 等于缓存里已有的位置数。清缓存就把计数器归零，归零就不必清缓存。

小结

• from_gguf 只认 llama.* 元数据键和 blk.{i}.attn_q.weight 这类名字；没有 output.weight 时共享 embedding；embedding 和 RmsNorm 权重反量化成 f32，投影保持量化。
• RoPE 约定按 general.architecture 选（0.11.0 起）；选错不报错，输出退化成重复。rotary 表固定 4096 个位置，超过就报 narrow 错误。
• mask 只在 seq_len > 1 时建，在主机上算、按 (seq_len, kv_len) 永久缓存；只返回最后一个位置的 logits。
• CPU 上一个 decode 步约一半的时间在注意力的十几个小 op 上：短上下文时是 softmax_last_dim 的 rayon 调度，长上下文时是 cat 和 repeat_kv 的复制；Metal 上 span 只量到编码时间。
• index_pos 是没人检查的计数器：decode 时传错静默算错，多 token 时传错报广播错误。index_pos == 0 会丢掉旧缓存，所以换 prompt 不清缓存也没事；清了缓存却不归零才出错。