最近在做一个基于 ReAct 范式的 Agent 项目，底层用 PyTorch 2.1 + torch.compile 想优化推理速度。但发现 Agent 每轮都要根据上一步的 observation 动态构造 prompt 重新推理，输入 shape 和计算图一直在变，torch.compile 好像一直在触发 recompile，日志里全是 guard failure。试过 dynamic=True 和 mark_dynamic，效果也不太理想，有时候甚至比 eager 还慢。想问下大家在实际 Agent 场景里是怎么处理这种动态性的？是干脆放弃 compile，还是有什么别的 trick？
楼主
25天前
PyTorch 2.x 的 torch.compile 在 Agent 多轮推理场景下真的能加速吗？
请 登录 后发表回复
全部回复
共 3 条
2楼
24天前
我踩过一样的坑，Agent 这种每轮 prompt 长度都在变的场景，torch.compile 基本就是在做无用功，guard 一直失效还不如不开。后来我的做法是把编译粒度缩小到模型内部的固定子模块，比如 attention 或 FFN 单独 compile，外层动态拼装的部分留给 eager。另外 dynamic=True 对变长 seq 帮助有限，真要提速可以试试 pad 到固定 bucket 长度再 compile。
3楼
19天前
Agent 多轮推理这个场景确实挺折磨 torch.compile 的，你这个 guard failure 刷屏基本是必然的，因为 ReAct 每轮 prompt 长度都在变，KV cache 的 shape 也跟着变，编译图根本稳不下来。我之前也踩过类似的坑，dynamic=True 只能救一部分，遇到控制流分支（比如工具调用走不同路径）还是会不停 recompile。后来我的做法是把 compile 只加在真正稳定的那块，比如固定的 system prompt 编码或者某个 shape 锁死的小模块，动态部分老老实实跑 eager，反而整体更快。另外你可以看看是不是每次都在用新的 input_ids 触发 padding 变化，把 padding 策略固定成 max_length 或者用 left padding 能减少一些重编译。还有一点值得试的是 torch.compile 的 mode="reduce-overhead" 配合 CUDA graph，但前提是你的 batch 和 seq len 得尽量固定，不然 CUDA graph 也白搭。说实话 Agent 这种高度动态的负载，我现在的态度是能不 compile 就不 compile，省下来的调优时间拿去优化调度和 KV cache 复用更划算。如果非要上，建议先 profile 一下到底哪段占大头，别把时间浪费在编译不稳定的小算子上。
4楼
14天前
Agent 场景图老变，compile 基本白搭，我后来直接关了换 vLLM 跑推理。