Lean4-TileLang 张量程序超级优化器：用形式化推理自动搜索 GPU Kernel，实现 Attention 4.08x 加速

项目概览

• 核心项目：Lean4-TileLang 张量程序超级优化器（开发中）
• 项目地址：https://leloykun.github.io/ponder/lean4-tilelang/
• 内容规模：5000 字，阅读约 18 分钟，播客约 25 分钟
• 技术关键词：Lean4 形式化推理、TileLang GPU Kernel 生成、张量程序超优化器、算子融合与显存优化、Attention 4.08x 加速

内容导读

本文旨在深入解读 Franz Louis Cesista 在博客文章《Lean4-TileLang Tensor Program Superoptimizer [WIP]》中阐述的核心思想，并联合 tile-ai/tilelang 仓库，全面剖析 Lean4、TileLang、张量计算图、GPU Kernel 优化以及超优化器之间的内在联系与协同工作机制。

原文揭示了一个正处于开发阶段的创新系统：该系统首先利用 Lean4 形式化地描述神经网络的计算图，然后自动搜索在数学上等价但执行效率更高的实现方案，最终将这些方案编译为 TileLang Kernel，并在 GPU 上高效执行。

本文将力求用最直观的语言，阐明该系统的工作原理、其存在的意义，以及其中几个典型 Kernel 的优化策略。

文章目录

• 一、整体结论先行
• 二、TileLang 的本质
• 三、Lean4 的引入价值
• 四、自动化优化的必要性
• 五、实验结果速览
• 六、Attention：逼近 FlashAttention 的优化路径
• 七、SwiGLU：矩阵乘与激活函数的深度融合
• 八、Matmul：基础 GEMM 的优化空间有限
• 九、RMSNorm：并非所有变换都能提速
• 十、RMSNorm-MLP：融合后收益重现
• 十一、案例分析揭示的规律
  • 1. 显存带宽比计算更关键
  • 1. 算子融合是性能提升的主要来源
  • 1. 基础算子的独立优化空间有限
  • 1. 自动搜索依赖精准的成本模型
• 十二、该系统的长期价值
• 十三、在 TileLang 生态中的定位
• 十四、一个形象化的比喻
• 总结

一、整体结论先行

根据《Lean4-TileLang Tensor Program Superoptimizer》的内容，可将整个系统概括为：

使用 Lean4 形式化地定义“需要计算什么”，借助超优化器自动搜索“如何等价地改写”，最后通过 TileLang 生成“能够在 GPU 上高效运行的 Kernel”。

换言之，这并非单纯地手写几个高性能 Kernel，而是试图构建一条自动化的优化链路：

神经网络算子的数学定义
↓
Lean4 计算图
↓
等价变换与搜索
↓
生成 TileLang Kernel
↓
JIT 编译
↓
在 GPU 上运行并基准测试

在此流程中，各组件的职责如下：

简而言之：Lean4 确保“计算正确”，TileLang 负责“运行高效”。

二、TileLang 的本质

tile-ai/tilelang 是一个专用于编写高性能 GPU/CPU Kernel 的领域专用语言（DSL）。

其设计目标是让开发者能够使用相对简洁的 Python 风格代码，编写出性能接近 CUDA、CUTLASS、TVM 等底层系统的 Kernel。

在 tile-ai/tilelang 仓库中，TileLang 支持的典型算子包括：

• GEMM / Matmul
• Dequant GEMM
• FlashAttention
• Linear Attention
• Convolution
• DeepSeek MLA
• 其他复杂 AI Kernel

TileLang 代码看起来类似于 Python，但它并非普通 Python。它会通过即时编译（JIT）转换为底层 Kernel。其常见写法及通俗解释如下：

可以将 TileLang 理解为：

一种比 PyTorch 更底层、比 CUDA 更易写的 Kernel 编程语言。

PyTorch 更擅长表达模型逻辑，而 TileLang 更擅长表达“这个算子在 GPU 上应该如何分块、如何搬运数据、如何使用共享内存、如何调用 Tensor Core”。

三、为什么需要 Lean4？

既然 TileLang 已经能够编写高性能 Kernel，为何还需要引入 Lean4？

关键在于：GPU Kernel 优化常常涉及数学上的等价变换。

在标准的 Attention 实现中，代码通常如下所示：

scores = Q @ Kᵀ
probs = softmax(scores)
out = probs @ V

然而，高性能的 FlashAttention 类实现并不会将完整的 scores 矩阵写入显存。相反，它采用分块计算策略，并借助在线 softmax 技术来维护中间状态。

这里必须确保一个关键点：

分块在线计算的结果

与以下计算的结果在数学上完全等价：

完整 scores + softmax + matmul 的结果

另一个例子是 RMSNorm 中的变换：

out = (X * S) @ W

在特定条件下，它可以被重写为：

out = (X @ W) * S

这里的 S 是按行广播的缩放因子。如果条件不满足，这个变换就会出错。

Lean4 的作用，正是为这些计算图和变换提供一个更加严谨的表达基础。

它能够帮助系统清晰地描述以下内容：

• 输入和输出的形状（shape）；
• 轴（axis）的约束；
• 归约（reduction）的语义；
• 广播（broadcast）的规则；
• 矩阵乘法的收缩轴（contraction axis）；
• 逐元素操作的细节；
• 哪些变换是合法的；
• 哪些计算图在逻辑上是等价的。

因此，在这个系统中，Lean4 扮演的是“数学与推理层”的角色，而 TileLang 则负责“高性能执行层”。

四、为何需要这种自动化优化？

如今，神经网络模型日益复杂，新结构层出不穷，例如新的 Attention 变体、MLP 结构、Norm 方法以及门控机制等。

如果每出现一种新结构，都需要手动编写一套 CUDA kernel，就会面临几个问题：

1. 开发成本极高；
2. 需要非常专业的 GPU kernel 编写经验；
3. 极易出错；
4. 难以快速比较大量候选结构；
5. 小规模实验结果可能无法反映大规模训练的真实表现。

正如博客原文开头所述，作者并不看好仅在小规模上进行“代码 fuzzing”式的自动化研究，因为小规模表现良好并不代表大规模下也能同样出色。要公平地比较不同架构，至少需要同时考虑：

• GPU / TPU kernel 是否优化充分；
• optimizer 是否选择得当；
• 参数化方式是否合理；
• 超参数如何随宽度、深度、batch size 和训练步数变化；
• 小规模规律能否迁移到大规模场景。

因此，自动生成高性能 kernel 只是更大目标中的一个环节：

用形式化系统定义神经网络结构
        ↓
自动生成高性能 kernel
        ↓
自动调整 optimizer / parametrization
        ↓
研究 scaling law
        ↓
更公平地比较不同架构

本文主要聚焦于解读其中“Lean4 → TileLang kernel”这一阶段。

五、实验结果速览

根据原文提供的 benchmark 数据，几个 workload 的结果如下：

从中可以得出几个直观结论：

1. Attention 的加速效果非常显著，大约提升了 4 倍；
2. SwiGLU 也有明显提升，大约为 2 倍；
3. Matmul 的提升相对有限，因为矩阵乘法本身已经得到了高度优化；
4. RMSNorm 单独测试时反而变慢了；
5. 但将 RMSNorm 与 MLP 融合后，性能又得到了提升。

这揭示了一个重要事实：

TileLang 和自动优化器的优势，更多地体现在“算子融合”和“减少显存读写”上，而不仅仅是把单个基础算子写得更快。

六、Attention：接近 FlashAttention 的优化思路

Attention 的数学形式如下：

scores = Q @ Kᵀ
probs = softmax(scores)
out = probs @ V

其对应的 Lean4 描述大致为：

let q ← input “Q” shapeHTqDh
let k ← input “K” shapeHTkvDh
let v ← input “V” shapeHTkvDh
let scores ← matmul q k “d_h” “d_h”
let probs ← softmaxSubdag scores “t_kv”
let out ← matmul probs v “t_kv” “t_kv”
output out

这表达的是标准的 Attention 计算图。

问题在于，scores 的规模会非常庞大。以原文中的配置为例：

h = 16
t_q = 4096
t_kv = 4096
d_h = 128

那么 scores 的形状将是：

16 × 4096 × 4096

如果完整生成这个矩阵，再写入显存，然后读回来做 softmax，最后与 V 相乘，显存流量将会非常巨大。

Derivation of top-1 kernel

高性能实现的关键在于，避免将完整的 scores 矩阵落地到显存中。

生成的 TileLang kernel 采用了类似 FlashAttention 的分块在线计算方式：

固定一块 Q
循环读取 K/V 的小块
每次计算局部 QKᵀ
做局部 softmax 统计
更新全局 softmax 状态
累加输出

TileLang 代码中有几个关键的状态变量：

state_pass0_m
state_pass0_l
state_pass0_o

它们的含义可以这样理解：

为什么需要这些状态？

因为 softmax 不能简单地分块各自计算。为了保证分块计算与完整 softmax 的等价性，必须维护全局的最大值和归一化因子。

这个在线 softmax 的过程可以理解为：

每来一块新的 K/V
就把它纳入当前 softmax 统计
同时修正旧的累积结果

这样就不需要存储完整的 scores 矩阵。

这正是 FlashAttention 系列算法的核心思想：

用更多的片上计算，来换取更少的显存读写。

因此，针对 Attention 这一场景，我们实现了大约 4 倍的性能提升，并且原文中特别指出，这一效果可与 Flash Attention 2 相媲美。

七、SwiGLU：将两个矩阵乘法与激活函数融为一体

SwiGLU 是大语言模型 MLP 中一种非常常见的结构。它可以简化为以下计算步骤：

up = X @ W_up  
gate = X @ W_gate  
out = up * silu(gate)  

其中，silu 函数的定义为：

silu(x) = x * sigmoid(x)  

在 Lean4 的计算图描述中，它被表示为：

let x ← input "X" shapeMD  
let wUp ← input "W_up" shapeDN  
let wGate ← input "W_gate" shapeDN  
let gatePre ← matmul x wGate "d" "d"  
let up ← matmul x wUp "d" "d"  
let gate ← silu gatePre  
let out ← mul gate up  
output out  

如果使用普通的框架来执行，可能会产生多个中间张量，例如：

gatePre  
up  
gate  
out  

这些中间张量每次都需要写回显存，然后再被读取出来进行下一步操作。

最优内核的推导

TileLang 内核的做法是将这些步骤全部融合在一起：

1. 在同一个内核中计算 X @ W_up；
2. 在同一个内核中计算 X @ W_gate；
3. 直接对 gate 执行 SiLU 操作；
4. 直接与 up 相乘；
5. 最后，只将最终结果写回显存。

其核心代码形态类似于：

T.gemm(input_0, input_4, matmul_8, clear_accum=False)    # up  
T.gemm(input_0, input_10, matmul_14, clear_accum=False)  # gate  

随后，融合 SiLU 和乘法操作：

matmul_14[i0, i1] = matmul_14[i0, i1] / (1.0 + T.exp2((-(matmul_14[i0, i1])) * scale))  
matmul_14[i0, i1] = matmul_8[i0, i1] * matmul_14[i0, i1]  

这里还运用了一个 fast math 技巧：

scale = 1.44269504  # log2(e)

原因在于：

exp(x) = exp2(x * log2(e))

在 GPU 上，exp2 通常更容易实现高效的计算。

SwiGLU 的性能提升主要来自两个方面：

1. 两个矩阵乘法（matmul）共享同一个输入 tile；
2. SiLU 和乘法操作不再需要单独将中间结果写入显存。

因此，该场景获得了大约 2.14 倍的加速。

八、Matmul：基础 GEMM 的优化空间相对有限

矩阵乘法（Matmul）是最基础的运算：

O = X @ W

其 Lean4 计算图非常简洁：

let x ← input "X" shapeMD  
let w ← input "W" shapeDN  
let out ← matmul x w "d" "d"  
output out  

最优内核的推导

生成的 TileLang 内核是一个标准的分块（tiled）GEMM：

把 X 的一块搬到 shared memory  
把 W 的一块搬到 shared memory  
用 T.gemm 做 tile-level 矩阵乘  
沿 d 维循环累加  
最后写回 O  

其核心代码类似：

for k_pass0 in T.Pipelined(T.ceildiv(d, block_d), num_stages=num_stages):  
T.copy(X[gy * block_m, k_pass0 * block_d], input_0)  
T.copy(W[k_pass0 * block_d, gx * block_n], input_4)  
T.gemm(input_0, input_4, matmul_8, clear_accum=False)  

它使用了典型的 GPU GEMM 优化手段，包括：

• 分块（tiling）；
• 共享内存（shared memory）；
• 片段累加（fragment accumulation）；
• 流水线循环（pipelined loop）；
• 交错（swizzle）；
• 张量核心 GEMM（tensor core GEMM）。

然而，矩阵乘法本身就是 GPU 上被优化得最为充分的算子之一。像 cuBLAS、CUTLASS、PyTorch 后端以及 torch.compile 等都已经进行了大量的优化工作。

因此，TileLang 版本仅比 torch.compile 快大约 1.29 倍，这并不令人意外。

由此可以得出一个重要结论：

自动生成 TileLang 内核的最大价值，并不一定在于单独击败成熟的 GEMM 库，而在于能够将 GEMM 与其他操作有效地融合起来。

九、RMSNorm：正确的变换不一定更快

RMSNorm 的简化公式如下：

S = 1 / sqrt(mean(X²) + eps)  
norm = X * S  
out = norm @ W  

在 Lean4 中，首先描述 RMSNorm 的子图：

let x2 ← square x  
let mean ← meanAlongSubdag x2 axis  
let eps ← epsConst  
let denom ← add mean eps  
let scale ← rsqrt denom  
mul x scale  

接着是一个矩阵乘法：

let out ← matmul norm w "d" "d"  

自动生成的 TileLang 实现将其拆分为两个内核。

第一个内核负责计算每一行的 scale 值：

S[gy * block_m + i0] = 1.0 / T.sqrt(row_state[i0] / d + eps)  

第二个内核执行矩阵乘法，并在最后乘上 scale：

T.gemm(X_tile, W_tile, acc_0, clear_accum=False)  

for i0, i1 in T.Parallel(block_m, block_n):  
acc_0[i0, i1] = acc_0[i0, i1] * S[gy * block_m + i0]  

这里运用了一个代数变换：

(X * S) @ W = (X @ W) * S  

只要 S 是一个按行广播的缩放因子，这个变换就是合理的。

然而，在性能上却出现了问题。

因为这种实现方式需要…

深度重写与降重版本

在执行一个 Kernel 计算 S 之后，需要将其结果写回显存。随后，必须再次启动一个新的 Kernel 来完成矩阵乘法（matmul）操作，并从显存中重新读取之前存储的 S 数据。最后，将这个读取到的 S 值乘到最终的输出结果上。

对于像 RMSNorm 这样结构相对简单的算子而言，额外增加一次 Kernel 启动（kernel launch）以及中间数据的读写操作，很可能完全抵消掉优化所带来的性能提升。

最终的测试结果也证实了这一点：

TileLang: 0.206 ms  
torch.compile: 0.174 ms  
Speedup: 0.842x  

最优 Kernel 的推导过程

这意味着，在这个具体案例中，TileLang 生成的版本反而比 torch.compile 更慢。

这揭示了一个关键问题：

Superoptimizer 虽然能够找到数学上完全等价的实现方案，但数学上的等价性，并不等同于在实际运行中拥有更优的性能。

真实的 GPU 性能会受到诸多复杂因素的共同影响，其中包括：

• Kernel 启动（launch）的开销
• 显存带宽与读写延迟
• 共享内存（shared memory）的利用效率
• 寄存器（register）资源压力
• 占用率（occupancy）水平
• 缓存（cache）的命中率
• 底层编译器优化的质量

十、RMSNorm-MLP：融合后的性能增益重新显现

RMSNorm-MLP 是一个更贴近实际模型结构的组合算子，其计算流程如下：

norm = rmsnorm(X)  
up = norm @ W_up  
gate = norm @ W_gate  
out = up * gate  

用 Lean4 图表示，可以理解为：

let norm ← rmsNormSubdag x "d"  
let up ← matmul norm wUp "d" "d"  
let gate ← matmul norm wGate "d" "d"  
let out ← mul up gate  
output out  

在 TileLang 的实现中，仍然需要首先为每一行计算 S 值：

S = scale_kernel(values['X'], eps_value)  

接着，在第二个 Kernel 中，同时执行两个矩阵乘法运算：

T.gemm(X_tile, W_gate_tile, acc_0, clear_accum=False)  
T.gemm(X_tile, W_up_tile, acc_1, clear_accum=False)  

然后，将计算出的 scale 值乘到结果上：

acc_0[i0, i1] = acc_0[i0, i1] * S[gy * block_m + i0]  
acc_1[i0, i1] = acc_1[i0, i1] * S[gy * block_m + i0]  

最后，执行逐元素乘法操作：

acc_0[i0, i1] = acc_0[i0, i1] * acc_1[i0, i1]  

尽管这里依然存在单独计算 S 的开销，但第二个 Kernel 融合了更多的计算任务：

• X @ W_gate
• X @ W_up
• 两次 RMSNorm 的 scale 操作
• 最终的逐元素乘法

这种融合方式显著减少了多个中间张量在显存中的读写次数，使得性能收益重新超过了开销。

因此，RMSNorm-MLP 获得了大约 1.91 倍的加速效果。

这个例子很好地说明了：

单独一个小算子可能并不值得进行专门的优化，但当多个算子融合在一起后，特化 Kernel 的优势会得到显著放大。

十一、从这些案例中能总结出哪些规律？

综合多个工作负载（workload）的测试结果，我们可以观察到一些通用的规律。

1. 显存读写往往比计算本身更为关键

Attention 和 SwiGLU 获得加速的主要原因，并非因为减少了多少浮点运算量（FLOPs），而是由于显著减少了中间张量在显存中的读写操作。

GPU 的计算速度很快，但显存带宽是有限的。因此，高性能 Kernel 的核心目标之一就是：

让数据尽量留在寄存器或共享内存中  
不要频繁写回 global memory  

2. 算子融合是性能提升的主要来源

在这些例子中，性能收益最大的部分通常来自于融合操作：

Attention: matmul + softmax + matmul  
SwiGLU: matmul + matmul + silu + mul  
RMSNorm-MLP: rmsnorm + matmul + matmul + mul  

融合得越充分，就越有机会省去中间张量的存储和读取。

3. 基础算子的单独优化空间有限

矩阵乘法（Matmul）本身已经有非常成熟的优化库。自动生成 Kernel 能够取得一点优势已经不错，但很难轻松实现大幅度的领先。

真正有价值的优化方向是：

围绕 Matmul 做图级融合  

而不是仅仅优化孤立的 GEMM 操作。

4. 自动搜索需要一个优秀的成本模型（cost model）

RMSNorm 变慢的案例说明，搜索到的“数学等价实现”并不等同于“最快实现”。

一个优秀的 superoptimizer 不仅要了解数学上的等价变换是否可行，还需要准确评估：

• 多启动一个 Kernel 是否值得；
• 多写入一个中间张量是否值得；
• 融合后寄存器的压力是否会过高；
• 共享内存是否足够使用；
• Tensor Core 是否能被充分利用；
• 块大小（block size）是否合适。

因此，形式化推理和性能建模这两者都至关重要。

十二、这类系统的长期价值

根据原文的设想，这套基础设施的目标并不仅仅是优化几个孤立的算子，而是要为更大规模的自动化 AI 研究流程服务。

未来，如果要比较大量新型模型结构，理想的流程可能是：

在 Lean4 中定义模型结构  
↓  
自动生成候选计算图  
↓  
自动验证等价变换  
↓  
自动生成 TileLang kernel  
↓  
自动 benchmark  
↓  
自动调整 optimizer 和超参数  
↓  
研究 scaling law  

这样做的重要意义在于：

1. 降低新架构的 Kernel 开发成本；
2. 避免因 Kernel 优化不到位而误判架构的优劣；
3. 更快地探索新的 Attention、MLP、Norm 设计方案；
4. 让实验从“小规模代码试错”升级为更系统的 scaling law 比较；
5. 为 AI 编译器引入形式化验证能力。

如果这个方向走向成熟，未来的研究者可能只需要描述数学结构，系统就能自动生成性能足够高的实现方案。

十三、结合 TileLang 来看它的定位

tile-ai/tilelang 仓库本身提供的是高性能 Kernel 编程和编译能力。它有点像一个“面向 AI Kernel 的底层执行语言”。

而 Lean4 superoptimizer 则更像是 TileLang 上方的一层“自动生成器”。

二者之间的协同关系可通过以下架构图清晰呈现：

Lean4 / Superoptimizer
职责范围：
– 计算图描述
– 等价变换推导
– 搜索 derivation 路径
– 选择实现策略

↓ 生成

TileLang
职责范围：
– shared memory tiling
– fragment accumulation
– pipelining
– swizzle
– T.gemm
– JIT 编译
– autotune
– 后端代码生成

↓ 编译运行

GPU / CPU / Accelerator

由此可见，TileLang 在这个体系中扮演的是后端执行者的角色：

它将上层搜索得出的计算方案，转化为真实可运行且具备高性能的 kernel。

unsetunset十四、一个形象比喻unsetunset

我们可以把整个系统类比为自动化烹饪过程。

传统 CUDA 优化的方式好比名厨亲自掌勺，每一步都依赖经验积累。

而 Lean4 + TileLang 的核心理念是：

用户只需描述想做什么菜，系统自动推理哪些步骤可以合并、如何切块、怎样下锅，最终生成一套高效的操作流程。

unsetunset总结unsetunset

基于 Lean4-TileLang Tensor Program Superoptimizer[3] 的研究成果，其核心思想可概括为：

利用 Lean4 严谨地描述张量计算，通过 superoptimizer 搜索等价实现方案，再借助 TileLang 生成高性能的 GPU kernel。

实验数据表明：

• Attention：约 4.08 倍加速；
• SwiGLU：约 2.14 倍加速；
• Matmul：约 1.29 倍加速；
• RMSNorm：反而比 torch.compile 更慢；
• RMSNorm-MLP：约 1.91 倍加速。

这些结果揭示了以下几点：

1. 算子融合是性能提升的主要来源；
2. 减少显存读写操作至关重要；
3. 单独优化基础算子的空间有限；
4. 自动搜索仍需更完善的性能模型；
5. Lean4 与 TileLang 的组合有望让 AI kernel 优化变得更自动化、更可靠。

最简洁的总结是：

Lean4 让系统理解“哪些计算变换是正确的”，TileLang 则确保这些计算“在 GPU 上能够高效运行”。

参考资料
[1] Lean4-TileLang Tensor Program Superoptimizer [WIP]: https://leloykun.github.io/ponder/lean4-tilelang/
[2] Lean4-TileLang Tensor Program Superoptimizer: https://leloykun.github.io/ponder/lean4-tilelang/
[3] Lean4-TileLang Tensor Program Superoptimizer: https://leloykun.github.io/ponder/lean4-tilelang/