Luminal 拆解一个“用搜索代替启发式”的推理编译器

01它要解决的问题

部署一个模型，要面对三个相乘的维度：模型结构（LLM、扩散、MoE、CNN……）、硬件（H100、B200、Apple GPU、AMD……）、数据类型（fp32、bf16、fp8、nvfp4……）。主流框架的做法是为高频组合手写 kernel（FlashAttention、FlashInfer、cuBLAS、CUTLASS 模板），再用编译器 pass 把模型“降”到这些 kernel 上。

Luminal 的 README 和 why.mdx 指出了两个问题：

• 组合爆炸：每加一个设备，就要把框架支持的每个算子都实现一遍；再乘上 dtype，代码量几乎呈指数增长。
• 破坏式重写的复杂度：XLA、TVM、torch.compile 的优化由大量“单向”改写组成。改写一旦执行，原图就没了，所以每条规则只能在确定有收益时才触发。判断“有没有收益”需要越来越多的特判和启发式，最后还是要靠 Triton / Pallas 这种 DSL 手工找回性能。

02核心思路：三根支柱

① RISC 式的小原语集（HLIR）

README 说一切最终都归结为 15 个原语：Log2 Exp2 Sin Sqrt Recip、Add Mul Mod LessThan、SumReduce MaxReduce Iota Gather Scatter Cast。注意几个刻意的选择：只有以 2 为底的 Exp2/Log2（exp(x)=exp2(x·log2e)），没有减法和除法（a-b = a + b·(-1)，a/b = a·recip(b)），矩阵乘也不是原语，而是“广播 → 逐元素乘 → 沿 k 求和”。

当前分支的逻辑层比 README 多：src/logical/op/ 下有约 36 个 Logical* 算子，增加了 Select、And/Or/Not、Floor/Ceil/Round/Trunc、TruncDiv/TruncRem、三种 scan 等。主要原因是对齐 PyTorch 的精确语义（整数取整方向、比特级 bool），见 associative-fold.md 里关于 prod = exp(sum(log x)) 在负数上出错的讨论。

② 非破坏式重写：e-graph（egglog）

所有优化都写成等价关系，而不是“替换”。e-graph 同时保存一个计算的所有等价写法，规则可以按任意顺序、反复触发，没有“先做哪个 pass”的问题。规则引擎是 egglog（Datalog + e-graph），Luminal 用自己的 fork：luminal-ai/egglog。

③ 以实测运行时间为代价函数的搜索

e-graph 里可能有天文数字种组合。Luminal 不写代价模型，而是采样候选、生成代码、在目标 GPU 上真跑，用遗传算法迭代。HN 上作者的话：“The best heuristic is no heuristic.”

另外两个配套原则：一切 AOT（整个网络是静态图，编译器有全局视野），以及原生的符号维度（(s, 4096)、(b, h, w+3)），用分桶（bucket）为不同的动态取值范围分别搜索。

03整体管线

spec 里写着当前分支的管线是 recorder → egglog_core → runtime 自己负责的 extraction/search → bufferize → runtime。展开如下：

04深入：e-graph 与 egglog

为什么“非破坏式”很关键

经典例子（来自 egg 论文）：表达式 (a×2)/2，有两条规则：强度削减 x×2 → x≪1，以及消去 (x×y)/y → x。用传统 pass 管线时，如果强度削减先跑，表达式变成 (a≪1)/2，消去规则就再也匹配不上了。这就是 phase ordering 问题：结果取决于 pass 的顺序，于是只能写启发式来决定顺序和触发条件。

一条真实的规则：把“广播-乘-求和”识别为矩阵乘

矩阵乘不是原语，它是一个模式。下面这条规则把这个模式和一个辅助节点 LogicalHelperMatrixMultiply 合并到同一个 e-class，后端规则只需要匹配这一个节点：

; src/logical/helper/matrix_multiply/recognize.egg
(rule
  ((= ?out  (LogicalReduceSum ?prod 0))
   (= ?prod (LogicalMul ?a_b ?b_b))
   (= ?a_b  (LogicalHelperBroadcastAxis ?a 1 ?n))
   (= ?b_b  (LogicalHelperBroadcastAxis ?bT 2 ?m))
   (= (shape-of ?prod) (ShapeLit [?m ?n ?k]))      ; 形状把两条腿绑到同一个 [m,n,k]
   (= (shape-of ?a)    (ShapeLit [?m ?k]))
   (= (shape-of ?bT)   (ShapeLit [?n ?k])))
  ((union ?out (LogicalHelperMatrixMultiply ?a (LogicalHelperMatrixTranspose ?bT))))
  :ruleset main_ruleset
  :name "helper/matrix_multiply/recognize: reduce-sum of broadcast muls is a MatrixMultiply")

关键字是 union，不是“替换”。另外有一条方向相反的 expand.egg，把矩阵乘展开回原语链。两个方向同时存在，正是非破坏式的意思。博客里提到一个好处：模型作者没有显式写 matmul 的地方，只要出现这个模式，也会被识别出来。

可交互：一个 e-graph，三种 genome

下图是上面矩阵乘的 e-graph。虚线框是 e-class（一组等价计算），实线框是 e-node（一种写法）。搜索里的 genome 就是“每个 e-class 选哪个 e-node”。点按钮看不同 genome 选出的实现：

05深入：布局层（逻辑值 × 物理布局）

Luminal 把“算什么”和“数据在内存里怎么摆”分成两件事。逻辑 IR 里的张量没有物理布局（spec：只有输入输出有）；在 e-graph 里，LayoutTensorLit(逻辑值, 布局) 把两者配对，同一个逻辑值可以有很多种布局写法，而且它们都在 e-graph 里互相等价。

后端规则只匹配硬件能直接用的那种写法。例如 CUDA 的 ReduceSum kernel 写的是 out[i]，所以规则要求输出布局是行主序紧凑，并且映射是单射：

; crates/luminal_cuda_lite/src/ops/reduce_sum/match_functional.egg（节选）
(= ?out_layout (RightMajorContiguousElementLayoutLit ?out_rm_shape ?out_rm_bits))
(= ?out_logical (LogicalReduceSum ?in_logical ?reduction_dim))
(!= ?value_dtype (Int))                         ; 整数单独一条规则：要处理回绕语义
(= (injectivity-of ?out_layout_tensor) (Injective))
; ⇒ union 成 LayoutTensorOpReduceSumGeneric（一个具体 CUDA kernel）

注释写道：“The e-graph joins equivalent spellings; the matcher matches the hardware's literal one.” 也就是说：等价性交给 e-graph，硬件约束交给匹配条件。转置、切片这类“视图”在这一层只是布局的改变，可以零拷贝；需要实体化时再插入 MaterializeLayoutCopy，是否插入同样由搜索决定。

06深入：搜索（遗传算法 + 真机实测）

为什么不像 egg 那样用代价函数做贪心提取？因为在 kernel 层面代价不可加：同一个 matmul，用 cuBLASLt 本身更快，但如果它要求的布局导致前后各多一次拷贝，整体可能更慢；融合、内存复用、launch 开销都是全局性质。Luminal 的选择是：不建模，直接测。

几个细节值得注意：

• 无环保证：e-graph 本身可以有环（x ≡ x+0），genome 的选择必须构成 DAG。采样和变异都维持这个不变式（mutate_genome 里的 flip_closes_cycle）。
• 相关采样：源码注释写道，独立地为每个位置抽选择，在大型重复图上“几乎永远抽不出一个连贯的实现”，因为每个候选里总会混进几条巨大的实体化路线。于是还有一种抽法：随机排一个“实现族”顺序，整个 genome 共用这个顺序（sample_genome_correlated）。
• “描述”与“决定”分离：core 只负责“把 genome 变成图”和“怎么抽 genome”，不做任何选择；怎么给候选定价、循环怎么跑，归每个 runtime 自己（lib.rs 顶部注释）。
• 搜索不负责正确性：spec 明确规定，所有可提取的图都必须等价，搜索只比快慢，不允许“绕开”一个已知错误的候选。错了就去修规则。

// crates/luminal_cuda_lite/src/search.rs
impl Default for CompileOptions {
    fn default() -> Self { Self {
        generations: 8, generation_size: 8, mutations: 2, trials: 3,
        seed: 0, keep_finalists: 4, candidate_timeout: None,
        device_budget_bytes: None, /* 默认取可用 CUDA 显存 */ .. } }
}

07深入：厂商库也是搜索空间里的普通候选

“不写手写 kernel”不等于“不用手写 kernel”。cuBLASLt 调用在 Luminal 里就是一种 e-node，和生成的通用 kernel 平起平坐，由搜索决定用不用。难点在于：库调用的参数空间（转置、主序、leading dimension、batch stride、精度、epilogue）很大，而且每种组合在什么条件下合法，必须被严格证明。

当前做法分四步，每步一组 egglog 规则（cublaslt/README.md）：

1. 算术识别：认出 D = α·op(A)·op(B) + β·C 加可选 epilogue（bias、ReLU、GELU）。这一步只是“命名”，不代表允许执行。
2. 布局证明：证明某个存储张量确实按某个 vendor 描述符（7 个字段：dtype、主序、行、列、ld、batch 数、batch stride）寻址。
3. 精度证明：只有作者显式写出的精确扩宽才允许直接读 BF16 / FP16；存储格式不能决定计算精度。
4. 准入：上面这些事实再和白名单里的某一条完整配置对上，才生成可执行节点。白名单有 2,176 条完整配置，文件 approved_whitelist.egg 有 21,787 行。空白名单就什么都不生成。

和 FlashInfer 的关系：README 的“Where are we”里写着 “Many kernel libraries supported in the search space (FlashInfer, cuBLASLt, etc.)”，这是旧管线的状态。BRINGUP.md 把旧 crate 描述为一个“零件库”（cuBLASLt / FlashInfer / MoE、CUDA Graph），并说明新管线是按 CL-3 阶段逐个迁回来的。目前迁回来的只有 cuBLASLt。环境变量 LUMINAL_FLASHINFER_DIR 等仍留在 README 中。

08深入：内存规划与动态形状

DPS + one-shot bufferization

提取出来的图里，中间值还没有存储。Luminal 先把每个算子改写成“目标传递”风格（DPS：输出缓冲区作为参数传入），再用 MLIR 风格的 one-shot bufferization 决定原地复用还是新分配。冲突引擎检查三件事：可写性（不能写只读绑定）、写后读（所有读必须在 DAG 上可证明地先于覆盖写）、跨操作数共享许可（只有算子自己声明了 Sharing::May 才允许）。被拒的原地复用一律“修复”成新分配加必要的拷贝，所以拒绝永远不是错误。

显存剪枝

搜索之前，CUDA 后端会对 e-graph 做一次强制的内存剪枝：在整个声明的动态维度范围内，物理容量超过显存预算的张量连同它们的生产者实现一起删掉，边界张量不能删（egraph_postpass.rs）。这样搜索根本不会抽到“需要 200 GB 中间结果”的路线。这属于“删掉不可行的候选”，不违反“不在 Rust 里做模式匹配和选择”的规定。

符号维度与分桶

形状是整数表达式（c + q、ceil(n/128)），egglog 前导程序里有一整套带值域的整数代数：截断除法与向上取整除法分成不同构造器，除法只在能证明除数非零时才参与合并。运行时给出 dyn_map（如 {q:1, c:812}），按所在 bucket 分发到对应编译结果。spec 要求：在代表点上测出来的最优计划，必须在整个声明范围内都合法。

09自动“发现” FlashAttention：说法与现状

这是 Luminal 最出名的一句话：不靠手写算子，靠搜索自动找到 FlashAttention 这类复杂优化。按作者在 Show HN 帖子中的解释，路径是：重写规则里既有循环 / 分块结构的改写，也有代数改写，其中关键一条是 softmax → online softmax。两者组合之后，搜索空间里就出现了“按 KV 分块、在块间重缩放、S 矩阵不落显存”的写法，也就是 FlashAttention 的结构。

这也揭示了这条路线的一个本质约束：搜索只能找到规则能到达的写法。SEARCH_DIAGNOSIS 原话：“Increasing generations/population explores combinations of existing producers; it does not extend the rewrite vocabulary.” 所以能力上限取决于规则库，而“该写哪些规则”本身仍然需要人的洞察。HN 讨论中也有人指出，等价保持的搜索很难找到 Strassen 这类“换算法”的优化。

10难点与解法总表

11现状：仓库自带的实测数字

以下数据来自仓库中的验证报告（2026-09-21，NVIDIA H200，batch 1，FP32，128 输入 / 128 输出 token，3 次热请求取中位数）：

估算一下差距：Llama3-8B 的 FP32 权重约 8.03 B × 4 B ≈ 32 GB，H200 的 HBM 带宽约 4.8 TB/s，所以 decode 每 token 的带宽下限约 6.7 ms。实测 221 ms，约为下限的 33 倍，带宽利用率约 3%。这是我的粗略估算，不是仓库给出的数字，也忽略了 KV cache 读取，但量级足以说明：新管线目前关注的是正确性和架构，性能还没跟上。

搜索本身也很贵。仓库的一次计时分析里，Llama3-8B 的搜索用了 347 s，其中 60.5% 花在 pinned 主机内存的分配、清零和释放上，并不是在 GPU 上计时。下面是其中一个候选的耗时分解（按比例绘制）：

正确性方面比较扎实：Metal 报告里，Llama3 / Qwen3 / Gemma3 / Qwen3-MoE 的小型实例与 ReferenceRuntime 的 logits 误差都小于 1e-4；CUDA 上四个模型 128 个 token 输出前后一致。README 中 “Q8 Llama 3 8B ~80% of theoretical max on H100”，以及 HN 帖子中 “M 系列 Mac 上 Q8 Llama 3 8B 15–25 tok/s”，对应的是旧管线（作者在 HN 上也承认当时还比不上 llama.cpp 这类手调实现）。

12与其他编译器 / 系统对比

13评价与借鉴

14源码阅读路线（约 3 小时）

1. 想法（20 分钟）：README → spec.md（注意它已过时，但契约部分仍然成立）→ AGENTS.md 的 Debugging 和 Rewrite Boundary 两节。
2. 一个算子的完整生命周期（40 分钟）：logical/op/add/ 的 10 个 .egg 文件 → CUDA 侧 cuda_lite/ops/add/ → KernelOp 怎么生成源码。
3. 模式识别（20 分钟）：matrix_multiply/ 的识别与展开 → cublaslt/README.md。
4. 提取与搜索（40 分钟）：Genome → sample_genome / mutate_genome → search_implementations。
5. 内存（30 分钟）：bufferize.rs 文件头的长注释（冲突引擎讲得很清楚）。
6. 诚实的现状（20 分钟）：SEARCH_DIAGNOSIS.md 和 BRINGUP.md。
7. 动手：cargo test -p luminal --lib（不需要 GPU）；在 5090 上可以试 cargo run --release -p luminal_cuda_lite --example llama3 --features device，加 LUMINAL_LOG=1 看搜索过程，加 LLIR_DUMP_DIR=… 导出选中的计划。