PyTorch 2.x 的 compile 和 Agent 工具调用混用时报 graph break，大家怎么绕？

这问题我去年踩过一模一样的坑，当时也是 ReAct 循环里动态拼 prompt，compile 完直接负优化。后来想明白一件事，graph break 的根因其实不是工具调用本身，而是你把「控制流」和「张量计算」混在同一个被编译的函数里了，Python 层的 if/while 对 dynamo 来说就是天敌，它没法把动态分支 trace 成静态图。我最后的做法是把 compile 的边界卡在 LLM 前向那一段，也就是 tokenize 之后、logits 出来之前，用 dynamic=False 加上 mark_dynamic 固定住输入 shape 的维度，工具路由和 prompt 组装全部留在 eager 里跑。KV Cache 对不上大概率是因为你在外层反复重建 cache，试试把 past_key_values 当成显式的输入输出传进编译区，别让它隐式捕获，这样每次 forward 的签名就稳定了。另外 2.3 的 mode="reduce-overhead" 配合 cudagraphs 对固定 shape 的小 batch 挺友好，但 batch 一变就重新 capture，Agent 场景里其实不太划算。我现在基本是 selective compile，只编最热的那几个 matmul 密集层，整体收益比全图 compile 稳得多。你要是延迟敏感，也可以看看 vLLM 或者 TensorRT-LLM 那套 paged attention，比在 torch.compile 上硬抠动态图省心。

我之前也踩过这个坑，compile 碰上 ReAct 那种动态拼 prompt 基本无解，graph break 是必然的。后来我干脆只 compile LLM 前向那一段，把工具路由和 prompt 拼接全放在 eager 里，KV Cache 手动管理反而更稳。动态控制流那块想静态化代价太大，不如接受部分加速。你可以试试用 cudagraph 只包住固定 shape 的 decode 步骤，收益比整体 compile 靠谱。

这个坑我也踩过，torch.compile 碰上 ReAct 这种动态拼 prompt 的循环基本就是灾难，graph break 一多，recompile 的开销直接把那点 kernel fusion 收益吃干净了。我的做法是把 compile 的范围缩到最小，只包住 LLM 前向那一段纯 tensor 计算，工具调用、prompt 拼接、解析这些全放在 eager 里，让 compile 的边界清清楚楚。KV Cache 对不上这个问题，通常是因为你把 cache 的更新也塞进了编译区，其实可以把 cache 当成外部传入的静态输入，每次只 compile 单步 decode，prefill 那块单独走一遍不编译。另外 dynamic control flow 那堆警告里，很多其实是 Python 层的分支，torch.compile 根本没法 capture，与其硬绕不如承认这部分本来就不该进图。真要压性能，可以看看 vLLM 或者 TensorRT-LLM 这类专门做 serving 的，它们对 KV Cache 和动态 batching 的处理比自己在 PyTorch 层面抠 compile 靠谱得多。我现在的组合是 Agent 逻辑用普通 Python 跑，LLM 推理丢给专门的 runtime，torch.compile 只在训练或者固定 shape 的场景用。纯为了 Agent 推理去跟 graph break 死磕，投入产出比是真的低。

我这边也踩过一模一样的坑，ReAct 那种每轮拼不同长度 prompt 的玩法，compile 基本等于白开。后来干脆只 compile LLM 前向里 shape 稳定的那一段，工具调用和 prompt 拼装全留在 eager，反而比硬 compile 快。KV Cache 对不上大概率是因为你把 prefill 和 decode 混在一个编译图里了，拆成两个独立入口会好很多。真要静态化，可以试试把工具路由改成固定长度的 slot 输入，别直接拼字符串。

我也踩过这个坑，torch.compile 碰上 ReAct 那种边跑边拼 prompt 的循环基本就是灾难，graph break 多到怀疑人生，最后性能还不如直接 eager 跑。我的做法是把 compile 的范围缩到最小，只包住真正吃算力的那段纯前向，也就是 input_ids 已经定下来之后的那部分，工具调用、prompt 拼接、KV Cache 的搬运全放在外面用 eager 管。KV Cache 对不上这个问题，其实是因为你把 cache 的分配和更新也塞进了被 compile 的函数里，把它挪出来当外部状态传进去会好很多。另外 dynamic control flow 那块可以考虑用 torch.cond 或者干脆按工具类型分几条静态分支，虽然写起来丑但能减少 break。还有个思路是别硬 compile 整个 LLM，用 vLLM 或者 TensorRT-LLM 这类专门做推理的框架，Agent 的动态逻辑留在 Python 侧，反而更省心。我现在基本就是 compile 只负责固定 shape 的 kernel 级加速，动态路由交给上层，你们要是追求端到端那确实挺难搞的。

这个问题我去年底也踩过，当时用 torch.compile 包整个 ReAct loop，结果编译出来的图碎得跟拼图一样，每次工具返回新 observation 就重编译一次，缓存直接爆掉。后来我的做法是把动态拼 prompt 和工具路由全部留在 Python 侧，只对 LLM 的 forward 做 compile，而且输入固定成 padded tensor 而不是变长 list，这样至少能吃到 inductor 的 kernel fusion。KV Cache 对不上是因为你拆出去之后 cache 的 shape 和 stride 变了，可以试试用 static cache 配合 mark_dynamic，或者干脆手动管理 past_key_values 再喂进去。说实话 Agent 这种控制流本质就是数据依赖的动态分支，指望 compile 全图吃掉不太现实。我现在更倾向于用 vLLM 或 TensorRT-LLM 跑 LLM 部分，Agent 编排逻辑用普通 Python，性能反而更稳。如果非要 compile，建议只 compile 单个 decode step，别碰整个循环。

我一般把动态拼 prompt 那段放 eager，只 compile 纯前向，graph break 少很多，KV Cache 也稳。

我也踩过这个坑，ReAct 那种循环里动态拼 prompt，compile 基本没法收敛。后来我的做法是只把 LLM 前向那段单独 compile，工具调用和 prompt 组装全留在 eager 里，KV Cache 用静态 shape 的 past_key_values 手动管理，别让 compile 去 trace 那部分。效果是有点提升，但没想象中大，毕竟瓶颈很多时候在工具延迟和调度上。你们如果工具返回特别长，可能还得考虑把 observation 截断或缓存，不然 dynamic shape 也会一直触发重编译。

这个坑我们也踩过，torch.compile 对 ReAct 这种每轮 prompt 结构都在变的场景确实很不友好，graph break 基本是必然的。后来我们的做法是把工具调用和 prompt 拼装都留在 eager 里，只对 LLM 前向那一段固定 shape 的部分做 compile，KV Cache 用手动管理的静态 buffer 去接。动态路由那块别硬静态化，收益不大，反而容易把缓存搞乱。真要压延迟，可能还得从 batching 和调度层面下手，比死磕 compile 划算。

这个思路不错，收藏了。