最近在做一个基于 LLM 的 Agent 项目，需要频繁调用小模型做工具选择和结果解析，延迟卡在 200ms 下不去。听说 PyTorch 2.x 的 torch.compile 能加速推理，但我试了下发现首次编译巨慢，而且 Agent 里输入长度变化频繁，动不动就触发重新编译，反而更慢了。是我用法不对还是这玩意儿根本不适合动态 shape 的 Agent 工作流？有没有大佬在实际 Agent 场景里成功用起来的，求指点下配置姿势。
楼主
17天前
PyTorch 2.x 的 torch.compile 在 Agent 推理场景下到底能不能打？
请 登录 后发表回复
全部回复
共 1 条
2楼
15天前
动态 shape 这块确实是 torch.compile 在 Agent 场景里最难受的地方，你遇到的重编译问题不是用法错，而是它默认对 shape 变化太敏感。可以试试开 dynamic=True 配合 mark_dynamic 把那些会变的维度标出来，能少触发很多次编译。不过说实话 Agent 里工具调用这种小模型，输入长度从几十到几百 token 来回跳，编译收益经常被重编译开销吃掉。我自己的做法是这种场景干脆退回 eager，或者用 ONNX / TensorRT 提前把固定几个常用 shape 编好。torch.compile 更适合那种 shape 相对稳定的批量推理，比如固定 max_len padding 之后跑。如果非要上，建议加个 shape 分桶，把长度归到几个区间里，别让它每个新长度都编一遍。