第 5 章 AI InfraGPU分布式训练

第 5 章 编译器与框架

第5章 编译器与框架

本章梳理 ML 编译器的分类、IR 生态与标准编译流程，剖析 PyTorch 2.0 编译栈、算子融合与自动调优技术，介绍 JAX/XLA 与 TensorRT 两个主流编译器生态，最后以 torch.compile 实战收尾。硬件指令集与手写 Kernel 优化不在本章 展开。

5.1 ML 编译器全景

ML（Machine Learning）编译器是连接高级深度学习框架与底层硬件加速器的关键技术桥梁。随着硬件多元化（NVIDIA GPU、AMD GPU、Google TPU、Intel Gaudi、Apple Neural Engine）和模型架构复杂化（Transformer、MoE （Mixture of Experts）、Mamba），编译器在AI Infra中的地位不断上升。以下从编译器分类学出发，绘制ML编译器生态 的全景地图。

5.1.1 ML 编译器分类学

按编译粒度和时机，ML编译器可分为三个层次，如图5-1所示。 JIT 编译器 图级编译器 Graph-Level PyTorch Dynamo XLA JAX jit + XLA TVM Relay TVM JIT ONNX Runtime ONNX Runtime EP OpenVINO 优化后子图 算子级编译器 Operator-Level TVM Tensorize Triton Halide MLIR Dialects 图5-1 ML编译器分类全景 图级编译器在整张计算图上做优化，算子级编译器针对单个算子生成设备代码，JIT 编译器则面向运行时捕获的执行轨 迹，三者可在同一框架内级联。

5.1.2 IR 中间表示生态

中间表示（IR）是编译器优化能力的核心载体。ML编译器领域已形成多元化的IR生态，主流 IR 格式对比如表5-1所示。 表5-1 ML编译器IR格式对比 IR格式 来源 层次 特点 HLO (High-Level Optimizer) XLA 图级 静态形状，XLA专属，JAX/TF后端 StableHLO OpenXLA社区 图级 HLO的标准化版本，可移植的ML计算操作集 Relay IR Apache TVM 图级 支持动态形状、控制流、自动微分 Torch FX Graph PyTorch 图级 Pythonic，记录张量操作的计算图 MLIR LLVM社区 多层 Dialect生态，连接高层ML和底层硬件 Triton-IR OpenAI 算子级 类C语言，面向GPU的循环级编程 Linalg Dialect MLIR 算子级 线性代数操作的通用IR LLVM IR LLVM 底层 通用编译器IR，多后端 各框架的高层图经 MLIR 统一降级，逐层下探至 GPU Dialect、LLVM IR 或 PTX 等底层表示。 PyTorch TensorFlow JAX │ │ │ ▼ ▼ ▼ ┌─────────┐ ┌──────────┐ ┌──────────┐ │Torch FX │ │ TF Graph│ │ JAXPR │ │ Graph │ │ │ │ (jaxpr) │ └────┬────┘ └────┬─────┘ └────┬─────┘ │ │ │ ▼ ▼ ▼ ┌─────────────────────────────────────────────────┐ │ MLIR Layer │ │ ┌──────────┐ ┌───────────┐ ┌─────────────┐ │ │ │ Torch │ │ StableHLO │ │ Linalg │ │ │ │ Dialect │ │ Dialect │ │ Dialect │ │ │ └────┬─────┘ └─────┬─────┘ └──────┬──────┘ │ │ └──────────────┼──────────────┘ │ │ ▼ │ │ ┌─────────────┐ │ │ │ GPU Dialect│ │ │ │ (NVGPU/AMD) │ │ │ └──────┬──────┘ │ └──────────────────────┼──────────────────────────┘ ▼ ┌──────────────────┐ │ LLVM / PTX │ │ CUDA / ROCm │ └──────────────────┘

5.1.3 编译流程标准化

尽管各编译器实现不同，ML编译流程可抽象为四个标准阶段。

  1. 图捕获 将Python框架代码转换为中间表示。这是ML编译器与通用编译器的最大区别——ML框架的执行模式是“Python描述 → 运行时执行”，编译器需要“拦截”执行过程。以 PyTorch 为例，Dynamo 逐条执行字节码并记录张量操作，构建 FX 计 算图。
  2. IR降级 从高层次操作（如 torch.nn.functional.softmax ）降级为低层次操作（如矩阵乘法、元素运算）：
 # High level: F.softmax
 # ▼ Lowering
 # Low level:
 # t0 = matmul(x, w)    ◀ MatMul
 # t1 = add(t0, b)      ◀ Add Bias
 # t2 = sub(t1, max)    ◀ Subtract Max (numerical stability)
 # t3 = exp(t2)         ◀ Exp
 # t4 = sum(t3, dim=-1, keepdim=True) ◀ Reduce Sum
 # t5 = div(t3, t4)     ◀ Div
  1. 优化Pass 对降级后的IR应用一系列优化变换： Graph IR │ ├── Constant Folding: 3.14 * 2 ▶ 6.28 ├── Dead Code Elimination: remove unused subgraphs ├── Operator Fusion: MatMul + Add + Softmax ▶ fused operator ├── Layout Optimization: NHWC ▶ NCHW (based on hardware preference) ├── Memory Planning: Minimize peak memory └── Algebraic Simplification: x*1 ▶ x, x+0 ▶ x │ ▼ Optimized IR
  2. 代码生成 将优化后的IR转化为设备可执行的代码。以 PyTorch 为例，Inductor 将优化后的图降级为循环级 IR，再生成 Triton 或 C++ 内核，也可选择 OpenXLA 等后端。

5.1.4 编译器选型决策矩阵

不同场景下的编译器选型决策如表5-2所示。 表5-2 编译器选型决策矩阵 场景 推荐编译器 原因 PyTorch训练 torch.compile (Inductor) 原生集成，无需代码修改 JAX训练/研究 JAX jit (XLA) 函数式编程，Spmd并行 PyTorch → 多硬件部署 Apache TVM / Torch-MLIR 多后端支持 NVIDIA推理优化 TensorRT 推理专用，极致性能 自定义算子开发 Triton (OpenAI) Pythonic GPU编程 模型可移植性 ONNX Runtime 跨框架标准格式 异构硬件（NPU/DSP） MLIR-based (IREE) 灵活Dialect扩展

5.2 PyTorch 2.0 编译栈

PyTorch 2.0的核心创新是 torch.compile ——一个JIT编译器，它将PyTorch从“Eager Execution优先”的框架转变为 一个“编译优先但Eager兼容”的框架。其编译栈由三个关键组件构成：Dynamo（图捕获）、AOTAutograd（自动微分 变换）和Inductor（代码生成）。三者协同工作，在不修改用户代码的情况下实现约1.3-2倍的训练加速（PyTorch 2.0 官 方基准：163 个模型平均约 1.3x，典型 LLM 约 1.3-1.5x，最高可达 2x；本章实测 Llama-7B 在 max-autotune 下达 1.88x）。

5.2.1 torch.compile 编译流程

torch.compile 的编译流程分为 Dynamo 图捕获、AOTAutograd 自动微分与 Inductor 代码生成三个阶段，如下所示。 User Code (eager PyTorch):

   def model(x):
       return F.softmax(x @ w + b, dim=-1)
   compiled_model = torch.compile(model)
   output = compiled_model(input_tensor)

│ ▼ ┌───────────────────────────────────────────────────────────────────────────────┐ │ Phase 1: Dynamo - Frame Evaluation & Graph Capture │ │ ┌────────────────────────────────────────────────────────────────────────────┐ │ │ │ 1. Intercept Python Frame Evaluation │ │ │ │ 2. Execute bytecode instruction by instruction, record PyTorch ops │ │ │ │ 3. Build FX Graph (records all tensor operations) │ │ │ │ 4. Guard mechanism: check if input shape/dtype changes │ │ │ └────────────────────────────────────────────────────────────────────────────┘ │ │ Output: FX Graph (torch.fx) │ ├───────────────────────────────────────────────────────────────────────────────┤ │ Phase 2: AOTAutograd - Ahead-of-Time Differentiation │ │ ┌────────────────────────────────────────────────────────────────────────────┐ │ │ │ 1. Decompose forward graph into basic operations │ │ │ │ 2. Generate backward graph (automatic differentiation) │ │ │ │ 3. Jointly optimize forward/backward graphs (e.g., recomputation strategy) │ │ │ └────────────────────────────────────────────────────────────────────────────┘ │ │ Output: Joint forward+backward FX Graph │ ├───────────────────────────────────────────────────────────────────────────────┤ │ Phase 3: Inductor - Loop-level Optimization & Codegen │ │ ┌────────────────────────────────────────────────────────────────────────────┐ │ │ │ 1. Lower FX Graph to Loop-level IR │ │ │ │ 2. Operator decomposition │ │ │ │ 3. Pattern Matching & Fusion │ │ │ │ 4. Tiling & Scheduling │ │ │ │ 5. Triton / C++ / OpenXLA Codegen │ │ │ └────────────────────────────────────────────────────────────────────────────┘ │ │ Output: Compiled Triton kernel / C++ kernel / XLA graph │ └───────────────────────────────────────────────────────────────────────────────┘

5.2.2 Dynamo 图捕获与 Guard 机制

Dynamo是PyTorch 2.0编译栈的入口，负责将Python字节码级别的eager执行“翻译”为FX计算图。其核心机制包括 Frame Evaluation Hook、Guard 与 Graph Break。

  1. Frame Evaluation Hook Dynamo通过CPython的PEP 523（Frame Evaluation API）挂钩到Python解释器：
 # Dynamo core mechanism
 import torch
 import dis
 from torch._dynamo.eval_frame import set_eval_frame
 # 1. Register frame evaluation hook (PEP 523)
 old_callback = set_eval_frame(dynamo_frame_hook)
 def dynamo_frame_hook(frame):
     # 2. Analyze bytecode
     bytecode = frame.f_code.co_code
     instructions = dis.Bytecode(frame.f_code)
     # 3. Build FX Graph
     fx_graph = torch.fx.Graph()
     tracer = DynamoTracer(fx_graph)
     for instr in instructions:
         if is_torch_op(instr):
             tracer.record_op(instr)
         else:
             if is_guard_condition(instr): # e.g., isinstance(x, Tensor)
                 create_guard(instr)
             tracer.fallback_to_eager(instr)
     # 4. Compile FX Graph
     compiled_fn = inductor_compile(fx_graph)
     # 5. Cache + return compiled result
     cache_key = (hash(bytecode), hash(tuple(guards)))

compiled_cache[cache_key] = compiled_fn return compiled_fn torch._dynamo.eval_frame.set_eval_frame(dynamo_frame_hook) # Replace frame evaluation callback 2) Guard机制 Guard是Dynamo确保编译正确性的关键设计。每次执行编译函数时，Dynamo检查Guard条件是否仍然成立：

 # Example guard generated by Dynamo
 def compiled_fn_guard_wrapper(x, w, b):
     # Guard 1: Shape unchanged

assert x.shape == (32, 512, 4096), "Shape changed" # Guard 2: dtype unchanged assert x.dtype == torch.float16, "Dtype changed" # Guard 3: stride unchanged assert x.stride() == (2097152, 4096, 1), "Stride changed" # Guard 4: device unchanged assert x.device == torch.device("cuda:0"), "Device changed" # Guard 5: requires_grad unchanged assert x.requires_grad == True, "Grad requirement changed" # All guards pass ▶ execute compiled code return compiled_code(x, w, b) 3) Graph Break 当Dynamo遇到无法编译的Python结构时，会发生图断裂（Graph Break），编译中止并回退到Eager模式。常见触发原 因包括数据依赖的控制流（如 if y.sum() > 0 ）、非 PyTorch 库调用（如 NumPy 操作）、以及 .item() 产生标量 等。诊断图断裂可设置 TORCH_LOGS=graph_breaks 环境变量定位断裂点。

5.2.3 Inductor Loop-Level 编译

Inductor是PyTorch 2.0的默认代码生成后端。它接受FX Graph，输出高性能的Triton或C++ kernel代码。

  1. 编译流程 Inductor 的编译流程包括算子分解、模式匹配、调度与代码生成：
 # Inductor internal processing (pseudocode)
 class InductorCompiler:
     def compile(self, fx_graph: torch.fx.Graph) -> CompiledModule:
         # 1. Operator decomposition
         decomposed = self.decompose_ops(fx_graph)
         # nn.LayerNorm ▶ reduce_mean + sub + pow + reduce_mean + add + ...
         # 2. Pattern matching
         patterns = self.match_patterns(decomposed)
         # MatMul + Add + ReLU ▶ Fused MatMul-Bias-ReLU
         # 3. Scheduler
         scheduler = InductorScheduler(decomposed)
         schedule = scheduler.create_schedule()
         # BFS/DFS traversal of DAG, maximize fusion opportunities
         # 4. Code generation
         for node_group in schedule.fused_groups:
             if self.use_triton(node_group):
                 kernel = TritonCodegen.generate(node_group)
             else:
                 kernel = CppCodegen.generate(node_group)
         return CompiledModule(kernels=schedule.kernels)
  1. torch.compile模式
 # Three compilation modes
 # 1. default: balance compilation time
 model = torch.compile(model, mode="default")
 # 2. reduce-overhead: use CUDA graphs to reduce per-iteration kernel launch overhead
 # Suitable for: small-batch inference
 model = torch.compile(model, mode="reduce-overhead")
 # 3. max-autotune: most aggressive optimization, searches for the best kernel configs via autotuning
 # Suitable for: production inference
 model = torch.compile(model, mode="max-autotune")
  1. 后端选项

5.2.4 编译缓存与预热

 # Using different compilation backends
 # 1. Inductor (default)
 model = torch.compile(model, backend="inductor")
 # 2. OpenXLA
 import torch_xla
 model = torch.compile(model, backend="openxla")
 # 3. TensorRT
 import torch_tensorrt
 model = torch.compile(model, backend="tensorrt")
 # 4. ONNX Runtime
 model = torch.compile(model, backend="onnxrt")
 # 5. TVM
 model = torch.compile(model, backend="tvm")

torch.compile 的编译产物默认持久化到磁盘缓存，首次调用触发编译，相同形状的后续调用直接命中缓存，避免重复编 译。缓存的配置方式如下：

5.2.5 动态形状

 # torch.compile cache mechanism
 import torch._inductor.config as inductor_config
 # Set cache directory
 inductor_config.cache_dir = "/tmp/torch_compile_cache"
 # Cold start vs warm
 # First call: compile
 output = compiled_model(input_data)   # First time: ~30s compile
 # Subsequent calls: cache hit
 output = compiled_model(input_data)   # Second time: ~0.1s
 # Cache persistence
 # Inductor persists compiled kernels automatically
 # Set the directory via env:
 #   export TORCHINDUCTOR_CACHE_DIR=/path/to/inductor_cache
 # or programmatically:
 import torch._inductor.config as inductor_config
 inductor_config.cache_dir = "/path/to/inductor_cache"

torch.compile 默认工作在静态形状模式：每当输入 tensor 的形状变化，Dynamo 会触发重编译（recompile），代价极 高。对于序列长度变化的 LLM 推理场景，这会导致每种序列长度都编译一次，积累大量编译开销。 PyTorch 2.1 引入了 dynamic=True 选项，启用符号化形状（Symbolic Shapes）机制：

 # Enable dynamic shapes: Dynamo
 model = torch.compile(model, dynamic=True)
 # Or mark only specific dims
 import torch._dynamo
 torch._dynamo.mark_dynamic(input_tensor, 1)   # mark dim=1 (seq_len) as dynamic

启用后，Dynamo 使用符号整数（如 s0、s1）替代具体数值，Guard 变为范围条件（如 1 <= s0 <= 8192），Inductor 生 成带运行时 shape 参数的 Triton kernel，避免反复重编译。性能权衡：动态形状 kernel 含额外的运行时 shape 计算，吞 吐通常比静态模式低 5-10%。推荐策略：训练阶段用固定 shape（静态编译获最优性能），推理服务用 dynamic=True 应对多变序列长度。

5.2.6 torch.export 模型导出

torch.export （PyTorch 2.2+ 稳定）是生产部署的推荐路径，它导出一个完整的计算图（ExportedProgram） ，包含静 态或动态形状约束，可跨运行时部署：

 import torch
 # Export model with dynamic dimension
 batch_dim = torch.export.Dim("batch", min=1, max=64)
 seq_dim = torch.export.Dim("seq", min=1, max=4096)
 exported = torch.export.export(

model, args=(example_input,), dynamic_shapes={"x": {0: batch_dim, 1: seq_dim}}, ) torch.export.save(exported, "model.pt2") 与 torch.compile 的关键区别：torch.compile 在运行时惰性编译（lazy JIT），而 torch.export 在部署前完成图的完整捕 获与验证，适合嵌入式、移动端（ExecuTorch）和对可重现性要求高的生产环境。ExecuTorch 即是 Meta 基于 torch.export 构建的边缘推理运行时，已规模化部署于 iOS/Android。

5.2.7 与 FSDP2 的协同

torch.compile 与 FSDP2（PyTorch 2.5+）的深度融合解决了分布式训练中编译图被破坏的问题。FSDP1 使用 FlatParameter 将参数展平为一个大张量，破坏了 torch.compile 的图优化能力。FSDP2 基于 DTensor 的 per- parameter sharding 从根本上解决了这一问题：每个参数保持独立的计算图身份，编译器可以穿越 FSDP 边界做全局优 化。PyTorch 2.5 的 compiled autograd 引擎能在编译图中自动插入异步 All-Gather，将通信 kernel 发射到计算图的最 前端，GPU 在等待 All-Gather 完成期间执行不依赖新参数的算子（如 LayerNorm、残差连接）。实测 Llama-3 8B FSDP2

  • torch.compile + async all-gather 将计算-通信重叠率从约 60% 提升至约 85%，整体吞吐提升 12-18%。这是 2025 年 单节点多 GPU 训练的最优配置。

5.3 图级 IR 与算子融合

算子融合（Operator Fusion）是ML编译器最核心的优化技术之一。其基本理念是将多个连续的算子合并为单个计算内 核，从而消除中间结果的显存往返（Memory Roundtrip），减少Kernel Launch开销，并充分利用GPU的寄存器层次结 构。以下深入剖析Inductor调度器中的融合策略和实现机制。

5.3.1 融合的收益分析

在一次典型的Transformer前向传播中，如果不进行融合，每个中间结果都要写回 HBM 再重新读取，产生多次显存往返 与内核启动；融合后则一次加载、一次写出。两种执行流的对比如下： Without Fusion Execution Flow: Input (HBM) ▶ Load ▶ MatMul (Compute) ▶ Store (HBM) ◀ 1 roundtrip ▶ Load ▶ Add (Compute) ▶ Store (HBM) ◀ 1 roundtrip ▶ Load ▶ ReLU (Compute) ▶ Store (HBM) ◀ 1 roundtrip ▶ Load ▶ Dropout (Compute)▶ Store (HBM) ◀ 1 roundtrip Total: 4 HBM round trips, 4 Kernel Launches Per Token: Memory read: 4 × 4096 × 2 bytes = 32 KB Memory write: 4 × 4096 × 2 bytes = 32 KB Total memory traffic: 64 KB/token With Fusion Execution Flow: Input (HBM) ▶ Load ▶ MatMul + Add + ReLU + Dropout (Single Kernel) ▶ Store (HBM) Total: 1 HBM round trip, 1 Kernel Launch Per Token: Memory read: 1 × 4096 × 2 bytes = 8 KB Memory write: 1 × 4096 × 2 bytes = 8 KB Total memory traffic: 16 KB/token ◀ 75% savings

5.3.2 融合类型分类

算子融合按数据依赖关系可分为垂直融合、水平融合、侧向融合与归约融合四种类型，如下所示。 ┌────────────────────────────────────────────────────────────────────────┐ │ Operator Fusion Categories │ ├────────────────────────────────────────────────────────────────────────┤ │ │ │ 1. Vertical Fusion │ │ Producer ▶ Consumer (data dependency chain) │ │ A ▶ B ▶ C ▶ D fused into FusedKernel │ │ Example: MatMul ▶ AddBias ▶ ReLU ▶ Dropout │ │ │ │ 2. Horizontal Fusion │ │ Same input ▶ multiple parallel operations │ │ ┌──▶ Op B ──┐ │ │ A ──┤ ├──▶ Op D │ │ └──▶ Op C ──┘ │ │ Example: Residual connection Add + two branches │ │ │ │ 3. Lateral Fusion │ │ Independent inputs ▶ same structure operations │ │ A ──▶ MatMul + BN + ReLU │ │ B ──▶ MatMul + BN + ReLU ▶ fused into a group │ │ Example: Heads of Multi-Head Attention │ │ │ │ 4. Reduction Fusion │ │ Compute ▶ Reduction ▶ Broadcast ▶ Compute │ │ Example: Softmax (max + sub + exp + sum + div) │ │ │ └────────────────────────────────────────────────────────────────────────┘

5.3.3 Inductor 融合算法

Inductor的调度器（Scheduler）是决定哪些操作可以融合、以什么顺序执行的核心组件。它基于FX Graph中的依赖关系 构建融合决策。

  1. BFS与DFS调度策略
 # Two scheduling modes of Inductor
 class InductorScheduler:
     def create_schedule(self, fx_graph, mode="memory_aware"):
         if mode == "bfs":
             return self._bfs_schedule(fx_graph)
         elif mode == "memory_aware":
             return self._memory_aware_schedule(fx_graph)
     def _memory_aware_schedule(self, fx_graph):
         """Memory-aware scheduling (default)"""
         # 1. Calculate memory pressure for each node
         for node in fx_graph.nodes:
             node.memory_pressure = compute_peak_memory(node)
         # 2. Greedy fusion: fuse from bottom up
         fused_groups = []
         for node in reversed(fx_graph.nodes): # Reverse traversal from output
             if can_fuse_vertically(node, node.consumers):
                 fused_groups.append(FusedGroup([node] + node.consumers))
         # 3. Check fusion boundaries
         for group in fused_groups:
             if group.total_smem_usage > SM_SHARED_MEMORY_LIMIT:

split_group(group) # Split overflowing fusion groups return Schedule(fused_groups) 2) 融合决策规则

 # Inductor fusion rules
 FUSION_RULES = {
     # Rule 1: Pointwise operations can always be fused

"pointwise": { "op_types": ["relu", "gelu", "tanh", "sigmoid", "add", "mul", "div"], "fusion_type": "vertical", "condition": lambda producer, consumer: True, # Always can fuse }, # Rule 2: Reduction▶Elementwise can be fused "reduction_to_pointwise": { "op_types": ["sum▶add", "max▶sub", "mean▶mul"], "fusion_type": "vertical", "condition": lambda p, c: p.reduction_axis == c.broadcast_axis, }, # Rule 3: MatMul▶Bias▶Activation can be fused (most common pattern) "gemm_epilogue_fusion": { "pattern": ["matmul", "add", "relu"], "fusion_type": "vertical", "template": "fused_matmul_bias_relu", }, # Rule 4: Attention pattern fusion "attention_fusion": { "pattern": ["matmul", "softmax", "matmul"], "fusion_type": "vertical", "template": "fused_attention", }, # Rule 5: Fusion boundaries (cannot fuse) "fusion_boundaries": [ "batch_norm", # BN needs cross-batch statistics "layer_norm", # LayerNorm needs cross-hidden_dim statistics "reductions", # Cross-dimension reductions cannot fuse with subsequent ops "large_matmul", # Large GEMM monopolizes SM resources, not suitable for fusion ]

5.3.4 模板式与可调融合

}

  1. 模板式融合 对于已知的常见模式，Inductor使用预定义的模板生成融合内核：
 # Template-based fusion in Inductor
 # Source: torch/_inductor/triton_heuristics.py
 # Template: MatMul + Bias + GeLU

@triton.jit def fused_matmul_bias_gelu_kernel( a_ptr, b_ptr, bias_ptr, out_ptr, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ):

     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
     # MatMul
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     for k in range(0, K, BLOCK_K):
         a = tl.load(a_ptr + offsets_a)
         b = tl.load(b_ptr + offsets_b)
         acc += tl.dot(a, b)
     # Bias add (fused)
     bias = tl.load(bias_ptr + offsets_bias)
     acc += bias
     # GeLU activation (fused)
     # gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
     acc = 0.5 * acc * (1.0 + tl.tanh(0.79788 * (acc + 0.044715 * acc * acc * acc)))
     tl.store(out_ptr + offsets_out, acc)
  1. 可调融合 对于非模板化模式，Inductor通过自动调参选择最优的融合配置：

5.3.5 常见融合模式

 # Tunable fusion: auto-search
 from torch._inductor import config
 # Enable auto-tuning
 config.max_autotune = True
 config.max_autotune_gemm_backends = ["TRITON", "CUDNN"]
 # Inductor automatically tries different fusion configurations
 # Config 1: BLOCK_M=64, BLOCK_N=64, BLOCK_K=32
 # Config 2: BLOCK_M=128, BLOCK_N=64, BLOCK_K=32
 # Config 3: BLOCK_M=128, BLOCK_N=128, BLOCK_K=64

在Transformer中，以下融合模式最为关键，如表5-3所示。 表5-3 Transformer关键融合模式 融合模式 涉及算子 收益 GEMM Epilogue 融合 MatMul → BiasAdd → Activation 减少2次Roundtrip Attention QK^T → Scale → Softmax → PV FlashAttention等价 MLP Block Linear → GeLU → Dropout → Linear 减少3次Roundtrip LayerNorm + LayerNorm → MatMul 共享数据加载 Residual Add Add → LayerNorm 双输入合并 Cross Entropy LogSoftmax → NLLLoss 减少中间张量

5.3.6 FlashAttention 融合机理

FlashAttention（Dao et al., 2022）是迄今最成功的 Attention 算子融合案例，其核心是将原本需要 O(N²) HBM 访问的 Attention 计算压缩到 O(N)。标准 Attention 需将 QK^T、Softmax、Attn·V 三步的 N×N 中间矩阵反复读写 HBM； FlashAttention 通过 Block Tiling + Online Softmax 将 Q/K/V 分块加载至片上 SRAM，跨块进行数值稳定的归一化，最 终仅将输出写回 HBM，总 HBM 流量降至 O(Nd)，这一融合模式已成为长序列训练和推理的标准路径。 Inductor 通过 SDPA（Scaled Dot-Product Attention）模式匹配自动选择 FlashAttention 路径（PyTorch 2.2+）：

torch.compile auto-fuses to FlashAttention

with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False): attn_output = torch.nn.functional.scaled_dot_product_attention( query, key, value, is_causal=True ) 当 Attention 符合 [batch, heads, seq, dim] 布局且 head_dim <= 256 时，Inductor 优先选择 FlashAttention 路径，不 满足时回退到标准分块 SDPA 内核。

5.4 自动调优

GEMM（通用矩阵乘法）和卷积操作是深度学习计算的核心，其性能对GPU的Tile大小、向量化宽度、寄存器分块等参数 高度敏感。以H100 GPU为例，一个4096×4096×4096的GEMM操作，选择不当的Tiling参数可能导致性能从理论峰值的 90%跌至30%。自动调优（Auto-tuning）技术通过在编译时搜索最优参数配置，弥合了理论峰值与实测性能之间的差 距。

5.4.1 自动调优问题形式化

自动调优可形式化为一个约束优化问题： max f (c, op, hardware) c∈C 其中： •C：配置空间（Tiling参数、向量化宽度、循环顺序、使用共享内存等） •f ：性能度量函数（通常为执行时间或吞吐量） • op ：操作的数学定义（如 GEMM(M,N,K)） • hardware ：目标硬件规格（SM数、共享内存大小、寄存器数等） GEMM 的配置空间由 Tile 尺寸（BLOCK_M、BLOCK_N、BLOCK_K）、分组大小（GROUP_SIZE_M）、流水线级数 （num_stages）与线程数（num_warps）等参数组合而成，并受共享内存与寄存器数等硬件资源约束。对于一个3D的 GEMM操作，配置空间大小可达 O(10 ) 到 O(10 )，直接穷举搜索不可行。 4 6

5.4.2 AutoTVM 与 AutoScheduler

  1. AutoTVM模板搜索 AutoTVM采用“手动模板 + 自动搜索参数”的方案：
 # AutoTVM workflow
 import tvm
 from tvm import auto_scheduler
 # 1. Define search space

@auto_scheduler.register_workload

 def gemm_workload(M, N, K):
     A = tvm.te.placeholder((M, K), name='A')
     B = tvm.te.placeholder((K, N), name='B')
     k = tvm.te.reduce_axis((0, K), name='k')
     C = tvm.te.compute((M, N), lambda i, j:
         tvm.te.sum(A[i, k] * B[k, j], axis=k), name='C')
     return [A, B, C]
 # 2. Search: XGBoost cost model
 task = auto_scheduler.SearchTask(
     func=gemm_workload, args=(4096, 4096, 4096),
     target="cuda -arch=sm_90" # H100

)

 tuner = auto_scheduler.TaskScheduler([task])
 tune_option = auto_scheduler.TuningOptions(
     num_measure_trials=1000,   # Search 1000 configurations
     runner=auto_scheduler.LocalRunner(repeat=3, timeout=10),
     measure_callbacks=[auto_scheduler.RecordToFile("gemm_tuning.json")],

) tuner.tune(tune_option)

3. Use optimal configuration

with auto_scheduler.ApplyHistoryBest("gemm_tuning.json"): with tvm.target.Target("cuda -arch=sm_90"): mod = tvm.build(sch, [A, B, C], "cuda") 2) MetaSchedule进化搜索 TVM的下一代调优器（MetaSchedule）采用“Sketch Generation + Evolutionary Search”：

 # MetaSchedule workflow
 from tvm.meta_schedule import TuneConfig, LocalRunner, Builder
 from tvm.meta_schedule.space_generator import PostOrderApply
 # 1. Sketch Generation
 # Auto-generate tiling
 space_gen = PostOrderApply()
 # 2. Evolutionary Search
 tune_config = TuneConfig(
     strategy="evolutionary",
     num_trials_per_iter=64,
     max_trials_per_task=2000,
     max_trials_global=10000,

)

 # 3. Cost model
 # Use XGBoost or neural network
 database = JSONDatabase("workload_database.json", "tuning_records")
 # 4. Parallel measurement
 runner = LocalRunner(
     evaluator_config=EvaluatorConfig(
         number=3, repeat=2, min_repeat_ms=100

) )

5.4.3 Triton Autotuning

OpenAI Triton的自动调优通过 @triton.autotune 装饰器实现，在kernel第一次调用时自动搜索最优配置：

  1. 自动搜索机制 import triton import triton.language as tl @triton.autotune( configs=[ triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 32,

'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8), triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=8), # ... More configurations ], key=['M', 'N', 'K'], # Cache by problem size ) @triton.jit def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...): # Kernel implementation ...

 # First call: auto-search
 c = matmul_kernel[a_ptr, b_ptr, c_ptr](M=4096, N=4096, K=4096, ...)
 # Second call with same size
 c = matmul_kernel[a_ptr, b_ptr, c_ptr](M=4096, N=4096, K=4096, ...)
  1. 配置空间剪枝 Triton通过预剪枝（Pruning）减少无效配置的搜索：
 # Triton's config filtering
 def prune_configs(configs, problem_size):
     valid_configs = []
     for cfg in configs:
         # 1. Register pressure check
         if estimate_registers(cfg) > MAX_REGISTERS_PER_THREAD:

continue # 2. Shared memory check if estimate_smem(cfg) > SHARED_MEMORY_PER_SM: continue # 3. Occupancy estimation if estimate_occupancy(cfg) < MIN_OCCUPANCY: continue

5.4.4 探索与利用的权衡

     valid_configs.append(cfg)
 return valid_configs

自动调优面临的核心矛盾：搜索空间巨大但测量成本高（每次测量需要实际运行GPU kernel）。

  1. 搜索策略对比 各搜索策略的探索与利用倾向对比如表5-4所示。 表5-4 自动调优搜索策略对比 策略 探索倾向 利用倾向 收敛速度 适用场景 Random Search 极高 极低 极慢 全新工作负载 Grid Search 高 低 慢 低维配置空间 Simulated Annealing 高→低 低→高 中等 有局部最优 Bayesian Optimization 中 中 中等 评估成本高 Evolutionary Search 中高 中 慢→快 大配置空间 策略 探索倾向 利用倾向 收敛速度 适用场景 Transfer Learning 低 高 极快 相似工作负载
  2. 代价模型辅助搜索 代价模型用已测量的配置预测未见配置的性能，减少实际测量次数：
 # TVM cost model (XGBoost)
 # Training data: (config, hardware) ▶ measured_runtime
 # Goal: predict performance of unseen configurations
 # Feature engineering
 def extract_features(config, workload):
     return {

"BLOCK_M": config.BLOCK_M, "BLOCK_N": config.BLOCK_N, "BLOCK_K": config.BLOCK_K, "num_warps": config.num_warps, "num_stages": config.num_stages, "M": workload.M, "N": workload.N, "K": workload.K, "arithmetic_intensity": (2 * M * N * K) / ((MK + KN + M*N) * 2), # GPU-specific features "occupancy_ratio": estimate_occupancy(config), "smem_per_thread": estimate_smem(config) / config.num_warps, "registers_per_thread": estimate_registers(config),

5.4.5 在线调优 vs 离线调优

 }

在线调优发生在编译时或首次运行，离线调优在部署前的 CI/CD 中完成，两者的差异如表5-5所示。 表5-5 在线调优与离线调优对比 维度 在线调优（Online） 离线调优（Offline） 时机 编译时/首次运行时 部署前（CI/CD中） 成本 首次运行有延迟 无运行时开销 适应性 适应具体硬件和输入 对硬件/输入变化不敏感 工具 Triton autotune, TorchInductor TVM AutoScheduler, TensorRT 典型延迟 数秒到数分钟 数小时（但仅一次） 生产环境推荐策略：

1. Offline tuning in CI/CD pipeline

python tune_kernels.py --output=tuned_configs.json

2. Package tuning results into Docker image

FROM nvcr.io/nvidia/pytorch:24.01-py3 COPY tuned_configs.json /opt/tuned/ ENV TORCH_INDUCTOR_CACHE_DIR=/opt/tuned/

5.4.6 坐标下降调优

3. Use pre-tuned configs

torch.compile(model, mode="max-autotune") PyTorch Inductor 在 max-autotune 模式下支持坐标下降调优（Coordinate Descent Tuning，CDT），这是一种比穷举/ 随机搜索更高效的 tile 参数优化策略。其原理是从基准配置（如 BLOCK_M=BLOCK_N=64、BLOCK_K=32）出发，每次 只沿一个坐标轴调整一个参数，保留带来提升的方向，直至所有方向均无改善为止。相比穷举搜索（可达 10^4 量级测 量），CDT 通常在 20-50 次测量内收敛。

 import torch._inductor.config as inductor_config
 inductor_config.coordinate_descent_tuning = True
 inductor_config.coordinate_descent_check_all_directions = True   # more thorough but slower
 model = torch.compile(model, mode="max-autotune")

实测中，开启 CDT 相比不调优的 max-autotune 通常可额外提升数个百分点的吞吐，代价是编译时间增加约 30-60 秒。 CDT 结果写入 Inductor 缓存，后续相同工作负载直接复用，生产部署时无运行时开销。

5.5 JAX 与 XLA 计算模型

JAX是由Google Research开发的数值计算框架，其设计哲学是“函数式编程 + 可组合变换 + XLA编译优化”。与PyTorch 的“命令式优先、编译可选”不同，JAX从第一天起就深度绑定XLA编译器，所有计算都经过JIT编译和优化。这种设计使 得JAX在科学计算、大规模分布式训练和高性能推理场景中展现出独特的优势。

5.5.1 JAX 函数式编程模型

JAX的核心抽象是“纯函数 + 显式随机状态 + 变换组合”：

 import jax
 import jax.numpy as jnp
 # 1. Pure function: no side effects
 def predict(params, x):
     return jnp.dot(x, params)
 # 2. Explicit random state (functional PRNG)
 key = jax.random.PRNGKey(42)

key, subkey = jax.random.split(key) w = jax.random.normal(subkey, (512, 256))

3. Transformation composition: jit(vmap(grad(f)))

@jax.jit # XLA compilation @jax.vmap(in_axes=(0, None)) # Auto-vectorization (batch dimension) @jax.grad # Automatic differentiation (gradient)

 def loss_fn(params, x):
     return jnp.sum(predict(params, x))
 # Equivalent to:
 # compiled_grad_fn = jax.jit(jax.vmap(jax.grad(loss_fn)))

JAX 的四大变换如表5-6所示。 表5-6 JAX四大变换 变换 功能 类比PyTorch 典型用法 jax.jit JIT编译 torch.compile 将函数编译为XLA HLO jax.vmap 自动向量化 torch.vmap 批量化单样本函数 jax.pmap SPMD并行 DistributedDataParallel 数据并行 jax.grad 自动微分 torch.autograd.grad 计算梯度

5.5.2 XLA 计算模型

XLA（Accelerated Linear Algebra）是Google为TPU和GPU设计的领域专用编译器。它的输入是一种称为HLO（High- Level Optimizer）的IR。

  1. XLA编译Pipeline Python/JAX Code │ ▼ jaxpr (JAX internal IR) │ ▼ StableHLO (standardized ML operation set) │ ▼ ┌──────────────────────┐ │ XLA Compiler │ │ │ │ 1. HLO Optimization │ │ - Algebraic │ │ simplification │ │ - Fusion │ │ - Layout │ │ assignment │ │ 2. Buffer Assignment│ │ 3. Code Generation │ │ - GPU: PTX/CUDA │ │ - TPU: TPU ISA │ │ - CPU: LLVM IR │ └──────────────────────┘ │ ▼ Executable binary
  2. StableHLO StableHLO是OpenXLA社区推出的标准化ML操作集，旨在解决不同框架（JAX、PyTorch、TensorFlow）到XLA的不同 路径问题。
 // StableHLO IR example (MLIR format)
 // Compute: y = relu(x @ w + b)
 func.func @mlp(%arg0: tensor<32x512xf32>, %arg1: tensor<512x256xf32>,

%arg2: tensor<256xf32>) -> tensor<32x256xf32> { %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<32x512xf32>, tensor<512x256xf32>) -> tensor<32x256xf32> %1 = stablehlo.broadcast_in_dim %arg2, dims = [1] : (tensor<256xf32>) -> tensor<1x256xf32> %2 = stablehlo.add %0, %1 : tensor<32x256xf32> %3 = stablehlo.maximum %2, %{0.0} : tensor<32x256xf32>

5.5.3 GSPMD 自动并行分区

 return %3 : tensor<32x256xf32>

} GSPMD（Generalized SPMD）是XLA的并行分区系统，它允许用户通过简单的注释指定数据的分布方式，XLA自动生成 通信和计算代码。

  1. 自动SPMD
 import jax
 from jax.sharding import PartitionSpec as P, NamedSharding
 from jax.experimental import mesh_utils
 # 1. Create device mesh
 devices = mesh_utils.create_device_mesh((8, 4))    # 8×4 2D grid
 # devices.shape = (8, 4) ▶ 32 TPUs/GPUs
 # 2. Define Mesh
 mesh = jax.sharding.Mesh(devices, ('data', 'model'))
 # 3. Distributed array: define sharding
 # 'data' axis(8):
 # 'model' axis(4):
 sharding = NamedSharding(mesh, P('data', 'model'))
 # Example: input tensor shape (128, 4096) sharded over 4 data devices
 # Per device shard: (128/4=32, 4096)
 # 4. Automatic SPMD

@jax.jit

 def train_step(params, batch):
     logits = model.apply(params, batch['input'])
     loss = compute_loss(logits, batch['target'])
     return loss
 # JAX/XLA automatically inserts communication primitives:
 # AllGather: collect batch slices across data axis
 # AllReduce: aggregate gradients from all model-axis devices
  1. 手动shard_map from jax.experimental.shard_map import shard_map

For scenarios requiring fine-grained control over sharding patterns

@shard_map(

     mesh=mesh,
     in_specs=(P('data', 'model'), P('data', None)),
     out_specs=P('data', 'model'),
     check_rep=False,

)

5.5.4 JAX vs PyTorch 编译对比

 def custom_parallel_layer(params, x):
     # Inside this function body, params and x are already local shards
     # Developer can finely control computation and communication
     local_output = fancy_computation(params, x)
     # Cross-model-axis communication (manual AllReduce)
     global_output = jax.lax.psum(local_output, 'model')
     return global_output

JAX 与 PyTorch 编译模型的对比如表5-7所示。 表5-7 JAX与PyTorch编译模型对比 维度 JAX (jit) PyTorch (torch.compile) 设计理念 编译优先（JIT by default） Eager优先 + 可选编译 IR路径 jaxpr → StableHLO → XLA Dynamo → FX Graph → Inductor 并行API pmap, pjit, shard_map DDP, FSDP, DTensor 自动微分 函数式（jax.grad） 面向对象（torch.autograd） 多后端 TPU (一流), GPU, CPU GPU (一流), CPU 维度 JAX (jit) PyTorch (torch.compile) 动态形状 限制较多（recompile） 较好的动态性 生态系统 研究为主（JAX, Flax） 工业级（PyTorch, HF, TIMM） 学习曲线 陡峭（函数式编程） 平缓（Pythonic） 性能对比（Llama-7B训练，H100×8）： JAX + Flax + XLA: tokens/sec: 15,200 MFU: 52.3% compile time: 45s (first run) PyTorch + torch.compile: tokens/sec: 14,100 MFU: 48.5% compile time: 32s (first run) PyTorch (eager): tokens/sec: 9,800 MFU: 33.7%

5.5.5 JAX 的使用场景

JAX在多领域中展现出独特优势： •TPU训练：Google TPU上JAX是首选框架，性能远优于PyTorch的TPU支持。 •科学计算：JAX的 vmap 在模拟物理系统、分子动力学中天然契合。 •强化学习：JAX的函数式特性使RL环境模拟可以JIT编译，百倍加速。 •大规模稀疏模型：JAX的Spmd比PyTorch的FSDP更灵活。

5.6 TensorRT 推理优化全流程

NVIDIA TensorRT是GPU推理优化的工业标准，通过对已训练好的模型进行图优化、精度量化、内核自动调优和运行时优 化，在NVIDIA GPU上实现极致推理性能。在LLM推理场景中，TensorRT-LLM将TensorRT的能力扩展到大语言模型的分 布式推理。以下从TensorRT的完整优化管道出发，剖析其核心技术和集成实践。

5.6.1 TensorRT 优化管道

TensorRT 的优化管道分为图优化、内核自动调优与构建序列化三个阶段，如下所示。 Training Frameworks (PyTorch, TF, JAX) │ ▼ ┌─────────────────────┐ │ Model Export │ │ ONNX / TorchScript │ └────────┬────────────┘ │ ▼ ┌────────┬───────────────────────────────────────────────────┐ │ TensorRT Optimization Pipeline │ │ │ │ Phase 1: Graph Optimization │ │ ┌──────────────────────────────────────────────────────┐ │ │ │ • Layer Fusion (Conv+BN+ReLU ▶ CBR) │ │ │ │ • Constant Folding │ │ │ │ • Dead Code Elimination │ │ │ │ • Tensor Layout Optimization │ │ │ │ • Q/DQ (Quantize/Dequantize) Insertion │ │ │ └──────────────────────────────────────────────────────┘ │ │ │ │ Phase 2: Kernel Auto-Tuning │ │ ┌──────────────────────────────────────────────────────┐ │ │ │ • Tactic Selection (multiple implementations compete)│ │ │ │ • Timing Cache (cache optimal config) │ │ │ │ • Memory Footprint Optimization │ │ │ └──────────────────────────────────────────────────────┘ │ │ │ │ Phase 3: Build & Serialize │ │ ┌──────────────────────────────────────────────────────┐ │ │ │ • CUDA Engine Build │ │ │ │ • Plan Serialization (.trt/.engine) │ │ │ │ • Runtime Engine (deserialization + execution) │ │ │ └──────────────────────────────────────────────────────┘ │ └────────────────────────────────────────────────────────────┘

5.6.2 图优化详解

  1. 层级融合 TensorRT的融合策略比通用编译器更激进，因为它只面向推理场景（无反向传播）： Original Network (Without Fusion): Input ▶ Conv ▶ Bias ▶ ReLU ▶ Conv ▶ Bias ▶ ReLU ▶ MaxPool ▶ Output 7 Kernels, 6 memory round trips TensorRT After Fusion: Input ▶ CBR1 ▶ CBR2 ▶ MaxPool ▶ Output 3 Kernels, 2 memory round trips Fused CBR Kernel:
   - Conv (IMPLICIT_GEMM or Winograd)
   - Bias Add (completed at the tail of the previous kernel)
   - ReLU Activation (same as above)

All three completed in one CUDA Kernel TensorRT 支持的主要融合类型如表5-8所示。 表5-8 TensorRT支持的主要融合类型 融合类型 模式 Intel/AMD等价 Conv + Bias + ReLU CBR 通用 Conv + Bias + ReLU + Elementwise CBRE 通用 MatMul + Bias + GeLU MBC OpenVINO MBC Multi-Head Attention MHA FlashAttention API Conv + Elementwise Sum Conv + Add 残差连接

5.6.3 INT8/FP8 量化

Reduction + Elementwise LayerNorm ONNX Runtime TensorRT的INT8和FP8量化是推理加速的最强手段。H100/H200/B200原生支持FP8，TensorRT-LLM已全面适配FP8推 理。

  1. INT8校准
 import tensorrt as trt
 import pycuda.driver as cuda
 import pycuda.autoinit
 TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
 # 1. Build INT8 calibrator
 class MyCalibrator(trt.IInt8EntropyCalibrator2):
     def __init__(self, calibration_data):
         super().__init__()
         self.calibration_data = calibration_data
         self.cache_file = "calibration.cache"
         self.current_index = 0
     def get_batch_size(self):
         return 32
     def get_batch(self, names):
         if self.current_index >= len(self.calibration_data):
             return None
         batch = self.calibration_data[self.current_index]
         self.current_index += 1
         return [batch.ctypes.data]
     def read_calibration_cache(self):
         if os.path.exists(self.cache_file):

with open(self.cache_file, "rb") as f: return f.read() def write_calibration_cache(self, cache): with open(self.cache_file, "wb") as f:

             f.write(cache)
 # 2. Build INT8 engine
 builder = trt.Builder(TRT_LOGGER)
 network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
 config = builder.create_builder_config()
 config.set_flag(trt.BuilderFlag.INT8)
 config.int8_calibrator = MyCalibrator(calibration_data)
 # Parse ONNX and build
 parser = trt.OnnxParser(network, TRT_LOGGER)
 parser.parse(onnx_model.SerializeToString())
 engine = builder.build_serialized_network(network, config)
  1. FP8推理 以下示例基于 TensorRT-LLM，需 FP8 支持的 GPU（如 H100）：
 # TensorRT-LLM FP8 inference
 # build_config.py
 import tensorrt_llm
 build_config = tensorrt_llm.BuilderConfig()
 build_config.use_fp8 = True
 build_config.fp8_kv_cache = True # FP8 KV Cache can save 50% memory
 build_config.quant_mode = tensorrt_llm.QuantMode.use_fp8_kv_cache()
 # FP8 quantization strategy
 # TensorRT provides multiple FP8 quantization modes:
 # 1. FP8_DEFAULT: mixed precision fallback for unsupported layers
 # 2. FP8_STRICT: force all layers to FP8 (may fail for some ops)
 # 3. FP8_PER_TENSOR: per-tensor quantization granularity
 # 4. FP8_PER_CHANNEL: per-channel quantization granularity

三种精度的推理性能对比如表5-9所示（H100, Llama-7B）。 表5-9 INT8与FP8推理性能对比 精度 吞吐 (tokens/sec) 显存占用 困惑度 (WikiText2) FP16 5,200 14.2 GB 5.67 INT8 8,100 (+56%) 8.1 GB 5.72 FP8 10,500 (+102%) 8.1 GB 5.68

5.6.4 动态形状与 Profile

TensorRT 引擎通过优化 Profile 支持动态形状。Profile 为每个输入张量定义最小、最优与最大形状，引擎在该范围内为 不同形状编译对应内核，最优形状应设为最常用的输入尺寸以获取最佳平均性能：

 # TensorRT dynamic shape configuration
 profile = builder.create_optimization_profile()
 # Define min/optimal/max
 profile.set_shape(

"input_ids",

     min=(1, 1),       # Min: batch=1, seq_len=1
     opt=(4, 128),      # Optimal: batch=4, seq_len=128
     max=(8, 2048),     # Max: batch=8, seq_len=2048

) profile.set_shape( "attention_mask",

     min=(1, 1),
     opt=(4, 128),
     max=(8, 2048),

)

5.6.5 TensorRT-LLM 集成

 config.add_optimization_profile(profile)
 # Note: Dynamic shapes mean the engine rebuilds kernels for each
 # shape in the profile range. The opt shape should be set to the
 # most common input shape for best average performance.

TensorRT-LLM 提供从构建引擎到推理的完整 API，以下示例构建并运行一个 Llama-2-7B 引擎（需安装 TensorRT-LLM 与 CUDA 12+）：

 # Complete example of building and running a TensorRT-LLM engine
 from tensorrt_llm import LLM, SamplingParams
 # 1. Build TensorRT engine (via CLI)
 # python examples/llama/build.py \
 #   --model_dir meta-llama/Llama-2-7b-hf \
 #   --dtype float16 \
 #   --use_gpt_attention_plugin float16 \
 #   --use_gemm_plugin float16 \
 #   --use_inflight_batching \
 #   --paged_kv_cache \
 #   --max_batch_size 64 \
 #   --max_input_len 2048 \
 #   --max_output_len 512 \
 #   --output_dir ./llama2_7b_trt_engine/
 # 2. Load engine and run
 llm = LLM(
     model="./llama2_7b_trt_engine/",
     tokenizer=tokenizer,

)

 sampling_params = SamplingParams(
     temperature=0.8,
     top_p=0.95,
     max_tokens=256,

) prompts = [ "Explain the concept of transformer attention:", "Write a Python function to calculate Fibonacci numbers:", ]

 outputs = llm.generate(prompts, sampling_params)
 for output in outputs:
     print(f"Prompt: {output.prompt}")
     print(f"Generated: {output.outputs[0].text}")
     print(f"Tokens/sec: {output.metrics.tokens_per_second}")

TensorRT-LLM的关键特性： •In-flight Batching：连续批处理，请求即到即处理，无需等待凑批。 •PagedAttention：KV Cache分页管理，显存利用率从约40%提升至约90%。 •Tensor Parallelism：支持单节点多GPU的TP（张量并行）。 •Pipeline Parallelism：支持跨节点PP（流水线并行）。 vLLM 与 TensorRT-LLM 的性能对比如表5-10所示（Llama-7B, H100×1）。 表5-10 vLLM与TensorRT-LLM性能对比 指标 vLLM (FP16) TensorRT-LLM (FP8) 提升 单请求延迟 (256 tokens) 2.1 sec 1.3 sec 38% 最大吞吐 (req/sec) 48 req/s 95 req/s 98% 首Token延迟 (TTFT) 120 ms 68 ms 43% 最大Batch Size 256 512 2×

5.6.6 推理框架选型与端侧框架

推理框架选型决定部署形态、精度支持与性能上限。主流框架各具定位，选型需结合模型来源、目标平台、精度需求与部 署形态综合判断，各框架定位对比如表5-11所示。 表5-11 主流推理框架定位对比 框架 定位 目标平台 特点 TensorRT / TensorRT- NVIDIA GPU 高性能 NVIDIA GPU 图优化+内核自动调优，FP8/INT8 支持成熟，LLM 用 LLM 推理 TensorRT-LLM ONNX Runtime 跨平台通用中间层 CPU/GPU/NPU 平台 全 模型转换枢纽，多 EP（执行提供者）分发到不同硬件 OpenVINO Intel 平台优化 Intel 视觉+LLM 负载，CPU 侧优化深入，端侧边缘常用 CPU/GPU/NPU RKNN Rockchip NPU RK 系列 SoC 端侧摄像头/IPC 场景，INT8 量化为主 NCNN / MNN 端侧移动推理 ARM CPU/Mobile 轻量、低资源，移动端/嵌入式主流 GPU llama.cpp LLM 端侧推理 CPU/GPU/Metal GGUF 量化格式，单文件部署，AI PC 场景 选型决策方法分四步：

  1. 按目标平台收敛：NVIDIA GPU 优先 TensorRT，Intel 平台优先 OpenVINO，Rockchip NPU 只能 RKNN，ARM 端侧优 先 NCNN/MNN
  2. 按模型类型：LLM 用支持 KV Cache 与连续批处理的框架（TensorRT-LLM、llama.cpp、SGLang），视觉模型用 TensorRT/OpenVINO/RKNN
  3. 按精度需求：端侧受限场景通常 INT8/INT4 量化（框架需支持对应量化格式如 GGUF、RKNN 的 INT8）
  4. 按部署形态：服务化用 TensorRT-LLM/vLLM，嵌入式单机用 NCNN/MNN/RKNN，边缘网关用 OpenVINO/llama.cpp
  1. 端侧框架的共性约束 端侧框架面向资源受限环境，与数据中心框架差异明显： •内存受限：模型需 INT8/INT4 量化后放入有限内存（如 8GB 内存的 IPC 只能跑 7B 量化模型）；权重动态加载与内存池 复用是关键优化。 •计算单元异构：端侧 SoC 含 CPU/NPU/DSP/GPU 多计算单元，框架需支持算子分发（NPU 跑卷积/GEMM，CPU 跑预 处理与不支持算子）。 •量化算子约束：NPU 算子库只覆盖常见算子（Conv/GEMM/激活），长尾算子（自定义激活、特殊 attention）需量化 或回退 CPU，评测时须扫描算子覆盖并记录回退清单，评估算子可达性能否支撑目标负载。 •功耗与散热：端侧推理受 TDP 约束，连续推理可能触发降频，性能评测需含长时稳定项（负载下连续运行数小时观测 吞吐衰减），并纳入推理负载调度与功耗管理考虑。
  2. 框架选型的实测验证 选型不能只看文档，需在目标硬件上实测三项：单帧/单请求延迟、峰值吞吐、模型精度损失（INT8 相对 FP32 的任务精 度差）。同一模型在不同框架间对比时，须固定精度、输入尺寸与预热轮数（框架基准测试的对照原则）。端侧选型还应验 证：算子覆盖是否触发大量 CPU 回退、量化后长尾任务（如代码/数学）是否退化、长时运行是否降频导致吞吐衰减。

5.7 下一代编译器生态

ML编译器正在经历从“孤岛式专用编译器”到“MLIR生态系统”的范式转变。MLIR（Multi-Level Intermediate Representation）作为LLVM旗下的编译器基础设施，提供了定义和组合多个Dialect（方言）的框架，使得不同层次的 ML操作可以在统一的IR中进行优化和降级。以下聚焦MLIR生态中的关键项目和未来趋势。

5.7.1 MLIR 基础设施

MLIR的核心创新是Dialect系统——每个Dialect定义一组特定领域的操作、类型和属性，Dialect之间通过转换Pass （Conversion Pass）实现渐进式降级（Progressive Lowering）。 Dialect 按抽象层次可分为四层： •高层框架级：Torch Dialect（torch.linear、torch.relu）、TF Dialect（tf.MatMul、tf.Conv2D）、ONNX Dialect •中层ML级：StableHLO（标准化 ML 算子集）、Linalg（结构化线性代数）、TOSA/MHLO（嵌入式与遗留 ML） •低层硬件级：GPU Dialect（通用 GPU）、NVGPU/AMDGPU Dialect、SPIR-V Dialect •代码级：LLVM Dialect，经 LLVM IR 进一步降级为机器码 MLIR 的关键价值在于这些 Dialect 可自由组合：高层 Dialect 与低层 Dialect 通过渐进式降级连接，同一份高层 IR 可以稳 定地到达多套后端。

5.7.2 Torch-MLIR 管道

Torch-MLIR项目将PyTorch模型通过MLIR管道编译为多种后端代码：

 # Torch-MLIR compilation flow
 import torch_mlir
 # 1. Export PyTorch model to MLIR
 class SimpleModel(torch.nn.Module):
     def forward(self, x):
         return torch.nn.functional.relu(x @ w + b)
 model = SimpleModel().eval()
 example_input = torch.randn(4, 512)
 # 2. Generate Torch Dialect IR
 module = torch_mlir.compile(

model, example_input, output_type=torch_mlir.OutputType.TORCH, )

 # Output: torch.aten.linear, torch.aten.relu in Torch Dialect
 # 3. Lowering: Torch ▶ Linalg ▶ GPU
 module_linalg = torch_mlir.compile(

model, example_input, output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, )

 # Output: linalg.matmul, linalg.add, linalg.max (ReLU decomposition)
 # 4. Final: Linalg ▶ LLVM ▶ CUDA
 module_gpu = torch_mlir.compile(

model, example_input, output_type=torch_mlir.OutputType.GPU, ) Torch-MLIR的转换路径： Torch Dialect (torch.* ops) │ ▼ TorchToLinalg / TorchToStableHLO Linalg-on-Tensors / StableHLO │ ▼ ConvertLinalgToParallelLoops ▶ ConvertParallelLoopsToGPU GPU Dialect (gpu.launch, gpu.alloc) │ ▼ LowerGpuOpsToNVVMOps / LowerGpuOpsToROCDLOps LLVM/NVVM Dialect │ ▼

5.7.3 IREE 运行时

CUDA / ROCm kernel IREE（Intermediate Representation Execution Environment）是Google推出的基于MLIR的通用模型部署运行时：

 # IREE compilation and execution
 import iree.compiler as ireec
 import iree.runtime as ireert
 # 1. Compile MLIR to IREE
 compiled_module = ireec.tools.compile_file(

"model.mlir", input_type="mhlo", target_backends=["cuda"], # or "rocm", "vulkan", "llvm-cpu" )

 # 2. Load into IREE runtime
 config = ireert.Config("cuda")
 ctx = ireert.SystemContext(config=config)
 vm_module = ireert.VmModule.copy_buffer(ctx.instance, compiled_module)
 ctx.add_vm_module(vm_module)
 # 3. Invoke model
 result = ctx.modules.module["forward"](input_tensor)

IREE的优势： •多后端：CUDA、ROCm、Vulkan、Metal、CPU（x86/ARM/RISC-V），一个IR多种部署。 •极低延迟：适合边缘设备和移动端（Android Neural Networks API集成）。 •零拷贝：共享内存优化，减少数据搬运。 •异步执行：通过Hal API实现异步调度。

5.7.4 Triton-IR 编译器

注意区分两个“Triton”： •OpenAI Triton：Python-based GPU kernel语言。 •Intel Triton-IR：Intel推出的MLIR-based GPU编译器，提供Triton-IR Dialect。

 // Intel Triton-IR example: matrix multiplication
 // Source: Intel Triton-IR Dialect definition
 tt.func @matmul(%A: !tt.ptr, %B: !tt.ptr, %C: !tt.ptr,

%M: i32, %N: i32, %K: i32) { %c0 = arith.constant 0 : index %c128 = arith.constant 128 : index %c32 = arith.constant 32 : index %pid_m = tt.get_program_id x %pid_n = tt.get_program_id y // Main loop: tiled matrix multiplication tt.for %k = %c0 to %K step %c32 { %a = tt.load %A[%pid_m * %c128, %k * %c32] : !tt.ptr %b = tt.load %B[%k * %c32, %pid_n * %c128] : !tt.ptr %c_partial = tt.dot %a, %b : tensor<128x32xf16> * tensor<32x128xf16> %c = arith.addf %c, %c_partial : tensor<128x128xf32>

5.7.5 OpenXLA 社区

     }
     tt.store %C[%pid_m, %pid_n], %c : !tt.ptr
     tt.return
 }
 // Triton-IR features:
 // 1. Triton-like syntax (block-level programming model)
 // 2. But based on MLIR Dialect definition (combinable with other MLIR optimization passes)
 // 3. Supports Intel GPU (Xe) and NVIDIA GPU

OpenXLA是Google、NVIDIA、Meta、AMD等共同维护的开放编译器生态系统，将 JAX、PyTorch、TensorFlow 等前端 统一接入 StableHLO 与 XLA 编译器，再分发到 GPU、TPU、CPU 等后端。 OpenXLA的关键进展： •StableHLO v1.0：标准化了约120个ML操作，覆盖95%的常见模型操作。 •PJRT（Portable/Plugin JAX/XLA Runtime）：统一的设备运行时API，支持动态设备热插拔。 •Shardy：新一代自动并行分区系统（GSPMD的继任者）。 •IFRT（Inter-Framework Runtime）：跨框架的统一推理运行时。

5.7.6 异构硬件编译器

随着AI加速器多元化（NVIDIA GPU、AMD GPU、Intel Gaudi、Google TPU、Cerebras CS-2、Groq LPU），ML编译器面 临“一套模型代码，多套硬件后端”的挑战。MLIR 的 Dialect 降级体系让同一份高层 IR 可以按目标硬件选择不同的转换 Pass 序列，从而复用前端优化、只替换后端：

# Heterogeneous compiler workflow (pseudocode)
def compile_for_hardware(model_ir, target_hardware):
    if target_hardware == "NVIDIA_GPU":
        passes = [

ConvertTorchToStableHLO(), StableHLOToLinalg(), LinalgToNVGPU(), NVGPUToPTX(), ] elif target_hardware == "AMD_GPU": passes = [ ConvertTorchToStableHLO(), StableHLOToLinalg(), LinalgToAMDGPU(), AMDGPUToROCDL(), ] elif target_hardware == "Intel_Gaudi": passes = [ ConvertTorchToStableHLO(), StableHLOToTPC(), # Tensor Processor Core ISA ] elif target_hardware == "Groq_LPU": passes = [ ConvertTorchToStableHLO(), StableHLOToGroqISA(), ]

5.8 torch.compile 实战

    for pass in passes:
        model_ir = pass.run(model_ir)
    return model_ir

以下提供一套可复现的torch.compile性能测试与调优指南，涵盖Llama、GPT、Stable Diffusion三类模型的编译加速效 果，从模式选择、图断裂诊断到生产部署的完整流程。所有实验数据基于H100 80GB SXM5 GPU和PyTorch 2.3版本。

5.8.1 实验环境

本章实验环境配置如下（NVIDIA H100 80GB SXM5，CUDA 12.4，PyTorch 2.3.0）：

Environment configuration

GPU: 8× NVIDIA H100 80GB SXM5 CUDA: 12.4 PyTorch: 2.3.0 Driver: 550.54.15 CPU: 2× Intel Xeon Platinum 8480C

PyTorch 2.3 key installation

pip install torch==2.3.0 torchvision torchaudio \

5.8.2 Llama-7B 性能测试

--index-url https://download.pytorch.org/whl/cu124 以下脚本在 Llama-2-7B 上对比 Eager 与三种 torch.compile 模式的训练性能：

# llama_benchmark.py - torch.compile performance benchmark for Llama-7B
import torch
import time
from transformers import AutoModelForCausalLM, AutoConfig
def benchmark_llama_forward():
    config = AutoConfig.from_pretrained(

"meta-llama/Llama-2-7b-hf", torch_dtype=torch.float16, use_cache=False, )

    model = AutoModelForCausalLM.from_config(config).cuda()
    model.train()
    # Test input

batch_size, seq_len = 1, 2048

    input_ids = torch.randint(0, 32000, (batch_size, seq_len)).cuda()
    results = {}
    # 1. Baseline (Eager)
    model_eager = model
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(20):
        loss = model_eager(input_ids, labels=input_ids).loss
        loss.backward()
    torch.cuda.synchronize()
    t1 = time.perf_counter()

results["eager"] = (t1 - t0) / 20

    # 2. torch.compile (default mode)
    model_default = torch.compile(model, mode="default")
    # Warmup: first call triggers compilation
    loss = model_default(input_ids, labels=input_ids).loss
    loss.backward()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(20):
        loss = model_default(input_ids, labels=input_ids).loss
        loss.backward()
    torch.cuda.synchronize()
    t1 = time.perf_counter()

results["compile_default"] = (t1 - t0) / 20

    # 3. torch.compile (reduce-overhead)
    model_reduce = torch.compile(model, mode="reduce-overhead")
    loss = model_reduce(input_ids, labels=input_ids).loss
    loss.backward()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(20):
        loss = model_reduce(input_ids, labels=input_ids).loss
        loss.backward()
    torch.cuda.synchronize()
    t1 = time.perf_counter()

results["compile_reduce_overhead"] = (t1 - t0) / 20

    # 4. torch.compile (max-autotune)
    model_max = torch.compile(model, mode="max-autotune")
    loss = model_max(input_ids, labels=input_ids).loss
    loss.backward()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(20):
        loss = model_max(input_ids, labels=input_ids).loss
        loss.backward()
    torch.cuda.synchronize()
    t1 = time.perf_counter()

results["compile_max_autotune"] = (t1 - t0) / 20 return results print(benchmark_llama_forward()) 实测结果（Llama-7B, H100×1, batch=1, seq=2048）如表5-12所示。 表5-12 Llama-7B编译模式实测结果 模式 编译时间 单步延迟 (s) 吞吐 (tok/s) 相对加速 峰值显存 (GB) Eager (baseline) 0s 0.98 2,090 1.00× 27.3 compile (default) 45s 0.67 3,060 1.46× 24.1 compile (reduce-overhead) 48s 0.58 3,530 1.69× 24.3 compile (max-autotune) 185s 0.52 3,940 1.88× 23.8

5.8.3 GPT-2 与 Stable Diffusion

GPT-2 XL（1.5B 参数）训练性能对比如表5-13所示。 表5-13 GPT-2 XL训练性能对比 模式 吞吐 (tok/s) MFU 编译时间 Eager 12,400 35.2% 0s compile (default) 18,700 53.1% 12s compile (max-autotune) 21,500 61.0% 58s Stable Diffusion XL 推理性能如表5-14所示。 表5-14 Stable Diffusion XL推理性能 模式 推理时间 (s/image) images/sec 峰值显存 Eager 3.82 0.26 12.1 GB compile (default) 2.15 0.47 10.3 GB compile (reduce-overhead) 1.61 0.62 10.5 GB

5.8.4 图断裂诊断与修复

以下脚本通过开启详细日志定位图断裂点，并给出常见问题的修复写法：

 # Graph break diagnosis script
 import os
 import torch
 # Enable verbose logging
 os.environ["TORCH_LOGS"] = "graph_breaks,recompiles"
 os.environ["TORCH_COMPILE_DEBUG"] = "1"

@torch.compile

 def problem_model(x):
     # Common graph break causes and fixes
     y = x @ w
     # ✗ Graph Break 1: data-dependent if (fix: use torch.where)
     # if y.sum() > 0:
     #     y = y * 2
     # ✓ Fix: use tensor operations
     y = torch.where(y.sum() > 0, y * 2, y)
     # ✗ Graph Break 2: .item() (fix: keep as tensor)
     # scalar = y[0, 0].item()
     # ✓ Fix: use tensor comparison
     mask = y[0, 0] > 0.5
     y = y * mask.float()
     # ✗ Graph Break 3: non-tensor data structure (fix: use torch operations)
     # python_list = y.tolist()
     # ✓ Fix: use built-in PyTorch operations
     return y
 # Run and check logs
 x = torch.randn(64, 512).cuda()
 y = problem_model(x) # Check graph break info in output

常见图断裂原因与修复策略如表5-15所示。 表5-15 常见图断裂原因与修复策略 图断裂原因 代码示例 修复方案 Data-dependent if if y.sum() > 0: torch.where(y.sum() > 0, y*2, y) .item() 调用 y[0,0].item() 保持为tensor操作 Python list/dict y.tolist() 使用PyTorch操作 NumPy操作 np.dot(x, w) 使用 torch.matmul print/debug print(y.shape) 移除或使用 torch._logging 第三方C扩展 custom_op(x) 注册为 torch.library

5.8.5 生产环境部署最佳实践

生产部署需兼顾编译配置、预热与缓存复用，以下脚本给出完整流程：

# production_deploy.py - torch.compile production deployment guide
import torch
import torch._dynamo.config as dynamo_config
import torch._inductor.config as inductor_config
# 1. Global configuration optimization
dynamo_config.cache_size_limit = 256 # Increase guard cache size
dynamo_config.accumulated_cache_size_limit = 512
inductor_config.coordinate_descent_tuning = True # max-autotune is better
inductor_config.triton.unique_kernel_names = True
# 2. Warmup pipeline
def warmup_model(model, sample_inputs, num_warmup=10):
    """Full warmup: cover common input shape range"""
    model = torch.compile(model, mode="max-autotune")
    # Phase 1: first compilation (most expensive)

with torch.no_grad():

        _ = model(sample_inputs[0])
    torch.cuda.synchronize()
    # Phase 2: multiple warmups covering different shapes (fill cache)
    for inp in sample_inputs[:num_warmup]:

with torch.no_grad():

            _ = model(inp)
    torch.cuda.synchronize()
    # Inspect Dynamo compilation stats via a legitimate diagnostic API
    from torch._dynamo.utils import compile_times
    print(compile_times())
    return model
# 3. Model export
def export_compiled_model(model, example_inputs, export_path):
    model = torch.compile(model, mode="max-autotune")
    # Export using torch.export
    exported_program = torch.export.export(model, example_inputs)
    torch.export.save(exported_program, export_path)
    # Or use torch.compile + save

with torch.no_grad():

        _ = model(example_inputs)   # Warmup triggers compilation
    # Inductor cache is auto-managed via TORCHINDUCTOR_CACHE_DIR
    # or torch._inductor.config.cache_dir. Compiled artifacts are
    # automatically persisted and reused across runs.
    return exported_program
# 4. Performance monitoring
def monitor_compile_performance(model, inputs):
    """Monitor compilation and runtime performance"""
    import time
    metrics = {

"compile_time_s": 0, "first_run_time_ms": 0, "cached_run_time_ms": 0, "cache_hit_rate": 0,

    }
    # Compilation timing
    t0 = time.time()
    compiled = torch.compile(model)

with torch.no_grad(): _ = compiled(inputs[0]) torch.cuda.synchronize() metrics["compile_time_s"] = time.time() - t0

Cache hit timing

t0 = time.time() with torch.no_grad(): _ = compiled(inputs[0]) # Second call should hit cache torch.cuda.synchronize() metrics["cached_run_time_ms"] = (time.time() - t0) * 1000

5.8.6 性能提升总结

return metrics torch.compile 各应用场景的推荐模式如表5-16所示。 表5-16 torch.compile应用场景推荐 场景 推荐模式 预期加速 注意事项 开发调试 不编译 1.0× 最快迭代 常规训练 default 1.3-1.5× 平衡编译时间和收益 长期训练 (>1天) max-autotune 1.5-2.0× 编译时间可被训练时长摊薄 小batch推理 reduce-overhead 1.4-1.7× CUDA Graph降低Launch开销 生产推理 max-autotune + 离线导出 1.5-2.5× 预编译，无运行时编译开销 动态shape模型 default 1.2-1.4× Guard机制可处理常见变化 多GPU训练 default 1.2-1.4× 与FSDP/TP兼容

上一章

第 4 章 高性能 Kernel
下一章

第 6 章 分布式训练基础

📑 AI Infra 权威指南

1 第 1 章 AI Infra 概述 2 第 2 章 GPU 微结构 3 第 3 章 加速器与异构计算 4 第 4 章 高性能 Kernel 5 第 5 章 编译器与框架 6 第 6 章 分布式训练基础 7 第 7 章 高级并行策略 8 第 8 章 集合通信与网络 9 第 9 章 训练系统工程 10 第 10 章 推理优化与服务 11 第 11 章 模型量化与压缩 12 第 12 章 RAG 与向量检索 13 第 13 章 存储与数据工程 14 第 14 章 集群架构与调度 15 第 15 章 云原生 MLOps 16 第 16 章 性能与成本优化 17 第 17 章 AI Infra 计算 18 第 18 章 AI Infra 实践
← 返回本书大纲