跳到主要内容

数据截至 (上游 commit d7a2074112d2)

04 · llama 库:从模型加载到逐 token 生成

这一章讲什么: src/ 目录下的 libllama——GGUF 之上、server 之下的推理库。一次 llama_decode 调用内部发生了什么,模型架构是怎么「写」进代码的,KV cache 怎么管,采样器链怎么把 logits 变成一个 token。

1. 它要解决的小问题​

ggml 只给你「张量 + 图 + 后端」,它不认识 transformer。从「一张图」到「一个会说话的模型」,中间还缺一整层:

  • 几百个权重张量要按架构拼成层、层拼成模型;
  • 自回归生成要记住历史(KV cache),还要管多序列、滑动窗口;
  • 每个架构(LLaMA/Qwen/DeepSeek……)的层结构都不一样;
  • logits 要变成 token,中间有十几种采样策略。

libllama 就是这一层。它对外的表面是一组 C API(include/llama.h):llama_model_load_from_file 载模型、llama_init_from_model 开上下文、llama_decode 算前向、llama_sampler_sample 采样。

2. 加载:三个对象的分工​

llama_model_load_from_file(src/llama.cpp:464)走出的加载路径,分三个角色:

角色干什么锚点
llama_model_loader开文件、解析 GGUF、建 mmap、给每个张量名定位「在哪个文件、偏移多少」src/llama-model-loader.cpp:532、:1348
llama_model读 hparams(层数/维度/上下文长)、读词表;按架构把权重张量组织成层src/llama-model.cpp:1196(load_hparams)、:1366(load_vocab)
各架构类声明自己需要哪些张量、什么形状,向 loader「点菜」src/models/llama.cpp:36(load_arch_tensors)

load_tensors(src/llama-model.cpp:1390)负责把层放到设备上:按 n_gpu_layers 和 tensor_split 算出每张卡的层分配——默认按各卡剩余显存比例切(src/llama-model.cpp:1437-1456),然后逐层 create_tensor 到对应后端的 buffer。CPU 上的层直接用 mmap 指针,GPU 上的层此时才真正上传。

加载产物 llama_model 是只读的:权重、hparams、词表。所有「生成到一半」的可变状态都在另一个对象 llama_context 里(llama_init_from_model,src/llama-context.cpp:3636)——同一模型可以开多个 context 并行生成。

3. 构图:每个架构一段代码​

llama.cpp 不用 ONNX 之类的中间表示——架构就是代码。llama_model::build_graph(src/llama-model.cpp:2652)按 arch 分发到 src/models/ 下对应文件,那里有一个类直接用手写代码把一层 transformer「搭」出来。

以 LLaMA 为例,构图入口 llama_model_llama::graph(src/models/llama.cpp:99),一层的内容(节选自 src/models/llama.cpp:130-200 一带):

// 摘自 src/models/llama.cpp(精简,符号保留)
cur = build_norm(inpL, layer.attn_norm, ..., il);     // RMSNorm
Qcur = build_lora_mm(layer.wq, cur);                  // Q 投影
Kcur = build_lora_mm(layer.wk, cur);                  // K 投影
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, ...);       // 旋转位置编码(src/models/llama.cpp:146)
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, ...);
cur = build_attn(layer.wo, ..., Qcur, Kcur, Vcur, ...);  // 注意力(写 KV cache)
cur = build_norm(ffn_inp, layer.ffn_norm, ..., il);   // 第二个 RMSNorm
cur = build_ffn(cur, layer.ffn_up, layer.ffn_gate, layer.ffn_down, ..., il); // SwiGLU FFN

每个 build_* / ggml_* 调用只是往 ggml 图里挂节点(第 1 章的两段式)。n_layer 层循环完,再接输出 norm 和 output 投影,得到 logits 节点。

这种写法的取舍很 llama.cpp: 支持一个新架构 = 手写一个文件 + 一个转换器,没有「自动从 ONNX 导入」的魔法;换来的是每个 op 的选择、精度、融合都精确可控(比如 GLM4 的 wo 投影强制 FP32 累加规避数值问题,src/llama-graph.cpp:2842-2848)。

4. KV cache:cell 记账,写入即图节点​

4.1 它解决什么​

自回归生成的第 t 步,注意力要看 t-1 步之前所有 token 的 K/V——重算一遍是平方级浪费。KV cache 把每步的 K/V 存下来复用。难点在管理:多序列并发、序列删除/复制、滑动窗口(SWA)只留最近 N 个、缓存命中续写……

4.2 cell:最小的记账单位​

llama_kv_cells(src/llama-kv-cells.h:41)把 cache 看成一排 cell,每个 cell 记:属于哪(几)个序列、逻辑位置 pos 是多少、有没有被占。一组 cell 可以同时属于多个序列——llama_kv_cache::seq_cp 复制序列(src/llama-kv-cache.cpp:452)只是给已有 cell 加一个序列引用,不拷数据。前缀相同的两个对话因此共享缓存,这就是 server 端 prompt cache 的地基。

llama_kv_cache 本体(src/llama-kv-cache.cpp)持有各层的 K/V 大张量 + cells 元数据。新 batch 进来先 find_slot(src/llama-kv-cache.cpp:899)找一段能放下的 cell;放不下就由 decode 主循环触发 cache 整理再重试(src/llama-context.cpp:1762-1777 的 LLAMA_MEMORY_STATUS_FAILED_PREPARE 分支)。SWA、混合注意力(Mamba/线性注意力混合)等变体各有一副子类实现(llama-kv-cache-iswa.cpp、llama-memory-hybrid.cpp 等)。

4.3 妙处:cache 写入是图里的节点​

build_attn 的 KV 版本(src/llama-graph.cpp:2786)里,写 cache 不是图外胶水:

// 摘自 src/llama-graph.cpp:2824-2827(精简)
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));  // K 写入 cache = 一个图节点
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
ggml_tensor * k = mctx_cur->get_k(ctx0, il);   // 注意力读的 K = cache 张量的视图
ggml_tensor * v = mctx_cur->get_v(ctx0, il);

cpy_k/cpy_v 生成 ggml_cpy 节点把新 K/V 写进 cache 的指定槽位;注意力算子再拿整个 cache 的视图去读。好处是:写入自动参与调度和跨设备切分,RoPE 甚至能融合成「边算边直接写进 cache」(注释见 src/llama-graph.cpp:2812)。

5. decode 主循环:一次前向的全过程​

llama_context::decode(src/llama-context.cpp:1643,C API llama_decode 只是它的薄包装,src/llama-context.cpp:4212)做四件事:

① 校验 batch、初始化 batch 分配器 balloc
      │
② memory->init_batch():向 KV cache 申请槽位,
      │   并把逻辑 batch 切成若干 ubatch(微批)
      ▼
③ 对每个 ubatch:
      process_ubatch():建图(或复用)→ 分配图内存 → graph_compute()
      │
④ 收集输出的 logits / embeddings

两个值得记住的细节:

  • n_batch vs n_ubatch。 逻辑批 n_batch(一次 llama_decode 提交的 token 上限)与物理批 n_ubatch(一张图实际算的 token 上限)分离(include/llama.h:361-362 的字段注释):预填充一大段 prompt 时,decode 内部切成 ubatch 大小的段逐段算,每张图都不大、显存峰值可控。
  • 图复用。 process_ubatch(src/llama-context.cpp:1333)先问 res->can_reuse(gparams):图拓扑由 (ubatch 形状、memory context、图类型) 唯一决定,相同则跳过建图直接换输入(src/llama-context.cpp:1348-1358)。逐 token 解码阶段图形状不变,于是热路径上连建图开销都没有。

算图本身委托给第 1 章的调度器:graph_compute(src/llama-context.cpp:2488)设好线程数后调 ggml_backend_sched_graph_compute_async(src/llama-context.cpp:2507)。

6. 采样:一条过滤器链​

llama_decode 给出 logits,变成 token 是采样器的事。设计是责任链:llama_sampler_sample(src/llama-sampler.cpp:895)把 logits 包成候选数组 llama_token_data_array,交给链上每个采样器依次加工(llama_sampler_chain_apply,src/llama-sampler.cpp:681),最后一个采样器(通常是 dist)真正抽签。

logits ─► [penalties] ─► [top-k] ─► [top-p] ─► [min-p] ─► [temp] ─► [dist 抽签] ─► token
            惩罚重复        截断       截断       截断      缩放      随机采样

单个采样器极小:top-k 的全部逻辑是「排序取前 k、其余砍掉」(llama_sampler_top_k_apply,src/llama-sampler.cpp:1449);greedy 就是取 argmax(src/llama-sampler.cpp:1053)。组合顺序由应用层决定——common_sampler_init(common/sampling.cpp:187)按参数组装链,顺序在 common/sampling.cpp:353-406:先是 DRY/top-k/top-p/min-p/typical 等截断类,再温度,再 penalties,最后 dist 或 mirostat 收尾;grammar(约束解码,如强制 JSON)排在链首附近,先把不合法的 token 屏蔽掉。

接口的妙处: 每个采样器是同构的 llama_sampler(一组函数指针:apply/accept/reset/clone),llama_sampler_accept 把「刚出了这个 token」回灌给链上需要记忆的采样器(penalties、grammar 都靠它更新内部状态)。

7. 原理演示:一个最小的生成循环​

把本章串起来,库用户的生成循环长这样:

# 示意,非源码 —— 对应 include/llama.h:1231-1252 的官方示例骨架
model = llama_model_load_from_file("model.gguf", model_params)
ctx   = llama_init_from_model(model, ctx_params)
tokens = llama_tokenize(vocab, prompt)          # BPE 切词(src/llama-vocab.cpp:4121)

llama_decode(ctx, batch_of(tokens))             # 预填充:整个 prompt 一张图
smpl = build_sampler_chain(top_k=40, top_p=0.95, temp=0.8)

for _ in range(max_new):
    t = llama_sampler_sample(smpl, ctx, -1)     # 链式过滤 + 抽签
    if t == eos: break
    llama_decode(ctx, batch_of([t]))            # 每步 1 个 token,靠 KV cache 看历史

重点看: 第一次 decode 吃整段 prompt(GEMM,快),循环里每步只吃一个 token(GEMV,带宽瓶颈)——这就是「预填充快、解码慢」的结构性来源。

8. 关键细节/坑​

  • 一个模型多个 context 是安全的,一个 context 多线程并发 decode 不是。 可变状态(KV cache、图结果)都在 context 里。
  • n_ubatch 别乱调小。 非因果注意力(如 embedding 模型)要求 n_ubatch >= n_tokens(断言见 src/llama-context.cpp:1721)。
  • 图复用失败是静默的。 can_reuse 不通过只是重建图,不报错——性能问题要开 n_reused 统计或日志才看得见(src/llama-context.cpp:1357)。
  • 采样器顺序即语义。 top-k 在温度前还是后,结果分布不同;换链序等于换采样算法,common 里的顺序(§6)是社区沉淀的默认。
  • 词表在模型文件里,不在库代码里。 换个 GGUF 就换了分词器(BPE/SPM/Unigram,src/llama-vocab.cpp 里各有一套 session 实现);「模型输出乱码」先查分词器 key 是否被转换器正确写入。

9. 代码地图​

主题文件路径符号名
C API 表面include/llama.hllama_model_load_from_file、llama_decode、llama_sampler_sample
模型加载src/llama.cpp / src/llama-model.cppllama_model_load_from_file_impl、llama_model_base::load_tensors
架构分发src/llama-model.cppllama_model::build_graph、build_arch_graph
LLaMA 构图src/models/llama.cppllama_model_llama::graph、load_arch_tensors
注意力/KV 写入src/llama-graph.cppllm_graph_context::build_attn、build_attn_mha
KV cachesrc/llama-kv-cache.cpp / src/llama-kv-cells.hllama_kv_cache::find_slot、llama_kv_cells、seq_cp
decode 主循环src/llama-context.cppllama_context::decode、process_ubatch、graph_compute
分词器src/llama-vocab.cppllama_vocab::tokenize、llm_tokenizer_bpe
采样器src/llama-sampler.cppllama_sampler_sample、llama_sampler_chain_apply、llama_sampler_top_k_apply
采样链组装(应用层)common/sampling.cppcommon_sampler_init、common_sampler_sample