AOT 与 JIT：预制菜与量体裁衣的编译对比

1. 先用一句人话理解 JIT

JIT（Just-In-Time，即时编译）就是：不急着在软件发布前把所有代码都编好，而是等程序运行、拿到一部分真实信息后，再编出适合这次任务的代码。

代码最终需要变成处理器能够执行的指令。AOT 和 JIT 的核心区别首先是编译发生在什么时候，而不是谁一定更快：

AOT（提前编译）：软件发布前就编好，好比提前印好说明书，用户拿到就能照着执行。
JIT（即时编译）：运行到某段代码时才编译，好比遇到具体问题再现场翻译；既可能多花一次编译时间，也有机会利用已经知道的设备和输入条件做优化。

需要先纠正一个容易误会的地方：AOT 不等于通用低性能，JIT 也不等于自动最优。 AOT 完全可以事先编译许多精心调优的 kernel，运行时再挑一个；JIT 也可能只是在运行时编出一份普通代码。真正影响性能的，是编译器掌握了哪些信息、生成了什么代码，以及选择了什么执行策略。

想边读边试，可以直接运行文中配套的两个小实验：JIT 第一次为什么慢 和 GEMM 为什么要按 shape 选配置。只需 Python 就能跑通demo，不要求 NVIDIA GPU；第3节还提供了可选的 Triton/CUDA 版本。

2. 为什么宁愿“现场翻（慢）”，也不总用“印好的册子（快）”？

拿 GEMM（矩阵乘法）举例：

M、N、K 可以取很多不同的整数；再考虑 FP32 / FP16 / BF16 / INT8 等数据类型、矩阵是否转置、内存布局、设备架构、tile 大小和线程配置，实际要适配的场景就会非常多。

但这里有个重要区别：场景很多，不等于“变体空间是连续的”，更不等于必须为每个 shape 编一个 kernel。

2.1 完全可以先写一个通用 kernel

比如写一份通用 GEMM 逻辑：

template <typename T>void gemm(const T* A, const T* B, T* C, int M, int N, int K);

M/N/K 是运行时参数；T 可以在源码层复用同一份算法逻辑。发布时，针对 float / half / bfloat16 等类型分别实例化或分发实现，便能覆盖大量 shape 和 dtype。同一份通用源码不一定意味着同一份机器码：不同 dtype 往往对应不同指令、累加精度和硬件路径。

一个通用 kernel 只要正确处理边界、stride 和数据类型，就能完成工作。对于不常见 shape、调用次数少的场景，它甚至可能是更划算的选择。

问题在于：“什么都能算”和“每种情况都算得快”不是一回事。

• M/N/K 很小：kernel launch 和调度开销可能比计算本身还显著。
• M/N/K 很大：更在意 Tensor Core/矩阵指令利用率、tile 复用、流水线与寄存器压力。
• 矩阵很“瘦长”：最适合方阵的分块方式可能留下大量空闲线程或无效计算。
• dtype、对齐、layout 不同：可能需要不同的向量化、load/store 或矩阵指令路径。
• 如果让一套固定的 tile / warps / stages 应对所有情况，往往只能在一部分输入上接近最优。

2.2 JIT 的优势，是在“已知条件”上做针对性优化

假设服务启动后知道这次计算是：M=128、N=4096、K=4096、FP16、NVIDIA GPU。此时可以：

• 特化：固定某些 tile、对齐或静态配置，让编译器展开循环、消除满足条件时不必要的分支，并使用合适的指令。
• 选择算法：在不同 tile、线程组织、流水线深度之间做选择；这一步可以靠启发式，也可以靠 benchmark/autotune。
• 缓存结果：下次遇到满足同样编译前提的请求，直接复用已经编好的产物。
  这里也不意味着必须把 M/N/K 的每个值都写死。实际系统经常采用少量特化变体 + 动态 shape + 运行时分发，覆盖一个 shape 家族，而不是每个输入生成独立机器码。

可以把几种策略放在一起看：

所以 JIT 真正解决的是：什么时候值得拿运行时已知的信息，换一份更优的代码

2.3 有公开的 GEMM 性能数据吗？

有！Trition官方的 Matrix Multiplication 教程 给了两个很直观的观察点。

观察一：同样是 GEMM，换一种分块调度就能更快。 官方教程提到，在 A100 上把 block 的执行顺序调整为更利于 L2 复用的 grouped ordering，在某些配置下吞吐量可由约 220 提升到 245 TFLOPS（约 +11%）。这不是因为 GEMM 数学公式变了，而是数据访问和 tile 调度更贴合硬件。它说明执行策略的选择确实有价值，但不能单凭这组数据证明 JIT 比 AOT 快——AOT 同样可以提前编译这套调度。

观察二：Triton 也不保证处处赢过 cuBLAS。 官方教程附带的一组 FP16 方阵 benchmark 示例数据如下（单位 TFLOPS）：

这些是官方教程展示的样例输出，并不是真机复测；性能受 GPU、版本、时钟、精度与 benchmark 方法影响。更重要的是：cuBLAS 本身就是高度优化、拥有大量实现和启发式选择的库，不能把它简单等同于“单一通用 kernel”，也不能把此表当作“JIT 对 AOT”的受控对照。

真正要验证“现场翻得更合身”，应该在同一 GPU、同一 dtype、同一组 shape、相同计时条件下，测 固定通用 tile、离线准备的多变体分发、JIT 特化/调优，再单独统计编译耗时和缓存命中率。

2.4 只看结论还不够，我们自己跑一次 GEMM

上面的 Triton 和 cuBLAS 数据来自官方公开教程。为了让「不同 shape 可能需要不同执行策略」这句话可以自己动手验证，再做一个不需要 GPU 的小实验。

对应源码：engineering-labs / gemm-specialization；查看本次保存的原始 JSON。

# 只需要 Python 3.11+；第一次使用时安装依赖
git clone https://github.com/YuleZhang/engineering-labs.git
cd engineering-labs
python -m pip install -r requirements.txt
 
# 对比不同 M×N×K 下的 tile 配置
python run.py gemm
 
# 单独检查数值正确性
python run.py gemm --verify

实验用 Numba 编译同一份 CPU GEMM 算法的三种 tile 实现：8/16/32。输入由固定种子生成，结果先与 NumPy A @ B 比较，再统计每个配置的多次稳态运行时间。下面是一次实际记录，单位为毫秒；粗体是这一轮测得最快的配置。

测试环境：Linux 6.18.57 / x86_64 / glibc 2.41；Python 3.13.5、NumPy 2.3.5、Numba 0.65.1；每个 tile 的 warm 执行取 5 次中位数。
拿最后一行看：在这一轮测试里，Tile 16 比固定使用 Tile 8 的耗时约低 37%；但中间那行又是 Tile 8 更快。由此能验证的是：一份固定 tile 的通用实现可以处理多个 shape，却不保证在这些 shape 上都是最快的。

注意这个实验的边界：三种 tile 全都由 Numba JIT 编译，比较的是调度/变体选择，而不是「AOT 对 JIT」。真正要比较 AOT 与 JIT，必须进一步增加等价的离线编译实现，并统一执行条件。

这就把第 2 节的论点落到了可观察的证据上：JIT 可以成为生成、特化或管理不同变体的手段，但性能收益来自选对实现，不来自「编译发生得晚」本身。

3. 来从一个 Triton kernel 角度看 JIT

Triton 很适合观察这条链路，因为一个短小的 Python kernel 就能触发编译、生成 GPU 机器码并执行。

import torch
import triton
import triton.language as tl
 
@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask, other=0)
    y = tl.load(y_ptr + offs, mask=mask, other=0)
    tl.store(out_ptr + offs, x + y, mask=mask)
 
n = 100_000
x = torch.randn(n, device="cuda")
y = torch.randn(n, device="cuda")
out = torch.empty_like(x)
 
# 第一次运行：相应变体没有命中缓存时，需要编译。
add_kernel[(triton.cdiv(n, 256),)](x, y, out, n, BLOCK=256)
 
# 再次运行相同变体：通常可以复用编译结果。
add_kernel[(triton.cdiv(n, 256),)](x, y, out, n, BLOCK=256)
torch.cuda.synchronize()
assert torch.allclose(out, x + y)

这里最值得看的是 BLOCK: tl.constexpr：它是编译期常量。比如把 BLOCK=256 改成 BLOCK=128，对应的向量范围、程序布局等就可能变化，需要另一份编译变体。

但并不是所有参数都会按具体数值特化。例子里的 n 可以继续作为运行时长度，负责边界 mask；Triton 还会根据参数类型、对齐和显式特化策略决定哪些信息进入编译前提。triton.jit API 甚至提供 do_not_specialize，用于限制不必要的值特化。

3.1 第一次慢是为什么？

对 NVIDIA 后端，典型 Triton 编译路径可以简化为：

@triton.jit 装饰的 Python 函数
     ↓  解析函数、结合 constexpr / 类型等信息特化
Triton IR (TTIR)
     ↓  优化、layout / GPU 映射等 lowering
Triton GPU IR (TTGIR)
     ↓  lowering
LLVM IR
     ↓  生成目标相关代码
PTX
     ↓  ptxas 汇编
cubin（GPU 机器码）
     ↓  加载、launch
GPU 执行

这条路径是典型 NVIDIA 路径，不是所有 GPU 后端都使用 PTX/cubin。AMD 等后端会走其他目标相关 IR 和机器码路径，且具体阶段会随 Triton 版本演进。

第一次调用可能慢在多个环节：编译本身、模块加载、CUDA 上下文初始化、首次内存分配、autotune。所以“第一次慢”不可以全部归因于 JIT，做实验时要把这些成本分开。

还要区分三个术语：

• Codegen（代码生成）：可以从 IR 生成另一种 IR、目标汇编，甚至直接生成机器码；不只等于“拼源码字符串”。
• Lowering / Pass：对 IR 做变换和优化，是编译器内部的处理环节，本身不等于 JIT。
• JIT：描述编译在运行期发生。它可以调用上述 codegen 与 pass，也可以在 CPU 上直接发射机器指令。

3.2 编译产物靠什么标识并缓存？

一个编译产物需要能回答：“这次请求，是否还能安全复用上次的机器码？”

以 Triton 为例，影响编译变体的因素一般包括：kernel 源码/定义、参数类型及参与特化的属性、constexpr、num_warps / num_stages 等编译配置、目标后端与设备能力，以及相关编译环境版本。

不要简单理解成“M/N/K 和所有运行时变量都必须进 key”。只有生成代码确实依赖的条件，才需要成为特化依据或有效性检查的一部分；否则新 shape 也可能安全复用同一个动态 kernel。

这里存在三种不同的“键”或判断：

key 过细，会反复编译、缓存膨胀；key 过粗或 Guard 不完整，可能执行不适用的代码，带来错误结果甚至内存安全问题。

3.3 特化和调优是一回事吗？

不是。继续拿 GEMM 来看：

# 示意：@autotune 从多个合法配置中选一个。
@triton.autotune(configs=[...], key=["M", "N", "K"])
@triton.jit
def matmul_kernel(...):
    ...

特化（specialization）：根据已知信息，决定这份代码长什么样，比如 tile 大小、常量折叠、循环展开。
调优（autotuning）：在多个正确候选里进行测量或选择，决定这一类输入用哪份代码。
JIT 不必依赖 autotune，autotune 也不要求在线 JIT：提前编好的几个 AOT kernel 一样可以在运行时测速、选择并缓存结果。

3.4 没有 NV GPU，也能看到「第一次慢」

前面的 Triton 是 GPU JIT 的真实实现，但要有 NV GPU 才能运行。因此我们另外提供了一个普通 CPU 也能跑的例子：使用 Numba @njit 编译数值循环。

如果已经按第 2.4 节克隆过仓库，直接执行：

python run.py
 
# 观察首次调用、重复调用、换 dtype
python run.py --verify
 
# 与纯 Python 结果核对
python run.py --size 100000 --repeats 21
 
# 自己调整输入规模

下面是实验仓库保存的一次真实运行，输入为 50,000 个 float64，warm 执行统计 11 次中位数：

测试环境：Linux 6.18.57 / x86_64 / glibc 2.41；Python 3.13.5、NumPy 2.3.5、Numba 0.65.1。计时包含 Python 调用和 Numba dispatcher；启用了 cache=True。
再看编译签名的数量，会比只看耗时更直观：

这里也能验证第 2 节提到的一个关键点：shape 变化不等于一定重编译。 Numba 默认数组类型签名主要关心元素类型、维度和布局等信息，而不是将每个数组长度写成一个新变体。相反，输入从 float64 变成 float32，会产生另一份类型特化。

我们能从这组数据得出三件事：首次调用有明显的额外成本；重复执行同一种已编译实现可能明显更快；不同输入类型可能需要新的特化。 但它比较的是「Numba 编译的循环」与「Python 解释执行循环」，不能证明 JIT 比优化过的 AOT C++ 更快。GPU 上的实际编译链路仍以第 3.1 节的 Triton 为准；对应 CUDA 代码是可选路径，目前尚没有经过 GPU 实机验证的结果。

4. JIT 到底可以发生在哪些地方？

JIT 不只用于一个 kernel。按优化对象看，可以从计算图一路下沉到机器指令；但下面这些是不同的编译/优化层次，不是每一项都属于 JIT。

4.1 再看一个不同的例子：VkFFT

VkFFT 支持多种 GPU API，会在运行时根据 FFT 长度、维度、精度、设备约束等做计划并编译 kernel。不同长度可能走不同 radix 分解；某些质数长度会使用 Rader 或 Bluestein 等算法。

这仍然不是“连续变体空间”：FFT 长度是离散整数，实际还受设备内存和实现上限约束。真正难的是有效组合数量很大，而且不同组合的最佳算法和内存调度可能明显不同，不适合把所有排列组合都提前编成二进制。

VkFFT 能清晰展示另一种 JIT 形态：运行时做算法规划 → 生成针对该计划的代码 → 交给后端编译器 → 执行。它和 Triton 的区别主要在表达形式与规划逻辑，而不是“一个是离散、一个是连续”。

4.2 有三个概念，最好别混用

• JIT ≠ Fusion（融合）：JIT 说的是编译时机；融合说的是如何减少算子边界和中间数据搬运。AOT 也能融合。
• JIT ≠ Autotune（自动调优）：JIT 决定何时编译；autotune 决定如何选择。两者可以配合，也可以独立存在。
• JIT ≠ GPU 专属能力：CPU 上也有 JVM、oneDNN、Xbyak/AsmJit 等 JIT。移动设备可能受代码签名、W^X 和动态执行权限限制；在手机或封闭设备上，JIT 未必是实际可部署的最优方案。

5. 图级 JIT 是怎么工作的？看看 PyTorch

如果一个网络要跑 matmul → add bias → activation，只加速其中一个 GEMM 还不够：中间张量的读写和多次 kernel launch 也会耗时。

torch.compile 就是在更大的计算图上寻找这些机会。典型路径可以概括为：

Python / PyTorch 模型
     ↓ TorchDynamo：捕获可编译的 FX 图，建立 Guard
FX Graph
     ↓ AOTAutograd / Dispatcher（训练时还可处理反向图）
可优化的前向/反向图
     ↓ TorchInductor：融合、调度、代码生成
Triton kernel / C++ 等后端代码
     ↓ 运行时编译、加载、执行与缓存
计算结果

• 图级优化：决定哪些操作融合、哪些中间结果不必写回显存。
• kernel 级优化：将融合后的工作交给 Triton、C++ 或其他合适的后端实现。
• Guard：判断新输入是否仍满足已捕获图的假设，不满足时可能重新编译或走回退路径。
  例如 batch size 从 8 变到 16，并不天然意味着必须重新生成所有 kernel。PyTorch 支持 dynamic shapes，允许某些维度在一个编译产物里保持动态；代价是放弃对应维度的一部分静态优化机会。

它与第 3 节的 Triton cache key 是相关但不同层级的问题：前者保障图级假设，后者保障具体 kernel 编译变体的复用。

6. 图级和 kernel 级一定都要上吗？

不一定。要看瓶颈。

• 如果瓶颈是大量碎 elementwise op、频繁 launch 和中间 tensor 访存，图级融合往往很有价值。
• 如果瓶颈是一颗 GEMM 或 attention kernel 的 tile 组织、矩阵指令利用率，kernel 级调优更直接。
• 如果算子已经由高度优化的 cuBLAS/cuDNN/专用 attention kernel 实现，强行全部改成 JIT 可能增加复杂度却没有收益。
• 如果是一次性离线任务、短生命周期进程或受限端侧部署，编译成本可能根本摊不回来。
  两级可以协同，也可以只做一层。先找到瓶颈，再选择编译粒度；不要为了上 JIT 而上 JIT。

7. 常见框架擅长什么，该怎么选？

oneDNN 展示了 CPU 侧运行时代码生成的价值；CUTLASS 则展示了即使不用在线 JIT，也可以利用模板、多变体与调优获得高性能。把两者放在一起，反而更容易理解“JIT 只是手段，不是性能保证”。

8. AI 编译器为什么会用到 JIT？

AI 编译器要做的事情，往往不只是把一段代码翻成机器指令，而是把模型运算映射到特定硬件、特定输入和特定执行计划。

以一个 LayerNorm → Linear → GELU 子图为例：

模型定义（PyTorch / JAX / 其他前端）
           ↓
      捕获/导出计算图
           ↓
   图优化（融合、常量传播、布局转换）
           ↓
     算子选择 / 调度 / 特化
           ↓
       生成设备代码
           ↓
         运行

这里既可以完全离线 AOT，也可以在运行期引入 JIT。后者通常适合下面几类情况：

不过，图优化发生在运行时，不代表最终机器码一定由 JIT 生成：编译器也可能只在运行时选择一颗已经 AOT 编好的 kernel。判断一个系统是不是 JIT，要看它是否真的在运行期间生成/编译了新的可执行代码，而不能仅凭“运行时做了优化”就下结论。

9. 大模型开发中，JIT 到底扮演什么角色？

大模型训练和推理非常依赖 GEMM、attention、归一化、激活与量化。JIT 在这里的价值，主要不是改变 Transformer 公式，而是减少执行这些公式的开销。

9.1 训练：把反复执行的子图编得更高效

典型场景是 PyTorch 训练用 torch.compile 优化前向/反向图，融合一些短小的 elementwise 操作、改进内存访问和 kernel 调度。长 Transformer 中相似的 block 反复出现，也可以考虑区域编译、复用编译结果，降低冷启动成本。

PyTorch 的 regional compilation 教程 就是针对“模型大、首次编译慢”这个真实问题。

9.2 推理：Prefill 和 Decode，需要的优化经常不同

要注意：PagedAttention、continuous batching、KV cache、CUDA Graph 都不等于 JIT。 它们分别解决存储/调度/执行复用的问题，可以和 JIT 一起使用，也可以与 AOT kernel 配合。

9.3 一个真实例子：vLLM 为什么会接入 torch.compile？

vLLM 的编译配置文档 描述了从直接 eager 执行、标准 torch.compile，到 vLLM 自定义编译后端的不同模式。自定义后端涉及图分段、缓存、shape specialization 和 Inductor pass；其 CUDA Graph 设计 则负责捕获与重放执行。

把它们拆开看，会更容易理解整套推理系统：

用户请求（prompt 长度、batch、模型配置）
               ↓
      调度器组织本轮 token 工作
               ↓
     选择已编译图 / kernel 变体
               ↓
   必要时编译，或直接复用缓存产物
               ↓
  CUDA Graph replay 或常规 kernel launch
               ↓
       Attention / GEMM 等 GPU 执行

关键不是“每生成一个 token 都重新 JIT”，而是尽量让绝大多数请求落在已经编好、预热好、能复用的路径上。真实 serving 系统会努力避免 token 长度或 batch 每变化一次，就引发昂贵的整图重编译。

9.4 什么时候反而应该优先选 AOT？

部署到移动设备、NPU 或严格控制启动延迟的在线服务时，经常更适合提前编译 + 少量变体 + 运行时分发。例如，固定模型/设备、已知 batch 范围的场景，把编译成本挪到构建和发布阶段，可能比把完整编译器打包进终端更合适。

PyTorch 也提供 torch.compile 的 AOT 编译与产物序列化路径。这进一步说明：图级优化、Triton 代码生成和 autotune，并不天生要求发生在用户第一次请求时。

10. JIT 什么时候真的划算？算一笔账

假设某个 AOT 通用实现每次执行需要 T_generic，特化实现每次需要 T_specialized，首次编译需要 T_compile，之后复用 H 次。

简化估算：

只有当 T_generic > T_specialized 且 H > T_compile / (T_generic - T_specialized) 时，JIT 的累计时间才可能更划算。

例如：编译花 2 秒、每次快 0.2 毫秒，要复用超过 10,000 次才能回本（尚未计入 autotune、额外内存、缓存 I/O 和服务冷启动）。对高频模型服务可能值得；对偶尔触发一次的算子，就未必。

这个公式是理解成本的简化模型。在真实部署中，还要考虑 AOT 多变体的性能、缓存跨进程持久化、设备升级、代码体积、编译器版本、并发编译，以及 P50/P99 延迟。

10.1 用刚才的 CPU JIT 数据算一次「回本」

第 3.4 节实际测到：纯 Python 一次约 4.982 ms，Numba 首次调用约 474.240 ms，之后每次 warm 调用约 0.063 ms。如果仅作为粗略估算，假定调用次数为 H、数据和机器状态不变，那么：

两者相等时约为 96 次调用，所以按这次测量和模型，累计到第 97 次左右才可能开始弥补首次调用成本。即便单次 warm 执行快很多，调用次数太少时仍未必划算。

这里有两个边界不能省略：第一，474.240 ms 是本进程首调用的总耗时，并非纯 JIT 编译耗时，且受到磁盘缓存状态影响；第二，这个「回本」比较的是 Numba JIT 与纯 Python 循环，不是 JIT 与 AOT。若换成已优化的 C++、AOT 多变体或不同的输入规模，分母、结论和回本次数都可能完全不同。

你可以自己运行 python run.py --json my-jit.json 取得本机数据，再代入上面的式子；这比直接套用本文的 97 次更有意义。

11. 真要在工程里使用 JIT，先检查什么？

• 基准要分开测：分别记录冷启动/首次编译、缓存命中后的 steady-state、autotune 耗时和端到端延迟，避免把初始化混入 JIT 结论。
• 先识别特化维度：shape、stride、dtype、layout、设备架构、静态配置；哪些必须固定，哪些应保留动态？
• Guard 要守住正确性：每个产物都要有清晰的适用条件，失效时能够重新编译或安全回退。
• 缓存要可解释：知道 key 的组成、编译器/驱动升级后的失效条件、存储位置、容量上限和淘汰策略。分层缓存可以存在，但边界要清楚。
• 预热要针对真实流量：从热点 shape 开始，不要漫无目的地穷举所有变体。
• 要有非 JIT 回退路径：图抓取失败、编译失败、硬件不支持或冷启动预算不足时，仍能保证功能正确。
• 先保证数值和内存安全：验证不同 dtype 的累加精度、shape 边界、别名、stride、并发与设备行为，再看 GFLOPS。
• 最后才判断收益：对比真实瓶颈、编译摊销和全链路吞吐，别只比较单颗 kernel 的峰值。

12. 常见疑问，快速对照

把文章和代码连起来

文章里的三个问题，现在都能在源码中找到对应观察点：

两个 CPU Demo 均可用 --verify 运行数值检查；仓库的 GitHub Actions 还在 Linux 与 Windows 的 Python 环境中运行自动化 smoke test。CI 通过说明示例可执行与正确性检查通过，不表示不同机器的性能数值一致。

只想理解原理，可以直接运行 CPU Demo；想进一步追踪 GPU JIT，则阅读第 3 节 Triton 代码和可选 CUDA 示例。后者需要相应 GPU 环境，不能拿 CPU 结果代替 GPU 实测。

继续阅读

Engineering Labs：JIT 首次调用与 GEMM 特化源码、数据和运行命令
Triton：Matrix Multiplication 教程与公开 benchmark
Triton：@triton.jit API
PyTorch：Dynamic Shapes
PyTorch：Dealing with Recompilations
PyTorch：编译缓存机制
vLLM：Compilation 配置
vLLM：CUDA Graph 设计
VkFFT：运行时 FFT kernel 编译