推理引擎的编译优化：从计算图到高效执行

当 PagedAttention 解决了显存碎片问题之后，真正决定单 token 推理延迟下限的，是计算图的编译质量。

一、计算图：推理引擎的中间表示

现代推理引擎（vLLM、TensorRT-LLM、llama.cpp、GGML）虽然架构各异，但都共享同一个计算模型：将有向无环图（DAG）作为推理任务的中间表示（IR），然后对 IR 施加一系列变换，最终生成可在 GPU 上执行的 kernel。

一个典型的 transformer 单层计算图包含以下算子节点：

Input (hidden_states)
    │
    ├──→ LayerNorm ──→ [分支A]
    │                      │
    │                      ▼
    │                 Linear(Q) ──→ RotaryEmbedding ──→ [Q branch]
    │                 Linear(K) ──→ RotaryEmbedding ──→ [K branch]
    │                 Linear(V) ──────────────→ [V branch]
    │                                               │
    │                                               ▼
    │                                          KV Cache (external state)
    │                                               │
    ▼                                               ▼
              Attention(Q, K_cache, V_cache) ──→ Softmax ──→ [scores]
                                                           │
                                                           ▼
                                                  Linear(O_proj) ──→ [output_A]
                                                                          │
                                               [分支A] ──→ Add ───────────┘
                                                           │
                                                           ▼
                                                     FeedForward:
                                                    Linear(up) → GELU → Linear(down)
                                                           │
                                                           ▼
                                                      Add ──→ Output

这张图只是逻辑视图。当 batch_size=1、sequence_length=1（decode 阶段）时，每个 Linear 的矩阵乘法退化为矩阵-向量乘法（GEMV），计算密度极低，瓶颈从算力转为显存带宽。此时编译优化的目标完全不同——不是最大化并行度，而是最小化显存搬运。

1.1 编译优化的核心矛盾

推理引擎编译面临三重矛盾：

矛盾维度	训练优化	推理优化
计算粒度	大矩阵乘法，batch 越大越好	GEMV 主导，算术强度低
显存策略	重计算（activation checkpointing）换显存	预分配 + 静态规划，避免碎片
并行模式	张量并行 + 流水线并行	张量并行 + 连续批处理（continuous batching）

理解这些矛盾是把握推理编译优化的前提。

二、算子融合：降低显存带宽瓶颈的第一法则

算子融合（operator fusion）是将多个相邻的计算节点合并为一个 kernel 的技术。其核心收益是避免中间结果写回 HBM，直接在寄存器或共享内存中流转。

2.1 为什么融合能提升性能

以 FFN 层的 Linear → GELU → Linear 为例，未融合时需要 3 次 kernel launch + 2 次 HBM 中间张量读写。对于隐藏维度 4096、intermediate 维度 11008 的 Llama-2-7B 模型，中间张量大小为 batch × seq × 11008 × 2 bytes ≈ 22KB/token。每个中间结果写入 HBM 再读回，消耗约 22KB × 2 × 152（单层数）/ 1.6TB/s ≈ 4.2ms。而融合的 kernel 只在首尾各访问一次 HBM，节省约 60% 的显存带宽。

2.2 融合的分类

推理引擎中的融合策略可分为三个层次：

垂直融合（Vertical Fusion）——将同一数据流上的连续算子合并：

# 原始：三个独立 kernel
silu_out = silu(x @ gate_proj)        # kernel 1
up = x @ up_proj                       # kernel 2
hidden = silu_out * up @ down_proj     # kernel 3

# 融合后：单个 kernel（SwiGLU-Fused）
def swiglu_fused(x, gate_w, up_w, down_w):
    gate = x @ gate_w                   # 共享内存中间结果
    up = x @ up_w
    hidden = silu(gate) * up @ down_proj
    return hidden

水平融合（Horizontal Fusion）——将共享输入的多个线性投影合并为一个大矩阵乘法。QKV 投影是最经典的例子：

# 未融合：三次 GEMM
Q = X @ W_q      # [1, 4096] × [4096, 4096]
K = X @ W_k      # [1, 4096] × [4096, 4096]
V = X @ W_v      # [1, 4096] × [4096, 4096]

# 融合后：单次 GEMM + split
W_qkv = concat(W_q, W_w, W_v)  # [4096, 12288]
QKV = X @ W_qkv                # 一次大矩阵乘
Q, K, V = split(QKV, dim=-1)

当 batch_size=1 时，这三次 GEMV 被合并为一个 GEMM，每个 SM 从只处理一个小矩阵变为处理一个规整的大矩阵，GPU 利用率提升 2-3 倍。

注意力融合（FlashAttention Fusion）——将 QK^T、softmax、dropout、@V 合并为单个 kernel，彻底避免 attention score 矩阵写回 HBM。FlashAttention-2 在此基础上进一步将 Q 的分块循环移到最外层，最大化共享内存的复用。

2.3 Llama.cpp 中的融合实践

Llama.cpp 的 ggml 后端定义了一组融合算子，位于 ggml-backend.h：

// 融合 silu + mul + matmul（SwiGLU FFN 的 ggml 定义）
struct ggml_tensor * swiglu = ggml_silu_mul(ctx, 
    ggml_mul_mat(ctx, gate_proj, x),      // silu(gate @ W_gate)
    ggml_mul_mat(ctx, up_proj, x)          // up @ W_up
);
// 完整的 FFN：融合为单个 GPU kernel
struct ggml_tensor * ff = ggml_mul_mat(ctx, down_proj, swiglu);

在 CMakeLists.txt 或编译宏控制下，这些融合会映射到具体的 CUDA kernel 实现（如 ggml-cuda/fn_*.cu）。

三、编译期常量折叠：消灭"已知计算"

常量折叠（constant folding）是指在编译阶段将已知结果的计算直接替换为常量。在推理场景中，权重矩阵、scale 因子、位置编码系数都是编译期已知量。

3.1 权重量化折叠

INT8/INT4 权重量化的反缩放因子可以折叠进相邻算子：

# 原始计算（per-channel INT8 量化）
scale = weight_scale[channel_idx]        # 运行时读取
y = (x @ (w_int8 / scale)) + bias       # 先反量化，再计算

# 折叠后：将 scale 融合进下一层的输入 scale
y_scale = input_scale * weight_scale     # 编译期预计算
y_int8 = (x_int8 @ w_int8) >> shift      # 全 INT8 计算

GPTQ、AWQ 等量化方法的编译期预处理本质上就是对量化参数（group scale、zero point）在计算图中做折叠和重排。

3.2 旋转位置编码融合（RoPE Fusion）

RoPE 对 Q/K 施加如下旋转变换：

[θ₁, θ₂, ..., θ_{d/2}]  ← 序列位置 m 决定的频率
Q' = [q₀·cos(θ₁) - q₁·sin(θ₁), q₀·sin(θ₁) + q₁·cos(θ₁), ...]

编译优化中，RoPE 可以有两种融合策略：

与 QKV 投影后融合：将 cos/sin 表预先计算并存入常量内存，在 kernel 中直接完成旋转。由于 θ 只与序列位置 m 和有关，可以按层预计算 [max_seq, head_dim/2] 大小的 LUT，推理时查表+乘加即可。

与投影前融合：更进一步，将旋转矩阵的 cos/sin 分量直接乘到 W_q、W_k 权重上——这等价于对权重做一个位置相关的线性变换。但这种融合牺牲了灵活性（无法跨位置共享权重），所以主流引擎选择第一种。

3.3 LayerNorm 参数折叠

LayerNorm 的 γ（scale）和 β（bias）可以部分折叠进相邻的线性层：

# 原始：y = γ * (x - μ) / √(σ²+ε) + β;  z = W @ y + b
# 折叠：将 γ 合并到 W 的列方向
W_fused = W * γ[num_col]   # 列方向缩放
b_fused = b + W @ β         # 偏置平移（若 β ≠ 0）
# 结果：y' = (x - μ) / √(σ²+ε);  z = W_fused @ y' + b_fused

这消除了推理时读取 γ/β 的显存访问，在 decode 阶段（batch=1）这种访存敏感场景收益可观。

四、静态内存规划：预分配的艺术

推理引擎不需要通用分配器——在引擎启动时，模型形状、最大序列长度、batch 容量都是已知的（或可配的）。这允许编译器做全生命周期的静态内存规划。

4.1 内存规划的基本模型

以 Llama-2-7B（hidden=4096, layers=32, heads=32, head_dim=128, intermediate=11008）为例，令 max_batch=256, max_seq=4096：

张量	形状	大小
权重（FP16）	全部参数	≈ 14GB
KV Cache（每层/每token）	[2, heads, head_dim] × 2 matrices	16 MB/token
KV Cache 总计	256 batch × 4096 seq × 16 MB	≈ 16 GB
激活（峰值）	取决于最大并行 batch+seq	≈ 2-4 GB
显存总计	—	≈ 34 GB

静态规划的核心是：确定每个张量在显存中的起始偏移，使得在推理过程中没有两个活跃期间不重叠的张量共享内存。

4.2 张量生命周期分析

编译器构建活跃区间：

Timeline ──────────────────────────────────────────►

Input Embedding    [──────────]
Attention Mask              [──────────]
KV Cache (layer 0)    [═══════════════════════════════]  持续到第一次使用
QKV Linear               [──]
RoPE                     [──]  
FlashAttention                 [──────────]
O Proj                       [──]
FFN                           [──────────]  
Add + LayerNorm                 [────]
KV Cache (layer 1)         [══════════════════════════]
... (重复 32 层)

[═══] 表示长生命周期（KV Cache），[──] 表示短生命周期（激活张量）。静态规划的做法是：

1. 按活跃区间排序：尽量让短生命周期张量共享内存池
2. KV Cache 和权重常驻，不可复用
3. 激活内存池大小 = 同时活跃的最大激活张量之和

4.3 vLLM 的显存预分配策略

vLLM 的 CacheEngine 在启动时一次性分配最大容量的 KV Cache 显存：

class CacheEngine:
    def __init__(self, model_config, cache_config):
        self.gpu_memory = self._profile_available_memory()  # 实际显存
        self.num_gpu_blocks = self.gpu_memory // self._block_size()
        # 一个大 tensor，内部用 block table 索引
        self.gpu_kv_cache = torch.zeros(
            num_layers * 2,           # K and V
            self.num_gpu_blocks, 
            block_size,
            num_heads, 
            head_dim,
            dtype=torch.float16
        ).cuda()

这里的 block_size（通常 16 tokens/block）和 num_heads、head_dim 在编译期确定，所以 KV Cache 的 shape 完全静态。真正的动态性来自 block 的分配/释放，由 BlockTable 管理，而非底层 tensor 大小变化。

4.4 CUDA Graph 捕获与静态地址

最极致的静态优化是 CUDA Graph：将一次 forward 的所有 kernel 调用序列录制为固定图，推理时直接回放。CUDA Graph 要求所有 buffer 地址固定、所有 kernel 参数确定。

# 录制阶段
with torch.cuda.graph(static_graph, stream=capture_stream):
    static_output = model(static_input)   # 地址固定的 buffer

# 推理阶段（每次 batch 拿不同的 input buffer 地址？）
# 实际上需要每次 copy 新数据到预录制的 static_input buffer
static_input.copy_(dynamic_input)    # 唯一的动态操作
static_graph.replay()                 # 零 launch overhead

TensorRT-LLM 的 inflight batching 中，encoder 部分用 CUDA Graph 录制（shape 固定），decoder 部分用动态 planning（shape 随 batch_size 变化）——混合策略兼顾了 encode 阶段的极致效率和 decode 阶段的灵活性。

五、调度与并行：让 GPU 永不空闲

编译优化不只停留在静态变换——运行时的调度策略同样可以用"编译思维"建模。

5.1 连续批处理与抢占调度

传统的静态批处理（static batching）等所有请求处理完才一起返回，新请求在排队等待。连续批处理（continuous batching / iteration-level scheduling）每次迭代都可以插入/驱逐请求：

def continuous_batch_schedule(self, running_batch):
    # 每轮 iteration 检查
    for req in running_batch:
        if req.is_finished():
            free_blocks(req)
        if req.out_of_blocks():
            preempt(req)  # 驱逐到 CPU，释放 GPU blocks

    # 从 waiting queue 中调度新请求
    while self.gpu_blocks.available > new_request.required_blocks:
        new_req = waiting_queue.pop(0(self.gpu_blocks.available))
        running_batch.add(new_req)

    return running_batch

这种调度的"编译"体现在用预计算的成本模型来决定预驱逐阈值。vLLM 的 scheduling policy（FCFS with preemption）可视为一个静态配置：内存压力超过 threshold → 触发驱逐，驱逐选择 last-served 策略。

5.2 张量并行的图切分

当模型太大单卡放不下时，推理引擎需要在多卡间切分计算图。切分策略对性能有决定性影响：

切分维度	切分方式	通信量	适用场景
hidden dim	列并行（Column Parallel）	AllReduce(QKV output)	FFN 列投影
hidden dim	行并行（Row Parallel）	AllReduce(O_proj output)	FFN 行投影
num_heads	头并行	AllGather(KV after attention)	Attention

TensorRT-LLM 的编译期会自动确定每个算子的并行策略，编译输出为多个 GPU kernel 和 NCCL 通信的交错序列。这种"自动并行"本质上就是编译器对计算图的切分 + 通信插入。

六、实战：从零构建融合 kernel

以下是一个完整的 CUDA kernel 示例，实现 SwiGLU + ResidualAdd 融合：

// swiglu_residual_kernel.cu
#include <cuda_fp16.h>

// FP16 SwiGLU: output = (silu(gate) * up) + residual
// 对应 llama.cpp 中 2/3 * hidden 维度的 SwiGLU
template <int BLOCK_SIZE = 256>
__global__ void swiglu_residual_kernel(
    const half* __restrict__ x,           // 输入 [seq_len, inter_dim]
    const half* __restrict__ residual,    // 残差 [seq_len, hidden_dim]
    half* __restrict__ output,            // 输出 [seq_len, hidden_dim]
    const half* __restrict__ down_weight, // [hidden_dim, inter_dim]
    int hidden_dim,
    int inter_dim
) {
    int row = blockIdx.x;   // token 维度
    int tid = threadIdx.x;  // hidden_dim 内的并行线程
    int half_inter = inter_dim / 2;  // gate 和 up 各一半

    __shared__ half gate_val[BLOCK_SIZE / 2];
    __shared__ half up_val[BLOCK_SIZE / 2];

    extern __shared__ half smem[];

    // 1. Gate 分支
    if (tid < half_inter) {
        float sum = 0.0f;
        for (int k = 0; k < hidden_dim; k += 32) {
            // 分块加载权重
            half w = down_weight[k * half_inter + tid];  // 简化的列主序
            half v = x[row * hidden_dim + k + (tid % 32)];
            sum += __half2float(__hmul(w, v));
        }
        // SiLU: x * sigmoid(x)
        float sig = 1.0f / (1.0f + expf(-sum));
        gate_val[tid] = __float2half(sum * sig);
    }

    __syncthreads();

    // 2. Up 分支
    // ... 与 Gate 类似（实际工程会用 VectorizedLoad）

    // 3. SwiGLU + Residual
    if (tid < hidden_dim) {
        float result = __half2float(output_load);
        // ... 完整的矩阵乘
        half res = residual[row * hidden_dim + tid];
        output[row * hidden_dim + tid] = __float2half(result + __half2float(res));
    }
}

// 启动配置
void launch_swiglu_residual(
    const half* x, const half* residual, half* output,
    const half* down_weight, int tokens, int hidden, int inter
) {
    int threads = 256;
    int blocks = tokens;
    int smem = (inter / 2) * sizeof(half) * 2; // gate + up
    swiglu_residual_kernel<<<blocks, threads, smem>>>(
        x, residual, output, down_weight, hidden, inter
    );
}

这段代码揭示了推理编译优化的精髓：算法、并行策略、显存层级三者必须协同优化。融合 SwiGLU + Residual 不只是写出代码，更需要考虑：

• 权重是按 [out, in]（行主序）还是 [in, out]（列主序）排布
• 共享内存够不够同时容纳 Q 块和 KV 块
• 如果 hidden_dim=4096，一个 warp 32 线程能覆盖一行还是只覆盖 32 列
• FP16 的尾数位只有 10-bit，累加过程中是否溢出

七、生态全景：主流推理引擎的编译策略

引擎	编译模型	融合策略	内存规划	并行支持
TensorRT-LLM	TensorRT 编译 + 自定义 plugin	极致融合（FlashAttention + GEMM + Residual）全匹配	静态预分配 + CUDA Graph	TP + PP
vLLM	PyTorch eager + PagedAttention（动态）	torch.compile（ inductor 后端）自动融合	Paged KV Cache（块级别稀疏）	TP
llama.cpp / GGML	手写融合 kernel per arch	硬编码融合（SwigLM + KV copy + FFN）	全静态预分配	TP（实验性））
MLC-LLM	TVM 编译优化	Ansor 自动调优融合算子	静态规划 + KV Cache 换出	TP
tgi (Hugging Face)	Python + Rust 混合	Selective fusion（FlashAttention + GEMM）	Paged KV Cache	TP + Continuous Batching

核心差异在于"自动" vs "手写"：TensorRT-LLM 和 MLC-LLM 依赖自动调优来搜索最优 kernel；vLLM 依赖 torch.compile 在运行时融合；llama.cpp 用工程师的手写 kernel 来保证每种硬件上的性能上限。不同路径的 trade-off 本质上是 编译灵活性 vs 峰值性能 的取舍。

八、前沿方向

8.1 自动算子融合搜索

TVM Ansor 和 Meta 的 AOT Inductor 正在做编译期自动融合：给定计算图和硬件描述，搜索最优的融合矩阵分块 + 线程映射策略。2024 年的进展已经能在 Hopper（H100）上自动匹配手写 kernel 的 95% 性能。

8.2 混合精度计算图编译

Blackwell 架构（B100/B200）引入了 microscaling 格式（MXFP4/MXFP6），推理编译器需要在图上自动决定哪些层用 FP8、哪些用 MXFP4、哪些保持 FP16，使得精度损失 < 0.1% 且计算速度最大化。这类似于硬件综合中的资源分配问题。

8.3 投机解码（Speculative Decoding）的编译优化

投机解码中，draft model 和 target model 的计算图需要合并调度。LLM.d（AuxLossFree）和 DeepMind Sequoii 等方法通过将 draft 的推理提前到 target forward 的流水线气泡中，实现编译期优化的指令级并行（ILP）。

8.4 WASM 边缘推理编译

WAMR 和 wasm Micro Runtime 正在探索用 WASI-NN API 将推理 engine 编译为 WASM。这意味着整个推理引擎的计算图编译 + kernel 融合要让位于 WASM IR 变换 ——一层新的"编译的编译"。

九、总结

推理引擎的编译优化不是某一项孤立技术，而是一个系统工程：

• 算子融合 消除显存带宽瓶颈
• 常量折叠 消除冗余计算
• 静态内存分配 消除运行时碎片
• CUDA Graph 消除 CPU launch overhead
• 自动并行插入 释放多卡算力

当 PagedAttention 解决了显存外部碎片之后，最后这张拼图就是"让每次推理的每一步都精确落在 GPU 的算力、带宽、延迟的最优平衡点上"。而这一点，靠的不是更快的硬件，靠的是编译器的计算图优化。

优秀的推理引擎不只是"跑模型"——它是一台编译器，把 GPU 的硬件语义精确映射到 transformer 的数学语义上。