02. 从0开始学习硅基智能 - buddy-mlir deepseek模型导入代码解析

对deeppseek_r1 模型导入代码的理解

一、文档目标与整体认识

本文分析以下脚本：

buddy-mlir/examples/BuddyDeepSeekR1/import-deepseek-r1.py

这个脚本是 DeepSeek-R1-Distill-Qwen-1.5B 示例的 AOT 模型导入器。它不负责训练模型，也不直接执行完整的文本生成，而是完成以下转换：

HuggingFace/PyTorch 模型
    ↓
TorchDynamo 捕获的 FX Graph
    ↓
Buddy Graph
    ↓
图优化、算子融合和子图分组
    ↓
高层 MLIR + 二进制权重文件

脚本生成的高层 MLIR 还需要经过 buddy-opt、mlir-opt 和 LLVM 工具链，才能变成最终的目标文件和可执行程序。因此，该脚本处在完整编译流程的"模型导入与高层 IR 生成"阶段。

按照职责，可以将脚本划分为四个主要功能块：

功能块	主要内容
加载模型、配置图导入器	参数解析、模型加载、创建两个 DynamoCompiler
捕获计算图	构造 Prefill/Decode 输入和 StaticCache，调用 importer()
图优化与融合	校验图、提取参数、消除转置、Attention 融合和子图分组
生成高层 MLIR 并导出文件	GraphDriver lowering、主图生成、权重和 MLIR 落盘

由于 f16 与 f32/bf16 分支中存在重复代码，部分功能块的行号并不连续。理解时应按职责阅读，而不是把每个精度分支看成完全不同的流程。

二、前置概念总结

2.1 AOT 模型导入器

AOT 是 Ahead-Of-Time 的缩写，表示在模型正式运行之前完成图捕获、转换和编译。

传统 PyTorch 推理依赖 Python 和 PyTorch 运行时动态执行模型；AOT 流程则提前把模型转换为静态中间表示：

PyTorch 模型
→ 计算图
→ 编译器中间表示
→ MLIR
→ LLVM IR
→ 机器码

AOT 的主要价值是：

• 运行时不再逐个解释 Python/PyTorch 算子；
• 可以提前执行算子融合、内存规划和向量化；
• 可以部署到没有完整 Python/PyTorch 环境的平台；
• 可以针对 CPU、RISC-V 等目标硬件生成代码。

AOT 的代价是导出结果通常针对固定输入规格。例如本脚本固定了 batch size、Prefill 长度、Decode 长度、Cache 最大长度和数据精度。改变这些条件时，通常需要重新导入和编译。

2.2 图捕获

图捕获是指：用示例输入执行模型，同时记录执行过程中出现的张量算子和算子之间的数据依赖。

例如：

def forward(x, weight):
    y = x @ weight
    return torch.relu(y)

可以捕获为：

x ──────┐
        ├→ MatMul → ReLU → Output
weight ─┘

计算图通常保存以下信息：

• 算子节点，如 MatMul、Add、Reshape；
• 张量之间的数据依赖；
• 张量形状和数据类型；
• 模型参数与运行时输入的区别；
• 模型当前实际走过的执行分支。

本脚本的底层图捕获能力主要由 PyTorch TorchDynamo 提供。Buddy 的 DynamoCompiler 接收捕获后的 FX Graph，再将其转换为 Buddy Graph。

2.3 "完整、稳定的计算图"的准确含义

本脚本得到的不是适用于所有输入情况的通用动态图，而是：

针对当前模型配置、输入形状、数据类型和执行分支的一张完整静态图。

本脚本主要特化到以下条件：

batch size          = 1
Prefill 序列长度    = 1024
Decode 序列长度     = 1
KV Cache 最大长度   = 1024
use_cache           = True
precision           = f32、f16 或 bf16 中的一种

示例输入中的 token 值为零，但这些值一般作为图的运行时输入，而不是写死在图里。真正被固定的主要是形状、类型和 Python 控制分支。

TorchDynamo 会为图建立 Guard 条件，例如输入形状、dtype、模型类型和 use_cache 配置。Guard 条件成立时，该图有效；条件发生变化时，普通 PyTorch 编译模式可能重新捕获，但 AOT 产物通常需要重新导出。

脚本中的：

assert len(graphs_prefill) == 1
assert len(graphs_decode) == 1

要求 Prefill 和 Decode 各自只形成一张 Buddy Graph。如果发生 graph break 并产生多张图，断言就会失败。

2.4 图变换、算子融合与 Lowering

这三个概念容易混淆：

图变换：在同一抽象层级重写计算图
算子融合：把特定算子组合替换为更高效的操作
Lowering：把高层表示逐步转换成更低层表示

例如：

图变换：Transpose → Reshape 变成单个 Reshape
算子融合：Q/K/V + Cache 更新 + Attention 变成 GQAAttentionFusedOp
Lowering：TOSA MatMul 变成 Linalg，再变成循环、Vector 和 LLVM

本脚本主要执行 Buddy Graph 层面的图变换和融合；CMake 中的 pass pipeline 负责后续 MLIR lowering。

2.5 TOSA 算子与 MLIR Dialect

TOSA 是 Tensor Operator Set Architecture 的缩写，是面向机器学习编译器的标准张量算子集合，也是 MLIR 中的一种 Dialect。

典型 TOSA 操作包括：

tosa.add
tosa.mul
tosa.matmul
tosa.reshape
tosa.transpose

Buddy 使用算子注册表把内部图节点映射为 TOSA、Linalg、Arith、Tensor 等 MLIR 操作：

PyTorch aten.mm
→ Buddy MatmulOp
→ TOSA/Linalg MLIR

TOSA 的作用是提供一种相对独立于框架和硬件的高层表达，后续再逐步降低到更接近硬件的 IR。

2.6 Prefill、Decode 与 KV Cache

因果语言模型的推理分为两个阶段。

Prefill 阶段一次处理整个 Prompt：

[1, 1024] 个 token
→ 计算所有层
→ 生成初始 KV Cache
→ 输出各位置的 Logits

Decode 阶段每次只处理一个新 token：

[1, 1] 个新 token + 历史 KV Cache
→ 计算新 token
→ 更新 KV Cache
→ 输出下一个 token 的 Logits

KV Cache 保存各 Transformer 层之前 token 的 Key 和 Value，Decode 时不必重复计算历史 token。

StaticCache 使用固定最大形状，例如：

[batch, kv_heads, max_cache_len, head_dim]
= [1, 2, 1024, 128]

每次 Decode 只更新 cache_position 指定的位置，因此整体形状稳定，适合 AOT 编译和静态 MLIR 函数签名。

2.7 各组件的职责

组件	主要职责
HuggingFace Transformers	提供模型结构、模型权重、StaticCache 和 Cache 更新逻辑
TorchDynamo	监控 Python/PyTorch 执行并捕获 FX Graph
AOTAutograd	对捕获图进行规范化，为前向编译器提供图
TorchInductor decomposition	将复杂 PyTorch 算子拆成基础 Aten/Prims 算子
Buddy DynamoCompiler	将 FX Graph 转换成 Buddy Graph，并收集参数
Buddy 图变换	消除冗余节点、执行 Attention 等融合
GraphDriver	根据分组生成计算子图和调用子图的主图
MLIR/LLVM 工具链	将高层 MLIR 降低并编译为目标代码

从完整 Buddy 编译器的角度，DynamoCompiler 是接收 PyTorch 模型的前端；从 TorchDynamo 的接口角度，它又是接收 FX Graph 的自定义 backend。两种称呼观察的是不同系统边界，并不矛盾。

三、功能块一：加载模型、配置图导入器

3.1 解析导出参数

脚本首先定义两个命令行参数：

parser.add_argument("--output-dir", type=str, default="./")
parser.add_argument(
    "--precision",
    type=str,
    default="f32",
    choices=["f32", "f16", "bf16"],
)

--output-dir 决定 MLIR 和权重文件的保存位置；--precision 决定模型加载精度、MLIR 数据类型和权重文件格式。

随后创建输出目录：

output_dir = args.output_dir
os.makedirs(output_dir, exist_ok=True)

exist_ok=True 表示目录已经存在时不会报错。

3.2 确定模型路径

model_path = os.environ.get("DEEPSEEKR1_MODEL_PATH")
if model_path is None:
    model_path = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"

脚本优先使用环境变量指定的本地模型目录。如果没有设置，则使用 HuggingFace 模型名称。

这里实际加载的是蒸馏版 Qwen 1.5B 模型，不是完整参数规模的 DeepSeek-R1。较小模型更适合示例编译和 CPU 推理。

3.3 按精度加载模型

f16 分支：

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

bf16 分支使用 torch.bfloat16 和 .bfloat16()，默认分支使用 torch.float32。

关键调用的含义为：

• AutoModelForCausalLM：根据配置自动加载因果语言模型；
• from_pretrained()：加载模型结构和预训练参数；
• dtype：设置加载权重的数据类型；
• .eval()：关闭 Dropout 等训练行为；
• .half()/.bfloat16()：确保参数和相关 buffer 使用目标精度。

脚本随后设置：

model.config.use_cache = False

这只是关闭模型配置中的默认值。捕获两条图时仍会显式传入 use_cache=True，从而明确选择带 KV Cache 的执行路径。

3.4 创建 Prefill 和 Decode 图导入器

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",
)

创建 DynamoCompiler 对象只是配置一个图转换器，此时还没有捕获模型。

三个参数分别表示：

• primary_registry：指定 Buddy Op 到 MLIR 的主要生成规则；
• aot_autograd_decomposition：指定复杂 PyTorch 算子的拆解规则；
• func_name：指定最终生成的 MLIR 入口函数名。

必须创建两个导入器，是因为两次捕获需要保存不同的图集合和函数名。真正让它们分别成为 Prefill、Decode 图的，不只是函数名，而是后面传入的输入形状和 Cache 参数。

本功能块的结果是：

一个已加载、处于推理模式的指定精度模型
+ 一个 Prefill DynamoCompiler
+ 一个 Decode DynamoCompiler

四、功能块二：捕获计算图

4.1 使用 torch.no_grad()

with torch.no_grad():
    ...

该上下文临时关闭自动求导，避免构建反向传播图和保存训练中间量。它不会代替 .eval()；两者分别控制梯度记录和模块训练行为。

4.2 构造 StaticCache 和示例输入

past_key_values_prefill = StaticCache(
    config=model.config, max_cache_len=1024
)
past_key_values_decode = StaticCache(
    config=model.config, max_cache_len=1024
)

两个 Cache 分开创建，防止 Prefill 和 Decode 捕获过程共享可变状态。

示例输入为：

data_prefill = {
    "input_ids": torch.zeros((1, 1024), dtype=torch.int64),
}
data_decode = {
    "input_ids": torch.zeros((1, 1), dtype=torch.int64),
}
cache_position = torch.tensor([200], dtype=torch.int64)

这里固定了两条图的输入规格：

输入	形状	含义
Prefill input_ids	[1, 1024]	一次处理完整 Prompt
Decode input_ids	[1, 1]	每次处理一个 token
cache_position	[1]	新 K/V 写入 Cache 的位置

这些 Tensor 的数值主要用于让模型可以实际执行，形状和 dtype 才是静态图生成的关键条件。

4.3 捕获 Prefill 图

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

调用 importer() 后，TorchDynamo 执行并观察模型，将当前路径中的张量运算捕获为 FX Graph。Buddy 随后把 FX 节点转换为 Buddy Graph 节点，并保存到 graphs_prefill。

参数含义：

• input_ids：确定 Prefill 输入签名；
• use_cache=True：要求模型输出并维护 KV Cache；
• cache_implementation="static"：选择静态 Cache 实现。

脚本虽然创建了 past_key_values_prefill，但向 Prefill importer() 传入它的代码被注释掉了。因此该变量当前没有实际参与 Prefill 捕获，初始 Cache 由 Transformers 的静态 Cache 路径处理。

4.4 初始化并捕获 Decode 图

在捕获 Decode 前，脚本先调用一次模型：

model(
    input_ids=data_decode["input_ids"],
    past_key_values=past_key_values_decode,
    use_cache=True,
    cache_implementation="static",
)

这次调用的主要作用是初始化 past_key_values_decode 内部各层的 K/V Tensor，使后续捕获看到确定的张量结构和形状。

随后捕获 Decode 图：

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

与 Prefill 相比，Decode 多了两个重要输入：

• past_key_values：历史 K/V Cache；
• cache_position：当前 token 在 Cache 中的位置。

因此 Decode 图表示：

新 token
+ 历史 K/V Cache
+ Cache 位置
→ 单 token Transformer 计算
→ 更新后的 K/V Cache
+ 下一 token 的 Logits

4.5 DynamoCompiler.importer() 内部发生了什么

核心过程可概括为：

dynamo.optimize(self._compile_fx)(model)
→ 执行带示例输入的模型
→ TorchDynamo 生成 FX Graph
→ AOTAutograd 规范化图
→ Inductor decomposition 拆解复杂算子
→ Buddy _compile_fx() 遍历 FX 节点
→ 区分参数、buffer、运行时输入
→ 构造 Buddy Graph
→ 保存 imported_graphs 和 imported_params

模型中的 Transformer 层循环次数由配置确定，TorchDynamo 可以将固定层数的循环展开进图。输入 token 的值参与 Embedding 等张量操作，但通常不会改变 Python 控制流，因此零值输入仍可以捕获完整张量计算路径。

f16 与 f32/bf16 分支中的捕获逻辑基本相同，主要区别来自模型 dtype。当前代码存在重复，可以在理解时视为同一套算法的两个精度分支。

本功能块的结果是：

graphs_prefill：Prefill Buddy Graph 列表
graphs_decode：Decode Buddy Graph 列表

五、功能块三：图优化与融合

5.1 校验图并提取参数

assert len(graphs_prefill) == 1
assert len(graphs_decode) == 1

graph_prefill = graphs_prefill[0]
graph_decode = graphs_decode[0]
params = dynamo_compiler_prefill.imported_params[graph_prefill]

两个断言要求每个阶段只有一张图。检查通过后取出 Buddy Graph，并从 Prefill 导入器中取得按占位节点顺序排列的参数列表。

Prefill 和 Decode 使用同一模型，因此只导出一份模型权重。这里隐含的约束是：两张图对共享模型权重采用一致的参数顺序和布局。

5.2 消除冗余转置

graphs_prefill[0].perform(
    [eliminate_transpose, eliminate_matmul_transpose_reshape]
)
graphs_decode[0].perform(
    [eliminate_transpose, eliminate_matmul_transpose_reshape]
)

perform() 按顺序调用变换函数，并直接修改 Buddy Graph。

eliminate_transpose 主要将：

Weight → Transpose → MatMul

改成：

预先转置的 Weight → MatMul

该变换不仅删除图节点，还会修改真实参数 Tensor 和参数形状。这样运行时不再执行权重转置，导出的 .data 文件中直接保存目标布局的权重。

eliminate_matmul_transpose_reshape 尝试消除不会改变有效数据顺序的：

Transpose/Permute → Reshape/View

当前实现只处理 f32 节点，因此在 f16、bf16 图中通常会跳过。

这些变换放在 lowering 之前，是因为此时可以同时访问 Buddy Graph 和 PyTorch 参数，而且简化后的图更容易匹配后续融合模式。

5.3 Prefill 融合

pattern_list_prefill = [
    simply_fuse,
    apply_classic_fusion,
    flash_attention_prefill,
]

执行顺序如下：

1. simply_fuse

   将所有非 PlaceholderOp 节点放入一个名为 subgraph0 的 CPU 分组。它主要执行子图划分，不是把所有节点变成一个算子。

2. apply_classic_fusion

   处理常规线性计算模式。例如将目标 addmm 拆成 matmul + add，再将 Permute(weight, [1, 0]) + MatMul 替换为直接处理右侧转置的融合 MatMul。

3. flash_attention_prefill

   将 ScaledDotProductFlashAttentionForCpuOp 替换为 FlashAttentionForCpuPrefillOp，使用适合长 Query 序列的 Prefill Attention 实现。

5.4 Decode 融合

pattern_list_decode = [
    simply_fuse,
    apply_classic_fusion,
    gqa_attention_fusion,
]

Decode 前两个步骤与 Prefill 相同，最后执行 gqa_attention_fusion。

该变换识别静态 Cache 更新和 GQA Attention 组合，例如：

K Cache：IndexPut → Unsqueeze → Expand → Clone → View
V Cache：IndexPut → Unsqueeze → Expand → Clone → View
Q/K/V：Scaled Dot-Product Attention

识别成功后替换为 GQAAttentionFusedOp，减少中间 Tensor 和独立操作，适合每次只处理一个 Query token 的 Decode 阶段。

融合通过以下代码执行：

graphs_prefill[0].fuse_ops(pattern_list_prefill)
graphs_decode[0].fuse_ops(pattern_list_decode)

fuse_ops() 会按列表顺序调用变换函数，并原地修改图。

5.5 重命名分组并指定设备

graph_prefill.op_groups["subgraph0_prefill"] = (
    graph_prefill.op_groups.pop("subgraph0")
)
graph_prefill.group_map_device["subgraph0_prefill"] = DeviceType.CPU

graph_decode.op_groups["subgraph0_decode"] = (
    graph_decode.op_groups.pop("subgraph0")
)
graph_decode.group_map_device["subgraph0_decode"] = DeviceType.CPU

重命名可以避免两个阶段都使用 subgraph0 而产生符号冲突。新名称随后成为 MLIR 子图函数名。DeviceType.CPU 表示这些分组面向 CPU lowering。

本功能块的最终结果是：

优化后的 graph_prefill
  └── subgraph0_prefill（CPU）

优化后的 graph_decode
  └── subgraph0_decode（CPU）

六、功能块四：生成高层 MLIR 并导出文件

6.1 使用 GraphDriver 构造子图

driver_prefill = GraphDriver(graphs_prefill[0])
driver_prefill.subgraphs[0].lower_to_top_level_ir()

driver_decode = GraphDriver(graphs_decode[0])
driver_decode.subgraphs[0].lower_to_top_level_ir()

GraphDriver 根据 op_groups：

• 找出分组的外部输入；
• 找出分组输出；
• 构造独立 Buddy 子图；
• 分析多个子图之间的调用依赖；
• 为主图生成函数声明和调用关系。

本脚本每个阶段只有一个子图，因此：

driver_prefill.subgraphs[0] = subgraph0_prefill
driver_decode.subgraphs[0]  = subgraph0_decode

6.2 lower_to_top_level_ir()

该方法遍历 Buddy 子图节点，使用算子注册表生成高层 MLIR：

Buddy AddOp       → TOSA Add
Buddy ReshapeOp   → TOSA Reshape
Buddy MatmulOp    → TOSA/Linalg MatMul
Buddy AttentionOp → TOSA/Linalg/Vector 等操作组合

生成的模块保存在：

driver_prefill.subgraphs[0]._imported_module
driver_decode.subgraphs[0]._imported_module

子图函数包含实际模型计算，并把每个模型权重视为独立 Tensor 参数。此时仍然是 TOSA、Linalg、Tensor、Arith 等高层 MLIR，不是最终 LLVM IR。

6.3 子图 MLIR 与主图 MLIR

子图通过以下方式直接写出：

print(driver_prefill.subgraphs[0]._imported_module, file=module_file)

主图通过以下方式构造：

driver_prefill.construct_main_graph(True)
driver_decode.construct_main_graph(True)

True 表示启用参数打包。主图接收连续参数 MemRef，再按固定偏移恢复每个权重的形状，然后调用计算子图：

连续参数 MemRef
→ 提取参数切片
→ 恢复参数形状
→ 调用 subgraph0_prefill/subgraph0_decode
→ 返回 Cache 和 Logits

两类 MLIR 的职责如下：

MLIR 文件	职责
subgraph0_*.mlir	执行真正的模型张量计算
forward_*.mlir	对外入口、参数拆包、调用计算子图

主图和子图在 CMake 中分别编译，最终链接到同一个推理程序。

6.4 导出 f32 和 f16 权重

基本代码为：

all_param = numpy.concatenate(
    [param.detach().numpy().reshape([-1]) for param in params]
)
all_param.tofile(output_path)

每一步的作用为：

detach()          取消参数与自动求导图的关联
numpy()           转换成 NumPy 数组
reshape([-1])     展平单个参数
concatenate()     按图参数顺序拼成一个连续数组
tofile()          写出原始二进制字节

.data 文件没有文件头、参数名称和形状信息。因此，主图中的参数提取顺序、偏移、dtype 必须与参数拼接过程严格一致。

Prefill 和 Decode 使用相同模型，只写出一份权重文件，两张主图共同使用。

6.5 导出 BF16 权重

NumPy 对 BF16 的直接支持有限，因此脚本先转为 float32：

all_param = numpy.concatenate(
    [param.detach().float().numpy().reshape([-1]) for param in params]
)
all_param_bf16 = numpy.frombuffer(
    all_param.astype(numpy.float32).tobytes(),
    dtype=numpy.uint16,
)[1::2]

一个 float32 占 32 位，BF16 保存其高 16 位。[1::2] 在当前常见的小端平台上选取每个 float32 的高半部分，然后写入 arg0-bf16.data。

该实现隐含小端字节序假设。如果需要跨不同端序平台生成权重，应显式处理字节序。

6.6 输出文件总结

精度	Prefill 子图	Prefill 主图	Decode 子图	Decode 主图	权重文件
f32	subgraph0_prefill.mlir	forward_prefill.mlir	subgraph0_decode.mlir	forward_decode.mlir	arg0.data
f16	subgraph0_prefill-f16.mlir	forward_prefill-f16.mlir	subgraph0_decode-f16.mlir	forward_decode-f16.mlir	arg0-f16.data
bf16	subgraph0_prefill-bf16.mlir	forward_prefill-bf16.mlir	subgraph0_decode-bf16.mlir	forward_decode-bf16.mlir	arg0-bf16.data

6.7 Python 脚本之后的编译流程

该脚本只生成高层 MLIR。examples/BuddyDeepSeekR1/CMakeLists.txt 会继续执行：

TOSA
→ Linalg/Tensor/Arith
→ Bufferization/MemRef
→ Affine/Vector
→ SCF/OpenMP
→ LLVM Dialect
→ LLVM IR
→ .o 目标文件

四个 MLIR 模块最终生成四个主要目标文件：

forward_prefill.o
subgraph_prefill.o
forward_decode.o
subgraph_decode.o

这些目标文件组成静态库，并与 C++ 推理代码链接。C++ 运行时加载 .data 权重，通过 MLIR C Interface 调用 forward_prefill 和 forward_decode。

七、完整流程总结

整个脚本可以压缩为以下执行链：

1. 读取输出目录和精度
2. 加载 DeepSeek-R1-Distill-Qwen-1.5B 模型
3. 创建 Prefill/Decode DynamoCompiler
4. 构造固定形状的输入和 StaticCache
5. TorchDynamo 捕获两张 FX Graph
6. Buddy 将 FX Graph 转成 Buddy Graph
7. 校验单图并取得模型参数
8. 消除转置和冗余布局操作
9. Prefill 执行 Flash Attention 融合
10. Decode 执行 GQA + KV Cache 融合
11. 分别构造 Prefill/Decode 计算子图
12. 生成子图高层 MLIR
13. 生成参数打包主图 MLIR
14. 拼接并写出模型权重
15. 后续 CMake/MLIR/LLVM 流水线编译为目标代码

学习该脚本时，需要始终抓住三条主线：

计算主线：PyTorch → FX Graph → Buddy Graph → MLIR
状态主线：Prompt → Prefill Cache → Decode Cache 更新
参数主线：PyTorch 参数 → 参数列表 → 连续 .data → 主图拆包

理解这三条主线后，脚本中的精度分支、融合列表、四个 MLIR 文件以及一份共享权重文件之间的关系就会变得清晰。