用 CUDA Tile 加速 Mamba-2 前向传播内
来自 NVIDIA Developer · 期号 2026-10-20
摘要
NVIDIA 的 CUDA Tile（cuTile）是一个 Python DSL，让开发者用 tile 级原语写 GPU 内核，编译器自动利用 Tensor Core 等硬件特性。 Jane Street 的机器学习工程师用 cuTile 实现了 Mamba-2 状态空间模型，并与 NVIDIA 合作优化了三个前向传播内核：chunk_cumsum 在 log 空间计算每个 token 的衰减累积和；fwd_pass_state 用 Tensor Core 矩阵乘法计算 chunk 间的隐藏状态跳转；compute_out 把携带状态项、衰减因果交互矩阵和残差连接融合成一个 tile 内核，中间矩阵留在寄存器里，转成 TF32 交给 Tensor Core 执行。 优化手段包括：用延迟提示建模 DRAM 流量、交换操作数让融合后的操作数走 Tensor Memory、用转置内存布局适配 TMA、以及一个自动调优器为每种输入形状搜索最优的 block 大小和 occupancy。 原文是技术详解，附了完整的内核代码。
深读
导读 NVIDIA 的 cuTile 是一个 Python DSL，让开发者用 tile 级原语编写 GPU 内核，编译器自动利用 Tensor Core 等硬件特性。Jane Street 的机器学习工程师用 cuTile 实现了 Mamba-2 状态空间模型，并与 NVIDIA 合作优化了三个前向传播内核。读完本文能理解 Mamba-2 的 chunk 计算如何转化为三个可并行的内核，以及 cuTile 在 Blackwell 架构上的具体优化手段。 关键信息 - Mamba-2 的核心递推公式为 h_t = exp(A·dt_t)h_{t-1} + dt_t·B_t·x_t，其中隐藏状态 h_t 是矩阵而非向量，大小为 S×H。 - 三个前向内核分别是：chunk_cumsum（在 log 空间计算每个 chunk 内 token 的衰减累积和）、fwd_pass_state（用 Tensor Core 矩阵乘法计算 chunk 间的隐藏状态跳转）、compute_out（将携带状态项、衰减因果交互矩阵和残差连接融合成一个 tile 内核）。 - compute_out 内核将完整的 [chunk_size, chunk_size] 交互矩阵 C@B.T 保留在寄存器中，避免写入全局内存。 - 优化手段包括：用延迟提示建模 DRAM 流量、利用 B^T A^T = (AB)^T 将融合后的操作数路由到 Tensor Memory、将 decay_cumsum 和 dt 转置为 [batch, head, token] 布局以启用 TMA、以及用 autotuner 搜索 block size 和 occupancy。 - 在单张 NVIDIA B300 SXM6（Blackwell Ultra）上，cuTile 内核比 Triton 参考实现快 1.14x 到 1.82x，加速比随序列长度和状态大小增加。 逐节讲解 Mamba-2 的计算权衡与 chunk 化 序列建模架构面临一个计算权衡：传统 softmax attention 训练快（矩阵乘法可用 Tensor Core），但推理时计算和内存需求随序列长度线性增长；RNN 训练慢且不可并行化，但推理时计算和内存恒定。Mamba-2 的目标是兼得两者：它有一个固定大小的隐藏状态（矩阵，而非向量），其循环更新可以按 chunk 用 Tensor Core 高效计算。 核心递推公式为 h_t = exp(A·dt_t)h_{t-1} + dt_t·B_t·x_t，输出 y_t = c_t^T h_t + D x_t。其中 dt_t 是每个 head 的标量（值为负，控制状态更新幅度），x_t 类似 attention 中的 value，B_t 类似 key，c_t 类似 query。 直接按 token…