【翻译】深入探讨 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策略来获得更多加速。

正文

请到「今天看啥」查看全文

TMA描述符包含张量的关键属性：

• 基指针
• 形状和块大小
• 数据类型

TMA描述符在kernel执行前在主机上创建，然后通过将描述符传递给torch张量来移动到设备上。因此，在Triton中，GEMM kernel接收一个指向张量映射的全局指针。

Triton Host Code

   desc_a = np.empty(TMA_SIZE, dtype=np.int8)
   desc_b = np.empty(TMA_SIZE, dtype=np.int8)
   desc_c = np.empty(TMA_SIZE, dtype=np.int8)

   triton.runtime.driver.active.utils.fill_2d_tma_descriptor(a.data_ptr(), m, k, block_m, block_k, a.element_size(), desc_a)

   triton.runtime.driver.active.utils.fill_2d_tma_descriptor(b.data_ptr(), n, k, block_n, block_k, b.element_size(), desc_b)

   triton.runtime.driver.active.utils.fill_2d_tma_descriptor(c.data_ptr(), m, n, block_m, block_n, c.element_size(), desc_c)
  
   desc_a = torch.tensor(desc_a, device='cuda')
   desc_b = torch.tensor(desc_b, device='cuda')
   desc_c = torch.tensor(desc_c, device='cuda')

这是在kernel调用函数中用于设置描述符的代码。

Triton Device Code

偏移量/指针算术:

   offs_am = pid_m * block_m
   offs_bn = pid_n * block_n
   offs_k = 0

Load:

  a = tl._experimental_descriptor_load(a_desc_ptr, [offs_am, offs_k], [block_m, block_k], tl.float8e4nv)
  b = tl._experimental_descriptor_load(b_desc_ptr, [offs_bn, offs_k], [block_n, block_k], tl.float8e4nv)

Store:

 tl._experimental_descriptor_store(c_desc_ptr, accumulator, [offs_am, offs_bn])

我们不再需要在 kernel 中为加载和存储函数计算指针数组。相反，我们只需传递一个描述符指针、偏移量、块大小和输入数据类型。这简化了地址计算并减少了寄存器压力，因为我们不再需要在软件中进行复杂的指针算术，也不需要专门分配CUDA Core来进行地址计算。

TMA 性能分析

下面，我们讨论Hopper架构上不同加载机制的PTX指令。

用于Load Tile 的 PTX（cp.async）- H100无TMA

# 这两行计算了共享内存中的目标地址。%r100可能是共享内存的基地址,%r8和%r9是偏移量。
add.s32  %r27, %r100, %r8;
add.s32  %r29, %r100, %r9;
# 这行根据条件%p18选择%r102或0,结果存入%r30。这可能用于控制是否执行复制操作。
selp.b32  %r30, %r102, 0, %p18;

# 这两行是关键的异步复制指令。它们从全局内存(%rd20和%rd21)复制数据到共享内存(%r27和%r29)。0x10表示复制16字节。%p1是一个谓词,控制是否执行这些指令。
@%p1 cp.async.cg.shared.global [ %r27 + 0 ], [ %rd20 + 0 ], 0x10, %r30;
@%p1 cp.async.cg.shared.global [ %r29 + 0 ], [ %rd21 + 0 ], 0x10, %r30;

# 这行提交之前的异步复制操作组,确保它们开始执行。
cp.async.commit_group ;

总的来说,这段代码实现了从全局内存到共享内存的异步数据复制。它使用了H100之前的cp.async指令,而不是新的TMA机制。这种方法需要更多的寄存器来计算地址,并且每个线程都参与了数据移动,这与TMA的轻量级、单线程触发的方式形成对比。

在这里，我们观察到较旧的cp.async指令负责全局内存复制。从下面的跟踪中我们可以看到，两次加载都绕过了L1缓存。

请到「今天看啥」查看全文