autotune

主线位置：算子主线里 CubeBackend::float_matmul 和 cubek 之间。burn-cubecl 后端默认用 MatmulStrategy::Autotune，于是每次 matmul 先进 matmul_autotune：算出调优键，查这个设备的调优表；命中就直接启动选好的那个 cubek 内核，没命中就当场把候选都测一遍。一次 matmul 数到的那一万八千次启动、第一次 4 秒多的耗时，都发生在这一页。

flowchart TB
  mm["matmul_autotune"] --> key["create_key 得调优键"]
  key --> exec["LocalTuner::execute"]
  exec --> hit{"内存表命中?"}
  hit -->|是| run["执行选中的候选"]
  hit -->|否| check["check_tune"]
  check --> disk{"磁盘缓存有且校验和对?"}
  disk -->|是| run
  disk -->|否| plan["按组和优先级排出一轮候选"]
  plan --> bounds["roofline 求时间上限"]
  bounds --> probe["测峰值和启动开销"]
  probe --> bench["逐个测候选"]
  bench --> save["写内存表和 default.db"]
  save --> run

本页源码：burn 0.22.1 的 burn-cubecl/src/kernel/matmul/tune/，cubek 0.3.1 的 cubek-matmul/src/tune_key.rs，cubecl 0.11.3 的 cubecl-runtime/src/tune/ 和 cubecl-std/src/throughput/。实验 crate rt-tune（burn 只开 metal、tracing，另外直接依赖 cubek、cubecl 以构造调优键），release 构建，M2 Max、macOS 27、rustc 1.99.0。每个计时都注明缓存是空（冷）还是已有（热）。

一个 op，三十多个候选

matmul_autotune 的骨架：

    static TUNER: LocalTuner<MatmulAutotuneKey, CubeTuneId> = local_tuner!();

    let tune_id = CubeTuneId::new(&lhs.client, &lhs.device);
    let tunables = TUNER.init(&tune_id, move || {
        const PRIORITY_MAX: i8 = 3;
        const PRIORITY_HIGH: i8 = 2;
        const PRIORITY_MEDIUM: i8 = 1;
        const PRIORITY_MIN: i8 = 0;
        const PRIORITY_NEVER: i8 = -1;
        // ... 定义 accelerated / unit / tma / gemv / cpu 五个组
        let mut set = TunableSet::new(create_key, matmul_input_gen);

        set = with_matmul_bounds(set);

        // First entry should always work, since it is considered the fallback.
        set = set.with(
            Tunable::new("matmul_naive", |(lhs, rhs, out)| {
        // ... 再加 gemv、unit、gemm、cpu、accelerated（cmma / mma / tma）各个策略
        set
    });

    TUNER.execute(&tune_id, &client, tunables, (lhs, rhs, output.clone()));

• LocalTuner 是一个 static，按设备 id（CubeTuneId：DeviceId 加运行时名）各存一个 Tuner 和一份 TunableSet。init 的闭包每个设备只跑一次：候选集合是对着具体设备建的，闭包里会问 client 支不支持某种矩阵单元。
• 每个候选是一个 Tunable：名字加一个闭包，闭包里是 launch_matmul(&strategy, lhs, rhs, out)，也就是带着某个 cubek 策略去启动（见 burn-cubecl 后端）。候选一共 33 个，是 cubek 各种策略和参数的组合。
• 候选属于一个或几个 TuneGroup。组的优先级和组内优先级都是「调优键 → i8」的函数，负数表示这个键下跳过。按文档注释：组按优先级从高到低试，「once any tunable in a group returns a valid result, no later groups are tried」；组内只有最高的那一档成为这一轮，低档是这一轮全部失败时的后备。

组的规则摘要（kind 和 scale 见下一节）：

降级而不是删除，是代码注释里写明的：「They keep the minimum priority rather than being discarded, so they remain a last resort and the tune plan can never end up empty」。在只有 WGSL 的 wgpu 构建上能看到它的效果：设备登记不了 CMMA，accelerated 组里 20 个候选全都降到最低档；这一组的优先级仍是 2，于是第一轮就是这 20 个，每个只用了约 0.05 ms，结果都是「The profiled window resolved no device timing」，没有一个测出成绩，整轮作废，才轮到 unit 组，最后选中 matmul_double_unit_min_tile_size（wgpu 与 Metal 有两种编译器的对照）。

调优键：形状分桶

create_key 把张量的形状、步长、分块、元素类型交给 cubek 的 MatmulAutotuneKey::generate。键分两半：

pub struct MatmulAutotuneKey {
    pub definition: MatmulProblemDefinition,
    pub analysis: MatmulAutotuneAnalysis,
}
// ...
pub struct MatmulProblemDefinition {
    #[autotune(anchor)]
    pub m: usize,
    #[autotune(anchor)]
    pub n: usize,
    #[autotune(anchor)]
    pub k: usize,
    pub lhs_pow2_factor: u8,
    /// Power of two that lhs strides are aligned to
    pub lhs_stride_factor: u8,
    pub rhs_pow2_factor: u8,
    /// Power of two that rhs strides are aligned to
    pub rhs_stride_factor: u8,
    pub elem_lhs: ElemType,
    pub elem_rhs: ElemType,
    pub elem_out: ElemType,
    pub matrix_layout_lhs: MatrixBatchLayout,
    pub matrix_layout_rhs: MatrixBatchLayout,
    // ... lhs_storage / rhs_storage：普通行存储还是分块存储
}

• m、n、k 分桶。 #[autotune(anchor)] 让 #[derive(AutotuneKey)] 生成的构造函数对这三个字段调 anchor：默认的 Balanced 档把数向上取到 2 的幂（底数 2，ceil(log₂ x)）。Minimal 档底数放大到 2.5（桶更宽），Extensive 缩到 1.5，Full 不分桶。
• 向量化因子 pow2_factor：取（分桶后的）k 或 n 末尾零的个数，最多 4。
• 步长对齐 stride_factor：除最内维以外各个步长按字节算能被 2 的几次方整除，最多 10。它不分桶，用的是真实步长：步长对齐决定异步拷贝（TMA 一类）内核能不能用。
• 布局 MatrixBatchLayout：Contiguous、MildlyPermuted { transposed, batch_swap } 或 HighlyPermuted。转置过的输入是另一个键，因为它要的内核不同，这正是 matmul 不先 into_contiguous 的原因。
• analysis：MatmulGlobalScale::from_size（三边都小于 512 为 Small，都小于 2048 为 Medium，否则 Large，用分桶后的值）和 MatmulKind（General、MatVec、VecMat、ScalarVec、VecScalar、InnerProduct、OuterProduct）。上一节的组优先级就是看这两项。

rt-tune keys 直接调 MatmulAutotuneKey::from_parts，看几个 f32 方阵的键（T 表示 rhs 是转置视图）：

 512  -> m=512 n=512 k=512 lhs_pow2=4 lhs_stride=10 rhs_pow2=4 rhs_stride=10 layout_rhs=Contiguous scale=Medium kind=General
 768  -> m=1024 n=1024 k=1024 lhs_pow2=4 lhs_stride=10 rhs_pow2=4 rhs_stride=10 layout_rhs=Contiguous scale=Medium kind=General
1000  -> m=1024 n=1024 k=1024 lhs_pow2=4 lhs_stride=5 rhs_pow2=4 rhs_stride=5 layout_rhs=Contiguous scale=Medium kind=General
1024  -> m=1024 n=1024 k=1024 lhs_pow2=4 lhs_stride=10 rhs_pow2=4 rhs_stride=10 layout_rhs=Contiguous scale=Medium kind=General
1025  -> m=2048 n=2048 k=2048 lhs_pow2=4 lhs_stride=2 rhs_pow2=4 rhs_stride=2 layout_rhs=Contiguous scale=Large kind=General
2048  -> m=2048 n=2048 k=2048 lhs_pow2=4 lhs_stride=10 rhs_pow2=4 rhs_stride=10 layout_rhs=Contiguous scale=Large kind=General
1024T -> m=1024 n=1024 k=1024 lhs_pow2=4 lhs_stride=10 rhs_pow2=4 rhs_stride=10 layout_rhs=MildlyPermuted { transposed: true, batch_swap: false } scale=Medium kind=General
key(768) == key(1024): true
key(1000) == key(1024): false

• 768 和 1024 是同一个键：都落进 1024 这个桶，行步长 768×4 = 3072 字节和 4096 字节的对齐都封顶到 10。
• 1000 也落进 1024 的桶，但 1000×4 = 4000 = 2⁵ × 125，步长对齐只有 5，是另一个键。
• 1025 进了 2048 的桶，规模从 Medium 变成 Large。512 不小于 512，算 Medium。

burn 自己的 CubeAutotuneKey 只有 Sum、Conv、ConvTranspose 三个变体，给 burn 里的求和与卷积调优用；matmul 用的是 cubek 的 MatmulAutotuneKey。

第一次遇到一个键

LocalTuner::execute 先算键、查内存表，命中就 operations.fastest(index).execute(inputs)，这是每次调用都走的快路径。没命中进 Tuner::check_tune：

1. 再查一次。 把磁盘缓存里还没读进来的条目读进内存表（sync_persistent）；读到的条目带着写入时的校验和，要和当前候选集合的 compute_checksum（所有候选名字拼起来的哈希）比对。候选增删或改名，旧结果就作废，因为下标对不上了。
2. 占位。 都没有就 mark_pending(key)；只有一个候选的直接记下下标 0 返回。
3. 排计划、求上限。 tunables.plan(key) 按上一节的规则排出轮次；tunables.bounds(key, inputs) 求这个问题的时间上限（下一节）。
4. 测。 默认走 tune_adaptive：轮流测这一轮的候选，BenchConfig::default 规定每个至少 3 次、至多 10 次，比当前最快慢 1.5 倍以上的提前淘汰；有候选的 2 次采样落进上限就「短路」，不再测别的。
5. 记下。 中位数最快的成为结果，写进内存表和磁盘缓存，然后执行它。

测候选用的输入是 matmul_input_gen 给的：lhs、rhs 原样，输出 out.copy() 一份，免得测试时写坏真正的输出。

屋顶线和那十几万次启动

第 3 步的上限来自 burn 的 with_matmul_bounds → cubecl 的 roofline_bounds：

pub fn roofline_bounds(
    client: &Client,
    compute_key: ThroughputKey,
    work: Work,
    thresholds: Thresholds,
) -> Bounds {
    Bounds {
        bounds: alloc::vec![
            compute_bound(client, compute_key, work, thresholds.compute),
            memory_bound(client, MemoryAccess::Copy, work, thresholds.memory),
        ],
        launch_overhead: measure_launch_overhead(client),
    }
}

这是屋顶线模型（roofline）：一个 matmul 的运算量除以设备的计算峰值，访存量除以带宽峰值，两者取大，再除以阈值（Thresholds::for_level：Balanced 是 0.8，即「达到峰值的 80% 就够好」），加上一次启动的开销，就是 time_limit。候选跑进这个时间就可以短路。

三个峰值都要在这台设备上实测，各有一个探测内核：计算峰值用 ComputeCmmaThroughput（有矩阵单元时）或 ComputeDirectThroughput，带宽用 MemoryDirectThroughput（按问题的访存量取不同的工作集大小，所以换了规模的问题会再测一次带宽），启动开销用 launch_overhead：

#[cube(launch_unchecked)]
pub fn launch_overhead<I: Numeric, N: Size>(
    input: &[Vector<I, N>],
    output: &mut [Vector<I, N>],
    #[define(I)] _dtype: ElemType,
) {
    if ABSOLUTE_POS == 0 {
        output[0] = input[0];
    }
}

一个几乎什么都不做的内核，用单个线程组反复启动。测法在 ThroughputBenchmarker：warmup 不断加大一次采样里的迭代次数，直到一次采样大约 TARGET_DURATION = 20 ms，等结果稳定（相差 3% 以内连续 3 次、且稳定期至少 250 ms），然后在 200 ms 预算内取最快的样本。对这个探测，每次迭代就是一次真实的启动，所以启动次数大约是「几秒钟 ÷ 每次启动的开销」。一次 matmul 在 debug 构建下数到 18 046 次；release 构建下每次启动更便宜，同样 20 ms 的窗口要更多次，本页几次冷启动数到 6.9 万到 25 万次不等。这些数只决定「启动开销」这一个值（M2 Max 上 release 测出约 3.6 µs）。

吞吐量的结果另有一个缓存，命名空间按运行时名、适配器名和显存大小命名（ThroughputCache::get_for_device），和 autotune 表分开存，同一台机器上的所有键共用。

autotune 每做一次决定，可以把完整记录写成 JSON：在 crate 目录放一个 cubecl.toml，写 [autotune.decisions] 和 stdout = true（AutotuneConfig::decisions，只加一个输出口，不改别的日志级别）。冷缓存下 1024² f32 那一条，摘出要点：

m=1024 n=1024 k=1024 stride=10 scale=Medium fastest_index=14 measured=6 errs={'Skip': 23, 'Unknown': 4}
   [14] matmul_simple_cyclic_cmma_multi_rows                                   median 1.327 ms
   [16] matmul_ordered_double_cmma_partition_k2_row_count4_rows_per_plane2     median 1.445 ms
   [12] matmul_simple_cyclic_cmma                                              median 1.761 ms
   [22] matmul_double_cyclic_cmma_specialized                                  median 1.846 ms
   [18] matmul_ordered_double_cmma_partition_k2_row_count8_rows_per_plane2     median 2.300 ms
   [20] matmul_double_cyclic_cmma                                              median 2.472 ms
    log_context {"bounds": {"bounds": [{"resource": {"amount": 2146435072, "peak_per_s": 9044321293800.54}, "threshold": 0.8}, {"resource": {"amount": 12582912, "peak_per_s": 319882855399.6339}, "threshold": 0.8}], "launch_overhead": {"secs": 0, "nanos": 3626}}, "limit": {"secs": 0, "nanos": 300281}, ...

（用 Python 把那一行 JSON 汇总成这个样子，原始输出保存在实验目录。）

• 运算量 2 146 435 072 = 1024² × (2·1024 − 1)（cubek 的 MatmulCost::compute_ops：每个输出元素 k 次乘、k − 1 次加），计算峰值测得 9.04 × 10¹² op/s；访存量 12 582 912 字节（读两个输入、写一个输出，各 4 MiB），带宽峰值 3.2 × 10¹¹ 字节/秒。计算一侧是 0.237 ms，除以 0.8 再加 3.6 µs，上限 0.300 ms。
• 最快的候选 1.33 ms，远在上限之外，没有短路，这一轮 10 个候选都上了场：6 个测出成绩，4 个（specialized 和三个 tma 变体）报「The profiled window resolved no device timing」被排除。其余 23 个是计划里没轮到的（Skip）。
• 这一轮一共启动了 33 次 MatmulEntry，fastest_index 14 写进缓存。

缓存：内存表和 default.db

每个 Tuner 有一个 TuneCache：内存里的 HashMap<键, 结果>，背后挂一个持久化的 Store。TuneCache::new 里：

            let config = crate::config::CubeClRuntimeConfig::get();
            let namespace = Namespace::scoped("autotune", format!("{device_id}/{name}"));
            let table = namespace.as_str().into();

            if config.autotune.disable_cache {
                return TuneCache {
                    in_memory_cache: HashMap::new(),
                    persistent_cache: None,
                    // ...
                };
            }

• 存在哪。 所有持久缓存共用一个 turso（SQLite）库，默认位置由 CacheConfig 的 Target 决定：从当前目录往上最外层的 Cargo.toml 旁边的 target/environment/default.db，找不到 Cargo.toml 时退到系统的用户缓存目录；也可以配成 local（当前目录）、global（用户缓存目录）或指定目录（[environment] path = ..）。default 是环境名，CUBECL_ENVIRONMENT 可以换一个，同一个仓库里存几套。
• 按什么分。 表名是 autotune/<cubecl 版本>/<设备 id>/<调优器名>，键是 PersistentCacheKey { key, checksum }，值里除了 fastest_index 还存了所有候选的成绩、上限和时间限制，「so a cache entry can be inspected after the fact」（读回时只用 fastest_index）。
• 开关。 cubecl-runtime 的默认 feature 就有 persistence；burn 的 persistence feature 只是把它显式转发下去。运行期可以用环境变量 CUBECL_AUTOTUNE_CACHE=0 关掉磁盘缓存，内存表不受影响，同一进程里一个键仍只调一次。配置文件是 cubecl.toml，或 burn.toml 里的 [cubecl] 一节（section_file_names），从当前目录往上找。CUBECL_AUTOTUNE_LEVEL（minimal / balanced / extensive / full）同时改分桶的底数和屋顶线的阈值。

跑完下一节的冷启动之后，把 rt-tune/target/environment/default.db* 复制出来，用 sqlite3 数各命名空间的条目：

autotune/0.11.3/device-100-16384-wgpu<msl>/burn_cubecl-kernel-matmul-tune-base|3
records/v1/autotune|3
records/v1/compilation|19
records/v1/sessions|1
throughput/0.11.3/probe-v2/wgpu-msl-_spirv-0-0_Apple-M2-Max_mem4294967292_par0|4

• autotune 表 3 条：1024、1000、2048 三个键（768 和 1024 共用一条）。表名里的 burn_cubecl-kernel-matmul-tune-base 是 local_tuner!() 取的 module_path!()，:: 换成了 -。
• 吞吐量 4 条：计算峰值、启动开销和两种工作集大小的带宽。命名空间里的 spirv-0-0 是 wgpu 运行时的设备指纹（spirv_{vendor}_{device}，在 Metal 上 vendor 和 device 都报 0），runtime.rs 的注释说它是这类测量缓存里唯一能区分不同适配器的东西。
• records/v1/.. 是 cubecl 的环境记录（每次调优、编译的经过），不参与查找。键和值是 CBOR 编码，sqlite3 直接看是乱码。

实验：冷与热，同桶与新桶

rt-tune shapes 1024 768 1000 2048 在一个进程里依次跑四个方阵：每个形状先 random 出两个 f32 矩阵并 sync，然后计时第一次 matmul 到 sync，再跑 20 次取最短。启动和编译由 tracing 的 span 数出来（~compile(integrate) 是新内核的编译次数）。先删掉 target/environment 跑一次（冷），再原样跑一次（热）：

## 进程 1：冷缓存
1024x1024 first:       3.670s  launches=75199 [ComputeCmmaThroughput=50, LaunchOverhead=75082, MatmulEntry=33, MemoryDirectThroughput=32, PrimeBuffer=1, UnaryNumeric=1, ~compile(integrate)=11]
1024x1024 next 20: min 2.021ms median 2.725ms  launches=20 [MatmulEntry=20]
 768x768  first:    117.258ms  launches=1 [MatmulEntry=1, ~compile(integrate)=1]
 768x768  next 20: min 1.110ms median 1.778ms  launches=20 [MatmulEntry=20]
1000x1000 first:       1.430s  launches=34 [MatmulEntry=33, UnaryNumeric=1, ~compile(integrate)=6]
1000x1000 next 20: min 1.964ms median 2.223ms  launches=20 [MatmulEntry=20]
2048x2048 first:       1.411s  launches=82 [MatmulEntry=34, MemoryDirectThroughput=46, PrimeBuffer=1, UnaryNumeric=1]
2048x2048 next 20: min 3.839ms median 10.055ms  launches=20 [MatmulEntry=20]
## 进程 2：热缓存
1024x1024 first:    149.478ms  launches=1 [MatmulEntry=1, ~compile(integrate)=1]
1024x1024 next 20: min 2.015ms median 2.759ms  launches=20 [MatmulEntry=20]
 768x768  first:    157.049ms  launches=1 [MatmulEntry=1, ~compile(integrate)=1]
 768x768  next 20: min 1.159ms median 1.838ms  launches=20 [MatmulEntry=20]
1000x1000 first:    155.066ms  launches=1 [MatmulEntry=1, ~compile(integrate)=1]
1000x1000 next 20: min 2.046ms median 2.484ms  launches=20 [MatmulEntry=20]
2048x2048 first:     13.358ms  launches=1 [MatmulEntry=1]
2048x2048 next 20: min 3.605ms median 4.045ms  launches=20 [MatmulEntry=20]

几点：

• 每个进程的第一次至少要编一次 MSL（100–160 ms），因为编译产物不跨进程缓存（wgpu 与 Metal）。autotune 缓存省掉的是秒级的测量。
• 2048² 在冷进程里调优 34 次却一个内核也没编：它的候选在 1024² 那轮已经全部编过，KernelId 相同，管线表直接命中。热进程里也一样，1024² 编出来的那个内核 2048² 直接用。反过来先跑 2048² 时，1024² 第一次只要 1.35 ms（rt-tune shapes 2048 1024，热缓存）。
• 「之后 20 次」的中位数在 Metal 上波动很大（同一形状 1.8–10 ms），GPU 空闲后降频是一个可能的原因，本页没有核实；最短值比较稳定，表里只引最短值。1024² 和 1000² 的最短值接近，不对齐的步长在这里没有明显代价。

动脑筋答案

1000×1000 重新调优，768×768 不会：

 768x768  first:    117.258ms  launches=1 [MatmulEntry=1, ~compile(integrate)=1]
1000x1000 first:       1.430s  launches=34 [MatmulEntry=33, UnaryNumeric=1, ~compile(integrate)=6]

• 768 和 1024 分桶后都是 1024，步长对齐都封顶到 10，键完全相同，内存表命中，只启动选中的那一个内核（多编一次是因为 768 的编译期参数不同）。
• 1000 的 m、n、k 也分到 1024，但行步长 4000 字节只能被 2⁵ 整除，lhs_stride_factor / rhs_stride_factor 是 5，键不同。33 次 MatmulEntry 是这一轮候选的测量，UnaryNumeric 是 matmul_input_gen 复制输出用的那次拷贝。峰值和启动开销已经在进程里，不再测。
• 分桶只管 m、n、k 的大小；对齐、布局、元素类型任何一项变了都是新问题。

小结

• matmul 的 33 个候选按组和优先级排成轮次；设备不支持的（Metal 上的 mma、WGSL 上的全部 cmma）降到最低档而不删除，保证计划不空。
• 调优键 = 分桶后的 m、n、k（默认向上取 2 的幂）+ 向量化因子 + 真实步长对齐 + 布局 + 元素类型 + 规模和形状分类。同桶不一定同键。
• 第一次遇到键：查磁盘缓存并核对候选列表的校验和 → 排计划 → 用屋顶线求时间上限 → 自适应地测（3–10 次，慢 1.5 倍淘汰，进上限就短路）→ 写回。
• 冷缓存的上万次启动是 LaunchOverhead 探测：每次采样凑够约 20 ms 的启动，次数随每次启动的开销变化（debug 1.8 万，release 7 万到 25 万）；它和计算、带宽峰值一起存进单独的吞吐量缓存。
• 结果存在 target/environment/default.db（turso），表名带 cubecl 版本和设备 id；Device::wgpu 和 Device::metal 是两张表。调试日志会改变测出的启动开销，看调优结果用 [autotune.decisions]。