简街机器学习工程师用cuTile将Mamba-2递归状态计算映射到张量核心

序列建模架构面临一个根本性计算权衡：标准softmax注意力训练时利用张量核心加速矩阵乘法，但推理时计算和内存需求随序列长度线性增长；循环神经网络推理时计算量和内存固定，但训练时无法并行化。Mamba-2属于试图兼得两者的状态空间模型家族——它拥有类似RNN的固定大小隐藏状态，但状态更新可以按块（chunk）并行计算。

简街机器学习工程师Sylvain Gugger和Axel Feldmann与英伟达的Yinuo Liu、Siddharth Samsi合作，将Mamba-2前向传播实现为英伟达cuTile（CUDA Tile）的Python DSL核函数，并在英伟达技术博客上完整展示了推导过程。cuTile将线程和warp级调度抽象为tile-centric原语，编译器自动将操作映射到张量核心等高级硬件特性。

核心机制是跳跃公式。忽略衰减项时，从时刻t的隐藏状态跳到t+T只需累加中间的外积：h(t+T) = h(t) + BX^T，其中B和X分别是将chunk内各token的b和x堆叠成的矩阵。这直接是一个矩阵乘法，因此可以用张量核心执行。加入衰减后，对数空间的累积和将连续乘法替换为加法，预乘指数后的行仍保持矩阵乘法形式。

这一推导分解为三个核函数。chunk_cumsum用ct.cumsum在chunk内计算每个token每个头的衰减累积和，结果存储供后续两个核函数复用。fwd_pass_state是chunk上的顺序循环，每次迭代用ct.mma执行跳跃公式推进隐藏状态，将衰减后的b预乘后与x的转置相乘。compute_out是最终的输出计算，将携带状态项、衰减且因果掩码的交互矩阵、跳跃连接三项融合为单一tile核函数，中间矩阵保持在寄存器中，最终用TF32精度通过张量核心执行。

compute_out的设计值得注意。直接PyTorch实现会分别计算三项并将完整的[B, n_chunks, chunk_size, chunk_size]交互矩阵写入全局内存。cuTile版本将三项融合，交互矩阵cb_tile = C @ B^T在Phase 1中用ct.mma逐步构建，Phase 2应用衰减和因果掩码。因果掩码通过ct.arange和ct.where构造下三角条件，无需实例化掩码张量。衰减项exp(G - G^T)中i < j的位置产生正指数，但因果掩码立即将其置零，保证数值稳定。

针对Blackwell第五代张量核心的硬件优化包括：操作数交换——张量核心的MMA操作数不对称，操作数A可以驻留在Tensor Memory而操作数B必须驻留在Shared Memory，利用恒等式B^T A^T = (AB)^T将融合操作数路由到A侧以缓解共享内存争用；内存布局转置——将decay_cumsum和dt从默认strided布局转为[batch, head, token]以确保连续访问并启用TMA；以及cuTile的autotuner对每个输入形状搜索block size和occupancy参数。

自报性能结果显示，优化后的cuTile前向核函数比开源Triton实现快1.14x至1.82x，加速比随序列长度和状态大小增大而增长。测试环境为单张NVIDIA B300 SXM6（Blackwell Ultra，1920 MHz最大SM时钟，1100 W功率限制，268 GiB HBM），对比基线为mamba_ssm 2.3.2.post1的Triton参考实现。

这一加速仅为核函数级别。端到端训练性能还取决于反向传播实现，作者预告后续将发布反向传播的cuTile实现。cuTile目前锁定英伟达GPU，且该实现通过TileGym库提供。

从更广泛的视角看，这篇文章的价值不仅在于性能数字，而在于完整展示了Mamba-2的数学结构如何映射到GPU硬件原语：对数空间累积和对应ct.cumsum，跳跃公式对应ct.mma，因果掩码对应寄存器内的条件选择。这为理解状态空间模型为什么能在保持RNN式常数推理内存的同时获得类注意力的训练并行性，提供了可检查的具体机制。