SuperAPI AI 工具集

kernel-tileir-optimization

编程与开发

查看项目

更新于 10-16 09:47

技能介绍

用途:为已有 Triton kernel 针对 NVIDIA TileIR 后端(Blackwell,sm_100+)做优化,不负责从零编写 kernel。TileIR 后端用 CGA 级 tile 表示生成 CUDA 代码,与 PTX 后端差异明显:num_warps 被忽略,num_stages 仅作成本提示,occupancy(1-32)成为关键调参项,num_ctas 支持 Blackwell 的 2CTA,TMA 为 dot 类 kernel 所必需,且块尺寸通常越大越好。

主要能力:五阶段工作流——兼容性检查、分类、变换、校验、基准测试。按 kernel 类型选择策略:dot 类(GEMM、Attention)用 TMA+2CTA+occupancy 加大块,预期 1.2-2.0x;norm 类(LayerNorm、Softmax)用高 occupancy,预期 2.0-5.0x;element-wise 用 occupancy+num_stages,预期 1.5-3.0x;reduction 用高 occupancy,预期 1.8-4.0x。

使用方式:先运行 python scripts/tileir_check.py 检查 nvtriton、ENABLE_TILE、Blackwell GPU 是否就绪;用 python scripts/classify_kernel.py --file kernel.py 分类,加 --apply-optimizations 可直接应用配置(输出 optimized_code 与 changes_applied)。TileIR 通过 ENABLE_TILE=1 激活;TileIR 专属配置需用 torch.cuda.get_device_capability()[0] >= 10 做门控,pre-Blackwell 上运行 TMA/2CTA 会崩溃。校验用 verify_kernel.py,基准测试用 triton.testing.do_bench() 对比 ENABLE_TILE=0 与 1,并在多种输入尺寸下测试。

限制:需已安装 triton;一次只能安装一个 triton 包(pytorch-triton、triton、nvtriton);无 Blackwell 硬件时仍可加配置但不做 ENABLE_TILE 测试与基准。TMA 必须传 pre_hook=tma_set_block_size_hook,GEMM 需用 b.T.contiguous() 并 tl.dot(a, b.T, accumulator)。2CTA 需调整 NUM_SMS。配置函数必须接受 pre_hook=None。不支持 tl.sigmoid,需用 1.0/(1.0+tl.exp(-x))。若兼容性、校验失败或无加速应中止并报告。