PyTorch 的 JIT 编译器，通过图捕获、算子融合和后端优化将 eager 模式代码编译为高效 kernel。

为什么需要 torch.compile

PyTorch eager 模式每次 op 都经过 Python 调度和 GPU kernel launch，引入不必要的开销。torch.compile（基于 TorchDynamo + TorchInductor）自动将模型捕获为计算图，进行水平/垂直算子融合、内存布局优化和自动调优，生成融合 kernel，在不修改模型代码的前提下显著提升执行效率。

核心原理

• TorchDynamo：通过 frame evaluation hook 捕获 Python 字节码，将 eager 执行转换为 FX Graph。
• TorchInductor：后端编译器将 FX Graph lowering 为 Triton kernel（GPU）或 C++ kernel（CPU）。
• 算子融合：自动将连续的 pointwise op（如 SiLU + multiply）融合为单个 kernel，减少 HBM 访问。
• 动态形状：Inductor 支持动态形状符号化（symbolic shapes），适应 LLM 推理中变化的序列长度。

在源码中的实现

编译后端与图切分

• vllm/config/compilation.py — CompilationConfig 控制 compile 级别、splitting_ops（按注意力/unified_kv_cache_update 边界切图）、cudagraph_mode（NONE/PIECEWISE/FULL）、use_inductor_graph_partition。
• vllm/compilation/backends.py — VllmBackend 是 torch.compile 后端：split_graph() 把 FX 图切成子模块，为每个建 PiecewiseBackend。
• vllm/compilation/piecewise_backend.py — PiecewiseBackend 按动态形状桶（compile_range）预编译，运行时按 batch size 分派。切分点落在不能进 CUDA Graph 的注意力边界，使每段都可被 CUDAGraphWrapper 捕获（PIECEWISE 模式）；另一条路径 breakable_cudagraph.py（VLLM_USE_BREAKABLE_CUDAGRAPH）在整图 stream-capture 中于 attention/KV-cache op 处「断流」eager 执行。
• vllm/v1/worker/gpu/model_runner.py — 模型加载后根据配置调用 torch.compile() 编译。

自定义 Inductor 融合 pass

vLLM 通过 post_grad_custom_post_pass 钩子注入 PostGradPassManager（compilation/passes/pass_manager.py），在 post-grad、codegen 之前 运行一批基于 Inductor pattern matcher 的融合 pass（基类链 InductorPass → VllmInductorPass → VllmFusionPatternMatcherPass，compilation/passes/fusion/）：

pass	开关	融合内容
RMSNormQuantFusionPass	fuse_norm_quant	RMSNorm + FP8 量化
ActivationQuantFusionPass	fuse_act_quant	SiluMul + FP8/NVFP4 量化
AttnQuantFusionPass / MLAAttnQuantFusionPass	fuse_attn_quant	注意力直写 FP8/NVFP4 输出
QKNormRoPEFusionPass	enable_qk_norm_rope_fusion	Q/K RMSNorm + RoPE
RopeKVCacheFusionPass / QkNormRopeKvCacheFusionPass	fuse_rope_kvcache	RoPE(+QK-norm) + paged KV-cache 写入
MLARoPEKVCacheCatFusionPass	fuse_rope_kvcache_cat_mla	MLA RoPE + concat_and_cache_mla_rope_fused
AllReduceFusionPass	fuse_allreduce_rms	TP all-reduce + 残差 + RMSNorm（「延迟 all-reduce」）
AsyncTPPass / SequenceParallelismPass	enable_sp	GEMM↔reduce-scatter、SP 边界通信融合

每个 pass 的 uuid()（源码哈希）叠加 pass_config 与 compile_range 作 Inductor 缓存键——任一 pass 实现或开关变化都会触发重编译。

相关概念

相关概念

• cuda-graph — torch.compile 优化后的模型可进一步通过 CUDA Graph 消除 launch 开销
• flash-attention — torch.compile 可融合注意力前后的 pointwise op
• lora — LoRA 动态权重切换需考虑 compile 后的静态图限制
• tensor-parallelism — TP 的集合通信 op 需被 compile 正确处理