跳到主要内容
读懂 CuTe

前言

前言

读 CuTe 代码时,最难的往往不是某个接口,而是把一连串布局操作还原成具体的数据流:一个元素原来在哪里,交给了哪个线程,又以什么形式进入矩阵乘法指令?

这本笔记围绕这些问题展开。我们从可以手算的矩阵和地址表出发,把坐标、线程、共享内存与寄存器放在一起看,再用源码核对每一步。希望读完一段代码后,你既能说出它做了什么,也能指出某个线程实际持有哪些值。

从具体元素理解抽象​

Layout 描述坐标怎样映射到位置;Tensor 将布局与存储放在一起;Copy 和 MMA 则进一步规定线程怎样参与搬运与计算。这些抽象各有用途,却始终围绕同一批数据。

因此,文中的推导尽量保留具体数字。遇到 partition,就列出线程拿到的坐标;遇到 retile,就比较前后的寄存器视图;遇到跨轮累加,就把旧状态与新贡献分别算出来。公式、矩阵图和程序输出用来相互核对。

讨论范围​

这里主要讨论 CUTLASS 中的 CuTe C++ 接口,关注 Layout、线程分工、数据搬运以及 Tensor Core 的输入输出。FlashAttention 用来展示这些机制怎样组成一次完整计算。

不同 GPU 架构的指令、寄存器组织和同步方式会有区别。每篇文章会说明所用配置;其中的具体坐标与线程数量,应与该配置一起理解。

阅读约定​

文中默认读者了解基本的 CUDA kernel、block 与 warp,并能阅读 C++ 模板、auto 和 constexpr。涉及 CuTe 的类型、布局与接口时,会结合当前例子解释。

除特别说明外,线程编号、矩阵行列与数组下标都从 0 开始。数学符号可能沿用论文的编号方式,相关章节会说明它与代码下标的对应关系。矩阵图中的坐标、存储偏移和寄存器槽位分别表示不同的量,推导时会区分它们。

示例与源码​

每篇文章保留相关示例和源码引用,方便把图中的数值带回程序。小矩阵说明计算关系,完整配置说明真实线程分工;手算结果不能代替实际 GPU 上的精度与性能测试。

这本笔记关心的是读懂机制。在一次搬运或计算中,先找到数据,再确认线程和指令,最后把它们连成完整过程。