点击下方卡片，关注「AI视界引擎」公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
状态空间模型（SSMs）如Mamba最近引起了广泛关注。与基于Transformer的大语言模型（LLMs）相比，Mamba实现了线性计算复杂度，并且在序列长度上表现出优越性能。
然而，由于激活异常分散和复杂的计算依赖性，Mamba难以加速，导致现有的LLM加速器效率低下。本文提出了一种LightMamba，其通过结合量化算法和FPGA加速器架构以实现高效的Mamba推理。
作者首先提出了一种适用于FPGA的PTQ算法，该算法采用了旋转辅助量化和状态空间量化为2的幂次，从而将大部分计算量减少到4位。进一步地，作者设计了一个部分展平Mamba计算的FPGA加速器，以平衡效率和硬件成本。
通过计算重排以及细粒度的切分和融合，加速器的硬件利用率和内存效率得到了显著提升。作者将在赛灵思Versal VCK190 FPGA上实现LightMamba，并实现比GPU Baseline
更高的能效。当在Alveo U280 FPGA上评估时，LightMamba达到了93 tokens/s，这是GPU Baseline 的
。
unsetunset1. INTRODUCTIONunsetunset
状态空间模型（SSMs），如Mamba [1], [2]，最近已被提出并逐渐成为一类有前景的基础模型架构。与现有的基于Transformer的大语言模型（LLMs）[3]-[6]相比，Mamba在输入序列长度增加时只需要线性计算复杂度，但在各种下游任务上展现了更优异的表现。
Mamba [2] 的基本架构如图1 所示。每个 Mamba 块主要由两个线性投影层构成，即输入投影和输出投影，一个一维卷积（conv1d）层以及一个 SSM 层。Mamba 的计算包括一个预填阶段，用于总结输入 Prompt 词，以及一个自回归解码阶段以生成输出 Token。Mamba 的解码仅生成并存储固定大小的隐藏状态，而不是像 Transformer 中那样随序列长度线性增长的键值缓存。尽管 Mamba 有不同的架构，但作者重点关注 Mamba2 [2] 架构，并且除非另有说明，后续提及的 Mamba 均指代 Mamba2。
尽管前景广阔，如何加速Mamba处理，尤其是在FPGA等可重构平台上，仍然存在诸多开放问题。作者注意到，针对基于FPGA的大语言模型加速已有工作[7], [8]由于以下挑战无法直接应用。首先，虽然基于FPGA的加速可以从低比特精度量化中受益，但Mamba量化比基于Transformer的大语言模型更具挑战性，因为激活异常值分散出现。具体来说，对于不同的输入Token，激活异常值出现在不同的通道中。简单地应用现有的量化算法[9]-[11]会导致显著的准确性下降。其次，SSM层涉及大量的元素级计算。而现有工作[12]-[14]通常忽略了对SSM层的量化，并产生了较高的硬件成本，直接对SSM层进行量化也会引入显著的重新量化开销。第三，SSM层涉及复杂的计算和数据依赖性，这阻止了通过空间展开SSM计算来提高加速吞吐量。
为了应对这些挑战，作者提出了LightMamba，这是首款基于FPGA的Mamba加速框架，该框架协同设计了量化算法和加速器架构。作者提出了一种旋转辅助的PTQ（PTQ）算法，用于消除Mamba中的离群值，从而使作者可以将模型量化至4比特，同时最大限度地减少精度下降，并显著提高与外部存储器之间的通信效率。对于SSM层，作者提出了一种适用于FPGA的2的幂次量化方案，以实现通过简单移位进行的再量化，从而提高计算效率。在FPGA加速器架构方面，作者设计了专门的旋转模块来实现PTQ算法，并进一步提出了计算重排序以及精细粒度的切分和融合，以提升硬件利用率。
作者的主要贡献可以总结如下：
作者提出了首个适用于整个Mamba模型的Post Training Quantization (PTQ) 算法。通过旋转辅助量化，作者将Mamba成功量化至4位，并保持合理的精度。作者还采用FPGA友好的等比量化方案对SSM层进行量化，以提高计算效率。作者提出了首个基于FPGA的Mamba加速器，其架构集成了与所提出的量化算法协同设计的定制化旋转模块。
此外，作者还提出计算重新排序和细粒度调度策略，以提高计算吞吐量并减少片上内存使用。作者在Xilinx Versal VCK190 FPGA上实现LightMamba，结果显示每秒处理7.21个Token，并且相较于GPU Baseline 具有4.65至6.06倍更高的能量效率。
unsetunsetII. RELATED WORKSunsetunset
A. LLM Quantization
量化将权重和激活从高精度浮点数（例如FP16）映射到低精度整数，例如INT8或INT4，从而减少内存和计算成本。现有算法主要利用PTQ（PTQ）来处理大语言模型，并发现其关键障碍来自权重和激活中的异常值。LLM.int8O [15] 使用混合精度量化方法并保留高精度的异常值，这引入了大量计算开销。SmoothQuant [9] 和 Outlier Suppression [10] [11] 发现异常值主要存在于特定通道中，因此计算出的通道尺度或平移因子在量化之前重新缩放异常值。
QuaRot [16] 和 SpinQuant [17] 通过与旋转矩阵相乘来去除异常值。然而，这些方法仅证明适用于基于Transformer的大语言模型。对于Mamba [12]，仅使用四舍五入（RTN）方法对线性层进行了量化。Mamba-PTQ [13] 直接在Mamba上应用了SmoothQuant。尽管它们仅量化了线性层，而SSM层仍保持在FP中，但仍遭受了较大的准确率下降。
B.FPGA-based LLM Accelerators
之前的基于FPGA的大语言模型加速器可以分为两大类：时间架构和空间架构。
时间架构构建了一个处理引擎（PE），主要用于执行各种任务，特别是在矩阵乘法（MM）[7], [8], [18]方面。这种架构并不适合处理Mamba中多样且复杂的元素级乘积（EM）。
空间架构针对不同的操作定制了PE，并支持多个PE在Pipeline中的并发处理[19]-[22]，从而降低了延迟。然而，由于定制的PE分散了资源，这种设计导致了MM的并行度较低。在本文中，作者采用了一种部分展开的空间架构，该架构展开了一个Mamba块，并协同设计了量化算法和架构以提高硬件利用率并减少延迟。这些架构的对比分析见表1。
unsetunsetIII. MOTIVATIONunsetunset
挑战 1：由于存在离群值，Mamba 难以被量化为低位精度。作者发现输出投影层的激活离群值呈现出分散的分布情况，如图2(c) 所示，这使得以往的大语言模型 (LLM) 量化方法 [9]-[11] 较为无效。这是因为，在基于 Transformer 的大语言模型中，离群值会在不同的 Token 的固定通道中持续存在。因此，在这些模型中，可以通过逐通道计算缩放或平移因子来减小离群值通道的幅度，并降低量化误差 [9]。
对于 Mamba 而言，不同输入 Token 的激活离群值可能会出现在不同的通道中，并且可以在所有通道中分散出现。因此，现有的方法，例如 SmoothQuant [9] 和 OutlierSuppression
[10] 在 FPGA 加速中的量化误差与基准的 RTN 量化方法相当甚至更大，如表 HI 所示。旋转基量化算法在基于 Transformer 的大语言模型中已被证明有效 [16]、[17]，并且不需要离群值在固定的通道中持续存在。然而，由于旋转基算法通常需要操作符融合并引入额外的计算开销，尚不清楚如何将其应用于 Mamba 在 FPGA 上的加速。
挑战 2：现有工作 [12]-[14] 往往选择不对 Mamba 中的 SSM 层进行量化。SSM 层在 FPGA 上进行浮点计算会带来较大的硬件成本。此外，直接将 SSM 层量化到较低位精度也会如图3 所示，产生较大的硬件成本。这是因为 SSM 层包含大量的 EM（元素乘法），其再量化（即将高位精度的输出重新转换为低位精度）开销非常大。
挑战3：SSM层的计算图复杂，这增加了基于FPGA的加速难度。一方面，SSM层的输入包括
，这些输入均由输入投影层生成。尽管可以在FPGA上沿空间方向展开Mamba计算，但由于输入投影层和SSM层的计算被迫顺序进行，导致硬件利用率仅为不到60%。另一方面，SSM中间激活在芯片上的存储需求很高，占总量URAM使用量的超过70%。
为应对这些挑战，作者提出了LightMamba，一种高效的基于FPGA的Mamba加速框架。具体而言，在第四节中，作者提出了一种FPGA友好的Mamba量化算法，该算法包括一种旋转辅助的伪量化（PTQ）算法以缓解分散的异常值，并采用基于PoT的SSM量化以最小化再量化开销。同时，在第五节中，作者提出了一种基于FPGA的加速器，该加速器通过计算重排序来提高硬件利用率，并通过细粒度的切分和融合来降低片上内存成本。
unsetunsetIV. FPGA-FRIENDLY QUANTIZATION ALGORITHMunsetunset
A. Rotation-assisted Linear Layer Quantization
旋转辅助量化方法首次在[16]中提出，用于基于Transformer的大语言模型。通过用正交矩阵
乘以激活
和权重
，即
，结果与
相同，同时从
和
中去除了异常值。然而，尚不清楚旋转方法是否适用于Mamba。因此，作者研究了Mamba中的旋转等效性，并提出了如图4a所示的旋转辅助方法。
作者观察到线性层和SSM层的激活值中存在大量离群点，而输出投影层的离群点在不同通道中呈现出散乱分布的情况。旋转有助于去除这些离群点，因为它可以通过与其他元素相协调来减轻大离群点的影响。对于输入和输出投影层，作者在图2中展示了通过应用旋转并去除离群点的方法。值得注意的是，在对输出投影层之前的激活值进行旋转时，作者在图4a中在其前插入了一个在线哈达玛变换，这可以通过作者在第五节中定制的旋转单元高效实现。然而，作者发现无法对SSM进行旋转，因为其不满足旋转等价性。具体而言，SSM的原始计算公式是式1a。假设作者可以旋转隐藏状态
以去除离群点，即乘以哈达玛矩阵
，作者能够推导出式1b和式1c。然而，式1c不能推导出式1d，因为在EM中不满足矩阵结合律。因此，作者无法从式1a推导出式1d，即SSM不满足旋转等价性。
为了减少计算开销，作者尽可能将旋转与相邻操作融合。如图4a所示，作者可以将第一个旋转与嵌入表（即
）融合，最后一个旋转与大语言模型 Head （即
）融合，以及在每个Mamba块中输出投影层前后的旋转（即
）。
对于紧邻第一个根均方归一化（RMSNorm）操作符的旋转（即
），为了保证计算不变性，作者需要首先分割RMSNorm的缩放因子，即
，然后再与输入投影的权重融合。而对于紧邻第二个RMSNorm操作符的旋转（即
），作者发现是否将RMSNorm的缩放因子与输出投影的权重融合不会影响计算不变性，但会引入较大的量化误差，如图4b所示。因此，作者在算法中选择不将第二个RMSNorm的缩放因子进行融合。在作者的算法中，只有旋转
需要在线计算，在作者定制的FPGA模块支持下，这种计算开销很小。
B. FPGA-friendly SSM Quantization
为了量化SSM以通过FP计算减少高昂的硬件成本，作者利用INT8按组量化在准确度和硬件效率之间取得平衡。
然而，直接对SSM进行量化会导致大量的重新量化开销，如图3所示，这是因为EM固有的重新量化开销比MM大，因为EM中不存在减少量化的过程。为此，作者提出使用PoT量化方法对SSM进行处理，通过这种方式，重新量化可以通过位移操作而不是乘法来实现，从而显著减少重新量化开销。
unsetunsetV. FPGA-BASED MAMBA ACCELERATIONunsetunset
A. Overall Hardware and Hadamard Transform Unit
LightMamba的整体架构如图5(a)所示。作者设计了一种部分展开的空间架构，并在FPGA上展开一个Mamba块的计算。模型参数存储在片外DRAM中。LightMamba由三个主要模块组成：
(1)矩阵乘法单元(MMU)，负责处理所有线性层的计算，包括输入和输出投影；(2)SSM单元(SSMU)，完全展开SSM以实现流水线执行；
以及(3)哈达玛变换单元(HTU)，这是一种定制设计，支持旋转辅助量化。
MMU 设计用于以时间复用的方式支持输入和输出投影层。它采用基于树结构的乘积累加器（MAC）架构，接收
维的向量作为输入。此外，MMU 还配备了
条通道，在一个时钟周期内执行
个 MAC 操作。总体而言，MMU 包含
个 MAC 单元，这些单元通过利用 DSP 包装技术 [23]，高效地实现了
个 DSP，如图5(b) 所示。
SSMU 特设了一种细粒度且完全流水线化的架构来计算 SSM 层。如图5(c) 所示，每个操作均由一个专门的单元实现，并通过先进先出缓冲器（FIFO）连接。作者针对每个操作优化了并行性，以确保数据流的均衡 Stream 并尽可能减少 FIFO 的深度。作者通过元素级乘法单元（EMUs），其中 EMU 由 DSP 组成，来在 SSM 中实现各个操作。对于不同的操作，作者设置了不同的并行度。
HTU 专门支持作者在量化算法中使用的 Hadamard 变换。它有两种变体：2 的幂次型和非 2 的幂次型。例如，在 Mamba2.7B 中，需要两种类型的 HTU，即 128 点 HTU 和 40 点 HTU，分别如图5(d) 和图5(e) 所示。128 点 HTU 是基于快速 Hadamard 变换（FHT）算法[24]实现的。它包含七个阶段，每个阶段包含一个 Butterfy 核心和两个 FIFO。在第一阶段，前 64 个元素被推入输入 FIFO。当接下来的 64 个元素到达时，它们通过 Butterfy 核心按对进行处理。输出要么发送到下一个阶段，要么在输出 FIFO 中缓存。与基于 MM 的 Hadamard 变换相比，该设计在相同硬件资源下将延迟降低了 72%。对于较小的 40 点 Hadamard 变换，作者直接使用简单的 MMU 实现，并固定 Hadamard 矩阵的一个输入为 1 和 -1。
B. Computation Reordering
由于局部自注意力模块(SSM)和输入投影层之间存在不同的计算模式以及数据依赖关系，它们不得不像图6(a)所示那样顺序执行。然而，作者发现SSM包含多个独立的 Head ，并提出了一种粗粒度的Pipeline设计来提高硬件利用率，如图6(b)所示。这种Pipeline设计依赖于作者提出的计算重排序方法，该方法改变了输入投影层中的数据生成序列。具体来说，
、
、
首先生成并存储在一个片上缓存中，而
和
则交替生成用于逐个计算SSM的各个 Head ，如图6(b)所示。这种重排序使得SSM的计算可以在MMU生成
、
、
之后立即开始。与传统的顺序实现相比，作者的方法将网络的总计算时间减少了32%，并将硬件利用率从58%提高到96%，如图6(b)所示。
C. Fine-grained Tiling and Fusion
鉴于当前的设计，作者观察到SSMU消耗了超过70%的片上内存，因为它需要存储所有中间激活值，例如
、
和
等，如图7(a)所示。作者提出了一种精细粒度的分区与融合策略。通过利用操作融合，作者可以直接将前一个操作的输出传递给下一个操作，从而消除片上通信和数据处理。此外，作者采用精细粒度的分区来减少缓冲区大小。作者在头维度和隐藏状态维度上进行分区，分区大小为
，如图7(b)所示。逐块实现细化了SSMU的执行，使其能够支持图6(c)中的精细流水线，这将SSMU的URAM使用量减少了4倍。这不仅降低了片上缓冲区的要求，还通过消除流水线中的停顿提高了硬件利用率。
unsetunsetVI. EXPERIMENTSunsetunset
A. Experiment Setup
算法 作者评估了所提出的量化算法在Mamba模型家族[2]上的性能。作者对图4a所示的整个模型进行了量化以提高硬件效率。作者使用通道级权重量化和Token级激活量化，进行8位权重和激活量化（表示为W8A8）。作者还使用组级权重和激活（组大小设为128）进行4位量化（表示为W4A4）。作者使用lm-eval-harness [31] 在六项任务上评估了困惑度和零样本准确性：LAMBADA [25]、HellaSwag [26]、PIQA [27]、Arc（Easy和Challenge）[28]、Winogrande [29] 和OpenbookQA [30]。
硬件评估中，作者使用了两个FPGA平台：赛灵思Versal VCK190和Alveo U280。作者在VCK190上使用Vitis HLS和Vivado Design Flow实现LightMamba。图8显示了作者在VCK190 FPGA上的实现布局。通过PYNQ框架测量板载吞吐量，并借助赛灵思BEAM工具测量功耗。对于U280的评估，作者开发了一个周期精确的模拟器，该模拟器已在Vitis 2023.2的HLS仿真中得到验证。作者选择NVIDIA RTX 2070和RTX 4090作为GPU Baseline 。作者使用NVIDIA系统管理接口进行功耗测量。表4展示了FPGA和GPU的硬件参数，以及LightMamba的详细硬件利用率。
B.Evaluation Result
算法评估 作者将作者的量化算法与以前提出的用于Transformer基础的大语言模型（LLM）的权重-激活PTQ方法SmoothQuant（SQ）[9]和Outlier Suppression
[11]进行了比较。请注意，这些方法只能应用于线性层。作者在Mamba上重新实现它们，并使用WikiText2 [32]数据集中128个随机样本作为校准数据。作者评估了只对线性层进行量化的LightMamba方法，以及对所有模块包括SSM进行全面量化的LightMamba方法。如表1001所示，对于W8A8量化，LightMamba和LightMamba相比于FP16模型几乎没有准确性损失。尽管
在Lambada数据集上的困惑度表现更好，但在W4A4量化下却表现不佳。而作者的方法LightMamba和LightMamba*在W4A4量化下则优于其他所有方法，分别将困惑度改善了1.78和1.91，相较于先前的方法SQ。
硬件评估 作者在图9(a)中比较了不同输出序列长度下GPU、现有加速器和作者所开发的LightMamba的解码吞吐量。由于现有加速器不支持Mamba，因此作者将它们在运行基于Transformer的大语言模型时的表现进行对比。鉴于这些研究未提供长序列长度的吞吐量数据，作者根据每篇论文中的参数进行了性能模拟。在VCK190上，LightMamba分别实现了W8A8和W4A4下的实际吞吐量为3.61和7.21 Tokens/s。作者也对U280上的LightMamba进行了模拟以进行公平比较。LightMamba在U280上的吞吐量达到了93 Tokens/s，平均而言，这比RTX 2070快
。LightMamba在长序列上实现了更为显著的加速效果，因为Mamba只记录固定大小的隐藏状态。
作者在图9(b)中比较了LightMamba和GPU在不同模型规模下的能量效率(Tokens/J)。
在VCK190上，LightMamba始终优于GPU，并分别在RTX 2070和4090 GPU上实现了平均
倍和
的节能效果。对于小型Mamba模型，LightMamba能够实现更多的节能，因为作者的设计减少了SSM层的开销，而对于小型模型而言，这一开销更为昂贵。
C. Ablation Study
作者现在进行消融研究，以展示LightMamba中不同技术对准确性和效率的影响。如图10所示，通过权重和激活量化，吞吐量可以从2.23 token/s 提高到5.32 token/s。
借助旋转辅助量化算法和定制化的HTU，作者能够在几乎相同的吞吐量下将量化Mamba的精度提升4.3%。作者的计算重排序技术进一步提高了硬件利用率，并将吞吐量提高到7.21 token/s。最后，通过精细粒度的切分和融合，片上内存消耗减少了4倍，从246减少到61。
unsetunsetVII. CONCLUSIONunsetunset
在本文中，作者提出了一种高效的基于FPGA的Mamba加速框架——LightMamba，该框架集成了量化和硬件协同设计技术。作者指出了在FPGA上加速Mamba所面临的三个挑战。
为了解决这些问题，作者提出了一个适应FPGA的旋转辅助权重量化算法，将Mamba量化到4位，同时最小化了精度损失。
此外，作者还提出了一种基于FPGA的加速器，该加速器实现了计算重排序和精细粒度的切分与融合。通过这些方法，LightMamba在VCK190 FPGA上的性能达到了每秒处理7.21个Token，并且相比于GPU Baseline ，其能量效率提高了
倍。
unsetunset参考unsetunset
[0]. LightMamba: Efficient Mamba Acceleration on FPGA with Quantization and Hardware Co-design .
点击上方卡片，关注「AI视界引擎」公众号