NVIDIA 与简街公布了三个 GPU 内核，即在 GPU 上执行特定运算的代码，用于加速 Mamba-2 从输入生成输出的前向计算。优化模型的开发者可使用这些以 Python 编写的开源代码，据该公司公布的数据，其速度为现有开源内核的 1.14~1.82 倍。

1. 将循环模型计算转为分块计算

Mamba-2 是一种循环模型，将先前处理的信息存入固定大小的内部记忆，即“隐藏状态”，并不断更新。在训练或提示词处理时，它不逐个计算词元，而是将多个词元组成的分块计算转为矩阵乘法，以利用 GPU 上专用于矩阵运算的 Tensor Core。

NVIDIA 的 cuTile 是用 Python 编写 GPU 内核的工具。开发者编写数据块（tile）的运算，将具体执行安排交给编译器。简街工程师用它实现了 Mamba-2，并与 NVIDIA 一同优化了内核。

2. 三个内核分别计算状态与输出

第一个内核 chunk_cumsum 在分块内累加并存储衰减率与时间间隔的乘积，衰减率表示先前信息减弱的程度。后两个内核会复用这些值。

第二个内核 fwd_pass_state 计算各分块起始时的隐藏状态，以及最后一个分块之后的状态。它保持分块顺序，同时并行处理相互独立的状态片段。

第三个内核 compute_out 根据分块起始状态生成词元输出。它不将中间矩阵单独写入容量较大的全局内存，而是融合计算，将其保留在靠近运算单元的小型存储空间——寄存器中。

博客中的累加内核如下展示了读取数据块、在分块内累加并存储的流程。

    dt_tile = ct.load(dt, index=(batch, chunk_idx, head_tile_idx),
                          shape=(1, chunk_size, BLOCK_H),
                          padding_mode=ct.PaddingMode.ZERO).astype(ct.float32)
    a_tile = ct.load(a, index=(head_tile_idx,),
                        shape=(BLOCK_H,),
                        padding_mode=ct.PaddingMode.ZERO).astype(ct.float32)

    d_tile = ct.cumsum(a_tile * dt_tile, axis=1) * LOG2E
    ct.store(output, index=(batch, chunk_idx, head_tile_idx), tile=d_tile.astype(output.dtype))

3. 针对硬件优化内存与矩阵乘法

该实现为内存读写附加预期延迟信息，让编译器调整指令顺序。它根据 Blackwell 第 5 代 Tensor Core 的输入布局特性，调换乘法输入的顺序，减少共享内存争用，即多个任务争抢存储空间的现象。改变衰减累积值和时间间隔的内存布局，使数据可连续读取，也是优化措施之一。

由于瓶颈随输入形状变化，该实现会自动搜索块大小（即一次处理的数据组大小）及 GPU 资源利用设置。这种设计在缩短运行时间的同时，也需要针对各输入形状搜索设置组合。

4. 据公司公布的内核测试，速度最高为 1.82 倍

NVIDIA 博客中的测试在单块 B300 SXM6 上使用 cuTile 1.5.0 进行。GPU 最高计算时钟频率为 1920 MHz，高带宽内存容量为 268 GiB。

对比对象是用 Triton 3.7.1 编写的开源内核，为 mamba_ssm 2.3.2.post1 中 NVIDIA Dynamo 系列的参考实现。

测试使用了 PyTorch 2.13.0 和 CUDA 13.3。

据该公司公布的数据，cuTile 内核的速度为 Triton 的 1.14~1.82 倍，序列长度和状态大小越大，加速幅度越大。图表涵盖 1K~256K 个词元，以及 64、128、256、512 的状态大小，但各柱形没有标注精确数值。

博客说明，这些结果是前向内核层面的比较，整体训练性能还取决于反向计算（从输出向输入回溯的计算）的实现。

5. 在 TileGym 中查看实现

代码与基准测试可在 NVIDIA 的 TileGym 仓库中查看。默认运行方式为 cuTile，可按照仓库的安装说明查看运算示例和性能测试。

仓库明确要求 cuTile 1.5.0 及以上版本，以及运行时所需的编译器。安装说明写明，Blackwell GPU 支持 CUDA 13.1 及以上版本，Ampere GPU 支持 CUDA 13.2 及以上版本。

此次发布通过代码展示了 Mamba-2 前向计算的矩阵化和针对硬件的调优如何带来实测收益。