AI 编译技术的下一个突破点：自动 Kernel 生成、稀疏计算支持与异构编译器统一-育师

AI 编译技术的下一个突破点：自动 Kernel 生成、稀疏计算支持与异构编译器统一

一、从手写 CUDA Kernel 到让编译器自己写

GPU 编程的现状是割裂的。调试一个 matmul kernel 三个月，换个显卡架构又得重来。团队人手不足以维护每个算子的多个变体。

AI 推理场景的算子数量正在指数级增长。FlashAttention 之后，PagedAttention、RadixAttention、TreeAttention 接踵而至。每个新注意力机制都需要底层 kernel 配合。手写速度跟不上论文产出速度。

问题的本质不在工程能力，而在抽象层次。当前 GPU 编程仍处于"手工汇编"时代——开发者直接控制 warp、shared memory 和寄存器。这不是生产力的正确方向。编译技术应当接管这些底层细节。

自动 Kernel 生成、稀疏计算编译支持和异构编译器统一，是下一代 AI 编译器的三个核心突破方向。

二、原理剖析：从搜索到生成的范式迁移

自动 Kernel 生成的核心是搜索空间与代价模型。TVM 的 AutoScheduler 将调度原语（tiling、binding、unrolling）建模为一个搜索问题。Triton 走向另一条路——通过 block-level 编程模型，让编译器自动处理 thread-level 映射。两者的共同点：将优化决策从程序员转移到编译器。

Triton 的关键设计在于其编程模型介于 CUDA 和纯自动生成之间。开发者只需描述 block 级别的计算逻辑，编译器自动展开为 warp 和 thread 级别的代码。这个取舍在生产力上取得了最佳平衡——既不需要像 TVM 那样完全自动搜索（搜索时间可能数小时），也不需要像 CUDA 那样手工管理每个线程。

稀疏计算编译是长期被忽视的领域。结构化稀疏（2:4 pattern）已被 Ampere 架构硬件支持，但编译工具链远未成熟。当前稀疏 kernel 仍以手工实现为主。SparseTIR 尝试将稀疏格式（CSR/BSR/ELL）的转换与优化集成到编译流程中，但还处于研究阶段。

稀疏计算的难点在于：稀疏模式在运行时动态变化。权重剪枝后的稀疏模式在编译时已知，但激活值的稀疏性是输入相关的。编译时无法预测。这要求编译器生成多种稀疏格式的代码路径，在运行时根据实际稀疏度做选择。

异构编译器统一的最大推手是 MLIR。MLIR 提供多层 IR（dialect），允许不同抽象级别的优化共存。Linalg dialect 处理线性代数，GPU dialect 处理设备代码生成，LLVM dialect 完成最终的机器码生成。这种分层设计使得从 PyTorch 到各种硬件的编译路径可以复用大量优化 pass。

IREE 在 MLIR 之上构建了完整的编译和运行时栈。其核心价值在于：同一套编译流水线可以输出到 CUDA、Vulkan、Metal 甚至 WebGPU。对于需要同时支持云侧 GPU 推理和端侧 NPU 推理的团队，这意味着不必维护两套编译器。

三、代码实践：用 Triton 实现可移植的 Flash Attention

// 基于 Triton 的 Flash Attention 简化实现 // triton::autotune 宏告诉编译器：以下 kernel 的参数空间由我自动搜索 // 设计原因：避免手工调 block_size 和 num_warps，让编译器在 {64,128,256} x {4,8} 空间中搜索最优配置 use triton_rs::*; #[triton::autotune( configs = "attention_configs", key = ["BLOCK_SIZE"], prune_configs_by = { "early_config_prune": true } )] fn flash_attention_kernel( q: &Tensor, // [B, H, N, D] k: &Tensor, // [B, H, N, D] v: &Tensor, // [B, H, N, D] o: &mut Tensor, // [B, H, N, D] sm_scale: f32, block_size: i32, ) { // 每个 program 处理一个 (batch, head) 组合 // 设计原因：在 block 级别做 online softmax，避免全局 reduction let pid = program_id(0); let num_blocks = (seq_len + block_size - 1) / block_size; let batch_idx = pid / (num_heads * num_blocks); let head_idx = (pid / num_blocks) % num_heads; let block_idx = pid % num_blocks; // 加载 Q 的当前 block 到 SRAM // 设计原因：SRAM 带宽是 HBM 的 10x+，block_size 需匹配 SRAM 容量 let q_block = load_block(&q, batch_idx, head_idx, block_idx, block_size); // Online softmax 状态 // 设计原因：避免 O(N²) 显存占用，每步只保留 running max 和 running sum let mut m_i = vec![-f32::INFINITY; block_size as usize]; let mut l_i = vec![0.0f32; block_size as usize]; let mut acc = vec![0.0f32; (block_size * head_dim) as usize]; // 分块遍历 K, V for start_kv in (0..seq_len).step_by(block_size as usize) { let k_block = load_block(&k, batch_idx, head_idx, start_kv, block_size); // 计算 QK^T * scale — 当前 block 的注意力分数 // sm_scale = 1/sqrt(d_k) 防止点积过大导致 softmax 梯度消失 let scores = matmul(&q_block, &k_block.transpose()) * sm_scale; let scores = scores.clamp_max(0.0); // causal mask // Online softmax 更新 let m_new = max(&m_i, &scores.row_max()); let p = (scores - &m_new).exp(); let alpha = (&m_i - &m_new).exp(); l_i = l_i * alpha + p.row_sum(); // 累积加权 value let v_block = load_block(&v, batch_idx, head_idx, start_kv, block_size); let correction = alpha.unsqueeze(-1); acc = acc * correction + matmul(&p, &v_block); m_i = m_new; } // 最终归一化并写回 HBM // 设计原因：只在 block 计算完成时才写回，减少 HBM 访问 let result = acc / l_i.unsqueeze(-1); store_block(&mut o, &result, batch_idx, head_idx, block_idx); }

这段代码的核心在于online softmax 的数值稳定性处理。传统 softmax 需要三次遍历（max→exp→normalize），online 版本通过维护 running max 和 running sum，将三次遍历合并为一次。m_new - m_old作为 correction factor 保证了数值精度。

在 Triton 中，这段代码编译后自动映射到 GPU 的 block grid。编译器负责将block_size分派到 warp、将load_block展开为 coalesced memory access。开发者不再需要手写threadIdx.x和blockIdx.x。