NVIDIA用CUDA Tile加速Mamba-2前向核，状态空间模型推理提速

NVIDIA 开发者博客近期介绍了用 CUDA Tile（cuTile）为 Mamba-2 编写高性能前向传播核的实践。Mamba-2 属于状态空间模型（SSM）家族，以线性复杂度处理长序列，正成为 Transformer 在长上下文场景下的重要替代方案。而要把它的理论优势兑现为实际吞吐，核函数层面的工程优化是绕不开的一关。

CUDA Tile 是什么，为什么适合写 Mamba-2 核

cuTile 是 NVIDIA 推出的基于 Python 的领域特定语言（DSL）与并行编程模型。它让开发者用接近数学表达的方式描述分块（tile）级计算，由编译器自动映射到 GPU 的线程层次、共享内存和异步拷贝等硬件能力上。换句话说，原本需要手写大量 CUDA C++ 样板代码的工作，被抽象成了更高层的块操作。

这对 Mamba-2 尤其关键。Mamba-2 的前向计算包含选择性扫描（selective scan）、状态递推以及若干逐元素的门控运算，既有沿序列维度的串行依赖，又有沿特征维度的并行空间。cuTile 的分块语义恰好能表达这种"块内并行、块间递推"的结构，让扫描过程以 tile 为单位流水推进，同时把状态驻留在片上存储中，减少对全局显存的往返访问。

领域特定语言（DSL） 是针对某一特定应用领域设计的编程语言或编程模型，与通用语言（如 C++、Python）相对。DSL 通过限定问题范围来换取表达力和性能：编译器可以假定特定的计算模式（例如分块矩阵运算、张量收缩），从而自动完成线程映射、内存布局优化和指令选择。在 GPU 编程中，cuTile、Triton、Halide 等都属于此类尝试，目标是让开发者不必手动管理线程块、共享内存和同步屏障，就能获得接近手写 CUDA 的性能。代价是灵活性受限，超出 DSL 表达范围的操作往往需要回退到更低层的实现。

选择性扫描（selective scan） 是 Mamba 系列 SSM 的核心运算。与经典线性时不变 SSM 不同，Mamba 的离散化参数 Δ、输入矩阵 B 和输出矩阵 C 都随每个 token 动态变化，使模型能够根据内容“选择”保留或遗忘历史状态。前向计算因此是一个沿序列维度的递推：h_t = A_bar_t · h_{t-1} + B_bar_t · x_t，其中 A_bar、B_bar 由 Δ 决定。由于递推存在严格的时间依赖，无法像注意力那样一次性并行计算全部位置，必须借助分块扫描或并行前缀和算法把串行链拆解为块内并行与块间传递两阶段，这也是核函数设计的主要挑战所在。

Mamba-2 前向核的优化重点

据 NVIDIA 博客的描述，这套前向核的核心思路是把序列切分为若干块，每块内部并行计算局部贡献，块与块之间通过状态传递串联。这种分块扫描策略与并行前缀和的思想一脉相承，但针对 SSM 的输入依赖参数（如 Δ、B、C 随 token 变化）做了专门处理。

工程上的收益主要来自三处：一是分块后计算密度提高，矩阵乘与逐元素操作可以更充分地喂饱张量核心；二是中间状态尽量保留在片上，降低带宽压力；三是借助 cuTile 的自动调度，异步数据搬运与计算能够重叠。对于推理场景，前向核的效率直接决定了首 token 延迟和长序列吞吐，因此这类优化有直接的落地价值。

并行前缀和（parallel prefix sum / scan） 是一种把满足结合律的递推运算并行化的经典算法。以累加为例，朴素做法需要 O(n) 步串行，而 Blelloch 扫描等算法可以在 O(log n) 步内完成，代价是额外的计算量。SSM 的状态递推本质上是仿射变换的复合，而仿射变换的复合满足结合律，因此可以套用扫描算法：先把序列切成若干块，每块内部独立计算局部前缀结果，再对块间的聚合结果做一次小规模扫描，最后把块间结果广播回各块修正。这一“两遍扫描”结构正是分块扫描核的理论基础，块大小的选择需要在并行度、片上存储占用和重计算开销之间权衡。

从 Python DSL 到生产级核函数

cuTile 的定位并不是替代 CUDA C++，而是降低高性能核函数的开发门槛。研究者可以用 Python 快速迭代算法变体，验证分块策略是否奏效，再把成熟方案下沉到更底层的实现。NVIDIA 选择用 Mamba-2 作为示例，某种程度上也是在示范这类新型架构如何在新编程模型上获得良好支持。

对做模型部署的团队来说，这意味着两件事：一是 SSM 类模型的算子库生态在补课，过去围绕注意力机制打磨的融合算子经验，需要重新迁移到扫描与递推结构上；二是 Python 优先的核开发流程有望缩短"论文复现到线上服务"的周期。当然，cuTile 目前仍在演进中，其性能上限、对最新硬件的覆盖程度以及与其他编译栈的协同，还需要更多实测数据来检验。

算子融合（operator fusion） 指把多个逐元素或轻量运算合并进同一个核函数，避免中间结果反复写入全局显存。在 Mamba-2 前向中，扫描之后通常紧跟门控、归一化和输出投影等操作，若逐个独立启动核函数，每次都要读写一遍张量，带宽开销会迅速累积。融合后中间值可以停留在寄存器或共享内存中直接参与下一步计算，既减少访存也降低核启动延迟。围绕注意力机制打磨的 FlashAttention 等融合算子经验之所以难以直接照搬，是因为 SSM 的融合边界由扫描的递推结构决定，需要重新划分计算阶段并安排状态在片上的生命周期。

小结

Mamba-2 的竞争力不只取决于架构设计，也取决于底层核函数能否把线性复杂度的理论优势转化为实测吞吐。NVIDIA 用 cuTile 编写前向核的尝试，展示了一条用高层 DSL 快速逼近硬件极限的路径。对关注长上下文推理成本、以及状态空间模型工程落地的读者，这一方向值得持续跟踪。