Mamba-2前向计算搬进Python内核：序列越长、状态越大，比Triton实现快1.14到1.82倍
核心概要
Jane Street的机器学习工程师与NVIDIA合作，用NVIDIA CUDA Tile（cuTile）这一Python端GPU编程模型重写了Mamba-2的前向计算，并针对Blackwell张量核心做了调度与访存优化。Mamba-2兼有循环网络的固定大小隐状态和注意力式的矩阵乘法训练路径，但前向计算仍需高效内核。在单块B300上，这组内核比Triton实现快1.14到1.82倍，加速幅度随序列长度和状态规模增大而提高；这是内核级结果，端到端训练还取决于反向传播实现。
深度剖析
Mamba-2的前向计算现在可以用Python写成的三个自定义内核完成，不再依赖手写底层GPU代码。 此前这类序列模型的前向计算通常由Triton等实现承担，cuTile作为较新的Python端瓦片级编程模型尚处早期采用阶段。 在单块NVIDIA B300 SXM6（Blackwell Ultra）上，与mamba_ssm 2.3.2.post1的Triton参考实现对比，这组cuTile内核快1.14到1.82倍，且加速随序列长度和状态规模增大而提高。
三个内核分别负责块内累积和、块起始隐状态和块输出，把循环依赖限制在单个程序实例内。 直接按循环逐token计算无法并行，而并行扫描需要物化巨大的隐状态张量，受内存带宽限制且用不上张量核心。 chunk_cumsum计算每个块内衰减的累积和，fwd_pass_state用跳跃公式沿块串行推进隐状态，compute_out把交互矩阵留在寄存器中融合计算输出；不同batch、head和状态切片的块在网格上并行。
实现借助cuTile的延迟提示、操作数换位、内存布局转置和自动调参来贴近硬件。 Blackwell第五代张量核心的矩阵乘操作数不对称，A可放在张量内存而B必须放在共享内存，默认的跨步访存也会退化为不友好的加载指令。 利用转置恒等式把融合后的操作数送到A位以缓解共享内存压力，把衰减累积和与dt转成按token连续的布局以启用TMA，并用cuTile的穷举搜索按输入形状挑选块大小与占用率。
启示与展望
结果限于Mamba-2前向计算的内核级加速，测试环境为单块NVIDIA B300 SXM6、计算能力10.3，软件为cuda-tile 1.5.0、Triton 3.7.1与PyTorch 2.13.0。实现通过TileGym提供，示例以batch 4、16384个token、8个head、状态规模128、块大小128的bfloat16配置运行。后续将另文介绍反向传播在cuTile中的实现。
加速数字是内核级测量，端到端训练结果还取决于反向传播实现。部分指数项在计算时可能为正，代码随即用因果掩码将其剔除以避免溢出影响结果。不同输入形状下瓶颈会移动，因此需要自动调参而非固定一套优化。