大模型推理平台编译与图捕获调试指南

一、背景：为什么需要分段图捕获

在大模型推理场景中，Prefill 阶段与 Decode 阶段的执行特征差异很大。Prefill 需要一次性处理可变长度的输入序列，计算图中会频繁出现动态形状；Decode 阶段则通常以固定批量、固定序列长度逐步生成token，计算图相对稳定，非常适合用图捕获技术消除算子调度与内核启动开销。

早期实践通常只对 Decode 阶段做图捕获，因为其输入形状稳定、执行路径单一。而 Prefill 阶段由于输入长度不固定，直接对整张前向图做捕获会面临静态形状约束无法满足的问题，导致捕获失败或运行时回退到非图模式。

为了兼顾执行效率与动态形状需求，业界逐渐形成一种“分段图捕获”思路：将完整的前向计算图按子模块边界拆分成若干片段，对形状稳定的片段做图捕获，对形状敏感的片段保留动态执行路径。这样既能减少调度开销，又不会因为局部动态性而牺牲整体优化空间。

二、技术原理：从动态编译到子图切分

分段图捕获的基础是动态编译框架。动态编译前端会追踪模型前向过程，记录执行轨迹，并根据算子边界与数据依赖关系将代码转换为中间表示。这个中间表示并不是为了做内核级融合，而是为了完成计算图的结构化重写与子图边界识别。

2.1 切分目标

切分的主要目标是将模型中形状不稳定的算子隔离出去。典型做法是：

• 将词嵌入、归一化、多层感知机等形状相对固定的模块作为可捕获子图；
• 将注意力计算等依赖可变序列长度的模块保留为动态执行；
• 每个可捕获子图独立编译、独立捕获、独立回放。

2.2 简化执行流

一个典型的执行流可以抽象为：

片段0: 词嵌入       -> 图捕获
片段1: 第0层注意力  -> 动态执行
片段2: 第0层前馈    -> 图捕获
片段3: 第1层注意力  -> 动态执行
片段4: 第1层前馈    -> 图捕获
...

这种交替结构使得每一层中形状敏感的注意力部分保持灵活，而其余部分享受图捕获带来的低开销收益。

2.3 启用方式

在推理平台中，分段图捕获通常通过服务启动参数开启。以通用配置为例，启用方式可描述为：

python -m inference_launcher \
  --model-path /path/to/model \
  --enable-piecewise-graph

不同平台的参数名称和默认值可能不同，具体需要参考对应平台的配置说明。但整体设计目标一致：将动态编译前端用于图重写，而非启用完整的内核级编译后端。

三、稳定运行时的正常表现

当分段图捕获配置正确且模型适配良好时，启动日志通常会出现两类关键信息。

第一类是子图编译耗时信息，用于提示每个动态形状子图在首次编译时花费的时间：

Compiling a graph for dynamic shape takes 0.72 s

第二类是捕获进度信息，表明运行时正在逐个捕获子图，并显示当前可用显存：

Capturing num tokens (num_tokens=4 avail_mem=11.25 GB): 100%

如果捕获过程能够顺利完成，并且发送简单推理请求后输出内容与预期一致，基本可以认为该模型在分段图捕获模式下运行正常。此时可以进一步执行标准基准测试，观察吞吐与延迟是否满足上线要求。

四、常见问题与排查路径

分段图捕获虽然能显著提升稳定场景的执行效率，但在实际适配新模型时，开发者经常会遇到三类问题：编译失败、图捕获异常、推理结果精度下降。

4.1 编译失败

编译失败通常发生在动态编译前端追踪模型前向过程时。常见原因包括：

• 模型中存在动态编译无法识别的自定义算子；
• 控制流依赖运行时数据，导致追踪轨迹不稳定；
• 部分算子绕过了框架的算子注册机制，无法被纳入中间表示。

排查编译失败的第一步是完整阅读报错堆栈。堆栈中通常会指出触发失败的具体模块与算子类型。开发者需要判断该算子是否属于已知的受支持算子集合，或者是否需要为该算子注册专门的解析规则。

4.2 图捕获异常

图捕获异常的典型表现是进程在捕获阶段崩溃或挂起。常见原因包括：

• 捕获期间发生了设备内存分配，而分配结果与捕获时记录的地址不一致；
• 子图内部存在指向外部缓冲区的隐式依赖，这些依赖在回放时无法正确恢复；
• 子图边界划分不合理，导致跨子图的数据依赖被割裂。

针对图捕获异常，一个实用方法是逐个子图禁用捕获，先将所有片段都切换为动态执行，观察推理是否恢复正常。如果恢复正常，再逐个开启捕获，以二分方式确定引发异常的具体片段。

4.3 推理结果精度下降

精度下降是最隐蔽、排查成本最高的问题。常见原因包括：

• 子图边界处存在数值舍入顺序变化，导致累计误差被放大；
• 图回放时使用了与捕获时不同的随机数种子或缓存状态；
• 某些算子在图模式下启用了不同的内核实现，导致数值行为与动态模式不完全一致。

排查精度问题时，建议先在同一输入上分别运行动态模式与分段图捕获模式，对比每一层输出的最大绝对误差。通常误差会从某一层开始突然增大，据此可以定位到引入差异的算子。

五、对拍与根因定位方法

对拍是定位精度问题最有效的手段之一。具体步骤如下：

1. 使用相同权重与相同输入，分别执行动态模式与图捕获模式；
2. 在关键模块边界处保存中间张量，例如词嵌入输出、每层注意力输出、每层前馈输出；
3. 计算两种模式下中间张量的差异，记录差异首次显著增大的位置；
4. 将该位置对应的算子作为根因候选，进一步检查其内核实现与数值行为。

对拍过程中需要注意，浮点累加顺序的微小变化可能导致误差在不同层级之间传播与放大。将差异首次超过阈值的层级作为起点，而不是只关注最终输出，是更高效的做法。

六、实践建议与工程化经验

在多次适配与调试过程中，以下经验具有较高的通用性：

6.1 控制子图粒度

子图切分并非越细越好。过于细碎的切分会导致捕获与回放调度本身的开销增大，抵消图捕获带来的收益。通常以完整的前馈模块或归一化模块为最小切分单位的做法在性能与灵活性之间取得了较好平衡。

6.2 早验证、早对拍

在正式运行基准测试之前，先用简单请求做对拍验证。这一步骤能尽早暴露精度问题，避免在后续长时间压测中才发现结果异常。

6.3 建立可复现的排查样例

对于每个新模型，保留固定的输入样例、预期输出和中间层误差阈值，便于后续回归排查与版本比对。可复现的排查样例能大幅提高团队内部协作效率。

6.4 关注形状敏感算子

注意力计算、变长序列处理、条件控制流等是动态形状的主要来源。在分析图捕获稳定性时，优先关注这些算子的输入形状变化范围，必要时为它们单独配置动态路径。

七、总结

分段图捕获通过在“完全动态执行”与“整图静态捕获”之间引入中间地带，为大模型推理提供了更灵活的性能优化手段。其核心挑战不在于图捕获本身，而在于如何准确识别模型中形状敏感的部分、如何划分清晰的子图边界，以及如何在出现异常时快速收敛到根因。

对于开发者而言，掌握编译前端追踪机制、熟悉图捕获的约束条件、建立规范化的对拍流程，是提升调试效率的关键。不同推理平台在具体实现与参数配置上可能存在差异，但背后的设计目标与排查逻辑是相通的。以通用技术视角去理解分段图捕获，能够帮助开发者在不同平台之间迁移经验，更从容地应对新模型、新架构带来的挑战。