`torch.compile` 不是“开了就快“：RTX 3090 上 12 组对照给出该开、该关、该换模式的边界

EulerBlind

477人浏览 · 2026-05-02 13:11:45

EulerBlind · 2026-05-02 13:11:45 发布

`torch.compile` 不是"开了就快"：RTX 3090 上 12 组对照给出该开、该关、该换模式的边界

我在 RTX 3090 上用一个 6 层的手写 Transformer，跑完 12 组 eager / compile(default) / compile(reduce-overhead) 对照后，最让我意外的不是 “batch=1 时提速 2.87x”，而是 batch=16, seq=128 这组——开了 torch.compile 反而比 eager 慢了 9%。如果你最近正在按官方博客那句"推理前把模型 torch.compile 一下"照做，接下来这份对照表会告诉你，哪种形状值得开，哪种形状是纯亏，以及为什么 reduce-overhead 不是所有场景下的更优解。
结论先说：torch.compile 的收益几乎完全来自消除 eager 模式的 CPU 侧算子派发开销。当你的推理形状足够小、核足够多、每个核跑得足够快时，这个开销占比大，compile 速度起飞；一旦 batch/seq 变大让 GPU 真正忙起来，CPU 派发开销被摊薄到看不见，compile 带来的收益趋近于零，甚至因为图捕获和内存 planning 的副作用倒亏。这一节不是结论的全部，第 5 节会给出更完整的使用判断。

本文默认读者能看懂 PyTorch 代码、了解 SDPA 和 CUDA 执行模型的大致概念，定位偏向"做推理部署或者需要压榨 RTX 20/30 系消费卡的工程师"。不会展开 TorchInductor 和 CUDA Graphs 的底层原理，只讲工程判断。

1. 我为什么不相信"开 compile 一定更快"

最近半年 PyTorch 官方、HuggingFace、vLLM、SGLang 都在反复强调一件事：torch.compile 能给推理带来"显著"提速。问题是**"显著"从来不是一个工程可用的词**。我在真实项目里遇到过至少三种"开了 compile 反而变慢"的场景：

推理服务 batch 变化剧烈，每次新形状都触发重编译，尾延迟飙升。
1. 模型本身算子已经是大矩阵乘，compile 把几个点积融合了，节省不到 3%，却让冷启动多了 3 秒。
1. 开了 mode="reduce-overhead" 后遇到 dynamic shape，CUDA Graph 捕获失败静默回落到默认路径。
  所以这次我干脆把它做成一个可复现的最小实验：一个 768d / 12 head / 6 层的手写 Transformer，3 种模式 × 4 组 shape × bf16 精度，在一张 RTX 3090（SM 8.6，显存 24GB）+ PyTorch 2.9.1+cu128 上跑出实测表。模型和代码在第 3 节给出，拷贝即可复现。

2. 一眼看懂结论：12 组对照表

先上数据，单位毫秒，p50（30 次稳定迭代的中位数）：

输入形状 (batch, seq)	eager	compile(default)	compile(reduce-overhead)	default 提速	reduce-overhead 提速
(1, 128)	1.837	0.697	0.641	2.63x	2.87x
(1, 1024)	3.818	3.111	3.468	1.23x	1.10x
(16, 128)	5.472	6.042	5.997	0.91x	0.91x
(16, 1024)	38.694	38.077	37.969	1.02x	1.02x

我特地把负收益那两行加粗了。这是第一条需要记住的事实：把 batch 推到 16 + seq 推到 128 以上，torch.compile 的收益就已经压到几乎为零。在 batch=16, seq=128 这一组里，三个候选里最快的是 eager。

第二条同样重要的事实来自"decode 风格"的补充实验（batch=1、seq 很短，逐 token 生成时的典型形状）：

seq	eager	compile(default)	compile(reduce-overhead)	default	reduce-overhead
1	1.600	1.321	0.466	1.21x	3.44x
4	1.707	0.533	0.529	3.20x	3.23x
16	1.808	1.517	0.555	1.19x	3.26x
64	0.625	0.612	0.614	1.02x	1.02x
256	1.794	1.576	0.964	1.14x	1.86x

注意 seq=16 这一行：compile(default) 只有 1.19x，而 compile(reduce-overhead) 是 3.26x。这两个模式的差距不是 10%，是 3 倍。后面第 6 节会解释为什么。

3. 实验环境和最小复现路径

硬件/软件：

GPU : NVIDIA GeForce RTX 3090 (Ampere, SM 8.6, 24 GB)
PyTorch : 2.9.1+cu128
dtype : bfloat16
Driver / CUDA : 590.44.01 / CUDA 13.1 (nvidia-smi)

模型：6 层、dim=768、12 heads，SDPA + causal mask。整份 benchmark 代码 130 行左右，核心 block 如下：

class TinyBlock(nn.Module):
    def __init__(self, dim, heads, ff_mult=4):
            super().__init__()
                    self.norm1 = nn.LayerNorm(dim)
                            self.qkv = nn.Linear(dim, dim * 3, bias=False)
                                    self.proj = nn.Linear(dim, dim, bias=False)
                                            self.norm2 = nn.LayerNorm(dim)
                                                    self.ff1 = nn.Linear(dim, dim * ff_mult)
                                                            self.ff2 = nn.Linear(dim * ff_mult, dim)
                                                                    self.heads, self.hd = heads, dim // heads
    def forward(self, x):
            b, t, d = x.shape
                    y = self.norm1(x)
                            qkv = self.qkv(y).view(b, t, 3, self.heads, self.hd).permute(2, 0, 3, 1, 4)
                                    q, k, v = qkv[0], qkv[1], qkv[2]
                                            a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
                                                    a = a.transpose(1, 2).contiguous().view(b, t, d)
                                                            x = x + self.proj(a)
                                                                    y = self.norm2(x)
                                                                            return x + self.ff2(F.gelu(self.ff1(y)))
                                                                            ```
计时用的是 CUDA Events（比 `time.perf_counter()` 更准，且没法漏掉异步执行）：

```python
def cuda_time(fn, n_warmup=5, n_iter=30):
    for _ in range(n_warmup): fn()
        torch.cuda.synchronize()
            starts = [torch.cuda.Event(enable_timing=True) for _ in range(n_iter)]
                ends   = [torch.cuda.Event(enable_timing=True) for _ in range(n_iter)]
                    for i in range(n_iter):
                            starts[i].record(); fn(); ends[i].record()
                                torch.cuda.synchronize()
                                    return [s.elapsed_time(e) for s, e in zip(starts, ends)]
                                    ```
编译模式分别这样创建，`fullgraph=True` 可以确保没有静默的 graph break（一旦 break，收益会被严重稀释）：

```python
m_e = build_model()                                       # eager
m_c = torch.compile(build_model(), mode="default",         fullgraph=True, dynamic=False)
m_r = torch.compile(build_model(), mode="reduce-overhead", fullgraph=True, dynamic=False)

官方文档里 reduce-overhead 的描述是"reduce Python overhead using CUDA Graphs"。这一点在后面非常关键，因为 CUDA Graphs 对输入形状极度敏感。

4. 两个容易被忽略但严重的代价

很多教程写到这里就停了：跑个表、拉条曲线、结束。但这套 benchmark 最有价值的反而是两个不在推理稳态数字里、却真实影响生产的代价。

4.1 首次调用（冷启动）要付 2.5–7 秒

先看第一次调用 compile 模型时，编译阶段到底多贵：

输入形状 (batch, seq)	eager 首次	compile(default) 首次	compile(reduce-overhead) 首次
(1, 128)	935 ms	7137 ms	2805 ms
(1, 1024)	4 ms	3066 ms	2534 ms
(16, 128)	5.6 ms	3007 ms	2800 ms
(16, 1024)	37 ms	3541 ms	3077 ms

最长的一组是 7.1 秒。对一个稳态 latency 只有 1.8ms 的模型来说，这意味着你需要至少跑 3800 次推理才能把冷启动的时间赚回来。如果你的服务冷启动路径就那么一两个请求，这账单连成本都打不平。

这也解释了一个常见的线上现象：Canary 发布时 compile 模型的第一批请求全部超时。因为 torch._inductor 在第一次遇到某个形状时需要生成、编译 Triton kernel，这个过程是阻塞的。

我的建议：生产部署一定要在服务启动阶段主动做 “warm-up 所有典型形状”。不做 warm-up 的 compile 模型等于在给用户看一次压测。

4.2 形状变化的重编译陷阱

下面这个实验比稳态数字更重要。同一个 compile(default, dynamic=False) 模型，我连续喂入 seq = 128、129、130、131、132、128 六个形状，记录每次调用的耗时：

seq= 128 first-call = 2014.08 ms   # 第一次：编译
seq= 129 first-call =  982.38 ms   # 新 shape：重新编译
seq= 130 first-call = 1334.09 ms   # 新 shape：重新编译
seq= 131 first-call = 1131.51 ms   # 新 shape：重新编译
seq= 132 first-call = 1101.22 ms   # 新 shape：重新编译
seq= 128 first-call =    1.81 ms   # 命中缓存：快

每个新长度都要交 1 秒的编译税。把同样的实验切到 dynamic=True：

seq= 128 first-call =  157.54 ms
seq= 129 first-call =    2.77 ms
seq= 130 first-call =    2.60 ms
seq= 131 first-call =    2.60 ms
seq= 132 first-call =    2.59 ms
seq= 128 first-call =    2.53 ms

只有第一次付 157ms 的编译成本，后续所有不同长度都直接走已生成的动态形状 kernel。这个差距在推理服务里是"整体拖死"和"毫秒级抖动"的区别。

我的判断：只要你的推理 shape 会变（会变的例子包括变长输入、变长输出、batch padding 策略、speculative decoding），dynamic=True 几乎是必选。它会让峰值吞吐稍微逊色于 dynamic=False 下的专用 kernel，但换来的是稳定的尾延迟。

5. 形状决定一切：一张判断表

把上面所有数据拉在一起，我总结了一张 3090 上的决策表（其它卡需要自己跑一遍，但判断原则一致）：

你的工作负载	推荐	理由
纯 decode 阶段，batch=1，seq 为 1–16 的单步	`mode="reduce-overhead"`	3x+ 提速，CUDA Graphs 消除了 launch 开销
Prefill 阶段，batch 小，seq 短（<256）	`mode="default"` 或 reduce-overhead	1.2–3x 提速，差别不大，但 reduce-overhead 更稳定
大 batch 或长 seq，已经 compute-bound	不要开 compile	收益 ≤2%，冷启动和重编译反而拖尾延迟
形状频繁变化的推理服务	开 compile 必须配 dynamic=True	否则每条不同形状请求付 1 秒编译税
训练或微调	值得尝试但要独立 benchmark	很多 fused optimizer / gradient 相关的 compile 收益和本文不同

两个反直觉点值得强调：

reduce-overhead 不是"更激进的 default"。它的底座是 CUDA Graphs，对形状、对输入是否被原位改写、对 host-device 拷贝都敏感。当每次调用的输入张量都是独立新建时，CUDA Graphs 需要额外 copy-in 环节。对 seq=1024、batch=16 这种一次本身就很大的形状，CUDA Graphs 的 launch 收益趋近于零，额外开销甚至让它慢于 default（我的数据里 (1,1024) 这一组 reduce 就输给了 default）。
- 在 batch 大到 GPU 利用率 >80% 时，任何 compile 都不会再给你免费午餐。你的瓶颈是 FLOPs，不是 kernel 派发。这种情况应该去看量化、FlashAttention 版本、SDPA backend，而不是 torch.compile。

6. 为什么 `reduce-overhead` 在小 seq 下压倒 `default`

回到这张补充表：

seq=  1 | eager 1.600 | compile 1.321 (1.21x) | reduce 0.466 (3.44x)
seq=  4 | eager 1.707 | compile 0.533 (3.20x) | reduce 0.529 (3.23x)
seq= 16 | eager 1.808 | compile 1.517 (1.19x) | reduce 0.555 (3.26x)

seq=16 这一行非常有意思：compile(default) 只有 1.19x，而 reduce-overhead 是 3.26x。同样的图、同样的 Triton 生成器，差距从哪里来？

关键在 CUDA Graphs。Eager 模式每次 forward 要从 Python 侧逐个启动大约 50–100 个 kernel（matmul、bias、add、layer_norm、scaled_dot_product_attention、gelu 等），每个 kernel 启动都有几微秒的 CPU/driver 开销。这些开销加起来在 seq=1 这种单 forward 不到 0.5ms 的负载里，占比可以轻松到 60%–70%。

compile(default) 做算子融合、生成更少的 Triton kernel，但每次调用仍然要走 PyTorch dispatcher + Triton launcher，launch 开销并没有消失。它拼不过 reduce-overhead 的原因就是——默认 compile 砍的是"kernel 数量"，不是"launch 开销"。

reduce-overhead 用 CUDA Graphs 把整张图提前录制，之后每次调用相当于在 GPU 上 replay 一次预编排好的指令流。这一下 launch 开销几乎归零，这就是 3x+ 提速的来源。

代价是：CUDA Graphs 对输入形状极敏感（形状一变就得重录制）、对输入张量是否是原地修改敏感、对某些带有 host-device 同步的算子（比如 torch.tensor(...).cuda()）直接失败回落。PyTorch 2.9 里这些回落是有警告的，但很多人容易忽略。

我的判断：做 LLM inference server 的 decode 阶段，如果你的 KV cache 管理得当、输入形状稳定，那么 reduce-overhead 是默认选项；其它场景老老实实用 default。

7. 如果你要在项目里用，这份清单先过一遍

确认瓶颈是 launch overhead 而不是 FLOPs。简单做法：不 compile，用 nvidia-smi dmon 或 nsys 看 GPU 利用率。如果稳态利用率已经 >80%，别指望 compile 帮你。
1. 算一笔冷启动账。如果 QPS 很低或者服务经常重启，每次冷启动交的 3–7 秒可能比稳态提速赚的还多。
1. **shape 会不会变？**变就加 dynamic=True，且第一次访问每个典型 shape 时的首调用时间要计入 SLO 预算。
1. fullgraph=True 最好加上。不加的话你可能不知道哪里 graph break 了——收益被悄悄吞掉。
1. reduce-overhead 要确认输入张量生命周期。CUDA Graphs replay 时会复用被捕获时的指针，如果你每次用新申请的 input tensor，实际有 copy-in 代价。生产里通常做法是把输入张量 pin 成固定 buffer 复用。
1. 不要同时开 compile 和 torch.no_grad()。改用 torch.inference_mode()，避免 autograd 元数据的额外开销被 compile 错误缓存。
1. 测自己的形状。本文的数字是 RTX 3090 + PyTorch 2.9 + bf16。H100、A10、V100 上的 launch 开销占比都不一样，compile 收益也不一样，上线前要自己跑一遍。

8. 我的收尾结论

torch.compile 是一个好东西，但在 2026 年把它当成"开了就快"的大招用，会让生产侧吃亏。它真正发光的场景只有一个：你的工作负载是 launch-overhead bound。这条件主要出现在 LLM 推理的 decode 阶段、batch 很小的在线推理、和小模型的高吞吐服务。

对大 batch / 大 seq / 已经 compute-bound 的负载，花在 torch.compile 上的调参时间，远不如去换更快的 FlashAttention 版本、用 FP8/INT4 量化、或者调 SDPBackend。

一句话判断：先拿 nvidia-smi dmon 看 GPU 利用率；利用率 <60%，试 compile；利用率 >80%，别折腾。

参考与延伸阅读

PyTorch 官方文档：torch.compile、torch.compiler.reset、dynamic shapes 相关小节。
- PyTorch 官方关于 CUDA Graphs 与 mode="reduce-overhead" 的说明。
- torch._inductor 源码中 compile_fx.py 和 Triton codegen 的实现。
- NVIDIA 文档：CUDA Graphs 的使用模型和形状约束。
- SDPA 相关 blog：torch.nn.functional.scaled_dot_product_attention 的 backend 选择。
  以上数据均为本人在 RTX 3090 上亲自测出，代码和日志可完整复现；任何和你手上的卡不一致的结论，请以自己跑出的数据为准。