NVIDIA CUDA Tile（cuTile）是一种基于Python的领域专用语言和面向NVIDIA GPU的并行编程模型，能够自动调用先进硬件能力，为Tile级别的GPU编程提供峰值性能。通过将线程级和warp级调度细节隐藏在高层次、以Tile为中心的原语之下，cuTile让开发者能够用Python编写GPU内核，同时编译器优化会自动针对Tensor Core等先进硬件特性进行加速。
作为cuTile的早期采用者，Jane Street的机器学习工程师用cuTile实现了当前最先进的序列模型Mamba-2，并与NVIDIA合作对内核进行了优化。Jane Street是一家技术驱动型的全球交易公司，利用深度学习推动高级量化研究的创新。本文将介绍支撑Mamba-2前向传播的三个自定义内核，以及背后的cuTile优化方法。
概述与前提条件
要使用cuTile，需确保开发环境满足cuTile Python快速入门指南中列出的要求。在本文展示的内核实现中，将该软件包导入为cuda.tile，并在全文中以ct表示。
序列建模的计算权衡
序列建模架构面临一种计算权衡：传统的softmax注意力机制借助Tensor Core加速的矩阵乘法可以快速训练，但在推理阶段，其计算量和内存需求会随序列长度线性增长。相比之下，循环神经网络（RNN）训练速度慢且难以并行化，但无论序列长度如何，其计算量和内存占用都保持恒定。
Mamba-2架构
Mamba-2架构属于一类试图兼顾两者优势的模型家族。与RNN类似，它拥有固定大小的隐藏状态，并且其循环更新可以借助Tensor Core以分块（chunkwise）的方式高效计算。
Mamba-2的核心递归公式为：
\(h_t = \exp(A \cdot dt_t) h_{t-1} + dt_t B_t x_t\)
\(y_t = c_t^\top h_t + D x_t\)
其中：
\(x_t\)、\(dt_t\)、\(B_t\)和\(c_t\)是Mamba-2层在Token t处的输入（通过对模块输入进行投影获得，类似注意力层中的query、key、value）
\(y_t\)是Token t处的输出
\(h_t\)是跨Token携带的固定大小隐藏状态。它不像RNN中那样是一个向量，而是一个矩阵（具有更多数值和更强的表征能力）
\(A\)和\(D\)是学习到的模型参数
直观理解：
\(dt_t\)（每个头一个标量）表示状态应被更新的程度。\(dt\)中所有值均为负数，由于前一状态会被乘以\(\exp(Adt_t)\)，当\(dt_t = 0\)时状态保持不变
\(x_t\)（大小为H的向量）的作用类似于注意力机制中的value
\(B_t\)（大小为S的向量）的作用类似于注意力机制中的key
\(c_t\)（大小为S的向量）决定如何混合状态以生成输出，作用类似于注意力机制中的query
状态\(h_t\)是一个大小为S、H的矩阵（\(B_t x_t^\top\)是外积）。为了便于书写计算过程，令\(a_t = Adt_t\)，并将乘积\(dt_t B_t\)记为\(b_t\)。方程由此简化为：
\(h_t = \exp(a_t)h_{t-1} + b_t x_t^{\mathsf{T}}, y_t = c_t^{\mathsf{T}}h_{t-1} + Dx_t\)
这种递归形式非常适合顺序推理。然而，训练和预填充（prefill）阶段需要比for循环更快的方法。虽然这种线性状态更新支持并行扫描算法，但该方法需要物化一个巨大的隐藏状态张量，导致其主要受限于内存带宽，无法利用Tensor Core。
另一种方法是通过矩阵乘法直接计算未来的任意隐藏状态，从而利用Tensor Core。为简化说明，先忽略\(\exp(a_t)\)项，从\(h_t\)到\(h_{t+T}\)的转换为：
\(h_{t+T} = h_t + b_{t+1}x_{t+1}^{\mathsf{T}} + b_{t+2}x_{t+2}^{\mathsf{T}} + \cdots + b_{t+T}x_{t+T}^{\mathsf{T}}\)
若将所有\(b_{t+i}\)堆叠构建出形状为(T, S)的矩阵\(B\)，并将所有\(x_{t+i}\)堆叠构建出形状为(T, H)的矩阵\(X\)，则上式即为\(h_{t+T} = h_t + BX^T\)。
借助这个跳跃公式，可以高效计算序列中各个分块“检查点”处的隐藏状态。下一步是找到一种方法，仅凭\(h_0\)计算出输出\(y_1\)、\(y_2\)、\(y_3\)、\(y_4\)。如果逐一计算\(h_1\)、\(h_2\)、\(h_3\)、\(h_4\)……这会使跳跃公式失去意义。
同样暂时去掉衰减项\(\exp(a_t)\)，目标是计算：
\(y_{t+i} = c_{t+i}^{\mathsf{T}}h_{t+i} + Dx_{t+i} = c_{t+i}^{\mathsf{T}}\left(h_t + b_{t+1}x_{t+1}^{\mathsf{T}} + \cdots + b_{t+i}x_{t+i}^{\mathsf{T}} \right) + Dx_{t+i}\)
其中借助之前的公式展开了隐藏状态。跳跃连接项\(Dx_{t+i}\)比较容易理解。
再次将一个分块内的\(x_t\)、\(b_t\)、\(c_t\)和\(y_t\)组合成矩阵，就能用对Tensor Core友好的矩阵乘法重新表达这一过程（如图3所示），其中\(\mathit{C}\)与\(\mathit{B}\)乘积中的三角形部分对应因果掩码（与因果softmax注意力中的用法相同）。
计算输出\(y_{t+i}\)时，需要用\(c_{t+i}\)与所有满足\(j \leq t+i\)的\(b_j\)做乘积，因此在计算完整矩阵乘积时，需要屏蔽掉处于未来位置（对角线以上）的部分。用数学语言表达（加回跳跃连接项）：
\(Y = Ch_t + \left((CB^{\mathsf{T}})\cdot M\right)X + DX\)
其中\(M\)为传统的因果掩码。
需要注意的是，这与注意力机制（去掉softmax后）非常相似，因为去掉衰减项后，Mamba-2本质上就是线性注意力（只是换了名称）。加回衰减项会让情况稍微复杂一些。展开递归后的结果如下：
\(h_{(t+T)}=\exp(a_{(t+1)}+a_{(t+2)}+\cdots+a_{(t+T)})h_t + \exp(a_{(t+2)}+\cdots+a_{(t+T)})b_{(t+1)}x_{(t+1)} + \exp(a_{(t+3)}+\cdots+a_{(t+T)})b_{(t+2)}x_{(t+2)} + \cdots + b_{(t+T)}x_{(t+T)}\)
通过转换到对数空间，可以用累加和代替连续乘法。只要预先用指数项对\(B\)（或\(X\)）的行进行乘法，这一步仍然是矩阵乘法。用张量\(G\)记录\(a_t\)的累加和，得到如下新等式：
\(h_{t+T} = \exp(G_{last})h_t + \left(\exp(G_{last} – G)\cdot B\right)X^{\mathsf{T}}\)
将新的隐藏状态跳跃公式代入输出计算中，可得：
\(Y = \exp(G)Ch_{kT} + \left(\exp(G – G^{\mathsf{T}})\cdot (CB^{\mathsf{T}})\cdot M\right)X + DX\)
其中矩阵\(\exp(G – G^{\mathsf{T}})\)在\((i, j)\)位置的值为\(\exp(g_i – g_j)\)。注意，所有\(i < j\)的项都会被因果掩码屏蔽，因此只需保留\(g_t\)累加的\(a_t\)项数多于\(g_j\)的情况。由于所有\(a_t\)值按构造均为负数，这使得\(g_i – g_j\)为负值。任何正数的指数运算结果都不会被保留，从而保证了数值稳定性。
这将转化为三个内核：
chunk_cumsum：计算每个分块内\(a_t\)的累加和。
fwd_pass_state：使用上述跳跃公式计算每个分块起始处的隐藏状态。
compute_out：使用上一步计算得到的分块初始隐藏状态，计算每个分块的输出。
内核一：计算累加和
第2步和第3步都需要用到\(a_t = Adt_t\)的累加和（即前文中的\(G\)）。为避免重复计算，第一个内核只计算一次\(a_t\)并存储结果。需要注意的是，这不是全局累加和，而仅是每个分块内部的累加和。
该内核的输入为每个头一个标量的衰减值a，以及每个Token一个时间步dt。头h在Token t处的有效衰减为a[h] * dt[t, h]，内核沿分块维度累加该乘积。这些累加和均处于对数空间，之后始终需要对其取指数（或某种线性组合后取指数）。由于计算exp2比计算exp成本更低，结果会乘以log2(e)进行缩放。
该内核引入了三个内核都会用到的cuTile原语：ct.bid、ct.load / ct.store，以及ct.cumsum。
完整内核代码十分简洁（原文代码略）。
该内核以(num_chunks, num_head_tiles, batch_size)为网格启动，第一行通过ct.bid确定在网格中的位置。a_tile（形状为[BLOCK_H]）与dt_tile（形状为[1, chunk_size, BLOCK_H]）的广播乘积，得到该Tile上每个Token、每个头的有效衰减值。ct.cumsum(..., axis=1)随后沿分块内Token维度进行扫描，计算结果按照fwd_pass_state和compute_out所需的布局进行存储。
内核二：fwd_pass_state
fwd_pass_state计算每个分块的初始隐藏状态以及最后一个分块之后的最终隐藏状态，它并不直接计算输出，而是准备好供compute_out后续投影为输出的携带状态。该内核的核心是前文所述的跨分块顺序循环。每次内核启动会加载隐藏状态中一个(batch, head, head_tile, state_tile)切片。在该切片内，hid_tile的形状为[BLOCK_HEAD, BLOCK_STATE]，将保存每个分块的初始隐藏状态Tile（原文代码略）。
在循环的每次迭代中，首先存储隐藏状态（原文代码略）。
循环的其余部分加载一个分块的x、b、dt和decay_cumsum，用于计算下一个分块的初始隐藏状态。如前所述，其数学表达为：
\(h_{t+T} = \exp(G_{last})h_t + \left(\exp(G_{last} – G)\cdot B\right)X^{\mathsf{T}}\)
ct.extract可用于从已加载的衰减Tile中取出最后一个Token的值（用于计算\(G_{last}\)）。出于性能考虑，如前所述使用exp2代替exp。ct.mma是用于执行矩阵乘加运算并调用Tensor Core的API（原文代码略）。
在最后一个分块处理完成后，hid_tile中保存的就是最终隐藏状态，并将其存储下来。在推理场景中，该值可作为下一次调用的init_state被回传；在训练场景中，该值通常会被丢弃（原文代码略）。
这种顺序依赖关系仅存在于单个网格点（程序实例）内部，而相互独立的(batch, head, head_tile, state_tile)切片则会在整个网格中并行执行。
内核三：compute_out
compute_out是前向传播中的最后一个内核。它接收fwd_pass_state生成的每个分块的初始隐藏状态，计算每个Token的输出，实现前文推导出的输出方程：
\(Y = \exp(G)Ch_{kT} + \left(\exp(G – G^{\mathsf{T}})\cdot (CB^{\mathsf{T}})\cdot M\right)X + DX\)
若直接用PyTorch实现，会分别计算这三项，并将完整的[B, n_chunks, chunk_size, chunk_size]交互矩阵C @ B.T物化到全局内存中。而compute_out将它们融合进一个分块化的内核，并将该矩阵始终保留在寄存器中。
每次内核启动负责一个(batch, chunk, head, head_tile)输出Tile，并分两个阶段进行计算。阶段一遍历各个状态Tile，使用ct.mma构建交互矩阵cb_tile = C @ B.T，同时累加携带状态项C @ H.T（原文代码略）。
阶段二对cb_tile应用衰减和因果掩码（通过ct.arange和ct.where构建下三角掩码M，因此无需物化任何掩码张量），随后将所有项相加得到最终输出Tile。两个累加器均以ct.full寄存器Tile初始化，最后的ct.mma会转换为ct.tfloat32类型，以接近FP32精度通过Tensor Core进行运算（原文代码略）。
如前所述，尽管在ct.exp2(d_tile[:, None] - d_tile[None, :])这一行中会计算出一些数值不稳定的指数项（其中部分为正数），但会立即将其屏蔽掉，以防止潜在的数值溢出污染计算结果。
Mamba-2性能优化
在用cuTile-Python DSL定义好Mamba-2工作负载之后，引入了若干关键性能优化：
延迟提示：cuTile允许为ct.load和ct.store操作标注延迟提示，以建模DRAM访问流量。编译器会利用这些提示来优化指令调度和流水线缓冲区深度。
交换A/B操作数：在NVIDIA Blackwell第五代Tensor Core上，MMA运算的操作数是非对称的：操作数A可以存放在Tensor Memory中，但操作数B必须存放在共享内存中。对于fwd_pass_state中类似SIMT运算加点积运算的输入融合模式，可利用恒等式\(B^{\mathsf{T}}A^{\mathsf{T}} = (AB)^{\mathsf{T}}\)，将融合后的操作数（b_dt_d）路由到操作数A，从而缓解共享内存的争用。
全局内存布局：decay_cumsum和dt默认的跨步内存访问方式会被编译为对缓存不友好的LDG指令。将其转置为[batch, head, token]布局可以确保访问连续，从而支持使用TMA。
cuTile-Python自动调优器：由于性能瓶颈会随输入形状变化，没有任何单一优化策略能适配所有模型配置。这些优化手段与block_size、occupancy等cuTile参数一起被整合进统一的搜索空间，并借助cuTile的exhaustive_search功能，针对每种输入形状自动找出最优配置。
Mamba-2前向传播性能基准测试
该Mamba-2前向传播内核的实现已通过TileGym开放（具体示例代码略）。
在以下软硬件配置下，对所实现的Mamba-2前向传播性能与一个用Triton编写的开源Mamba-2内核实现进行了对比评测：
软件环境：
mamba_ssm 2.3.2.post1（NVIDIA Dynamo-Triton参考实现），Triton 3.7.1
cuTile：cuda-tile 1.5.0
TileIR汇编器（tileiras）13.4
PyTorch 2.13.0，CUDA 13.3（驱动和运行时API，NVCC V13.3.33），驱动版本595.58.03
cuda-python / cuda-bindings 13.3.1
Python 3.12.3，Ubuntu 24.04.4 LTS，内核版本6.8.0-107
硬件环境：
单张NVIDIA B300 SXM6，计算能力10.3（NVIDIA Blackwell Ultra架构）
SM最高时钟频率1920 MHz，功耗限制1100 W，HBM显存268 GiB
测试结果显示，经过优化的cuTile前向传播内核速度提升了1.14倍至1.82倍，且随着序列长度和状态规模的增加，加速效果更为明显。需要说明的是，这些仅是内核级别的加速比，端到端训练效果还取决于反向传播的具体实现。
开始使用cuTile
cuTile让开发者能够用Python编写清晰易读且性能卓越的内核。敬请期待后续文章，介绍反向传播在cuTile中的具体实现方式。
欢迎开始使用cuTile，并查阅相关文档、GitHub代码仓库与示例代码。
Q&A
Q1：什么是NVIDIA CUDA Tile（cuTile）？
A：cuTile是一种基于Python的领域专用语言和并行编程模型，面向NVIDIA GPU，能够自动利用先进硬件能力，让开发者用Python编写高性能的Tile级GPU内核，编译器会自动调用Tensor Core等硬件特性进行优化。
Q2：Mamba-2架构相比传统注意力机制有什么优势？
A：Mamba-2兼具RNN和注意力机制的优点，拥有固定大小的隐藏状态，推理时计算量和内存占用恒定，同时支持分块并行计算，可借助Tensor Core加速训练，解决了传统softmax注意力推理成本随序列长度线性增长的问题。
Q3：用cuTile实现的Mamba-2前向传播性能提升了多少？
A：在NVIDIA B300 SXM6硬件上，与Triton实现的开源Mamba-2内核相比，cuTile优化后的前向传播内核速度提升了1.14倍至1.82倍，且序列长度和状态规模越大，加速效果越明显。
