Mamba-2炸裂升级！大模型推理飙升20%，效率翻2.5倍！

在当今人工智能飞速发展的时代，大型语言模型（LLMs）的进步令人瞩目。这些模型在处理长序列数据时，往往面临巨大的计算和内存挑战。Mamba-2模型作为序列模型领域的新星，凭借其独特的状态空间双重（SSD）框架，成功地将结构化状态空间模型（SSM）与基于注意力机制的Transformer相结合，展现出处理超长序列的卓越能力。它不仅是Mamba模型的优化升级，更在实际应用中，尤其是在高达128K甚至更长序列长度的场景下，展现出远超传统自注意力机制的效率。

例如，国际商业机器公司（IBM）近期推出的Granite 4.0模型系列，便巧妙地采用了混合架构，将Mamba-2模块与Transformer模块融合。在Granite 4.0的设计中，每使用一个注意力层来处理长距离上下文时，会配备九个Mamba-2层，这充分说明了Mamba-2在长序列处理上的高效性。随着Mamba-2在这些主流模型中扮演越来越核心的角色，对其性能进行深度优化，以实现更快的推理速度，已成为行业关注的焦点。

Mamba-2模型的核心计算部分是SSD模块，它有效地替代了传统Transformer层中的注意力机制。然而，在原始的Mamba-2 SSD实现中，我们观察到其性能瓶颈主要集中在内存带宽和延迟上，这涉及到大量中间数据的重复写入和读取。针对这一关键问题，有巨大的优化空间。本次，我们聚焦于加速SSD的预填充（Prefill）操作，并通过创新的融合内核技术，取得了令人振奋的成果。

图1. 融合SSD Triton内核在英伟达A100和H100上的加速效果

Mamba-2核心操作与融合的必要性

Mamba-2模块的完整运行涉及一系列复杂的操作。其中，有五项SSD核心操作，它们虽然在概念上构成了一个完整的SSD流程，但在最初的实现中却是作为独立的图形处理器（GPU）内核（Kernel）顺序执行的。这些关键操作包括：

1. 层归一化（Layernorm）：负责稳定数值。
2. 输入投影（In Projection）：将输入数据映射到SSD的通道或维度。
3. 深度卷积（Depthwise Convolution）：混合最近的若干个标记（token）。
4. SSD分块累加（SSD Chunk Cumsum）：计算每个标记的dt值以及分块内的累积衰减。
5. SSD分块状态（SSD Chunk State）：独立计算每个分块末尾的状态。
6. SSD状态传递（SSD State Passing）：计算每个分块末尾的全局状态。
7. SSD批矩阵乘法（SSD BMM）：计算每个输入x分块如何影响对应的输出y分块。
8. SSD分块扫描（SSD Chunk Scan）：根据输入x分块和前一个分块的全局状态计算当前的输出y分块。
9. 输出投影（Out Projection）：将输出映射回模型的隐藏维度。

在Mamba-2的预填充阶段，这五项SSD核心操作在图形处理器上以管道流的形式顺序执行。然而，频繁启动多个独立的GPU内核会引入显著的开销，并且阻碍图形处理器在不同阶段之间高效地重用数据。为了克服这一限制，内核融合技术应运而生，它带来了多重显著优势：

首先，消除内核启动开销。通过将多个内核融合为一个，图形处理器只需进行一次启动，从而大幅减少中央处理器（CPU）与GPU之间的同步和调度延迟。

其次，提升缓存局部性。在融合后的内核中，前一阶段生成的数据可以直接由下一阶段在同一个线程块（threadblock）内使用，大大增加了缓存命中率，减少了全局内存的传输负担。

最后，实现计算重叠。融合内核的不同部分可以在彼此独立的情况下并行执行，从而更充分地利用GPU的计算资源。

我们的创新解决方案便是将这五项SSD核心操作融合到一个统一的Triton内核中，确保一个Mamba-2层的所有SSD预填充计算都能在一次GPU启动中完成，极大地提升了效率。

高效内核融合技术揭秘

与简单的矩阵乘法（matmul）与激活函数融合不同，SSD的融合要复杂得多，因为它涉及多步骤的计算和复杂的依赖关系。在原始实现中，各个独立内核之间的隐式同步是其正常运行的基础。然而，当我们尝试将所有内核融合到一个任务中时，这种隐式的全局同步机制便不复存在，这就要求我们必须显式地管理分块内部和分块之间的所有依赖。

最初，Mamba-2 SSD的五个步骤——分块累加、批矩阵乘法、分块状态、状态传递和分块扫描，是作为独立的内核运行的，它们处理固定大小的标记分块。其中，状态传递步骤在分块之间存在依赖关系，原始的状态传递内核通过在线程块内循环处理分块，并将状态的通道（channels）在线程块之间进行分割以实现并行化，从而处理了这些依赖。这种设计配合内核启动时的隐式全局同步，保证了所有依赖都能得到妥善处理。

然而，将这五个内核融合到一次启动中，我们便失去了原有的隐式全局同步。这意味着，我们必须在代码层面精细管理分块内部和分块间的依赖。对于分块状态、状态传递和分块扫描这三个大型内核，由于大部分依赖都发生在同一分块的不同步骤之间，我们本可以通过让同一个线程块处理分块的所有步骤来解决分块内部依赖，并利用寄存器或L1缓存来保存中间数据。但这种方法在实际操作中存在挑战，因为原始状态传递内核的线程块网格（grid）与其他内核不匹配。

为了实现融合，我们采取了巧妙的策略：将状态传递循环的迭代过程，在不同分块之间分割到独立的线程块中，从而使线程块网格得以匹配。通过原子操作（atomics）对这些线程块进行排序，我们保证了计算的正确性。虽然这看起来可能导致串行化（serialization），效率不高，但通过与其他两个部分的计算重叠，这种开销可以被有效缓解。例如，如果8个分块并行处理，虽然状态传递可能面临约8倍的局部减速，但由于融合后的状态传递在总计算中占比很小，且不再需要从全局内存读取状态，实际的减速远低于理论值。英伟达Nsight Compute的基准测试结果显示，因状态传递同步导致的指令停滞（warp stalls）不到3%，这表明串行化延迟被有效地隐藏了。

对于批矩阵乘法和分块累加这两个速度较快的步骤，我们通过启动独立的线程块来处理。当这些线程块完成任务后，其他处理分块状态、状态传递和分块扫描的线程块会通过原子操作等待它们的依赖，确保数据一致性。当一个线程块开始执行内核时，它会被分配处理分块累加，如果所有分块累加任务都已分配，则会转向批矩阵乘法任务。只有在所有这些快速步骤都被分配后，后续的线程块才会开始处理分块状态，然后是状态传递，最后是分块扫描，从而实现完整的SSD预填充流程。

新媒网跨境了解到，除了上述融合核心机制，还有一系列辅助优化技术：例如，重新调整线程块的执行顺序以隐藏串行化延迟；为数据加载/存储添加缓存提示，优先处理重用数据；将特殊情况的处理分离出融合内核以降低寄存器压力；调整某些中间数据类型；精细调整分块大小；以及重新组织操作以减少延迟。这些共同构建了高效的融合方案。

性能瓶颈的深入分析

为了进一步提升性能，我们使用英伟达Nsight Compute工具对优化后的融合SSD内核进行了细致的瓶颈分析，评估了其最终利用率、停滞模式和资源权衡。

图4. A100 Nsight Compute概述

图5. H100 Nsight Compute概述

从宏观上看，融合SSD内核的计算利用率约为40%至50%，内存利用率约为65%至75%。尽管由于初始加载/存储延迟及其他开销，我们无法实现100%的利用率，但通常在经过良好优化的内核中，达到80%或更高是可行的。与Mamba-2中使用的H100和A100矩阵乘法计算（计算利用率可达85%至96%）相比，SSD内核的利用率仍有提升空间。由于计算和内存利用率均未达到理想状态，这表明瓶颈并非单一地由内存带宽或计算吞吐量引起，而是更为复杂。

我们可以观察线程束（warp）的状态统计数据，了解线程束停滞的原因。其中，“Selected”表示线程束执行了新指令，而“Stall Long Scoreboard”和“Stall Barrier”则指示线程束在等待L2缓存/显存（VRAM）数据或进行同步时处于空闲状态。

图6. H100融合SSD内核的线程束状态统计

为了减少这些停滞的影响，提升计算或内存利用率，我们可以从几个方面着手：

1. 提高占用率（Occupancy）：现代英伟达GPU的每个线程束调度器通常拥有12至16个线程束。如果每个调度器中只有一个线程束，那么每当该线程束停滞时，就会浪费周期。若每个调度器中有16个线程束，即使某个线程束在15/16的时间内处于停滞状态，硬件也不会空闲。占用率表示可用线程束槽位被激活线程束填充的比例。提高占用率有助于隐藏内存和指令延迟，从而提升GPU的利用率。
   
   图7. H100融合SSD内核的占用率

   目前，该融合内核的占用率仅为25%，这主要是受限于寄存器和共享内存。尽管我们可以通过增加线程束数量和减少每个线程的寄存器使用来提高占用率，但在实践中，这可能会因同步成本增加和寄存器压力增大而导致性能下降。

2. 增强指令级并行性（Instruction-Level Parallelism）：通过优化代码，减少指令之间立即依赖的程度，使得即使前一条指令尚未完成，线程束也能执行后续指令。这与提高占用率一样，能够隐藏延迟，且无需增加线程束数量。

3. 减少同步和数据传输：由于线程束通常停滞于加载/存储内存或等待屏障，我们可以通过减少屏障数量，或通过更优的缓存策略、不同的块大小来减少总数据传输量，从而提升性能。

值得注意的是，这三类优化技术有时会相互冲突，需要进行权衡。例如，GPU中每个流多处理器（SM）的寄存器和共享内存是有限的，若每个线程块使用过多资源，则会降低占用率。通过分阶段加载数据可以增加指令级并行性，但这又需要更多的寄存器和共享内存，反而降低了占用率。调整块大小有助于减少数据传输总量或提高缓存命中率，但这同样会消耗更多资源，降低占用率。因此，融合内核的内存和计算利用率未能达到极高水平，正是这些复杂权衡的结果。

图8. H100融合SSD内核的内存图表

从内存图表中可以看出，报告的65%至75%内存利用率主要来源于通过L2缓存进行的读取。这些读取可能包括：能够完全放入L2缓存的张量；在多个线程块之间重复使用的张量；线程块之间传递的状态数据；以及自然通过L2缓存的显存读取。由于L1缓存是每个SM私有的且在线程块之间不连贯，因此将流量转移到L1缓存不可行。同样，绕过L2缓存进行显存访问也无益，因为所有全局内存访问都必须经过L2缓存。

这份内存图表表明，除了次优的内存利用率，该内核更受L2缓存限制而非显存限制。进一步的优化可能需要：增加内存利用率；调整块大小/配置；或者进行根本性的算法改变。

英伟达Nsight Compute的逐行分析显示，融合内核中的大部分线程束停滞都发生在数据加载、同步和计算中，而原子操作和分块间同步造成的开销较小。例如，在状态传递步骤中，初始状态的读取和最终状态的写入虽然概念上简单，但由于其条件逻辑，增加了寄存器压力并降低了融合内核的速度。为此，我们可以将这些特殊情况的处理移到融合SSD内核之外。通过在状态张量中增加一个维度，将初始状态复制到第0个分块，并将最终状态从最后一个分块复制出来，利用PyTorch的切片赋值语法完成这些复制操作，从而在引入极小运行时开销和启动开销的情况下，优化了寄存器使用和性能。此外，对于一些计算，例如在分块扫描中对B应用A衰减，使用fp16代替fp32进行计算，可以在减少类型转换指令的同时带来约16%的速度提升。

Triton要求线程块中张量块的维度必须是编译时已知的2的幂。这导致所有的存储和加载操作都作用于2的幂的块，可能无法精确地整除目标张量。因此，我们使用掩码来覆盖整个张量，同时避免读取或写入越界数据。这些掩码与张量块的维度相同。然而，这些掩码并非总是必需的，因为模型维度（如headdim）通常可以被块大小整除，并且不会在不同输入之间发生变化。Triton支持tl.constexpr编译时参数和通过@triton.heuristics根据其他参数进行设置。因此，我们可以根据headdim是否能被块大小整除，在运行时自动启用或禁用掩码的headdim维度。虽然这发生在“运行时”，但实际上只在针对该模型的内核进行初始即时（JIT）编译时发生一次。

Mamba-2 SSD算法的计算复杂度与序列长度呈线性关系（每个标记的计算量渐近常数），但其基本情况是在某个分块大小内进行二次计算。分块之间使用线性算法，而分块内部使用二次算法。最佳分块大小是更高计算和所需资源与更高硬件利用率和更少中间状态之间的权衡。对于原始的未融合内核，Mamba-2 2.7B模型的最佳分块大小是256。然而，对于新的融合内核，相同模型的最佳分块大小现在是128。这个较小的分块大小还具有降低寄存器压力的额外优势，使得内核对微小变化（如启用掩码或使用更高精度进行中间结果计算）不那么敏感。

在分块状态中，我们可以等效地将A衰减应用于X而不是B，因为要缩放的维度是X和B的矩阵乘法的内维度。这通常更快，可能是因为更相似的内存布局减少了寄存器数据移动。

实测数据惊艳亮相

新媒网跨境获悉，我们对优化后的Triton内核在典型的推理场景下进行了全面的基准测试，包括批量大小从1到32，序列长度从1K到256K个标记，以及采用fp16状态。这些测试结果清晰地展示了我们的融合内核相对于原始未融合内核所带来的显著加速效果。

图9. 英伟达A100融合内核加速图

图10. 英伟达H100融合内核加速图

数据显示，融合SSD内核在SSD部分的性能比未融合实现快了1.50倍至2.51倍。在序列长度较低（尤其是在批量大小为1时）的情况下，内核启动开销的减少对融合内核的加速效果最为明显。但随着序列长度的增加，这些固定成本被分摊，融合内核的数据移动优势在缓存抖动增加的情况下变得更加突出。

将SSD部分的加速效果转化为整体模型性能，对于像Mamba-2 2.7B这样的模型，在英伟达A100和H100图形处理器上，当批量大小为1且序列长度为128K时，端到端推理速度大约提升了8%至13%。而在更短的序列长度下，端到端加速甚至可以达到20%左右，这很可能得益于内核启动开销的显著降低。

精度与性能的平衡艺术

在追求极致性能的同时，我们同样重视模型的准确性和稳定性。融合内核在整体上表现出良好的准确性和正确性，但在与参考解决方案的输出之间，确实存在微小的差异。这些差异的产生，既与所运行的GPU平台有关，也与部分计算的精度设置密切相关。为了实现约16%的速度提升，融合内核在某些内部计算中采用了fp16浮点精度，而原始内核则使用的是fp32精度。此外，我们的加速报告主要针对fp16状态，但融合内核同样支持相同中间数据类型和fp32状态。

下表详细说明了在不同数据类型配置下，输出y张量与原始内核输出匹配的百分比。

表2. H100精度表

由于浮点加法不具有完美的结合律，我们不能期望输出张量的所有元素在0容差下完全匹配。即使是不同的Triton启动配置，也可能导致相同内核输出之间产生极其微小的差异。对于“精确数据类型”（包括fp16和fp32状态），输出在所有实际应用中几乎是相同的，因此在对精度最敏感的模型中，该内核也能很好地工作。

对于“宽松数据类型”（我们加速图中所采用的配置），大约有三分之一的元素与原始内核的输出未能完全匹配。然而，若允许1e-3的严格容差，超过99.7%的输出元素仍然能够匹配。更值得一提的是，在常用的1e-2（即1%）容差下，所有配置的准确率都达到了99.9995%以上，这实际上可以视为100%。从实际应用的角度来看，我们预计“宽松数据类型”所带来的精度差异几乎可以忽略不计。

图11. H100 fp32 vs fp16精度对比图

图11展示了当状态从fp16变为fp32时，我们的加速效果如何变化。当状态为fp32时，融合内核和原始内核在chunk_size=256的情况下都更快。这代表了用更高的计算量换取更小的状态张量的权衡。对于fp32状态，融合内核的加速效果不如fp16状态，这可能是因为计算和数据移动的平衡发生了变化。

未来展望与应用拓展

新媒网跨境认为，本次的SSD内核融合优化技术，其应用范围并非仅限于Mamba-2模型本身。它同样可以直接应用于线性注意力机制，因为当参数A等于1时，SSD公式会退化为线性注意力的更新方式。在这种特殊情况下，融合内核甚至可以进一步简化和优化，从而实现更高的性能。

展望未来，我们的融合SSD内核尚未充分利用新一代GPU的最新功能，例如英伟达Hopper架构GPU上的张量内存加速器（TMA）和线程块集群（thread block clusters），以及Blackwell架构GPU上的张量内存（Tensor Memory）。这些新特性有望显著降低寄存器压力，从而进一步提升SSD的运行速度，并可能实现更优的Triton配置（例如，更大的块大小）。特别是线程块集群，在处理SSD内核中跨多个头共享的C、B和CB矩阵的广播加载时，有望发挥巨大作用，进一步提高新一代GPU上的性能。

我们还在探索更深层次的融合，比如将卷积层和层归一化（Layernorm）也纳入融合范围。这些操作在SSD模块之前和之后，如果能够融合，将分别省去一次完整的内核间读写操作。虽然我们已经尝试过层归一化的融合实验，但目前来看效果有限。原因在于，若想实现显著提速，需要解决线程块间的同步等待和数据访问模式优化等复杂问题。但这方面的尝试仍在进行中，一旦取得突破，模型的整体效率将再次迈上新台阶。

新媒网跨境预测，随着Mamba-2预填充计算成本的显著降低，Mamba-2层在模型设计中的性价比将进一步提升。这意味着在构建未来的大型语言模型时，增大Mamba-2层的规模和数量，可能会成为实现模型性能与效率最佳平衡的关键。例如，当前融合内核在最快分块大小下，计算利用率仍偏低，这表明我们有能力在不过度增加总运行时间的情况下，承担稍微更复杂的计算操作。同时，状态精度对性能的影响也不容忽视，尤其是在状态传递步骤中，保持低精度（如fp16）有助于显著提升融合内核的效率，这对于隐藏串行化延迟至关重要。此外，融合内核在L2带宽利用率上高于显存带宽利用率，这意味着在一些性能高度依赖小规模组的架构中，即使增加显存读取，对性能的负面影响也可能小于原始内核。

与vLLM的深度融合

为了在不进行填充（padding）的情况下支持可变长度序列和初始状态，vLLM引入了“伪分块（pseudo chunks）”的概念。任何包含多个序列标记的分块都会有多个伪分块，每个序列对应一个。融合内核通过在当前分块中循环处理所有伪分块来实现对vLLM的支持。vLLM的Chunk Scan根据伪分块在实际分块中的起始位置来偏移其读写操作。而我们的融合内核则采用基于序列索引的掩码（masking）方法，虽然偏移和掩码在运行时读写相同的数据量，但掩码可能对编译器而言更具可预测性、对齐更优或更为简洁，从而带来性能提升。目前，vLLM融合内核的集成工作正在进行中，并已展现出相似的加速效果。

总结与期许

综上所述，通过将Mamba-2 SSD预填充阶段的五个Triton内核成功融合为一个，我们在SSD本身实现了2倍的性能提升，进而为Mamba-2模型带来了8%至20%的端到端推理速度提升。这一成果显著增强了采用Mamba-2层的模型的吞吐量。我们期待将这些内核优化成果融入开源项目，让广大开发者社区能够轻松享受到Mamba-2模型更快推理所带来的便利。请大家密切关注，该融合SSD内核将很快登陆Mamba代码库，并集成到vLLM等主流推理框架中，共同推动人工智能技术的发展。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/mamba-2-upgraded-20-inference-25x-faster.html