【翻译】深入探讨 Hopper TMA 单元在 FP8 GEMM 运算中的应用（文末送3本大模型书）

主要观点总结

本文探讨了Hopper TMA单元在FP8 GEMM运算中的应用，介绍了张量内存加速器（TMA）的工作原理，以及其在PyTorch和Triton中的实现和使用。文中详细解释了TMA如何帮助提升GEMM运算性能，特别是针对H100架构的GPU。此外，还讨论了Triton和CUTLASS之间的实现差异，以及这些差异如何影响性能。文章还包含了一些性能分析和未来工作计划的讨论。

关键观点总结

关键观点1: TMA的工作原理

TMA是H100硬件的一个新增功能，允许异步且双向地在GPU全局内存和共享内存之间传输1D-5D张量。TMA非常轻量级，只需一个单独的线程就可以启动TMA传输，这避免了早期GPU中使用寄存器在不同内存空间之间移动数据的要求。

关键观点2: TMA在Triton中的使用

在Triton中，TMA被用于执行从全局内存到共享内存的加载操作，它通过创建一个TMA描述符，该描述符包含张量的关键属性，如基指针、形状和块大小、数据类型等。在kernel调用函数中，我们只需传递一个描述符指针、偏移量、块大小和输入数据类型，这简化了地址计算并减少了寄存器压力。

关键观点3: 性能分析和比较

文章通过详细的性能分析，展示了利用TMA对H100 kernel的重要性，以及在Triton中构建支持TMA的FP8 GEMM kernel所能获得的性能提升。我们还比较了Triton和CUTLASS之间的实现差异，这些差异可能解释了在使用TMA时报告的性能回归。

关键观点4: 未来工作

在未来的研究中，我们计划将CUTLASS架构的TMA加载方式整合到Triton中，并研究FP8 GEMM的Cooperative kernel，以进一步改善结果。此外，我们还将探索在TMA GEMM kernel中利用SplitK策略来获得更多加速。

文章预览

深入探讨 Hopper TMA 单元在 FP8 GEMM 运算中的应用 PyTorch博客资料：https://pytorch.org/blog/hopper-tma-unit/ PyTorch实现和使用Demo：https://github.com/pytorch-labs/applied-ai/blob/main/kernels/triton/inference/fp8/tma_gemm.py 在本文最后也添加了这个代码的解释 摘要 Hopper（H100）GPU架构被称为"第一款真正的异步GPU"，它包含了一个新的、完全异步的硬件复制引擎，用于在全局内存和共享内存之间进行大规模数据移动，这个引擎被称为张量内存加速器（TMA）。虽然CUTLASS通过其异步流水线范式内置了对TMA的支持(https://github.com/NVIDIA/cutlass/blob/56b46e2d13875b46b8f6a03f9f5ac91e2bfdc01a/include/cute/arch/copy_sm90_tma.hpp)，但Triton则通过一个实验性API (https://github.com/triton-lang/triton/blob/538556a66ee49630e1cb0b239f93e63b968b2478/python/triton/tools/experimental_descriptor.py#L25) 来提供TMA支持。 在这篇文章中，我们将深入探讨TMA的工作原 ………………………………