用途:指导智能体用 OpenAI Triton 在 Python 中编写 GPU 内核,适合需要比纯 PyTorch 更快、但不想写原生 CUDA 的场景,如自定义 PyTorch 算子、融合内核原型(softmax+scale+bias)、用 autotuner 比较块大小与 warp 数、把 NumPy 风格逐元素运算移植到 GPU、以及系统化基准测试内核变体。
主要能力:@triton.jit 装饰器定义内核;tl.program_id 取块索引;tl.arange 生成块内索引向量;tl.load/tl.store 支持 mask 与 other 参数处理尾部元素;块指针(Triton 2.x+)支持结构化二维访问;tl.atomic_add 等原子操作(应少用,优先块内归约后每块单次原子);BLOCK: tl.constexpr 作为编译期常量以分配共享内存并展开循环;@triton.autotune 首次运行时基准测试配置并按 key 形状缓存最优解。
使用方式:用 grid lambda 调用内核并传 BLOCK 等参数;用 triton.testing.Benchmark 做系统对比,或用 triton.testing.do_bench 快速计时;通过 torch.library.custom_op 加 register_fake 元函数接入 PyTorch 与 torch.compile;调试用 tl.debug_barrier 同步块内线程,用 TRITON_PRINT_AUTOTUNING=1 与 TRITON_CACHE_DIR 检查编译产物。
限制与常见问题:BLOCK 过大导致 OutOfResources;缺 mask 导致尾部结果错误;块大小不优时慢于 PyTorch;类型不一致引发编译错误;未初始化掩码通道产生 NaN(需 other=0.0);torch.compile 缺少 fake kernel 会失败。JIT 函数中应避免 Python 副作用。资料未提供具体安装命令、版本要求与许可信息。