[Hopper 架构特性学习笔记 Part1] Distributed Shared Memory

主要观点总结

本文主要介绍了NVIDIA Hopper架构的新特性——分布式共享内存（Distributed Shared Memory）和线程块簇（Thread Block Clusters），以及在CUDA上如何使用分布式共享内存。文章还讨论了在一些特定场景下，如GEMM SplitK类（包括Flash Decoding）算子，使用分布式共享内存和线程块簇进行优化的情况。

关键观点总结

关键观点1: NVIDIA Hopper架构的分布式共享内存和线程块簇介绍

Hopper架构通过引入线程块簇来解决单个线程块处理数据规模有限的问题，提供了更大粒度的线程组。Distributed Shared Memory是线程块簇的一个重要特性，它允许不同线程块访问各自共享内存的集合，提高了处理大规模数据的能力。

关键观点2: 如何在CUDA中使用Distributed Shared Memory

在CUDA中，使用Distributed Shared Memory需要通过Cooperative Groups模块提供的接口来实现。包括获取当前cluster、获取其他Thread Block的Shared Memory地址、Cluster级别的同步等关键操作。

关键观点3: Thread Block Cluster Kernel的启动方式

由于传统的Kernel Launch方式无法传递Cluster Size这个参数，所以Thread Block Cluster Kernel的Launch方式需要使用cudaLaunchKernelEx这个函数，通过设置cudaLaunchConfig_t以及cudaLaunchAttribute来设置grid、block以及cluster配置。

关键观点4: 同步优化

为了减少cluster.sync()的同步开销，CUDA提供了拆分cluster.sync()的能力，通过arrival_token和barrier_wait两个接口实现同步点与等待点的设置，从而在等待期间插入与Shared Memory无关的操作，提升性能。

关键观点5: 总结和示例应用

文章总结了Hopper架构的Distributed Shared Memory和Thread Block Clusters的使用场景和优势，并以Flash Decoding Kernel为例，说明了如何使用这些特性进行优化。

文章预览

原文：https://zhuanlan.zhihu.com/p/708645371 最近 FlashAttention3 爆火，看了下其中的技术报告，发现大量使用了 NV Hopper 架构的新特性。Hooper 架构围绕影响并行程序性能的两个关键因素进行设计，分别是 数据局部性 以及 异步执行 。所以 Hopper 架构（https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth）针对这两个关键因素，在软件层提供了两方面的编程能力： 新线程、显存层次 ：通过新增 Thread Block Cluster 这一线程层次，提供跨 Thread Block 的 Shared Memory 访问。开发者可以基于 Thread Block Cluster ，利用  Distributed Shared Memory  实现高效的多 Thread Block 的协同运行； 访存计算异步执行： Hopper 在硬件层提供了  TMA 单元 ，在软件层可以通过 cuda::memcpy_async 使用 TMA 单元实现异步的 Global Memory 和 Shared Memory 之间的拷贝。 本文主要学习 Thread Block Clusters，研究如何基于 Di ………………………………