Hardware-Aware Algorithm

定义

Hardware-Aware Algorithm 是 Mamba 为高效计算 selective-state-space而设计的 GPU 优化算法。核心思想是利用 GPU 内存层次结构（SRAM vs HBM），通过 kernel fusion 和并行扫描（parallel associative scan）避免将扩展状态写入慢速 HBM。

为什么需要

S6 的选择机制使 SSM 变为时间变化 → 不能用卷积（卷积要求 LTI）。唯一的计算方式是循环模式：

h_t = A_bar_t * h_{t-1} + B_bar_t * x_t   （需要 scan）

朴素循环是 O(n) 顺序的，对训练不可接受。Mamba 通过以下方案解决。

关键技术

1. 并行关联扫描（Parallel Associative Scan / Blelloch Scan）

将循环计算的序列依赖展开为前缀和（prefix sum）形式，利用关联性在 GPU 上并行完成：

思路: h_t = A_t * h_{t-1} + B_t * x_t
      = f(f(...f(h_0, x_1), x_2), ..., x_t)
      
展开为关联操作，用 Blelloch scan 在 O(log n) 并行步完成

2. IO 感知 Kernel Fusion

位置	操作
HBM（慢）	存储输入 x、参数
SRAM（快）	加载 → 离散化 → scan → 写回
HBM（慢）	存储输出 y

关键：扩展后的状态 A_bar, B_bar 维度为 (B, L, D, N)，不在 HBM 中物化——在 SRAM 中计算完毕即写回压缩结果。

3. 重计算（Recomputation）

反向传播不保留前向的中间状态，在 backward pass 中从输入和参数重新计算 → 减少内存占用。

效果

比所有基于卷积的 SSM（S4、H3、Hyena）快 3×（A100 GPU）
理论复杂度：O(BLDN) vs 卷积 SSM 的伪线性
使选择性 SSM 的训练在实践中可行

参考

gu-mamba (Gu & Dao, 2024) Section 3.3
FlashAttention (Dao et al., 2022) — 同作者的 IO 感知优化风格

2.3 KiB Raw Blame History Unescape Escape