修复 DeepSeek-V4 Auto-Functionalized 问题学习

问题描述

在使用 torch.compile 运行 DeepSeek-V4 模型时，出现 AssertionError: auto_functionalized was not removed 错误。

根因分析

  1. 自定义算子: fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert 是 DeepSeek-V4 MLA attention 使用的自定义算子
  2. In-place 修改: 该算子对 q 和 k_cache 张量进行了 in-place 修改
  3. torch.compile 行为: 在 torch.compile 过程中，AOTAutograd 将这个算子包装为 auto_functionalized 节点
  4. 功能化传递未处理: FixFunctionalizationPass 没有处理这个算子，导致节点残留在计算图中，触发 PyTorch 的 assertion 失败
  5. 修复的本质: splitting_ops 在 FixFunctionalizationPass 之后执行，所以这个修复只是解决了 compile 阶段的问题（让 graph 能通过编译）。在运行时，这些算子仍然走 eager 路径

torch.compile 的 auto_functionalized 机制

什么时候会包装 auto_functionalized？

  1. 纯函数 Ops: 不会包装，可进行 DCE（Dead Code Elimination）、Fusion
  2. In-place 操作: 需要创建副本，会被包装为 auto_functionalized

为什么需要 FixFunctionalizationPass？

标准流程

  1. FX Graph 用函数包装，产生处理副本的节点
  2. Inductor 清除副本节点，还原
  3. Dispatch 到自己的 kernel

本问题的场景

PyTorch 不能很好地去掉复制操作，vLLM 通过手动处理，在步骤 1、2 之间介入，手动 defunctionalize

修复方案

在 FixFunctionalizationPass 中添加对该算子的处理：

elif at_target == torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.default:
    mutated_args = {1: "q", 3: "k_cache"}
    self.defunctionalize(graph, node, mutated_args)

torch.compile 执行模式

Piecewise 模式

  1. 编译器生成 FX Graph
  2. Split FX Graph 成为 Sub Graph
  3. CUDA capture / eager

缺点: 切分逻辑复杂，需要支持随意中断和恢复（如中间加入无法 trace 的函数）

BreakGraph 模式

  1. 执行到 @eager_break_during_capture 装饰的函数自动 eager
  2. CPU 接管程序，调用无法 capture 的 kernel

什么场景下 kernel 需要 BreakGraph

判断一个 kernel 要不要用 BreakGraph，基本看几个条件：

  1. 依赖运行时动态上下文 - 比如需要从外部获取 batch 状态、metadata 之类的，编译期拿不到 → Break
  2. 访问全局共享状态 - 不是通过参数传进来，而是直接读写全局变量 → Break
  3. 有 host 端副作用 - 除了 GPU 计算，还在 CPU 端干了点事（比如打印、发请求） → Break
  4. 纯计算 - 即使有 inplace 修改，只要不涉及上面几条 → Piecewise + defunctionalize 就行

BreakableCUDAGraphWrapper

启发与思考

  1. Piecewise 切分逻辑确实不合理，BreakGraph 方案更合理
  2. auto_functionalized 是 torch.compile 处理 in-place 操作的核心机制，理解它对于调试编译问题很重要
  3. vLLM 的 FixFunctionalizationPass 是一个很好的自定义编译传递案例，展示了如何在框架层面解决 PyTorch 的局限性