State Space Models (Mamba / S4): 解释 Mamba 的硬件感知实现（并行扫描 + 核融合）。

Core Concept: 把离散化、选择性扫描、输出投影融合进一个 kernel，中间数据留在 SRAM，避免 HBM 往返（同 Flash Attention 思想）。

数学机理：<strong>朴素实现的瓶颈</strong>——Mamba 的 SSM 层包含多步计算：输入投影 → 计算 Δ、B、C（输入依赖）→ 离散化（算 Ā、B̄）→ 选择性扫描（递归 h_t=Ā_t h_{t−1}+B̄_t x_t）→ 输出投影。若每步都是独立 kernel，则中间张量（尤其<strong>扩展状态</strong>：把隐维度 d 扩展到 d×N 的'状态张量'）需反复写入/读取 HBM，产生巨大带宽开销——<strong>这正是 RNN/SSM 传统实现的性能瓶颈</strong>。<strong>硬件感知实现（Mamba 的核心工程）</strong>：(1) <strong>并行扫描（parallel/associative scan）</strong>——递归 h_t=Ā_t h_{t−1}+B̄_t x_t 是<strong>线性递归</strong>，满足结合律，故可用<strong>关联扫描</strong>（Blelloch scan）在 O(log L) 深度内并行计算；这提供了训练所需的并行性。(2) <strong>核融合（kernel fusion）</strong>——把整个 SSM 层的计算（投影、离散化、扫描、输出）<strong>融合进一个 kernel</strong>：输入从 HBM 读入 SRAM，在 SRAM 内完成所有中间计算（包括 d×N 的扩展状态），<strong>只把最终输出写回 HBM</strong>。这样 HBM 访问量从 O(L·d·N)（每步读写状态）降到 O(L·d)（只读写输入输出），与 Flash Attention 的'不物化中间矩阵'完全同构。(3) <strong>重计算（recomputation）</strong>——反向时不保存中间状态，而是重新计算（类似 Flash Attention 的反向），进一步降低显存。<strong>效果</strong>——Mamba 在长序列上比朴素实现快数倍，且显存 O(L·d)（不随状态维度 N 增长）；论文报告 Mamba 的推理吞吐随序列长度<strong>线性</strong>增长（而非 Transformer 的平方）。

深度剖析与工程权衡：① <strong>'核融合'是通用范式</strong>——Flash Attention（注意力）、Mamba（SSM）、以及各种 fused 算子（fused Adam、fused LayerNorm）都遵循'把多步计算融合、让中间数据留在 SRAM'的原则；这是 memory-bound 时代的核心优化手段。② <strong>并行扫描的常数开销</strong>——虽然并行深度是 O(log L)，但关联扫描需要多次数据搬运（up-sweep 与 down-sweep），故其<strong>常数因子</strong>较大；在短序列上，串行扫描可能更快。这解释了'为什么 SSM 在短序列上未必优于注意力'。③ <strong>扩展状态的显存</strong>——d×N 的状态张量（N 常为 16~256）在朴素实现中占用大量显存；核融合通过'不物化它'解决（只在 SRAM 中分块计算）。④ <strong>与 chunked scan 的折中</strong>——另一种实现是'分块 + 块内并行'（把序列分块、块间用串行、块内用并行扫描），在并行度与常数开销间折中。⑤ <strong>与硬件的耦合</strong>——核融合的效果依赖 SRAM 大小与带宽；不同 GPU 上需重新调优（与 Flash Attention 同理）。⑥ <strong>面试要点</strong>——被问'Mamba 为什么快'，应给出'<strong>并行扫描（并行性）+ 核融合（不物化扩展状态、只读写输入输出）</strong>'，并指出这与 Flash Attention 的思想一致；能说明'并行扫描的常数开销导致短序列上优势不明显'是深度理解的标志。