[Paper] tritonBLAS：基于 Triton 的 GEMM 核参数选择分析方法

Source: arXiv - 2512.04226v1

概览

本文介绍了 tritonBLAS，一种确定性的分析模型，能够自动为 GPU 上的 GEMM（通用矩阵乘）内核选择高性能参数。通过利用缓存层次结构和数据布局等架构细节，tritonBLAS 可以在无需大多数库依赖的昂贵运行时自动调优的情况下，生成接近最优的内核。

主要贡献

• 分析性能模型：将 GPU 微架构（缓存大小、共享内存布局等）映射到 GEMM 阻塞参数。
• 仅基于 Triton 的实现：轻量级 GEMM 框架，无需手写 CUDA 内核或外部库。
• 零运行时自动调优：模型在编译时预测最优配置，达到 >95 % 的最先进自动调优方案的性能。
• 广泛评估：在多种矩阵形状和现代 GPU（如 NVIDIA Ampere、Hopper）上进行实验，展示了一致的加速和低开销。
• 开源潜力：该方法可集成到现有基于 Triton 的项目中，或作为生产流水线中 cuBLAS/rocBLAS 的即插即用替代方案。

方法论

1. 架构剖析 – 作者从目标 GPU 中提取关键硬件参数（L1/L2 缓存容量、每个 SM 的共享内存、warp 数量等）。
2. 分析阻塞模型 – 他们构建了一组方程，将这些参数与 GEMM 切分选择（块大小、线程块形状、寄存器使用）关联。模型捕获了三层阻塞：
   • 寄存器切分（内部微内核）
   • 共享内存切分（中间层）
   • 缓存层切分（外层）
3. 参数选择算法 – 通过轻量级搜索（通常为闭式解），评估可行的切分配置，并在满足内存带宽和占用约束的前提下，选取能够最大化估计算术强度的方案。
4. 在 Triton 中生成内核 – 选定的参数被填入通用的 Triton GEMM 模板，Triton 编译器将其特化为具体的 GPU 内核。
5. 验证 – 将生成的内核在一系列矩阵尺寸（从小的 64×64 到大的 8192×8192）上与自动调优库（如 cuBLAS、CUTLASS）进行基准测试。

结果与发现

• 性能一致：在方阵、长瘦矩阵和短宽矩阵上均表现稳定。
• 零调优时间：分析模型仅需几微秒，而传统自动调优器的搜索阶段需数分钟。
• 内存效率：所选切分遵循 L2 缓存复用，导致相较于朴素内核更低的 DRAM 流量。
• 可扩展性：该方法可扩展至多 GPU 环境，因为模型完全本地于每个设备，无需跨设备剖析。

实际意义

• 更快的部署周期 – 团队可以在不等待每台机器进行自动调优的情况下，直接发布新的 GEMM 密集型工作负载（如 transformer 推理、科学模拟）。
• 可预测的性能 – 由于模型是确定性的，性能回归更易追踪和调试，胜过随机的自动调优结果。
• 降低云成本 – 消除长时间的自动调优阶段直接转化为按需 GPU 实例的计算费用下降。
• 易于集成 – TritonBLAS 可通过一次 import 嵌入现有 Triton 代码库（例如用于扩散模型的自定义内核），即提供高性能的 GEMM 原语。
• 可移植性 – 分析模型会自动适配新一代 GPU；仅需更新硬件参数提取步骤，即可面向未来架构保持有效。

局限性与未来工作

• 模型精度 – 虽然 >95 % 的自动调优性能已相当出色，但在极度不平衡的矩阵形状上仍存在小幅差距。
• 非 GEMM 内核 – 当前框架聚焦于稠密矩阵乘；将分析方法扩展到卷积或稀疏内核仍是未解挑战。
• 动态工作负载 – 对于运行时会改变矩阵大小的任务，需要对模型进行轻量级重新评估；作者计划缓存并复用先前的选择。
• 硬件多样性 – 本研究主要针对 NVIDIA GPU；将相同方法应用于 AMD 或 Intel GPU 需要对架构模型进行相应调整。

总体而言，tritonBLAS 证明了精心构建的分析模型能够取代昂贵的 GEMM 自动调优，为开发者提供一种快速、可靠且可移植的方式，以实现接近最优的 GPU 性能。

作者

• Ryan Swann
• Muhammad Osama
• Xiaohu Guo
• Bryant Nelson
• Lixun Zhang
• Alex Brown
• Yen Ong
• Ali Yazdani
• Sean Siddens
• Ganesh Dasika
• Alex Underwood

论文信息

• arXiv ID: 2512.04226v1
• 分类: cs.DC
• 发表时间: 2025 年 12 月 3 日
• PDF: Download PDF