一个可定制的编译器，用于为AI模型生成高效的融合GPU内核 [P]
摘要
作者介绍了一款用 Python 编写、高度可定制且易于修改的 ML 编译器。该编译器通过多级 IR 流水线将 LLMs 转换为优化的 CUDA 内核，在特定操作上实现了与 PyTorch 相当甚至更优的性能。文章详细阐述了该编译器的优化过程、降级规则以及用于生成高效融合 GPU 内核的 CLI 用法。
现代 ML（LLM）编译器栈非常复杂繁重。TVM 包含超过 50 万行 C++ 代码。PyTorch 则将 Dynamo、Inductor 和 Triton 层层叠加。我从头构建了一个可 Hack 的 LLM 编译器，并正在记录这一过程。它接收小型模型（TinyLlama、Qwen2.5-7B），并通过六层 IR（中间表示）将其降阶为一系列 CUDA 核函数。目前，在 RTX 5090 上，生成的 FP32 核函数运行速度为 **几何平均 PyTorch eager 的 1.11 倍** 以及 **torch.compile 的 1.20 倍**，在 seq=128 时于 TinyLlama-128 和 Qwen2.5-7B 上实现全块性能持平。在小型规约 / SDPA / KV 投影上取得优势（最高达 4.7 倍）；在 seq=512 的稠密矩阵乘法上略有落后。[第一部分](https://medium.com/data-science-collective/a-principled-ml-compiler-stack-in-5-000-lines-of-python-17f2db9549d4) 端到端地走通了 RMSNorm 层，并详细梳理了该管线的前半部分。本文作为第二部分将填补这一空白，并深入讲解 Tile IR、Kernel IR 及其相关的降阶规则。完整文章：[用 5000 行 Python 构建一个原则性 ML 编译器栈](https://medium.com/data-science-collective/a-principled-ml-compiler-stack-in-5-000-lines-of-python-2337d208f278)
本文重点介绍如何为以循环嵌套形式编写的操作生成 GPU 调度（Loop IR）。RMSNorm 示例：
```python
v0 = reciprocal(2048)
for a0 in 0..32: # free
for a1 in 0..2048: # reduce
in2 = load x[0, a0, a1]
v1 = multiply(in2, in2)
acc0 <- add(acc0, v1)
v2 = multiply(acc0, v0)
v3 = add(v2, 1e-06)
v4 = rsqrt(v3)
for a2 in 0..2048: # free
in3 = load x[0, a0, a2]
in4 = load p_weight[a2]
v5 = multiply(in3, v4)
v6 = multiply(v5, in4)
merged_n0[0, a0, a2] = v6
```
该编译器栈模拟了 CUDA 工程师在优化核函数时会执行的一系列优化步骤：将输入排期（stage）到 smem、减少内存银行冲突、提高占用率（occupancy）等。
```diff
LoopOp │ ▼
[001] tileify — 将外层自由循环提升为线程轴
[002] chunk_matmul_k — 将 K 维规约分块为 K-outer × K-inner（CTA 内部）
[003] split_matmul_k — 将 K-outer 分块循环提升为网格维度
[004] cooperative_reduce — 协同规约：让多个线程共享一次规约；使用 Combine 进行树形归并
[005] blockify_launch — 配置启动块：选择块维度；将自由轴划分为 BLOCK 和 THREAD
[006] chunk_reduce — 分块规约：对非矩阵乘法的规约操作进行分块，使其加载操作适配共享内存
[007] stage_inputs — 输入排期：将热输入数据块提升（hoist）至 Stage 节点
[008] register_tile — 寄存器分块：复制内层分块，使每个线程独占一个寄存器块
[009] permute_register_tile — 寄存器分块重排：重排寄存器条带，使可能发生银行冲突的加载操作落在较远的列上
[010] double_buffer — 双缓冲：将 K-outer Stage 提升为 BufferedStage（乒乓缓冲）
[011] tma_copy — TMA 拷贝：将符合条件的 BufferedStage 细化为 TmaBufferedStage (sm_90+)
[012] split_inner_for_swizzle — 为交错分割内层：拆分 TmaBufferedStage 的内层缓存轴以进行交错（swizzle）
[013] async_copy — 异步拷贝：将剩余部分细化为 AsyncBufferedStage (cp.async, sm_80+)
[014] pad_smem — 共享内存填充：填充共享内存步幅以打破银行冲突
[015] pipeline_k_outer — K-outer 流水线化：将 K-outer 循环轮转为前导（prologue）/稳态（steady-state）/尾声（epilogue）阶段 (cp.async + TMA)
[016] mark_unroll — 标记展开：为小型内层循环添加 #pragma unroll 注释
│ ▼
TileOp (fully scheduled)
```
每个阶段都可以通过 CLI 命令复现。例如，`stage_inputs` 通道会在条件允许且能带来性能收益时（即输入在 CTA 内被多次读取），将输入缓冲区排期到 smem。要查看此效果，可使用以下命令：
```bash
deplodock compile \
-c "torch.nn.RMSNorm(2048)(torch.randn(1,32,2048))" \
--ir tile -vv \
| awk '/^>>> t:007/,/^<<< t:007/'
```
```diff
>>> t:007_stage_inputs
@@ matched at rms_norm (in-place) @@
@@ -2,6 +2,7 @@
v0 = reciprocal(2048)
Tile(axes=(a0:256=THREAD, a1:32=BLOCK)):
+ x_smem = Stage(x, origin=(0, a1, 0), slab=(a2:2048@2))
StridedLoop(a2 = a0; < 2048; += 256): # reduce
- in2 = load x[0, a1, a2]
+ in2 = load x_smem[a2]
v1 = multiply(in2, in2)
acc0 <- add(acc0, v1)
@@ -11,5 +12,5 @@
v4 = rsqrt(v3)
StridedLoop(a2 = a0; < 2048; += 256): # free
- in3 = load x[0, a1, a2]
+ in3 = load x_smem[a2]
in4 = load p_weight[a2]
v5 = multiply(in3, v4)
<<< t:007_stage_inputs
```
RMSNorm 层最终生成的 CUDA 核函数：
```bash
deplodock compile \
-c "torch.nn.RMSNorm(2048)(torch.randn(1,32,2048))" \
--target sm_120 --ir cuda
```
```c
extern "C" __global__ __launch_bounds__(256) void k_rms_norm_reduce(
const float* x,
const float* p_weight,
float* rms_norm) {
float v0 = 1.0f / 2048.0f;
int a1 = blockIdx.x;
int a0 = threadIdx.x;
int lane = threadIdx.x & 31;
int warp = threadIdx.x >> 5;
float acc0 = 0.0f;
__shared__ float x_smem[2048];
for (int x_smem_flat = a0; x_smem_flat < 2048; x_smem_flat += 256) {
float x_smem_v = x[a1 * 2048 + x_smem_flat];
x_smem[x_smem_flat] = x_smem_v;
}
__syncthreads();
for (int a2 = a0; a2 < 2048; a2 += 256) {
float in2 = x_smem[a2];
float v1 = in2 * in2;
acc0 += v1;
}
float acc0_w = acc0;
acc0_w = acc0_w + __shfl_xor_sync(0xffffffff, acc0_w, 16);
acc0_w = acc0_w + __shfl_xor_sync(0xffffffff, acc0_w, 8);
acc0_w = acc0_w + __shfl_xor_sync(0xffffffff, acc0_w, 4);
acc0_w = acc0_w + __shfl_xor_sync(0xffffffff, acc0_w, 2);
acc0_w = acc0_w + __shfl_xor_sync(0xffffffff, acc0_w, 1);
__shared__ float acc0_smem[8];
if (lane == 0) {
acc0_smem[warp] = acc0_w;
}
__syncthreads();
for (int s = 4; s > 0; s >>= 1) {
if (warp < s) {
acc0_smem[warp] = acc0_smem[warp] + acc0_smem[warp + s];
}
__syncthreads();
}
float acc0_b = acc0_smem[0];
float v2 = acc0_b * v0;
float v3 = v2 + 1e-06f;
float v4 = rsqrtf(v3);
for (int a2 = a0; a2 < 2048; a2 += 256) {
float in3 = x_smem[a2];
float in4 = p_weight[a2];
float v5 = in3 * v4;
float v6 = v5 * in4;
rms_norm[a1 * 2048 + a2] = v6;
}
}
```
相似文章
KForge：面向AI加速器的LLM驱动跨平台内核生成
KForge是一个跨平台框架，利用两个协作的基于LLM的智能体，自动生成和优化适用于多种AI加速器的高性能计算内核，在NVIDIA B200和Intel Arc B580硬件上实现了显著的加速效果。
数据并行内核的微型编译器
一篇博客文章，描述了一个微型编译器，演示了如何通过将for循环转换为带有通道和掩码的向量化循环来降低数据并行内核，实现代码约180行Python。
使用CUDA内核重写模型推理：瓶颈不仅仅是GEMM [P]
作者描述了构建FlashRT的过程，这是一个以CUDA为核心的推理运行时，通过使用C++/CUDA内核重写模型推理路径，来解决小批量/实时工作负载中超出GEMM的瓶颈，在Jetson Thor和RTX 5090上实现了显著的延迟改进。文章讨论了关于精度的经验（FP8有帮助，FP4好坏参半）以及绕过通用运行时进行实时推理的必要性。
Kuma：将PyTorch模型编译为自包含的WebGPU可执行文件 [P]
Kuma是一个编译器/运行时，它将导出的PyTorch模型编译成自包含的WebGPU可执行文件，从而实现无需Python或服务器依赖的直接浏览器推理。
@aryanvs_: 在外行人看来，这也许只是噪音。但这里蕴藏着数月来编写编译器的努力，最好的部分…
一位开发者分享了数月来构建一个编译器的成果，该编译器在A100 GPU上的矩阵乘法性能超越了cuBLAS，并附带了可视化效果。