Read OSS

llama.cpp 如何将模型权重转化为计算

高级

前置知识

  • 第 1 篇:架构概览与代码导航
  • 理解 Transformer 前向传播(自注意力、FFN、层归一化、KV 缓存、RoPE)

llama.cpp 如何将模型权重转化为计算

在第 1 篇中,我们看到 model.build_graph() 是每次解码调用的核心。但对于 120 多种模型架构来说,"构建计算图"究竟意味着什么?答案藏在 llama.cpp 最优雅的设计模式之一里:一套可复用的构建模块工具包。模型开发者用这些模块拼装前向传播过程,就像从同一套积木中组合出不同形状。

本文将深入探讨这套系统。我们会分析 llm_graph_context 基类及其工具方法,逐步拆解一个具体的模型构建器(LLaMA),理解将架构路由到对应构建器的分发机制,了解计算图复用如何避免重复开销,以及 Mamba、RWKV 等非 Transformer 架构如何融入同一套框架。

图上下文工具包

llama.cpp 中的每个模型构建器都继承自 llm_graph_context,定义在 src/llama-graph.h。这个基类提供两方面能力:一是一批已初始化的成员变量(来自超参数、上下文参数和当前 ubatch),二是一套构建器方法,封装了构建计算图片段时的重复性工作。

核心构建器方法如下:

方法 用途
build_inp_embd() 创建 token embedding 查找输入
build_inp_pos() 创建位置输入张量
build_norm() 应用 LayerNorm 或 RMSNorm
build_ffn() 构建前馈网络(支持 SiLU/GELU/ReLU 及门控)
build_moe_ffn() 构建混合专家(MoE)FFN
build_attn() 构建带 KV 缓存读写的注意力(5 个重载)
build_lora_mm() 矩阵乘法,支持可选的 LoRA 和逐张量缩放
build_cvec() 应用控制向量以引导生成
build_inp_out_ids() 输出过滤(仅为请求的 token 计算 logits)

这些方法定义在 src/llama-graph.h。值得注意的是,build_attn() 单独就有五个重载,分别对应不同的注意力模式:无缓存(BERT 风格)、KV 缓存、仅 K 缓存、交错滑动窗口(iSWA)和交叉注意力。

classDiagram
    class llm_graph_context {
        +arch: llm_arch
        +hparams: llama_hparams&
        +cparams: llama_cparams&
        +ubatch: llama_ubatch&
        +n_embd, n_layer, n_head...
        +ctx0: ggml_context*
        +gf: ggml_cgraph*
        +res: llm_graph_result*
        +build_inp_embd()
        +build_inp_pos()
        +build_norm()
        +build_ffn()
        +build_moe_ffn()
        +build_attn() ×5
        +build_lora_mm()
        +build_cvec()
    }
    class llm_build_llama {
        Constructor builds full graph
    }
    class llm_build_qwen2 {
        Constructor builds full graph
    }
    class llm_build_bert {
        Constructor builds full graph
    }
    llm_graph_context <|-- llm_build_llama
    llm_graph_context <|-- llm_build_qwen2
    llm_graph_context <|-- llm_build_bert

这套设计的核心理念是"约定优于配置"。模型构建器无需手动管理 KV 缓存槽位索引、注意力掩码构建或 LoRA 应用——这些细节都由工具包在内部处理。模型开发者只需专注于前向传播的拓扑结构:用哪些归一化层、哪些投影、哪些激活函数,以什么顺序排列。

具体案例:LLaMA 计算图构建器

我们来逐步分析 src/models/llama.cpp 中的参考实现。完整的前向传播过程都在 llm_build_llama 的构造函数中完成:

flowchart TD
    EMB["build_inp_embd(tok_embd)"] --> POS["build_inp_pos()"]
    POS --> ATTN_INP["build_attn_inp_kv()"]
    ATTN_INP --> LOOP_START["For each layer il = 0..n_layer"]
    
    LOOP_START --> NORM1["build_norm(attn_norm, RMS)"]
    NORM1 --> QKV["Q/K/V projections via build_lora_mm"]
    QKV --> ROPE["ggml_rope_ext (Q and K)"]
    ROPE --> ATTN["build_attn(inp_attn, wo, Q, K, V)"]
    ATTN --> ADD1["residual add"]
    ADD1 --> NORM2["build_norm(ffn_norm, RMS)"]
    NORM2 --> FFN_CHECK{MoE layer?}
    FFN_CHECK -->|No| FFN["build_ffn(SiLU, PAR)"]
    FFN_CHECK -->|Yes| MOE["build_moe_ffn(...)"]
    FFN --> ADD2["residual add + build_cvec"]
    MOE --> ADD2
    ADD2 --> NEXT["Next layer"]
    
    NEXT --> FINAL_NORM["build_norm(output_norm, RMS)"]
    FINAL_NORM --> LM_HEAD["build_lora_mm(output, cur)"]
    LM_HEAD --> DONE["ggml_build_forward_expand(gf, cur)"]

构建器首先创建 embedding 输入和位置张量(第 13–16 行),然后初始化注意力输入,包括 KV 缓存索引和掩码张量。模板参数 <embed> 让同一套代码既能服务于生成模式,也能服务于嵌入模式——这是 C++ 模板的一个简洁用法。

在层循环内部(第 31–153 行),每次迭代遵循标准的 Transformer 流程:

  1. 对输入做 RMSNorm
  2. 通过 build_lora_mm() 完成 Q/K/V 投影(透明支持 LoRA)
  3. 对 Q 和 K 施加 RoPE 旋转编码
  4. 通过 build_attn() 完成注意力计算,内部处理 KV 缓存写入、掩码应用和 flash attention
  5. 残差连接
  6. 根据 ffn_gate_inp 是否存在,执行 FFN 或 MoE FFN第 107–143 行
  7. 第二次残差连接,并应用控制向量

所有层处理完毕后,输出经过归一化,再通过语言模型头投影到词表大小(第 156–171 行)。

提示: 代码中随处可见的 cb(cur, "attn_norm", il) 调用并非只是调试标签——它们会触发图回调,后端调度器依据这些回调做算子放置和 offload 决策。

架构分发

llama_context::process_ubatch() 调用 model.build_graph() 时,最终会进入 src/llama-model.cpp 中一个庞大的 switch 语句:

ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const {
    std::unique_ptr<llm_graph_context> llm;

    switch (arch) {
        case LLM_ARCH_LLAMA:
            llm = std::make_unique<llm_build_llama<false>>(*this, params);
            break;
        case LLM_ARCH_QWEN2:
            llm = std::make_unique<llm_build_qwen2>(*this, params);
            break;
        // ... 120+ cases
    }
}

arch 字段是一个 llm_arch 枚举值,在模型加载时通过读取 GGUF 文件中的 general.architecture 键来设置。GGUF 字符串名称到枚举值的映射定义在 src/llama-arch.cppLLM_ARCH_NAMES 中:

static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
    { LLM_ARCH_LLAMA,  "llama"  },
    { LLM_ARCH_QWEN2,  "qwen2"  },
    { LLM_ARCH_MAMBA,  "mamba"  },
    // ...
};

llm_arch 枚举 目前定义了 125 个架构标识符,而且这是整个代码库中迭代最频繁的部分之一——新的模型架构还在持续增加。

flowchart LR
    GGUF["GGUF file\ngeneral.architecture = 'llama'"] --> PARSE["llm_arch_from_string()"]
    PARSE --> ENUM["LLM_ARCH_LLAMA"]
    ENUM --> SWITCH["build_graph() switch"]
    SWITCH --> BUILDER["llm_build_llama<false>"]
    BUILDER --> GRAPH["ggml_cgraph"]

计算图复用优化

构建计算图并非没有代价——它涉及分配 GGML 张量、连接算子,以及运行后端调度器的内存分配流程。在自回归生成过程中,连续的解码调用往往会产生相同的图拓扑(相同的批大小、相同的序列结构、相同的模型配置)。llama.cpp 通过计算图复用优化来利用这一规律。

process_ubatch() 中,系统会先检查是否可以复用上一个计算图:

if (!graph_reuse_disable && res->can_reuse(gparams)) {
    n_reused++;
} else {
    res->reset();
    gf = model.build_graph(gparams);
    ggml_backend_sched_alloc_graph(sched.get(), gf);
}
res->set_inputs(&ubatch);
graph_compute(res->get_gf(), ...);

llm_graph_result 上的 can_reuse() 方法会检查新的 llm_graph_params 是否会产生与上次相同的图拓扑。复用条件定义在 allow_reuse() 中:ubatch 的维度必须匹配(n_tokens、n_seqs、n_seq_tokens),输入模式(token 还是 embedding)、输出数量以及配置标志也必须一致。

复用成功时,只会调用 set_inputs() 来更新张量数据,图结构、后端分配和张量内存均保持不变。每个 llm_graph_input_i 子类也有自己的 can_reuse() 检查,用于确认特定于该输入的状态(如 KV 缓存索引)没有发生导致图失效的变化。

提示: 设置环境变量 LLAMA_GRAPH_INPUT_DEBUG=1 可以输出调试信息,显示计算图复用成功或失败的情况,有助于定位性能瓶颈。

输入系统

计算图需要数据:token ID、位置、注意力掩码、KV 缓存索引。这些都通过 llm_graph_input_i 继承体系来管理,定义在 src/llama-graph.h

class llm_graph_input_i {
public:
    virtual void set_input(const llama_ubatch * ubatch) = 0;
    virtual bool can_reuse(const llm_graph_params & params) { return false; }
};

每个具体的输入类负责创建一个或多个 GGML 张量,并知道如何从 ubatch 中填充它们。主要的输入类型包括:

  • llm_graph_input_embd — token ID 或原始 embedding
  • llm_graph_input_pos — 位置索引(支持具有多个位置维度的 M-RoPE)
  • llm_graph_input_attn_kv — 标准 Transformer 的 KV 缓存槽位索引和注意力掩码
  • llm_graph_input_attn_kv_iswa — 同上,但用于交错滑动窗口注意力
  • llm_graph_input_attn_no_cache — BERT 风格模型的完整注意力掩码
  • llm_graph_input_rs — Mamba/RWKV 的循环状态复制索引
  • llm_graph_input_mem_hybrid — 混合模型(如 Jamba)的注意力与循环输入组合
classDiagram
    class llm_graph_input_i {
        <<interface>>
        +set_input(ubatch)*
        +can_reuse(params)*
    }
    class llm_graph_input_embd {
        tokens: ggml_tensor*
        embd: ggml_tensor*
    }
    class llm_graph_input_attn_kv {
        self_k_idxs: ggml_tensor*
        self_v_idxs: ggml_tensor*
        self_kq_mask: ggml_tensor*
    }
    class llm_graph_input_rs {
        s_copy: ggml_tensor*
    }
    class llm_graph_input_mem_hybrid {
        inp_attn: llm_graph_input_attn_kv
        inp_rs: llm_graph_input_rs
    }
    llm_graph_input_i <|-- llm_graph_input_embd
    llm_graph_input_i <|-- llm_graph_input_attn_kv
    llm_graph_input_i <|-- llm_graph_input_rs
    llm_graph_input_i <|-- llm_graph_input_mem_hybrid

这一设计将图拓扑与输入数据的生命周期解耦。模型构建器调用 build_attn_inp_kv() 创建注意力输入对象时,该对象既会创建占位符张量,也会将自身注册为后续 set_inputs() 调用时需要填充的对象。

非 Transformer 架构

llama.cpp 在架构层面的一大成就,是将非 Transformer 模型纳入同一套框架。关键在于一组特化基类,声明在 src/models/models.h 中:

llm_build_mamba_base 提供 build_mamba_layer()build_mamba2_layer() 方法,用于构建 SSM(状态空间模型)的计算过程:输入投影、卷积、选择性扫描和输出投影。Mamba 构建器 与 Transformer 构建器极为相似——同样是遍历各层、归一化、应用核心块、添加残差——区别只在于用 build_rs_inp() 替代了 build_attn_inp_kv(),用 build_mamba_layer() 替代了 build_attn()

llm_build_rwkv7_base 为 RWKV 架构的线性注意力变体提供 build_rwkv7_time_mix()build_rwkv7_channel_mix() 方法。

llm_build_delta_net_base 为 Delta Net 线性注意力机制同时提供分块和自回归两种实现。

Jamba 等混合架构则将两种模式结合在一起。它们使用 llm_graph_input_mem_hybrid,同时封装注意力 KV 输入和循环状态输入,在逐层循环中检查当前层是注意力层还是循环层,并做相应处理。

classDiagram
    class llm_graph_context {
        <<base>>
        build_norm()
        build_ffn()
        build_attn()
    }
    class llm_build_mamba_base {
        <<base>>
        build_mamba_layer()
        build_mamba2_layer()
    }
    class llm_build_rwkv7_base {
        <<base>>
        build_rwkv7_time_mix()
        build_rwkv7_channel_mix()
    }
    class llm_build_delta_net_base {
        <<base>>
        build_delta_net()
    }
    class llm_build_mamba {
        Constructor builds SSM graph
    }
    class llm_build_rwkv7 {
        Constructor builds RWKV graph
    }
    class llm_build_jamba {
        Constructor builds hybrid graph
    }
    llm_graph_context <|-- llm_build_mamba_base
    llm_graph_context <|-- llm_build_rwkv7_base
    llm_graph_context <|-- llm_build_delta_net_base
    llm_build_mamba_base <|-- llm_build_mamba
    llm_build_rwkv7_base <|-- llm_build_rwkv7
    llm_graph_context <|-- llm_build_jamba

内存系统(将在第 4 篇中详细介绍)与这套继承体系保持一致——llama_memory_recurrent 负责 SSM 状态,llama_memory_hybrid 将 KV 缓存与循环状态结合,create_memory() 工厂函数则根据架构类型自动选择合适的实现。

下一步

我们已经看到 llama.cpp 如何借助可组合的构建模块,将模型架构转化为 GGML 计算图。但 GGML 究竟是什么?一个惰性求值的张量库如何运作,又如何在 CPU、CUDA GPU、Metal 和 Vulkan 上执行同一张计算图?下一篇文章将深入剖析 GGML 的核心抽象、后端 vtable 系统,以及 GGUF 文件格式。