下载视频您将生成这些视频，使用本文中构建的B200注意力内核生成。由使用在本帖中实现的Blackwell B200注意力内核的视频模型生成。见下面的顶点项目。覆盖在14内核进程和最后的小优化中的同一运行的股票FA4的百分比。每种实现都在其首选的连续布局上计时；参见基准校准。这个博客的主题在这个博客中，我们从头开始用CUDA和一点PTX构建一个密集的B200注意力内核，从基线到FA4论文中使用的4K、8K和16K形状的FlashAttention-4性能的94.4%。主要贡献是视觉指南：围绕60幅图解构建的初学者友好的进度。我们首先建立对简单内核如何工作的直观理解，然后一次添加一个优化，附有详细的图解、简明的解释和代码。对于顶点项目，我们将最终内核插入到视频生成模型中。这里的重点不是要挤出每一分性能，这将是我下一个博客的重点。它是目前存在的最难的内核之一，运行在最新硬件上，所以很有趣。研究：这旨在为您在最新硬件上进行自己的GPU内核研究打下基础。我们专注于B200注意力，但许多概念和思维模型超出了此内核。到最后，您将理解此内核并能够提出自己的优化思路。前提条件：我包含了一个初学者入门。如果您是完全的初学者，或者觉得您的CUDA基础不牢固，请在继续之前遵循此脚注1。我仍然假设您对CUDA有基本了解，但对Blackwell没有先前知识。新概念以可视化方式逐步引入，只有在必要时才引入。因此，如果Blackwell对您来说是新事物，请继续阅读。您应该能跟上这个进度。我假设您知道什么是注意力。我还假设您对在线softmax背后的想法有一些熟悉。这部分不是Blackwell特定的，还有很多资源可供参考。您在家中不需要有B200才能跟随，我也没有。代码：所有代码都可以在我的B200注意力库中获得。在基线之后，每一章增加一个主要的优化。源代码是组织良好的，因此相邻的内核大部分容易比较差异。所以您可以看到每个优化更改的内容，一次一个。顶点项目：最后，我们将把您理解的内核插入到视频生成模型中，并生成美丽的视频。请参阅主章节后的顶点项目。DSL与CUDA：原始的FA4是用CuTe编写的，我个人发现原始的CUDA加一点PTX更容易理解（抽象层更少），所以我们将用CUDA来实现我们的内核。我们不会将FA4的CuTe实现翻译成CUDA语法，而是通过干净而直观的进度来理解一个快速的B200注意力内核。然而，FA4是这项工作的主要参考之一，大多数优化思路都改编自它（而FA4本身是基于cutlass和cute-dsl内核的，见致谢）。现有资源：有优秀的资源解释H200和B200上的优化matmul（见致谢）。但对于B200注意力，我没有找到我想要的深入解释。如果一些资源覆盖最终产生的内核，却不解释进程或大多数优化背后的低级动机。其他人则停留在高层和表面，比如总结管道和warp角色，但跳过了实际理解内核所需的大部分工作和引导。没有一个能给我我想要的深入解释。范围：我们即将优化的内核是密集的，头维度128，非因果，BF16。如果你仅仅了解AI，我们拥有与您一样的AI，可能在使用上更出色 – tomcr00se 第一部分 — 基础 第一章 – 基线 Blackwell B200注意力 i. 本章的路线图 我们首先研究每个CTA做什么工作，以及如何将工作分配给B200的不同CTA（工作并行化部分）。然后我们将进行一次初学者的可选旁路，涵盖逻辑与物理表示、指针和瓷片matmul。接着我们将深入探讨我们的b200注意力内核，并讨论每个CTA内部发生了什么。所有这些将在后面会以更详细的视觉形式解释。现在只是展示一下布局。然后我将链接代码（直接映射到我们的图解中）。基线注意力内核 让我们开始理解第一个内核。我们后来的内核大多使用相同的数学和此处介绍的相同tcgen05概念。因此，在本章中，我将涵盖我们将使用的基础知识。这就是为什么第一章比后面的章节更长。第一个内核存在作为起点，产生正确的数值结果，但效率远不及我们后来的内核。这个第一个内核已经不简单，并且使用了许多Blackwell特定的特性，我们将逐步覆盖它们。主要瓶颈