@leloykun：[进行中] 关于 Lean4-to-TileLang 张量程序超级优化器的博文：

X AI KOLs Following 2026/05/12 05:48 工具

摘要

一篇技术博文介绍了一种 Lean4-to-TileLang 张量程序超级优化器，能自动生成优化的 GPU/TPU 内核与超参数缩放规律，展示了相较 torch.compile 的性能提升。

[进行中] 关于 Lean4-to-TileLang 张量程序超级优化器的博文：https://t.co/9dQ8xyxAlj

查看原文

查看缓存全文

缓存时间: 2026/05/12 12:53

[WIP] Lean4到TileLang的张量程序超级优化器的博客文章在此：https://t.co/9dQ8xyxAlj

Lean4-TileLang 张量程序超级优化器 [进行中]

来源：https://leloykun.github.io/ponder/lean4-tilelang/

1. 概述

我对那些仅仅进行代码模糊测试的 AI 研究中的“自动研究”方法持看空态度。它们通常优化错了方向：在小规模下表现良好的性能很少能迁移到更大规模（例如，去掉权重衰减）。相反，我们应该更相信“苦涩教训”（bitter-lesson-pilled），去比较缩放规律。

但在进行架构搜索时，只有当我们至少同时优化 GPU/TPU 内核、优化器选择、超参数等，并弄清楚如何在规模扩大时调整它们，进行比较缩放规律才是公平的。这是昂贵且耗时的部分，而且目前我们必须为每个新架构都这样做！有些内核在权重和激活满足某些约束时也更稳定（有时甚至更快），而这些约束我们通常可以通过优化器选择和参数化来强制执行。因此，理想情况下，所有这些都应该联合优化。

通过这个项目，我现在拥有了形式化的基础设施，使我（或我的智能体）可以在 Lean4 中定义神经网络架构，然后自动获得：

经过优化的 IO 感知加速器内核。
能够使超参数在不同宽度和深度之间迁移的优化器选择和参数化。更多细节，请参阅我之前的博客文章。
超参数缩放规律，告诉我们如何在批量大小、训练时长、数据集大小等规模变化时调整超参数。更多细节，请参阅steepest-descent-lean。
优化器的低秩代理，以加速小规模下的超参数调优，并将其迁移到全秩情况（我们有一篇即将发表的论文，敬请期待）。

2. 结果

工作负载	案例	TileLang (ms)	torch.compile (ms)	加速比	与先前工作的等效性？
attention	h16_tq4096_tkv4096_dh128	0.531712	2.168960	4.079x	Flash Attention 2
swiglu	m1024_n2048_d2048	0.056797	0.121536	2.140x
matmul	m1024_d4096_n4096	0.125324	0.162112	1.294x
rmsnorm	m1024_n4096_d4096	0.206554	0.174016	0.842x	FlashNorm
rmsnorm_mlp	m1024_n1024_d1024	0.037851	0.072384	1.912x	-

3. 每个工作负载的 Top-1 内核

3.1. Attention

3.1.1. Lean4 源代码

def softmaxSubdag (scores : NodeId) (axis : Axis) : DagM NodeId := do
  let weights ← exp scores
  let denom ← red .sum weights axis
  div weights denom

def attentionGraph : Graph :=
  buildGraphWithOutput "attention" do
    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

3.1.2. Top-1 内核的推导

3.1.3. Top-1 内核的 TileLang 代码

def build_kernel(
    d_h: int = 128,
    h: int = 16,
    t_kv: int = 4096,
    t_q: int = 4096,
    block_t_kv: int = 128,
    block_t_q: int = 128,
    threads: int = 256,
    num_stages: int = 2,
    enable_swizzle: bool = True,
    enable_autotune: bool = False,
    autotune_warmup: int = 10,
    autotune_rep: int = 10,
    autotune_timeout: int = 100,
):
    dtype = T.float16
    accum_dtype = T.float32
    scale = 1.44269504  # log2(e)
    fast_math_pass_configs = {tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}
    jit_decorator = tilelang.jit(out_idx=[-1], pass_configs=fast_math_pass_configs)
    if enable_autotune:
        def decorate(fn):
            return tilelang.autotune(configs=get_configs(), warmup=autotune_warmup, rep=autotune_rep, timeout=autotune_timeout, skip_check=True)(jit_decorator(fn))
    else:
        def decorate(fn):
            return jit_decorator(fn)

    @decorate
    def attention_c2349_jit(d_h: int = d_h, h: int = h, t_kv: int = t_kv, t_q: int = t_q, block_t_kv: int = block_t_kv, block_t_q: int = block_t_q, threads: int = threads, num_stages: int = num_stages, enable_swizzle: bool = enable_swizzle):
        @T.prim_func
        def main(
            Q: T.Tensor((h, t_q, d_h), dtype),
            K: T.Tensor((h, t_kv, d_h), dtype),
            V: T.Tensor((h, t_kv, d_h), dtype),
            O: T.Tensor((h, t_q, d_h), dtype),
        ):
            with T.Kernel(T.ceildiv(t_q, block_t_q), h, threads=threads) as (gx, gy):
                input_0 = T.alloc_shared((block_t_q, d_h), dtype)
                input_4 = T.alloc_shared((block_t_kv, d_h), dtype)
                matmul_8 = T.alloc_fragment((block_t_q, block_t_kv), accum_dtype)
                red_max_9 = T.alloc_fragment((block_t_q,), accum_dtype)
                input_15 = T.alloc_shared((block_t_kv, d_h), dtype)
                matmul_19 = T.alloc_fragment((block_t_q, d_h), accum_dtype)
                state_pass0_o = T.alloc_fragment((block_t_q, d_h), accum_dtype)
                red_sum_22 = T.alloc_fragment((block_t_q,), accum_dtype)
                state_pass0_l = T.alloc_fragment((block_t_q,), accum_dtype)
                state_pass0_m = T.alloc_fragment((block_t_q,), accum_dtype)
                scale_old_state_pass0_m = T.alloc_fragment((block_t_q,), accum_dtype)
                scale_tile_state_pass0_m = T.alloc_fragment((block_t_q,), accum_dtype)
                cast_lhs_19 = T.alloc_fragment((block_t_q, block_t_kv), dtype)
                T.annotate_layout({
                    input_0: make_swizzle_layout(input_0),
                    input_4: make_swizzle_layout(input_4),
                    input_15: make_swizzle_layout(input_15),
                })
                T.use_swizzle(panel_size=10, enable=enable_swizzle)
                T.fill(state_pass0_m, -T.infinity(accum_dtype))
                T.clear(state_pass0_l)
                T.clear(state_pass0_o)
                T.copy(Q[gy, gx * block_t_q, 0], input_0)
                for k_pass0 in T.Pipelined(T.ceildiv(t_kv, block_t_kv), num_stages=num_stages):
                    T.copy(K[gy, k_pass0 * block_t_kv, 0], input_4)
                    T.gemm(input_0, input_4, matmul_8, clear_accum=True, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
                    T.reduce_max(matmul_8, red_max_9, dim=1, clear=True)
                    for i0, i1 in T.Parallel(block_t_q, block_t_kv):
                        matmul_8[i0, i1] = matmul_8[i0, i1] - red_max_9[i0]
                        matmul_8[i0, i1] = T.exp2((matmul_8[i0, i1]) * scale)
                    T.reduce_sum(matmul_8, red_sum_22, dim=1, clear=True)
                    T.copy(V[gy, k_pass0 * block_t_kv, 0], input_15)
                    T.copy(matmul_8, cast_lhs_19)
                    T.gemm(cast_lhs_19, input_15, matmul_19, clear_accum=True, policy=T.GemmWarpPolicy.FullRow)
                    for i0 in T.Parallel(block_t_q):
                        scale_old_state_pass0_m[i0] = state_pass0_m[i0]
                        state_pass0_m[i0] = T.max(scale_old_state_pass0_m[i0], red_max_9[i0])
                        scale_old_state_pass0_m[i0] = T.exp2((scale_old_state_pass0_m[i0] - state_pass0_m[i0]) * scale)
                        scale_tile_state_pass0_m[i0] = T.exp2((red_max_9[i0] - state_pass0_m[i0]) * scale)
                        state_pass0_l[i0] = state_pass0_l[i0] * scale_old_state_pass0_m[i0] + red_sum_22[i0] * scale_tile_state_pass0_m[i0]
                    for i0, i1 in T.Parallel(block_t_q, d_h):
                        state_pass0_o[i0, i1] = state_pass0_o[i0, i1] * scale_old_state_pass0_m[i0] + matmul_19[i0, i1] * scale_tile_state_pass0_m[i0]
                for i0, i1 in T.Parallel(block_t_q, d_h):
                    state_pass0_o[i0, i1] = state_pass0_o[i0, i1] / state_pass0_l[i0]
                T.copy(state_pass0_o, O[gy, gx * block_t_q, 0])
        return main

    return attention_c2349_jit()

3.2. SwiGLU

3.2.1. Lean4 源代码

def swigluGraph : Graph :=
  buildGraphWithOutput "swiglu" do
    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

3.2.2. Top-1 内核的推导

3.2.3. Top-1 内核的 TileLang 代码

def build_kernel(
    d: int = 2048,
    m: int = 1024,
    n: int = 2048,
    block_d: int = 128,
    block_m: int = 128,
    block_n: int = 128,
    threads: int = 256,
    num_stages: int = 2,
    enable_swizzle: bool = True,
    enable_autotune: bool = False,
    autotune_warmup: int = 10,
    autotune_rep: int = 10,
    autotune_timeout: int = 100,
):
    dtype = T.float16
    accum_dtype = T.float32
    scale = 1.44269504  # log2(e)
    fast_math_pass_configs = {tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}
    jit_decorator = tilelang.jit(out_idx=[-1], pass_configs=fast_math_pass_configs)
    if enable_autotune:
        def decorate(fn):
            return tilelang.autotune(configs=get_configs(), warmup=autotune_warmup, rep=autotune_rep, timeout=autotune_timeout, skip_check=True)(jit_decorator(fn))
    else:
        def decorate(fn):
            return jit_decorator(fn)

    @decorate
    def swiglu_c517_jit(d: int = d, m: int = m, n: int = n, block_d: int = block_d, block_m: int = block_m, block_n: int = block_n, threads: int = threads, num_stages: int = num_stages, enable_swizzle: bool = enable_swizzle):
        @T.prim_func
        def main(
            X: T.Tensor((m, d), dtype),
            W_up: T.Tensor((d, n), dtype),
            W_gate: T.Tensor((d, n), dtype),
            O: T.Tensor((m, n), dtype),
        ):
            with T.Kernel(T.ceildiv(n, block_n), T.ceildiv(m, block_m), threads=threads) as (gx, gy):
                input_0 = T.alloc_shared((block_m, block_d), dtype)
                input_4 = T.alloc_shared((block_d, block_n), dtype)
                matmul_8 = T.alloc_fragment((block_m, block_n), accum_dtype)
                input_10 = T.alloc_shared((block_d, block_n), dtype)
                matmul_14 = T.alloc_fragment((block_m, block_n), accum_dtype)
                T.annotate_layout({
                    input_0: make_swizzle_layout(input_0),
                    input_4: make_swizzle_layout(input_4),
                    input_10: make_swizzle_layout(input_10),
                })
                T.use_swizzle(panel_size=10, enable=enable_swizzle)
                T.clear(matmul_8)
                T.clear(matmul_14)
                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_up[k_pass0 * block_d, gx * block_n], input_4)
                    T.gemm(input_0, input_4, matmul_8, clear_accum=False)
                    T.copy(W_gate[k_pass0 * block_d, gx * block_n], input_10)
                    T.gemm(input_0, input_10, matmul_14, clear_accum=False)
                for i0, i1 in T.Parallel(block_m, block_n):
                    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]
                T.copy(matmul_14, O[gy * block_m, gx * block_n])
        return main

    return swiglu_c517_jit()

3.3. Matmul

3.3.1. Lean4 源代码

def matmulGraph : Graph :=
  buildGraphWithOutput "matmul" do
    let x ← input "X" shapeMD
    let w ← input "W" shapeDN
    let out ← matmul x w "d" "d"
    output out

3.3.2. Top-1 内核的推导

3.3.3. Top-1 内核的 TileLang 代码

def build_kernel(
    d: int = 4096,
    m: int = 1024,
    n: int = 4096,
    block_d: int = 128,
    block_m: int = 128,
    block_n: int = 128,
    threads: int = 256,
    num_stages: int = 2,
    enable_swizzle: bool = True,
    enable_autotune: bool = False,
    autotune_warmup: int = 10,
    autotune_rep: int = 10,
    autotune_timeout: int = 100,
):
    dtype = T.float16
    accum_dtype = T.float32
    scale = 1.44269504  # log2(e)
    fast_math_pass_configs = {tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}
    jit_decorator = tilelang.jit(out_idx=[-1], pass_configs=fast_math_pass_configs)
    if enable_autotune:
        def decorate(fn):
            return tilelang.autotune(configs=get_configs(), warmup=autotune_warmup, rep=autotune_rep, timeout=autotune_timeout, skip_check=True)(jit_decorator(fn))
    else:
        def decorate(fn):
            return jit_decorator(fn)

    @decorate
    def matmul_c23_jit(d: int = d, m: int = m, n: int = n, block_d: int = block_d, block_m: int = block_m, block_n: int = block_n, threads: int = threads, num_stages: int = num_stages, enable_swizzle: bool = enable_swizzle):
        @T.prim_func
        def main(
            X: T.Tensor((m, d), dtype),
            W: T.Tensor((d, n), dtype),
            O: T.Tensor((m, n), dtype),
        ):
            with T.Kernel(T.ceildiv(n, block_n), T.ceildiv(m, block_m), threads=threads) as (gx, gy):
                input_0 = T.alloc_shared((block_m, block_d), dtype)
                input_4 = T.alloc_shared((block_d, block_n), dtype)
                matmul_8 = T.alloc_fragment((block_m, block_n), accum_dtype)
                T.annotate_layout({
                    input_0: make_swizzle_layout(input_0),
                    input_4: make_swizzle_layout(input_4),
                })
                T.use_swizzle(panel_size=10, enable=enable_swizzle)
                T.clear(matmul_8)
                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)
                T.copy(matmul_8, O[gy * block_m, gx * block_n])
        return main

    return matmul_c23_jit()

3.4. RMSNorm

3.4.1. Lean4 源代码

def meanAlongSubdag (x : NodeId) (axis : Axis) : DagM NodeId := do
  let s ← red .sum x axis
  let n ← sizeConst axis
  div s n

def rmsNormSubdag (x : NodeId) (axis : Axis) : DagM NodeId := do
  let x2 ← square x
  let mean ← meanAlongSubdag x2 axis
  let eps ← epsConst
  let denom ← add mean eps
  let scale ← rsqrt denom
  mul x scale

def rmsnormGraph : Graph :=
  buildGraphWithOutput "rmsnorm" do
    let x ← input "X" shapeMD
    let w ← input "W" shapeDN
    let norm ← rmsNormSubdag x "d"
    let out ← matmul norm w "d" "d"
    output out

3.4.2. Top-1 内核的推导

3.4.3. Top-1 内核的 TileLang 代码

def build_kernel(
    d: int = 4096,
    m: int = 1024,
    n: int = 4096,
    block_d: int = 128,
    block_m: int = 128,
    block_n: int = 128,
    threads: int = 256,
    num_stages: int = 2,
    enable_swizzle: bool = True,
    enable_autotune: bool = False,
    autotune_warmup: int = 10,
    autotune_rep: int = 10,
    autotune_timeout: int = 100,
):
    dtype = T.float16
    accum_dtype = T.float32
    scale = 1.44269504  # log2(e)
    fast_math_pass_configs = {tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}
    scale_jit_decorator = tilelang.jit(out_idx=[2], pass_configs=fast_math_pass_configs)
    compute_jit_decorator = tilelang.jit(out_idx=[3], pass_configs=fast_math_pass_configs)
    if enable_autotune:
        def decorate_scale(fn):
            return tilelang.autotune(configs=get_configs(), warmup=autotune_warmup, rep=autotune_rep, timeout=autotune_timeout, skip_check=True)(scale_jit_decorator(fn))
        def decorate_compute(fn):
            return tilelang.autotune(configs=get_configs(), warmup=autotune_warmup, rep=autotune_rep, timeout=autotune_timeout, skip_check=True)(compute_jit_decorator(fn))
    else:
        def decorate_scale(fn):