03. 从0开始学习硅基智能 - buddy-mlir DeepSeekR1模型导入编译与调度总结本文总结前面两篇文章。

Buddy-MLIR DeepSeek 模型导入、编译与推理调度入门

本文以 examples/BuddyDeepSeekR1 中的 DeepSeek-R1-Distill-Qwen-1.5B、f16 独立示例为唯一主线，按照程序实际执行顺序，说明一个 Hugging Face PyTorch 模型如何变成 MLIR、目标文件、静态库，并最终被 C++ 程序用于逐 Token 推理。

源码基准：本文以 commit bc5210de927302e960b2ff699f1ab4c9c6eed109 为准。所有“源码位置”链接均指向该 commit 的 GitHub 永久地址。正文优先给出文件路径、类/函数名或关键语句，commit 行号只用于辅助核对。

1. 文档目标与范围

本文面向已经了解基本 Python、PyTorch 和 C++，但尚不了解 TorchDynamo、Buddy Graph 和 MLIR 的读者。读完后，应能回答以下问题：

import-deepseek-r1.py 为什么要分别捕获 Prefill 和 Decode？
PyTorch 模型如何依次经过 FX/ATen Graph、Buddy Graph 和高层 MLIR？
四个 MLIR 文件和一个权重文件分别做什么？
buddy-opt、mlir-opt、mlir-translate 和 llc 各自处在什么位置？
C++ 程序如何加载权重、管理 KV Cache、调用 Prefill/Decode 并选择 Token？

本文只分析下列 f16 独立示例，不把 f32、bf16、权重量化或新的配置驱动模型流水线混入主线：

examples/BuddyDeepSeekR1/import-deepseek-r1.py
examples/BuddyDeepSeekR1/CMakeLists.txt
examples/BuddyDeepSeekR1/buddy-deepseek-r1-f16-main.cpp

需要注意，脚本默认加载的是蒸馏模型 deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B，不是完整规模的 DeepSeek-R1。模型选择代码见 model_path 选择逻辑（import-deepseek-r1.py:67）`。

三个主要文件的职责如下。

阶段	主要文件	输入	输出	核心职责
模型导入	`import-deepseek-r1.py`	Hugging Face 模型、示例输入	4 个高层 MLIR、1 个权重文件	捕获 Prefill/Decode、构造 Buddy Graph、执行图变换
模型编译	`CMakeLists.txt`	高层 MLIR	4 个 `.o`、`libDEEPSEEKR1_F16.a`	运行 MLIR/LLVM lowering 和目标代码生成
推理调度	`buddy-deepseek-r1-f16-main.cpp`	用户文本、词表、权重、静态库	逐 Token 文本	分词、内存管理、KV Cache 交接、调用 Prefill/Decode

一句话概括三者边界：Python 导入器生成“计算表示”，编译流水线生成“可链接代码”，C++ 程序组织“运行时数据和调用顺序”。

2. 端到端流程总览

先看完整流程，再逐段进入代码：

Hugging Face DeepSeek-R1-Distill-Qwen-1.5B
                    │
                    ▼
       import-deepseek-r1.py --precision f16
       加载 FP16 权重，构造固定形状示例输入
                    │
                    ▼
              TorchDynamo 捕获
                    │
                    ▼
          FX Graph / ATen-Prims Graph
                    │
                    ▼
        DynamoCompiler 转换为 Buddy Graph
                    │
                    ▼
      转置消除、MatMul 融合、Attention 融合
                    │
                    ▼
              GraphDriver 组装
         ┌──────────┴──────────┐
         ▼                     ▼
   Prefill 子图/主图       Decode 子图/主图
         └──────────┬──────────┘
                    ▼
  4 个高层 MLIR 文件 + arg0-f16.data 权重文件
                    │
                    ▼
 buddy-opt → mlir-opt → buddy-opt → mlir-translate
                    │
                    ▼
              llvm-as → llc
                    │
                    ▼
  4 个目标文件 → 静态库 libDEEPSEEKR1_F16.a
                    │
                    ▼
      buddy-deepseek-r1-f16-main.cpp
                    │
                    ▼
  读取文本 → 分词 → 加载权重 → 分配 KV Cache
                    │
                    ▼
        forward_prefill → 首个 Token
                    │
                    ▼
       循环 forward_decode → 后续 Token
                    │
                    ▼
               反分词并输出

这个流程同时包含三条需要始终区分的“线”：

代码线：PyTorch 函数变成 MLIR，再变成目标代码。
权重线：PyTorch 参数被连续打包为 arg0-f16.data，运行时再读入内存；权重并不被编进静态库。
状态线：Prefill 生成 KV Cache，C++ 负责保存和交接，Decode 每轮读取并更新它。

3. 最小前置概念

本节只解释后续立即会用到的概念。

3.1 Prefill 与 Decode

自回归语言模型推理分为两个形状和工作负载明显不同的阶段。

Prompt: [t0, t1, ..., tn]
             │
             ▼
          Prefill
一次处理整个 Prompt，生成初始 KV Cache 和最后位置 logits
             │
             ▼
        选择第一个新 Token
             │
             ▼
           Decode
每次只输入一个 Token，复用 KV Cache，生成下一个 Token

本示例在导入时把 Prefill 固定为 [1, 1024]，把 Decode 固定为 [1, 1]，所以需要两次图捕获和两套编译结果。它们不是同一个函数的两个别名，而是面向不同输入形状和 Cache 数据流的两张静态图。

3.2 FX Graph、ATen Graph、Buddy Graph 与 MLIR

这几个名称对应同一计算在不同阶段的表示：

PyTorch Python 模型
        │ TorchDynamo 观察一次执行
        ▼
FX Graph
        │ AOTAutograd/Inductor decomposition
        ▼
以 aten.* / prims.* 为节点目标的 FX Graph
        │ DynamoCompiler._ops_map
        ▼
Buddy Graph（AddOp、MatmulOp、PlaceholderOp 等）
        │ GraphImporter + ops_registry
        ▼
TOSA/Linalg/Func 等高层 MLIR

FX Graph 是图的数据结构；ATen/Prims 是图中使用的规范化 PyTorch 算子集合；Buddy Graph 是 Buddy 前端自己的节点表示；MLIR 则是后续优化和 lowering 使用的编译器 IR。

3.3 Graph 与 GraphDriver

Graph 表示一张 Buddy 计算图，保存节点、输入、参数、shape、dtype 和 lowering registry。Graph.lower_to_top_level_ir() 会调用 GraphImporter 生成高层 MLIR，见 graph.py:540。

GraphDriver 不负责从 PyTorch 捕获图，也不负责选择最优分组策略。它消费已经写入 graph.op_groups 的分组，识别子图边界，构造子图，再生成按依赖顺序调用子图的 main graph，见 GraphDriver.__init__（graph_driver.py:46）和 graph_driver.py:200。本示例最终将所有计算放进一个 CPU 子图，因此 Prefill 和 Decode 各有一个子图。

3.4 MemRef 与 C 接口

MLIR lowering 后的张量接口使用 MemRef 描述符。Buddy C++ 的 MemRef<T, N> 保存数据指针、offset、N 个 sizes 和 N 个 strides，布局用于匹配编译后函数的 C ABI，定义见 MemRef 定义（Container.h:58）`。

例如：

MemRef<uint16_t, 4> kv({1, 2, 1024, 128}, 0);

表示一个四维连续缓冲区。在本 f16 示例中，uint16_t 负责承载 16 位浮点的原始位模式，而不是把数值当作普通无符号整数做模型计算。

4. 阶段一：导入 PyTorch 模型

4.1 本阶段负责什么

导入阶段负责把 Hugging Face 模型转换成 Prefill/Decode 两张 Buddy Graph，完成前端图优化与分组，然后输出高层 MLIR 和原始权重文件。它不生成目标机器码，也不执行最终文本生成循环。

4.2 输入是什么

导入器接收：

模型路径：优先读取 DEEPSEEKR1_MODEL_PATH，否则使用 Hugging Face 模型名。
精度参数：本文固定为 --precision f16。
输出目录：由 --output-dir 指定。
用于捕获的示例张量：Prefill [1, 1024]，Decode [1, 1]。

CMake 调用脚本的实际命令定义在 CMakeLists.txt:35：

python import-deepseek-r1.py \
  --output-dir <build/examples/BuddyDeepSeekR1> \
  --precision f16

4.3 执行流程图

解析 --precision f16 和 --output-dir
                    │
                    ▼
从本地路径或 Hugging Face 加载模型
dtype=torch.float16 → eval() → half()
                    │
                    ▼
创建 DynamoCompiler(forward_prefill)
创建 DynamoCompiler(forward_decode)
                    │
                    ▼
构造 input_ids、StaticCache、cache_position
                    │
        ┌───────────┴───────────┐
        ▼                       ▼
捕获 Prefill                 预热 Decode Cache
input [1,1024]                  │
        │                       ▼
        │                  捕获 Decode
        │                  input [1,1]
        └───────────┬───────────┘
                    ▼
          两张 FX/ATen Graph
                    │
                    ▼
          两张 Buddy Graph
                    │
                    ▼
转置优化 → 经典融合 → Attention 专用融合
                    │
                    ▼
    Prefill/Decode 各形成一个 CPU 算子组
                    │
                    ▼
GraphDriver 构造子图，GraphImporter 生成高层 MLIR
                    │
                    ▼
     写出 4 个 MLIR + 1 个权重文件

4.4 关键代码

4.4.1 加载 FP16 模型

脚本在 if args.precision == "f16"（import-deepseek-r1.py:73）按 FP16 加载模型：

model = (
    AutoModelForCausalLM.from_pretrained(model_path, dtype=torch.float16)
    .eval()
    .half()
)
model.config.use_cache = False

输入：Hugging Face 模型目录或名称。
动作：权重按 torch.float16 加载，切换推理模式，再确保浮点参数和适用 buffer 为 FP16。
输出：用于图捕获的 FP16 nn.Module。

model.config.use_cache = False 关闭配置中的默认行为；后面的具体捕获调用仍显式传入 use_cache=True，因此 Prefill/Decode 图仍包含 Cache 路径。

4.4.2 创建两个 DynamoCompiler

dynamo_compiler_prefill/decode（import-deepseek-r1.py:91）创建两个相同配置、不同函数名的前端实例：

dynamo_compiler_prefill = DynamoCompiler(
    primary_registry=tosa.ops_registry,
    aot_autograd_decomposition=inductor_decomp,
    func_name="forward_prefill",
)
dynamo_compiler_decode = DynamoCompiler(
    primary_registry=tosa.ops_registry,
    aot_autograd_decomposition=inductor_decomp,
    func_name="forward_decode",
)

输入：TOSA lowering registry 和 TorchInductor decomposition 表。
动作：配置 ATen 算子分解、Buddy 节点映射和 MLIR lowering 策略。
输出：分别记录 Prefill/Decode 图和参数状态的两个导入器。

DynamoCompiler.__init__() 会合并 math、linalg、TOSA、func 和用户传入的 registry，并建立 _ops_map，见 DynamoCompiler.__init__（frontend.py:75）。这里要区分两张表：

_ops_map：例如把 mm.default 映射为 MatmulOp，决定“创建哪种 Buddy 节点”。
_ops_registry：把 MatmulOp 等 Buddy 节点降低为 MLIR，决定“如何生成 MLIR”。

4.4.3 构造固定形状输入

f16 分支在 with torch.no_grad()（import-deepseek-r1.py:106）进入 torch.no_grad()，避免建立训练反向图，然后构造：

用途	PyTorch 张量或对象	shape/配置	dtype
Prefill Token	`data_prefill["input_ids"]`	`[1, 1024]`	`torch.int64`
Decode Token	`data_decode["input_ids"]`	`[1, 1]`	`torch.int64`
Decode 写入位置	`cache_position`	`[1]`，示例值 `200`	`torch.int64`
KV Cache	`StaticCache`	`max_cache_len=1024`	由模型路径建立

“f16 导出”不表示所有输入都是 FP16。Token ID 和 Cache 位置是索引，必须保持整数；模型权重、激活和 KV 浮点数据才是主要的 FP16 对象。

脚本创建了 past_key_values_prefill，但 Prefill 调用中对应参数被注释，没有直接传入该对象；Prefill 通过 cache_implementation="static" 请求模型建立静态 Cache。这个事实可在 StaticCache 构造（import-deepseek-r1.py:108）和 graphs_prefill = ...importer(...)（import-deepseek-r1.py:124）对照看到。

4.4.4 捕获 Prefill 与 Decode

Prefill 调用长度为 1024 的输入，见 graphs_prefill = ...importer(...)（import-deepseek-r1.py:124）：

graphs_prefill = dynamo_compiler_prefill.importer(
    model,
    input_ids=data_prefill["input_ids"],
    use_cache=True,
    cache_implementation="static",
)

Decode 捕获前，脚本先用显式 past_key_values_decode 实际调用模型一次，使 Cache 内部张量完成初始化；随后使用单 Token、Cache 位置和已有 Cache 捕获 Decode，见 # Initialize past_key_values...（import-deepseek-r1.py:131）：

model(..., past_key_values=past_key_values_decode, use_cache=True)
graphs_decode = dynamo_compiler_decode.importer(
    model,
    input_ids=data_decode["input_ids"],
    cache_position=cache_position,
    past_key_values=past_key_values_decode,
    use_cache=True,
    cache_implementation="static",
)

输入：两套固定 shape 的示例数据。
动作：真实执行触发 TorchDynamo，分别捕获长序列和单 Token 路径。
输出：graphs_prefill 和 graphs_decode 两个 Buddy Graph 列表。

脚本随后断言两个列表长度都为 1，见 if args.precision == "f16" 导出分支（import-deepseek-r1.py:188）。这意味着当前示例要求每条路径都被捕获成一张完整图；如果 TorchDynamo graph break 产生多张图，脚本不会继续静默导出。

4.4.5 DynamoCompiler 内部发生什么

DynamoCompiler.importer() 的入口位于 frontend.py:1204。核心调用是：

model_opt = dynamo.optimize(self._compile_fx)(model)
model_opt(*args, **kwargs)

其实际数据流为：

TorchDynamo 捕获 GraphModule
        │
        ▼
_compile_fx(gm, inputs)
        │
        ▼
aot_module_simplified(..., decompositions=inductor_decomp)
        │
        ▼
内部 _compiler 收到 ATen/Prims FX Graph
        │
        ▼
区分参数、buffer、运行时输入
        │
        ▼
读取 shape、dtype、用户关系和多返回值信息
        │
        ▼
_ops_map[算子名] 创建 Buddy Op
        │
        ▼
保存 imported_graphs 和 imported_params

_compile_fx() 从 frontend.py:887 开始。它先识别运行时输入、参数和 buffer，再多次遍历 FX 节点；_create_node() 在 frontend.py:818 记录参数、父子关系、shape 和 dtype。转换结果仍然是 Buddy Graph，还不是最终 LLVM IR。

4.4.6 参数提取与前端图变换

脚本从 Prefill 导入器获取参数：

params = dynamo_compiler_prefill.imported_params[graph_prefill]

代码位置为 params = ...imported_params[...]（import-deepseek-r1.py:194）。Prefill 和 Decode 来自同一个模型，示例只导出一份共享权重。

两张图首先执行：

graph.perform([eliminate_transpose, eliminate_matmul_transpose_reshape])

Graph.perform() 只是按顺序调用传入函数，见 graph.py:525。两个 Pass 的实际行为是：

eliminate_transpose：对仅被一个转置节点使用的权重参数，直接转置参数数据、更新 shape、重连消费者并删除运行时转置，见 eliminate_weight_transpose.py:28。
eliminate_matmul_transpose_reshape：寻找 transpose/permute 与 reshape/view 等组合并尝试消除。

这里存在一个重要的实现边界：第二个 Pass 在 eliminate_matmul_transpose_reshape.py:59 明确要求 TensorDType.Float32，因此虽然 f16 分支调用了它，当前实现对 f16 节点会直接跳过。不能把它描述成当前 f16 导出真正生效的优化。

4.4.7 融合、分组和 Attention 专用节点

Prefill 与 Decode 使用不同的融合列表：

pattern_list_prefill = [
    simply_fuse,
    apply_classic_fusion,
    flash_attention_prefill,
]
pattern_list_decode = [
    simply_fuse,
    apply_classic_fusion,
    gqa_attention_fusion,
]

代码位置为 pattern_list_prefill/decode（import-deepseek-r1.py:202）。这些函数按顺序原地修改 Graph：

变换	实际作用
`simply_fuse`	将所有非 Placeholder 节点放入一个 CPU `subgraph0`；它主要做分组，不做复杂模式融合
`apply_classic_fusion`	分解特定 `addmm`，并把 `permute([1,0]) + matmul` 融合为转置 MatMul 专用节点
`flash_attention_prefill`	将 CPU Scaled Dot Product Flash Attention 节点替换为 Prefill 专用节点
`gqa_attention_fusion`	检测 Decode 中 GQA Attention 与 KV Cache 更新模式，替换为融合节点

对应实现见 fuse_ops.py:246、 fuse_ops.py:271、 fuse_ops.py:292 和 fuse_ops.py:326。每个列表最后都重新生成一个 CPU subgraph0。脚本再把它们重命名为 subgraph0_prefill 和 subgraph0_decode，见 graph_prefill.op_groups[...]（import-deepseek-r1.py:216）。

4.4.8 GraphDriver 生成子图和主图

GraphDriver 构造时读取已有的 op_groups，分析每组外部输入、输出和依赖，再创建真正的子图 Graph。实现入口见 graph_driver.py:73。

graph_prefill.op_groups["subgraph0_prefill"]
                    │
                    ▼
GraphDriver(graph_prefill)
                    │
                    ├── subgraphs[0]
                    │   真正的计算函数 subgraph0_prefill
                    │
                    └── construct_main_graph(True)
                        参数解包 + 调用 subgraph0_prefill
                        对外入口 forward_prefill

子图通过 lower_to_top_level_ir() 生成 Tensor 形式的高层 MLIR。main graph 则使用 construct_main_graph(True)；其中 True 打开参数打包，把同 dtype 的大量权重表示为一个连续 MemRef，并在 main graph 中通过 param.extract 取出各个参数。参数打包实现见 graph.py:836 和 graph.py:928。

4.4.9 写出 MLIR 与权重

f16 输出代码位于 # Save the generated files...（import-deepseek-r1.py:275）。子图文件直接打印 subgraphs[0]._imported_module，主图文件打印 construct_main_graph(True) 的返回模块。

权重处理为：

all_param = numpy.concatenate(
    [param.detach().numpy().reshape([-1]) for param in params]
)
all_param.tofile(os.path.join(output_dir, "arg0-f16.data"))

输入：按图参数占位顺序排列的 FP16 Tensor 列表。
动作：脱离 autograd，转换 NumPy，逐个展平，按顺序拼成一维连续数组。
输出：无文件头、无 shape 元数据的原始 FP16 二进制流。

shape、offset 和参数顺序由 MLIR main graph 固化；C++ 只按既定元素数量读入，因此权重文件和 MLIR/静态库必须来自同一次兼容导出。

4.5 输出是什么

导入阶段产生五个文件：

文件	内容
`subgraph0_prefill-f16.mlir`	Prefill 实际计算子图
`forward_prefill-f16.mlir`	Prefill 对外入口、参数解包、子图调用
`subgraph0_decode-f16.mlir`	Decode 实际计算、Attention 与 KV 更新子图
`forward_decode-f16.mlir`	Decode 对外入口、参数解包、子图调用
`arg0-f16.data`	Prefill/Decode 共用的连续 FP16 权重

4.6 如何验证

不必先运行完整 C++ 推理，可以逐层检查导入结果：

ls -lh build/examples/BuddyDeepSeekR1/*-f16.mlir \
       build/examples/BuddyDeepSeekR1/arg0-f16.data

rg "func.func.*forward_prefill" \
  build/examples/BuddyDeepSeekR1/forward_prefill-f16.mlir

rg "func.func.*subgraph0_decode" \
  build/examples/BuddyDeepSeekR1/subgraph0_decode-f16.mlir

还应检查 MLIR 中浮点张量类型为 f16，Token 和位置类型为 i64。如果断言图数量失败，应先调查 TorchDynamo graph break，而不是跳过断言拼接多图。

5. 阶段二：优化和编译 MLIR

5.1 本阶段负责什么

编译阶段把 Python 导入器输出的高层 MLIR 逐步降低为 LLVM Dialect、LLVM IR 和目标机器码，再把四个目标文件归档为一个静态库。这个过程由 CMake 组织，buddy-opt 只是其中一部分。

5.2 输入是什么

输入是阶段一生成的四个 MLIR 文件。arg0-f16.data 不参加代码编译，它在运行时由 C++ 读取。

forward_prefill-f16.mlir
subgraph0_prefill-f16.mlir
forward_decode-f16.mlir
subgraph0_decode-f16.mlir

CMakeLists.txt 在开头定义 TOSA 标准 lowering pipeline：

set(TOSA_PIPELINE
  "builtin.module(
    func.func(tosa-to-linalg-named),
    func.func(tosa-to-linalg),
    func.func(tosa-to-tensor),
    func.func(tosa-to-arith))")

实际定义见 TOSA_PIPELINE（CMakeLists.txt:4）。

5.3 执行流程图

高层 MLIR（TOSA / Tensor / Func 等）
                    │
                    ▼
buddy-opt
示例专用 TOSA 简化
                    │
                    ▼
mlir-opt -pass-pipeline ${TOSA_PIPELINE}
TOSA → Linalg / Tensor / Arith
                    │
                    ▼
buddy-opt
Bufferization、内存优化、向量化、循环/OpenMP 转换
                    │
                    ▼
LLVM Dialect
                    │
                    ▼
mlir-translate -mlir-to-llvmir
                    │
                    ▼
LLVM IR 文本
                    │
                    ▼
llvm-as
                    │
                    ▼
LLVM Bitcode
                    │
                    ▼
llc -filetype=obj -O3
                    │
                    ▼
目标文件 .o
                    │
                    ▼
CMake add_library(... STATIC ...)
                    │
                    ▼
libDEEPSEEKR1_F16.a

5.4 关键代码

5.4.1 不是所有 MLIR 使用完全相同的优化参数

Prefill main graph 的完整命令从 CMakeLists.txt:313 开始。它依次执行 buddy-opt、mlir-opt、第二次 buddy-opt、mlir-translate、llvm-as 和 llc。

其中第二段 buddy-opt 大致完成：

Tensor → MemRef Bufferization
→ 内存释放与拷贝优化
→ MatMul/BatchMatMul 向量化
→ Linalg → Affine/SCF
→ SCF → OpenMP
→ Vector/Arith/MemRef/Func → LLVM Dialect

子图含有真实计算，所以比 main graph 多出 convert-elementwise-to-linalg、MemRef 布局和 MatMul 专用优化。Prefill 子图命令见 CMakeLists.txt:360。

Decode 子图还使用了针对单 Token 形状的专用向量宽度：

matmul-vectorization-decode=vector-size=32
batch-matmul-vectorization-decode=vector-size=128
batchmatmul-transpose-b-vectorization=vector-size=16

对应代码见 CMakeLists.txt:460。这再次说明为什么 Prefill 和 Decode 要独立导出：它们不仅输入 shape 不同，后端优化策略也不同。

5.4.2 C ABI 包装函数

lowering pipeline 中的 llvm-request-c-wrappers 为带有 llvm.emit_c_interface 的函数生成 C ABI 包装入口。C++ 不直接调用内部 MLIR 符号，而是声明：

_mlir_ciface_forward_prefill
_mlir_ciface_forward_decode

5.4.3 目标文件和静态库

四条命令分别产生：

高层 MLIR	目标文件
`forward_prefill-f16.mlir`	`forward_prefill-f16.o`
`subgraph0_prefill-f16.mlir`	`subgraph_prefill-f16.o`
`forward_decode-f16.mlir`	`forward_decode-f16.o`
`subgraph0_decode-f16.mlir`	`subgraph_decode-f16.o`

CMake 在 add_library(DEEPSEEKR1_F16 ...)（CMakeLists.txt:1510）将四个目标文件归档：

add_library(DEEPSEEKR1_F16 STATIC
  forward_prefill-f16.o
  subgraph_prefill-f16.o
  forward_decode-f16.o
  subgraph_decode-f16.o)

该静态库包含 Prefill/Decode 包装函数和实际子图代码，但不包含 arg0-f16.data 权重。C++ 可执行程序还链接 mlir_c_runner_utils 和 OpenMP，配置见 CMakeLists.txt:1624。

5.5 输出是什么

阶段二输出：

forward_prefill-f16.o
subgraph_prefill-f16.o
forward_decode-f16.o
subgraph_decode-f16.o
libDEEPSEEKR1_F16.a

随后 CMake 用 buddy-deepseek-r1-f16-main.cpp 创建 buddy-deepseek-r1-f16-run，并链接 DEEPSEEKR1_F16，见 add_executable(buddy-deepseek-r1-f16-run ...)（CMakeLists.txt:1559）和 CMakeLists.txt:1665。

5.6 如何验证

ls -lh build/examples/BuddyDeepSeekR1/*-f16.o
ls -lh build/examples/BuddyDeepSeekR1/libDEEPSEEKR1_F16.a

llvm-nm build/examples/BuddyDeepSeekR1/libDEEPSEEKR1_F16.a \
  | rg "_mlir_ciface_forward_(prefill|decode)"

如果 MLIR 已生成但 .o 失败，应查看具体是哪一段 pipeline 报错：TOSA lowering、bufferization、向量化、LLVM Dialect 转换和 llc 属于不同故障层，不应都归因于 buddy-opt。

6. 阶段三：C++ 推理调度

6.1 本阶段负责什么

C++ 程序不重新实现 Transformer 内部矩阵计算。它负责准备运行时数据、调用编译后的函数，并围绕函数构建自回归生成循环：

读取用户文本并分词；
从 arg0-f16.data 加载参数；
分配 logits、KV Cache 和位置缓冲区；
调用一次 Prefill；
从 logits 通过贪心 argmax 选择 Token；
将 Prefill Cache 交给 Decode；
循环调用 Decode、更新位置并判断 EOS；
把 Token ID 还原为文本。

6.2 输入是什么

运行时输入包括：

输入	来源
用户 Prompt	标准输入 `getline()`
词表	`examples/BuddyDeepSeekR1/vocab.txt`
FP16 权重	`build/examples/BuddyDeepSeekR1/arg0-f16.data`
编译后函数	链接进可执行程序的 `libDEEPSEEKR1_F16.a`

路径由 CMake 注入的 DEEPSEEKR1_EXAMPLE_PATH 和 DEEPSEEKR1_EXAMPLE_BUILD_PATH 拼出，使用位置见 buddy-deepseek-r1-f16-main.cpp:402。

6.3 执行流程图

启动 buddy-deepseek-r1-f16-run
                    │
                    ▼
getline() 读取用户 Prompt
                    │
                    ▼
tokenizeDeepSeekR1(vocab.txt, 1024)
                    │
                    ▼
读取 arg0-f16.data 到 ParamsContainer
                    │
                    ▼
分配 Prefill logits 和 56 个 KV MemRef
                    │
                    ▼
_mlir_ciface_forward_prefill(...)
                    │
                    ▼
Prefill 返回 KV Cache + [1,1024,vocab] logits
                    │
                    ▼
取 Prompt 最后有效位置 logits → argmax → 首个 Token
                    │
                    ▼
复制有效 Prefill KV 区间到 Decode Cache
设置 cache_position = Prompt Token 数
                    │
                    ▼
┌──────────────── Decode 循环 ────────────────┐
│ 准备单 Token 输入和当前位置                 │
│                  │                          │
│                  ▼                          │
│ _mlir_ciface_forward_decode(...)            │
│                  │                          │
│                  ▼                          │
│ 读取 logits → argmax → 新 Token             │
│                  │                          │
│          EOS? ───┴─── 是 → 结束             │
│            │ 否                             │
│            ▼                                │
│ 写入下一轮 input，cache_position += 1        │
└────────────┴─────────────────────────────────┘
                    │
                    ▼
revertDeepSeekR1() 反分词并输出文本

6.4 关键代码

6.4.1 运行时常量和返回结构

f16 main 在 ParamsSize/MaxVocabSize/MaxTokenLength（buddy-deepseek-r1-f16-main.cpp:32）固定参数数量、词表大小和最大长度：

constexpr size_t ParamsSize = 1777088064;
constexpr size_t MaxVocabSize = 151936;
constexpr size_t MaxTokenLength = 1024;
constexpr size_t HiddenSize = 128;
constexpr size_t HeadNum = 2;

PrefillReturns 包含 56 个四维 KV MemRef 和一个三维 logits，定义见 buddy-deepseek-r1-f16-main.cpp:44。 DecodeReturns 还包含 cache_position_out 和 27 个一维 dummy 返回槽，定义见 buddy-deepseek-r1-f16-main.cpp:105。这些结构的字段顺序必须与编译后 MLIR 多返回值的 ABI 完全一致，不能随意重排。

代码中的常量名 NUM_LAYERS = 56 容易误导。就返回结构而言，实际存在的是 56 个 KV 字段，即按 K/V 成对组织的 Cache 张量；阅读时不要仅凭该变量名断言模型有 56 个 Transformer Block。

6.4.2 分词和权重加载

tokenizeInput() 调用 Text<size_t, 2>::tokenizeDeepSeekR1()，见 tokenizeInput（buddy-deepseek-r1-f16-main.cpp:306）。Text 继承自 MemRef，所以既保存 Token ID，也可作为编译后函数的输入描述符。

loadParameters() 把权重文件直接读入 MemRef<uint16_t, 1>：

paramFile.read(reinterpret_cast<char *>(params.getData()),
               sizeof(uint16_t) * params.getSize());

实现见 loadParameters（buddy-deepseek-r1-f16-main.cpp:322）。这里没有解析文件头，因为 Python 输出本来就是连续的裸数据。

6.4.3 分配 KV Cache

main() 创建输入、参数和位置容器，然后通过 makeKV() 创建 56 个：

MemRef<uint16_t, 4>({1, HeadNum, MaxTokenLength, HiddenSize}, 0)

每个 Cache shape 为 [1, 2, 1024, 128]，并初始化为零，代码见 buddy-deepseek-r1-f16-main.cpp:412。C++ 在这里负责内存的分配和生命周期；具体 K/V 数值由编译后的模型函数计算。

6.4.4 调用 Prefill 并选择首个 Token

完成分词和权重加载后，程序调用：

_mlir_ciface_forward_prefill(
    &prefillRet, &ParamsContainer, &inputContainerPrefill);

位置见 _mlir_ciface_forward_prefill(...) 调用（buddy-deepseek-r1-f16-main.cpp:500）。调用结果包含完整 Prefill KV Cache 和每个位置的 logits。

程序使用 tokenIndex = tokenCnt - 1 定位 Prompt 最后一个有效 Token 对应的 logits，再调用 findMaxIndex() 做贪心选择。模型函数并不直接返回“自然语言 Token”；它返回 vocab 维度分数，C++ 才把最大分数的索引当作下一 Token ID。

6.4.5 Prefill Cache 交给 Decode

buildPrefillKVPtrs() 和 buildDecodeKVPtrs() 把结构中 56 个字段整理为指针数组，避免手写 56 次复制。随后 copy_kv_by_cache_position_block() 只复制 Prompt 已占用的序列区间：

复制长度 = min(Prompt Token 数, 1024)
每个 KV、每个 Head 复制：复制长度 × HiddenSize × 2 字节

实现见 buddy-deepseek-r1-f16-main.cpp:370，调用点见 copy_kv_by_cache_position_block(...) 调用（buddy-deepseek-r1-f16-main.cpp:614）。复制完成后，外部 cachePosition 被设置为 Prompt Token 数。

6.4.6 Decode 循环

Decode 循环从 for (int i = 1; ...)（buddy-deepseek-r1-f16-main.cpp:624）开始。每轮执行：

把当前 cachePosition 写入 27 个 dummy 位置槽。
调用 _mlir_ciface_forward_decode()，传入权重、单 Token、位置以及所有 KV Cache。
从 [1, 1, vocab] logits 中取 argmax。
Token ID 等于 151643 时停止。
否则把 Token 写成下一轮输入，并执行 cachePosition += 1。

接口调用见 buddy-deepseek-r1-f16-main.cpp:656，EOS 和状态更新见 buddy-deepseek-r1-f16-main.cpp:692。当前调度代码没有采用随机采样、temperature 或 top-k，而是确定性的贪心 argmax。

6.4.7 当前 f16 文件中的实现不一致

按照“以实现为准”的原则，需要明确指出当前文件中的几个不一致：

文件首行和 main 区域标题仍写着 bf16，运行时标题也输出 BF16 Inference，见 buddy-deepseek-r1-f16-main.cpp:1 和 buddy-deepseek-r1-f16-main.cpp:393。这与文件名和 f16 构建目标不一致，属于遗留命名/文案。
decode_f16() 使用 uint16_t << 16，这是 BF16 位模式扩展成 FP32 的方式，不是 IEEE FP16 到 FP32 的常规转换，见 decode_f16（buddy-deepseek-r1-f16-main.cpp:350）。而 Python 导出明确使用 torch.float16，MLIR 类型也为 f16。因此当前 f16 logits 的 argmax 解码存在实现风险；正式依赖数值结果前应修正或用参考实现验证。
DecodeReturns 含有 cache_position_out，但主循环维护的是外部 cachePosition 并手动递增，没有读取该返回字段作为下一轮位置。理解当前行为时应以主循环代码为准。

这些问题不改变导入、编译和调度的整体架构，但会影响 f16 示例的显示信息和潜在数值正确性。

6.5 输出是什么

C++ 程序逐轮打印 Token、耗时和最终文本，并输出 Prefill/Decode 吞吐。实际模型计算结果首先是 logits 和 KV Cache，Token 文本是调度程序后处理的结果。

6.6 如何验证

可以分三层验证：

接口层：静态库中能找到两个 _mlir_ciface_forward_* 符号
内存层：权重元素数、KV shape、返回结构字段顺序与 MLIR 一致
数值层：Prefill/Decode logits 与 PyTorch 参考结果比较

仅看到程序输出 Token 不足以证明 ABI 完全正确。特别是当前 f16 解码函数存在上述位模式疑点，建议加入 logits 抽样对比。

7. Prefill 与 Decode 的衔接

两条图在导入时独立，在运行时通过“首 Token + KV Cache + cache_position”连接：

Prompt Token IDs
       │
       ▼
forward_prefill
       │
       ├── 56 个 Prefill KV Cache
       └── Prompt 各位置 logits
                     │
                     ▼
             最后有效位置 argmax
                     │
                     ▼
               首个生成 Token
                     │
       ┌─────────────┴─────────────┐
       ▼                           ▼
复制有效 KV 到 Decode Cache    作为 Decode input_ids
       │                           │
       └─────────────┬─────────────┘
                     ▼
               forward_decode
                     │
                     ├── 更新后的 KV Cache
                     └── 下一 Token logits

项目	Prefill	Decode
主要目标	编码整个 Prompt	每轮生成一个新 Token
Token 输入 shape	`[1, 1024]`	`[1, 1]`
Token dtype	`int64/i64`	`int64/i64`
Cache 输入	脚本请求 Static Cache；C++ 首轮无历史 Cache	56 个已有 KV Cache + 位置/辅助槽
主要浮点计算	FP16	FP16
logits shape	`[1, 1024, 151936]`	`[1, 1, 151936]`
Cache 输出	56 个 `[1,2,1024,128]` KV 张量	更新后的 56 个 KV 张量
后端优化	Prefill Flash Attention、长序列向量化	GQA/Cache 融合、Decode 专用向量宽度
调用次数	每个 Prompt 一次	最多剩余序列长度次

Prefill 和 Decode 共用同一份 arg0-f16.data，但使用不同入口和不同目标代码。静态库必须同时包含 main wrapper 和 subgraph object，否则 wrapper 中对子图符号的调用无法解析。

8. KV Cache 的创建、传递和更新

KV Cache 不是只由 Python、MLIR 或 C++ 某一层独占完成。不同阶段有明确分工。

8.1 生命周期流程

导入期 Python StaticCache
用于固定图捕获时的 Cache shape 和数据流
             │
             ▼
导出后的 MLIR 函数签名
显式包含 KV 输入/输出
             │
             ▼
C++ 分配 KV MemRef 内存
             │
             ▼
Prefill 模型计算 K/V，写入 PrefillReturns
             │
             ▼
C++ 复制 Prompt 有效区间到 Decode Cache
             │
             ▼
Decode 模型读取历史 K/V，计算当前 K/V 并更新 Cache
             │
             ▼
C++ 保存 Cache，并在下一轮再次传回 Decode

8.2 所有权和操作对照

阶段/组件	拥有或管理什么	执行什么操作	不负责什么
Python `StaticCache`	导入时的 PyTorch Cache 对象	帮助捕获固定长度 Cache 路径	不会被序列化成 C++ 对象
Prefill 编译函数	Prefill 算子和 KV 输出逻辑	根据 Prompt 计算各层 K/V 和 logits	不负责长期保存 C++ Cache 容器
C++ `main()`	KV MemRef 内存、指针数组、`cachePosition`	分配、初始化、复制、跨轮保存和传参	不计算 Attention 的 K/V 数值公式
Decode 编译函数	单 Token Decode 与 KV 更新计算	读取历史 Cache，计算当前 K/V，更新并产生 logits	不选择自然语言 Token
C++ Decode 循环	当前 Token、EOS、轮次状态	argmax、写入下一 Token、位置递增、结束判断	不执行 Transformer 内部 MatMul/Attention

这种边界解释了为什么 C++ 接口参数很多：静态 Cache 已经从 Python 对象变成显式 MemRef，所有层的 K/V 状态都要通过 ABI 传递。

9. 中间文件和接口对照

9.1 五个导出文件

文件	是否参与代码编译	是否运行时读取	用途
`forward_prefill-f16.mlir`	是	否	参数打包入口，调用 Prefill 子图
`subgraph0_prefill-f16.mlir`	是	否	Prefill 实际模型计算
`forward_decode-f16.mlir`	是	否	参数打包入口，调用 Decode 子图
`subgraph0_decode-f16.mlir`	是	否	Decode、GQA 和 KV Cache 更新计算
`arg0-f16.data`	否	是	连续 FP16 模型权重

“forward” 文件与“subgraph”文件分开，是因为 GraphDriver 把对外 ABI/参数解包与实际计算函数拆开。前者适合作为稳定入口，后者可以使用针对具体计算图的 lowering 和优化。

9.2 C 接口

接口	主要参数	主要结果
`_mlir_ciface_forward_prefill`	packed FP16 参数、`[1,1024]` Token	56 个 KV Cache、`[1,1024,vocab]` logits
`_mlir_ciface_forward_decode`	packed 参数、`[1,1]` Token、`[1]` 位置、56 个 KV、27 个辅助位置槽	更新位置/辅助槽、更新后 KV、`[1,1,vocab]` logits

多返回值通过第一个 PrefillReturns* 或 DecodeReturns* 参数承载。接口声明必须与 MLIR lowering 后的结果顺序一致，而不仅仅是元素类型一致。

9.3 dtype 对照

语义	Python/PyTorch	Buddy dtype	高层/低层 MLIR	f16 C++ 示例
权重、激活、KV、logits	`torch.float16`	`TensorDType.Float16`	`f16`	`uint16_t` 保存原始 16 位模式
Token ID	`torch.int64`	`TensorDType.Int64`	`i64`	`size_t` Text 或 `long long` MemRef
Cache 位置	`torch.int64`	`TensorDType.Int64`	`i64`	`long long`
Bool/掩码条件	`torch.bool`	`TensorDType.Bool`	`i1`	经 lowering 后按 ABI/算子需要处理

PyTorch 到 Buddy dtype 的转换实现见 frontend.py:695，Buddy dtype 到 MLIR type 的转换见 graph.py:799。MLIR 的 i64 是 signless integer，因此 C++ 接口中 size_t 与 long long 的选择主要依赖 64 位位宽和容器用途；在非 64 位平台上不能默认 ABI 一致。

10. 构建、运行与验证

10.1 配置环境

根据示例 README，先构建 Buddy Python 包并设置 PYTHONPATH，见 README.md:37：

cd /home/king/buddy-mlir
cmake -G Ninja -S . -B build \
  -DMLIR_DIR=$PWD/llvm/build/lib/cmake/mlir \
  -DLLVM_DIR=$PWD/llvm/build/lib/cmake/llvm \
  -DLLVM_ENABLE_ASSERTIONS=ON \
  -DCMAKE_BUILD_TYPE=RELEASE \
  -DBUDDY_MLIR_ENABLE_PYTHON_PACKAGES=ON \
  -DBUDDY_DEEPSEEKR1_EXAMPLES=ON \
  -DPython3_EXECUTABLE="$(which python3)"

export PYTHONPATH=$PWD/build/python_packages:${PYTHONPATH}
export DEEPSEEKR1_MODEL_PATH=/path/to/DeepSeek-R1-Distill-Qwen-1.5B

如果不设置 DEEPSEEKR1_MODEL_PATH，脚本会尝试从 Hugging Face 加载默认模型；这要求网络和对应依赖可用。

10.2 构建 f16 可执行程序

cmake --build build --target buddy-deepseek-r1-f16-run -j

该目标会沿依赖关系自动完成：

运行 import-deepseek-r1.py
→ 生成 f16 MLIR/权重
→ 编译四个 MLIR
→ 创建 DEEPSEEKR1_F16 静态库
→ 编译并链接 f16 main

运行方式见 README.md:95：

./build/bin/buddy-deepseek-r1-f16-run

10.3 分阶段验证清单

检查阶段	建议检查
模型加载	`model.dtype`/参数 dtype 是否为 FP16，模型路径是否正确
图捕获	Prefill/Decode 是否各一张图，输入 shape 是否固定
Buddy Graph	未支持 ATen 算子是否导致 `_ops_map` 查找失败
MLIR 导出	四个文件是否存在，函数名、shape、dtype 是否正确
权重	文件大小是否等于打包元素数乘 2 字节，顺序是否与 main graph 一致
后端编译	四个 `.o` 和静态库是否生成，C 接口符号是否存在
C++ ABI	返回结构字段顺序、MemRef rank/shape、元素位宽是否匹配
数值正确性	Prefill 和前几轮 Decode logits 与 PyTorch 基准对比

建议先比较 logits，再比较 Token。argmax 可能在部分错误情况下偶然选到相同 Token，不能替代数值验证。

11. 常见问题

11.1 `import-deepseek-r1.py` 是否已经完成模型编译？

没有。它完成 PyTorch 前端捕获、Buddy Graph 变换和高层 MLIR 生成。机器码由后续 MLIR/LLVM 工具链产生。

11.2 是否可以说“buddy-opt 把 MLIR 编译成库”？

不够准确。完整链路是 buddy-opt → mlir-opt → buddy-opt → mlir-translate → llvm-as → llc；它们生成目标文件，CMake 再用 add_library(... STATIC ...) 创建静态库。

11.3 静态库是否包含模型权重？

不包含。libDEEPSEEKR1_F16.a 包含计算代码；权重单独保存在 arg0-f16.data，由 C++ 在运行时读取。

11.4 为什么同一模型要导出四个 MLIR 文件？

Prefill 和 Decode 形状、数据流和后端优化不同，所以需要两套代码。每套又拆成 main wrapper 与实际计算子图，因此共四个 MLIR 文件。

11.5 `DynamoCompiler.importer()` 是否直接返回 MLIR？

不是。它主要返回 Buddy Graph 列表并记录参数。脚本之后通过 GraphDriver 和 lower_to_top_level_ir() 生成 MLIR。

11.6 “所有算子放入 subgraph0”是否等于所有算子都融合成一个 MLIR Op？

不等于。simply_fuse 的关键作用是把节点放到同一个子图分组；节点仍可能降低成许多 MLIR Op。只有特定模式 Pass 才会把多个 Buddy 节点替换成专用融合节点。

11.7 f16 是否意味着 Token、位置和所有中间值都是 f16？

不是。权重和主要浮点计算使用 f16；Token ID、Cache 位置仍是 i64，条件可能是 i1。

11.8 Token 是模型函数直接生成的吗？

模型函数产生 logits。当前 C++ main 用 findMaxIndex() 做贪心 argmax 得到 Token ID，再进行 EOS 判断和反分词。

11.9 KV Cache 到底由谁更新？

C++ 负责分配、保存、复制和传递 Cache 内存，也维护外部位置；编译后的 Prefill/Decode 函数负责根据 Attention 计算真正的 K/V 数值并写入结果。两者共同完成 Cache 机制。

11.10 为什么 f16 脚本还调用 `eliminate_matmul_transpose_reshape`？

脚本复用了不同精度的变换流程，但该 Pass 当前实现只接受 Float32，因此对 f16 节点会跳过。这是“函数被调用”与“优化实际生效”之间的区别。

11.11 当前 f16 main 能否直接作为数值正确性基准？

不建议未经验证就作为基准。文件中存在 BF16 遗留标题，而且 decode_f16() 的位转换实现与 IEEE FP16 不一致。架构流程仍可学习，但数值结果应先与 PyTorch logits 对照。

12. 关键源码索引

主题	稳定定位标识	基准源码位置
f16 模型加载	`if args.precision == "f16"`	`import-deepseek-r1.py:73`
两个 DynamoCompiler	`dynamo_compiler_prefill`、`dynamo_compiler_decode`	`import-deepseek-r1.py:91`
Prefill/Decode 捕获	`with torch.no_grad()`	`import-deepseek-r1.py:106`
图优化与融合	`if args.precision == "f16"` 导出分支	`import-deepseek-r1.py:188`
MLIR/权重写出	`# Save the generated files`	`import-deepseek-r1.py:275`
DynamoCompiler 初始化	`DynamoCompiler.__init__`	`frontend.py:75`
FX → Buddy Graph	`DynamoCompiler._compile_fx`	`frontend.py:887`
Dynamo importer	`DynamoCompiler.importer`	`frontend.py:1204`
Buddy Graph → 高层 MLIR	`Graph.lower_to_top_level_ir`	`graph.py:540`
参数打包	`GraphImporter._pack_params`	`graph.py:836`
main graph import	`GraphImporter.import_main_graph`	`graph.py:928`
GraphDriver 子图构造	`GraphDriver.build_subgraph_by_group`	`graph_driver.py:73`
GraphDriver main graph	`GraphDriver.construct_main_graph`	`graph_driver.py:200`
权重转置消除	`eliminate_transpose`	`eliminate_weight_transpose.py:28`
经典融合	`apply_classic_fusion`	`fuse_ops.py:246`
f16 MLIR 编译流水线	`OUTPUT forward_prefill-f16.o`	`CMakeLists.txt:313`
f16 静态库	`add_library(DEEPSEEKR1_F16 ...)`	`CMakeLists.txt:1510`
C ABI 声明	`_mlir_ciface_forward_prefill/decode`	`buddy-deepseek-r1-f16-main.cpp:232`
分词	`tokenizeInput`	`buddy-deepseek-r1-f16-main.cpp:306`
权重加载	`loadParameters`	`buddy-deepseek-r1-f16-main.cpp:322`
C++ main 与内存分配	`main`	`buddy-deepseek-r1-f16-main.cpp:396`
Prefill 调用	`_mlir_ciface_forward_prefill(...)`	`buddy-deepseek-r1-f16-main.cpp:500`
KV Cache 交接	`copy_kv_by_cache_position_block(...)` 调用	`buddy-deepseek-r1-f16-main.cpp:614`
Decode 循环	`for (int i = 1; ...)`	`buddy-deepseek-r1-f16-main.cpp:624`

结语

沿着一次真实执行看，Buddy-MLIR 的 DeepSeek f16 示例不是单一“模型转换工具”，而是三个清晰衔接的阶段：

Python 导入器确定计算和参数表示
        → MLIR/LLVM 工具链生成目标代码
        → C++ 调度程序管理运行时状态并驱动生成循环

掌握这条边界后，再阅读其他精度、量化版本或新的配置驱动模型流水线时，可以继续用同样的问题定位代码：输入是什么、这一阶段改变了什么表示、输出交给谁、运行时状态由谁保存。

03. 从0开始学习硅基智能 - buddy-mlir DeepSeekR1模型导入编译与调度总结