mcpField
T
tilegym-cutile-autotuning

tilegym-cutile-autotuning

github.com/nvidia开发3.8k1.5k 次安装今天更新
代码仓库

关于 tilegym-cutile-autotuning

tilegym-cutile-autotuning 是一项为 CuTile GPU 内核添加自动调优配置的技能,面向开发和优化内核的工程师。它指导搜索空间设计、调优缓存、正确性测试及性能验证,同时限定改动范围为调优配置。

主要功能

  1. 01按内核类型选择搜索维度
  2. 02设计并筛选调优配置
  3. 03缓存调优结果并直接启动
  4. 04检查原地写入等常见问题
  5. 05运行正确性和性能对比测试

适用场景

  1. 01为新 CuTile 内核添加自动调优
  2. 02优化已有内核的调优配置
  3. 03排查搜索空间为空或编译过慢
  4. 04验证调优结果是否优于固定配置

安装

复制下面的命令,在终端中运行。

$npx skills add https://github.com/nvidia/skills --skill tilegym-cutile-autotuning

安装后在 Claude Code 中输入/tilegym-cutile-autotuning确认已加载。

包含的文件

除 SKILL.md 外,安装时一起下载的参考资料和脚本。

BENCHMARK.md
skill-card.md
skill.oms.sig
assets/examples/01_rmsnorm_occupancy_only/
autotuned_launch.py
fixed_launch.py
assets/examples/02_matmul_full_search/
autotuned_launch.py
fixed_launch.py
assets/examples/03_rope_inplace_splitbuffer/
autotuned_launch.py
fixed_launch.py
evals/
evals.json
references/
api-reference.md
hardware-constraints.md
kernel-type-templates.md
parameter-space-design.md
pitfalls.md
search-strategies.md
workflow.md

详细介绍

CuTile 自动调优

使用 exhaustive_search API,并采用“仅调优一次/缓存/直接启动”模式,为 CuTile 内核添加自动调优。

说明

按照决策树对内核进行分类,设计搜索空间,实现“仅调优一次/缓存/启动”模式,并验证性能。

  1. 分类 — 使用决策树确定搜索维度(仅占用率搜索或完整平铺搜索)
  2. 设计搜索空间 — 从 references/kernel-type-templates.md 中选择匹配的模板;通过架构筛选条件将最终代码中的配置数量精简至 ≤ 30 个(定向探索探测可能暂时超过此数量 — 参见设计理念)
  3. 实现 — 按照分步工作流添加 exhaustive_search + 缓存 + ct.launch;必要时使用分离缓冲区处理原地写入
  4. 测试 — 分别在启用自动调优和设置 DISABLE_AUTOTUNE=1 时运行正确性测试
  5. 验证 — 与已知最佳固定配置进行 A/B 基准测试;参见 references/search-strategies.md
  6. 精简 — 剔除从未胜出的冗余配置,目标是每种架构 ≤ 8 个配置,以尽量减少编译成本(步骤 10)

任务导航 — 直接跳转到所需内容

你要做什么?前往
为新内核添加自动调优(最常见)下方的快速参考 → 工作流:添加自动调优 → references/kernel-type-templates.md(按内核类型选择:T1=逐元素、T2=原地、T3=矩阵乘法、T4=持久化、T5=FMHA、T6=FP8、T7=分组 GEMM、T8=变长注意力、T9=双 GEMM 融合)
调试:首次运行后数据损坏/结果错误陷阱 #1(原地内核)
调试:自动调优耗时 5 分钟以上陷阱 #2(编译超时)
调试:搜索空间生成器返回零个配置先查看陷阱 #5;同时检查架构筛选条件、大小保护条件和 num_ctas 约束
优化现有自动调优配置工作流:优化现有配置

快速参考 — 仅调优占用率(调优一次/缓存/启动)

大多数 CuTile 内核(逐元素、归约、LayerNorm)只需调优占用率。请复制以下模式:

from types import SimpleNamespace
from cuda.tile.tune import exhaustive_search
import cuda.tile as ct
import torch

def _my_autotune_configs():
    for occ in [1, 2, 4, 8]:
        yield SimpleNamespace(occupancy=occ)

# Module-level cache: tune once, launch fast forever after
_autotune_cache = {}

def my_op(x, output):
    stream = torch.cuda.current_stream()
    NUM_SM = torch.cuda.get_device_properties(x.device).multi_processor_count

    # Cache key: anything that affects optimal config (use str() for device)
    cache_key = (x.shape, x.dtype, str(x.device))

    if cache_key not in _autotune_cache:
        configs = list(_my_autotune_configs())
        result = exhaustive_search(
            configs,
            stream,
            grid_fn=lambda cfg: (min(NUM_SM * cfg.occupancy, M), 1, 1),
            kernel=my_kernel,
            args_fn=lambda cfg: (x, output, ...),
            hints_fn=lambda cfg: {"occupancy": cfg.occupancy},
        )
        best_cfg = result.best.config
        tuned_kernel = my_kernel.replace_hints(occupancy=best_cfg.occupancy)
        _autotune_cache[cache_key] = (best_cfg, tuned_kernel)  # cache BOTH

    cfg, tuned_kernel = _autotune_cache[cache_key]
    grid = (min(NUM_SM * cfg.occupancy, M), 1, 1)
    ct.launch(stream, grid, tuned_kernel, (x, output, ...))

关键规则:

  • 调优一次、缓存并直接启动 — exhaustive_search 只会在每种形状首次调用时运行;后续调用会使用缓存的配置并通过 ct.launch 启动,零开销
  • 对于原地内核,搜索时使用分离缓冲区(输入/输出张量分开)
  • 最终代码中最多保留 30 个配置(有关临时定向探测,请参阅「设计理念」)
  • exhaustive_search 要求传入 Sequence(列表/元组)——请使用 list() 转换生成器
  • 搜索空间必须包含原始固定配置 — 这样可确保自动调优不会降低性能

何时使用此模式:内核具有固定的块大小(无法调优 tile 大小)。包括:逐元素运算(SwiGLU、GeGLU)、归约(RMSNorm、LayerNorm)、RoPE,以及使用启发式块大小的持久内核(分组 GEMM)。

对于复杂内核(具有 tile 大小的矩阵乘法、FMHA、使用 num_ctas 的 FP8),请阅读下方完整指南以及 kernel-type-templates.md。

⚠️ 三个几乎人人都会踩的坑——提交前请检查:

  • 热路径上调用了 replace_hints? → 从 exhaustive_search 缓存配置和内核对象。每次调用都执行 replace_hints() 会重新编译(慢 100–500 倍)→ 坑 #7
  • 原地内核(写回输入张量)?→ 搜索期间必须使用分离缓冲区模式 → 坑 #1
  • 搜索空间为空? → 检查架构筛选条件和 num_ctas 约束 → 坑 #5

最低覆盖要求:在 sm100+ 上,FMHA/矩阵乘法/变长搜索空间必须同时包含 num_ctas=1 和 num_ctas=2。对于核心维度(tile 大小、占用率),即使不确定哪个更好,也至少保留 2 个不同的值——让 exhaustive_search 来决定。

何时停止调优:平均加速比为 [0.98, 1.02],意味着你当前的搜索空间没有帮助——但这并不代表没有任何配置能带来提升。停止前,请检查是否已覆盖此类内核的关键维度(参阅 references/kernel-type-templates.md)。如果搜索空间已涵盖模板建议的维度,而最佳结果仍处于噪声底限,则停止调优——进一步的微调不会有帮助。如果缺少关键维度(例如,双 GEMM 内核从未尝试过 num_ctas=2),请扩展搜索空间,而不是放弃。

正确性测试通过且自动调优后的内核相较于固定配置基线有所提速后,就此停止——不要为了「确认」而再次运行。 由于时钟缩放和操作系统调度,GPU 内核计时在不同调用间会有 ±5–10 % 的波动;后续计时变慢并不意味着代码有问题。

要提升速度,只修改自动调优搜索空间(配置、tile 大小、占用率、num_ctas)。不要为了追求速度而修改其他代码(Python 封装、流管理等)——内核性能取决于配置选择,而非主机端代码。

阅读指南

  • 仅受占用率影响的内核(逐元素、归约、固定块大小的持久化内核):只需阅读快速参考和陷阱清单,跳过 references/ 文档。对于原地内核,还需阅读陷阱 #1。
  • 复杂内核(可调块大小的 matmul、FMHA、使用 num_ctas 的 FP8):快速参考 → 决策树 → API 参考 → 分步工作流 → 相关的 references/ 文档。

5 步概要:对内核分类 → 设计搜索空间(parameter-space-design.md)→ 使用模板实现(kernel-type-templates.md)→ 使用 A/B 测试验证 → 检查陷阱清单。

阅读参考文档:只阅读与你的内核类型相关的参考文档——例如,对于 FMHA,阅读 references/kernel-type-templates.md 中的模板 5 部分;对于硬件限制,只阅读目标架构对应的部分。如果有针对性的查阅已足够,就不要从头到尾阅读所有参考文档。

设计理念

自底向上构建小而精确的搜索空间,而不是先构建大空间再精简。 CuTile 的编译开销远高于 Triton(每个配置约 0.5–1 秒),因此最终代码应包含 ≤ 30 个配置。具体做法是:先对内核类型分类,然后只为该类型和架构构建相关配置。

开发期间的定向探索:如果初始模板配置带来的加速比 < 1.0,可以通过 bash + python3 -c 运行一个临时的较大规模探测(30–100 个配置),以确定哪些维度会产生影响——但该探测必须是有方向的,而不是盲目地计算笛卡尔积。利用内核类型分类来决定要改变哪些维度(例如,对于双 GEMM,在固定块大小的情况下探测 num_ctas × occupancy;对于 FMHA,在固定 TILE_N 的情况下探测 TILE_M × num_ctas)。探测确定最佳范围后,将最终代码的搜索空间限制为 ≤ 8 个最佳候选。不要把大型探测写进源文件——它只是一次性诊断工具。

决策树:此内核需要搜索哪些维度?

所有内核都应添加自动调优。问题不在于是否自动调优,而在于搜索哪些维度:

What type of kernel is this?
├── Compute-bound (matmul, GEMM, FMHA) → Does it have multiple tunable dimensions (tile sizes)?
│   ├── YES → Is it a fused multi-GEMM kernel (dual-GEMM, e.g. Linear+GLUAct)?
│   │   ├── YES → Template 9: low occupancy (1–2), conservative tiles (2× SHMEM/register pressure)
│   │   └── NO  → Full search: TILE_M × TILE_N × (TILE_K) × occupancy × num_ctas
│   │             (see matmul/FMHA templates in kernel-type-templates.md)
│   └── NO  → Occupancy-only search: [1, 2, 4, 8]
│             (see Quick Reference above)
├── Balanced (LayerNorm, reduction + compute) →
│   Occupancy-only search: [1, 2, 4, 8]
│   Expected benefit: 2-15%
└── Memory-bound (CE Loss, pure elementwise) →
    Occupancy-only search: [1, 2, 4, 8]
    Expected benefit: 0-15% (varies by kernel; zero-cost after tuning)

为什么内存受限内核只搜索占用率(而不搜索 num_ctas 或块大小):

  • num_ctas 毫无益处:num_ctas > 1 会启用 TMA 多播,让多个 CTA 在共享内存中共享块数据(例如,多个 CTA 复用 matmul 的 A/B 块)。内存受限内核使用逐元素的 ct.gather/ct.scatter,没有块复用——多 CTA 协作只会增加开销,却无法带来数据共享收益。
  • 块大小已预先确定:内存受限内核的 BLOCK_SIZE 由离线扫描确定(例如,在 B200 上,对 [256, 512, 1024, 2048, 4096, 8192] 进行扫描后,1024 是全局最优值)。这是一个常量,而不是运行时可调参数。
  • 占用率是唯一有效的调节项:更高的占用率可以让 GPU 在一个 CTA 因内存请求而停滞时切换到另一个 CTA,从而隐藏内存延迟。

证据——CE Loss 实验:在 Cross-Entropy Loss 上进行的 12 配置搜索(占用率 × num_ctas)只带来了 2.5% 的提升(相较 Triton,从 0.79x 提升到 0.81x)。num_ctas 维度没有任何贡献;由于编译成本超过了微小收益,因此撤销了该改动。仅搜索占用率(4 个配置)即可取得相同结果,编译时间减少至三分之一。

内存受限内核的注意事项:添加仅针对占用率的自动调优始终值得,因为:

  • 首次调用后,“调优一次/缓存/启动”的模式不会产生运行时开销
  • 搜索空间很小(4 个配置,编译约需 2–4 秒)
  • 即使提升很小,规模化后也有价值

占用率选择指南

占用率控制每个 SM 并发运行的 CTA 数量。设计占用率搜索空间时,可从以下范围开始:

占用率范围适用场景示例内核
1–4计算密集型(大量数学运算)复杂变换、矩阵乘法
4–8均衡型(GEMM、TMA)矩阵乘法、FMHA
8–16内存密集型(归约)Softmax、LayerNorm
16–32极轻量型(复制、类型转换)类型转换、逐元素运算

使用这些范围来初始化搜索空间。对于仅使用占用率的内核,[1, 2, 4, 8] 可涵盖大多数情况,参见上文的快速参考。

exhaustive_search API 参考

完整的 exhaustive_search API 内容请参阅 references/api-reference.md,包括当前签名、TuningResult、一次调优/缓存/启动模式、replace_hints、内核提示、search_space 设计和 grid_fn 模式。

分步工作流

端到端工作流请参阅 references/workflow.md,包括为新内核添加自动调优、处理现有的多架构配置、与 torch.autograd.Function 集成、跨后端配置迁移(Triton → CuTile),以及优化现有配置。

常见陷阱检查清单

常见陷阱的完整列表请参阅 references/pitfalls.md,包括原地数据损坏、编译超时、冷缓存导致的性能偏差、NCU 分析干扰、search_space 生成器耗尽、FP8 精度损失,以及热路径上 replace_hints 导致重新编译。

范围与边界

本技能仅涵盖自动调优配置:搜索空间设计、调用 exhaustive_search、缓存,以及使用调优后的提示调用 ct.launch。它不会修改内核代码。

范围内(自动调优配置):

  • 搜索空间生成器函数
  • exhaustive_search() 调用及结果处理
  • 用于应用调优后提示的 kernel.replace_hints()
  • 缓存逻辑(键设计、字典管理)
  • 使用调优后内核的 ct.launch()
  • DISABLE_AUTOTUNE 回退路径

范围外(内核代码修改——请勿进行以下更改):

  • 数学标志(flush_to_zero、rounding_mode)
  • 性能提示(slice_hint、buffer_depth、copy_config)
  • 内存访问模式(将 2D gather/scatter 转换为 1D)
  • 代码生成优化(safe_offs → padding_value)
  • 算法更改(K 循环拆分、负载均衡)

进一步的优化建议

添加自动调优后,以下内核级优化可能带来额外收益。这些优化不在本技能的范围内——可将其作为后续潜在步骤告知用户,但不要将其作为自动调优的一部分来实现:

  • 数学标志:flush_to_zero=True + rounding_mode=APPROX 可使 FMHA 类内核的性能提升 34–72%(通过环境变量 TILEIR_ENABLE_FTZ=1 TILEIR_ENABLE_APPROX=1 设置,或在内核代码中设置)。因果链:由于次正规数处理开销,较大的 tile 起初会使性能下降 18–43%;启用 FTZ+APPROX 可消除这一影响,使结果转为提升 34–72%。因此,数学标志是 FMHA 类内核上大 tile 配置发挥效果的前提条件。
  • 性能提示:slice_hint、buffer_depth、copy_config——需要修改内核 IR 代码
  • 内存访问模式:使用 TMA 加载(ct.load)而不是 ct.gather;在安全时移除不必要的边界检查(check_bounds=False)
  • 代码生成质量:使用 padding_value 参数而不是手动通过 ct.where 进行掩码;移除 safe_offs
  • 算法重构:K 循环拆分、负载均衡、代数化简

与 Triton Autotune 的差异

主要差异:Triton 使用 @triton.autotune 装饰器和 Config(...) 对象;CuTile 使用 exhaustive_search()、SimpleNamespace 配置、单独的缓存以及 ct.launch。CuTile 没有 num_warps/num_stages(由编译器决定),只有 tile 大小、occupancy 和 num_ctas。CuTile 编译开销更大(最终代码中配置数量应 ≤30)。CuTile 缓存由用户在内存中管理(不会自动持久化)。CuTile 将 args_fn(内核参数)与 hints_fn(编译器提示)分开。

参考文档

类别文档内容
API 参考api-reference.mdexhaustive_search 签名、TuningResult、一次调优/缓存/启动模式、replace_hints、内核提示、search_space 设计、grid_fn 模式
工作流workflow.md端到端工作流:为新内核添加自动调优、多架构配置、torch.autograd.Function 集成、Triton→CuTile 迁移、优化现有配置
陷阱pitfalls.md常见陷阱:原地操作导致数据损坏、编译超时、冷缓存偏差、NCU 干扰、search_space 耗尽、FP8 精度、replace_hints 重新编译
参数设计parameter-space-design.md各类内核的参数空间、跨架构模式、grid_fn 模式、剪枝规则
搜索策略search-strategies.md穷举搜索、A/B 测试方法、DISABLE_AUTOTUNE 模式
模板kernel-type-templates.md8 种内核类型的可复制粘贴自动调优模板
硬件hardware-constraints.md各架构的约束、tile 大小范围、num_ctas 规则、TMA 要求

源代码参考

关键文件:ops/cutile/matmul.py(matmul 自动调优)、ops/cutile/attention.py(FMHA 自动调优)、suites/unsloth/cutile/ct_ops.py(共享的 autotune_configs(),occupancy=[1,2,4,8])、suites/unsloth/cutile/swiglu.py(逐元素操作示例)、suites/unsloth/cutile/rope_embedding.py(分离缓冲区模式)、suites/unsloth/cutile/grouped_gemm.py(持久化 GEMM,仅使用 occupancy)。

实例详解

每个示例都展示改造前 → 改造后模式:fixed_launch.py(硬编码的 ct.launch)和 autotuned_launch.py(重构为一次调优/缓存/启动)。

目录内核自动调优模式复杂度重点说明
assets/examples/01_rmsnorm_occupancy_only/RMSNorm(归约)仅调优 Occupancy [1,2,4,8]低最常见的模式——不调优 tile,只寻找最佳 occupancy。网格 = NUM_SM * cfg.occupancy。不是原地操作。
assets/examples/02_matmul_full_search/GEMM C=A@B完整搜索:TILE_M/N/K + occupancy + num_ctas(sm90+)高计算密集型内核,包含多个可调优维度。args_fn 将 tile 大小作为 ct.Constant[int] 传递。grid_fn 依赖 cfg。配置数量 ≤30。
assets/examples/03_rope_inplace_splitbuffer/RoPE 嵌入(原地操作)仅调优 Occupancy,使用分离缓冲区中原地内核在搜索期间必须使用分离缓冲区,以避免数据损坏。搜索会写入暂存区;最终的 ct.launch 使用真实的原地参数。

评论

支持 Markdown