Mamba 状态空间模型推理完全拆解：没有 KV Cache 的模型，服务系统要替它重新记一本账

SSM 的标准形式

连续时间状态空间模型写成一对微分方程：

$h$ 是隐状态，$A$ 描述状态如何自我衰减或混合，$B$ 描述输入如何进入状态，$C$ 描述状态如何读出。把序列当作对连续信号的采样，用零阶保持（ZOH）离散化，得到递推形式：

其中 $\bar{A} = \exp(\Delta A)$、$\bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I)\Delta B$，$\Delta$ 是采样步长。对角 $A$ 时两个式子都是逐元素计算：$\bar{a}_i = \exp(\Delta a_i)$ 就是一个标量指数，$\bar{B}$ 也按同一方式逐元素展开。

S4 系列模型的矩阵 $A, B, C$ 都是时间不变的：与输入无关。这带来一个极好的性质——把上式展开，可以写成

即整个序列的处理等价于与一个长度为 $T$ 的卷积核做卷积。卷积核可以预先算好，然后用 FFT 或线性递归方式 O(T log T) 地一次性处理整段序列，训练完全并行。

代价是同样的时间不变性带来的问题：模型对每个 token 的处理方式相同，无法根据内容决定是否忽略一个输入（离散模态如文本不像连续信号那样平滑，选择性极重要），也无法在需要时保持状态不变。论文用一个合成任务说明了这一点：标准复制任务里输入与输出间隔固定，时间不变的模型可以靠卷积核”数间隔”完成；而”选择性复制”任务里间隔是随机指定的，只有能读懂内容、按内容决定写入哪些位置和何时读出状态的模型才能完成。时间不变的 S4 在这种任务上失败，选择性模型能做到接近 100%，而 S4 只有 18% 左右（论文图 4 的选择性复制实验结果）。

下图（Mamba 论文图 1）是标准 SSM 的结构：每个通道独立地经过一个更高维的隐状态，状态维度为 $N$、通道数为 $D$，所以有效状态大小是 $D \times N$。Mamba 论文特意指出：此前的 SSM 之所以能绕开”显式展开这个 $DN$ 大小的状态”，靠的正是时间不变性——同一套卷积核可以在序列维上并行计算，从不需要把每个位置的中间状态写回显存。一旦状态转移随输入变化，这条路被封死，就必须面对”每层、每通道、每状态维都要逐 token 推进”的现实。这也解释了下面为什么需要一套专门的硬件友好算法。

选择性：让参数随输入变化

Mamba 做了三件事：

第一，把 $B, C, \Delta$ 都变成输入的函数。输入先经过投影得到每个 token 专属的 $B_t, C_t$ 和步长 $\Delta_t$。$B_t$ 决定当前内容向状态里写什么，$C_t$ 决定从状态中读什么，$\Delta_t$ 决定写入强度和忘记程度——$\Delta$ 大时状态更多地被当前输入覆盖（类似”重视这个输入”），$\Delta$ 很小时状态几乎不变（类似”忽略这个输入”）。这正是”选择性”的具体含义。

第二，把 $A$ 简化为对角矩阵，并用一个标量先验 $A = -\exp(A_{\text{log}})$ 参数化。对角 $A$ 意味着各维度独立演化，$\bar{A}$ 可以逐元素计算，硬件上容易做。$A$ 的取值决定了每个状态维的衰减时间常数——负得越多，记忆越短。这条线索在数学上就是矩阵的谱，站内 特征值与特征向量 一文里对”特征值刻画长期行为”有更一般的讨论。

第三，在 SSM 层前加一个一维因果卷积（conv1d）和 SiLU 激活，让局部信息先混合，再由 SSM 做长程处理。这一层很关键：$B_t, C_t, \Delta_t$ 的投影只看当前 token，如果没有任何局部上下文，模型只能靠单点信息决定取舍。卷积给它一个几 token 的窗口。

下图是 Mamba 块的整体结构：输入先分成两路，一路依次经过输入投影、conv1d、SiLU、SSM 与门控 $\sigma$，另一路经过线性投影作为门，两路相乘后输出。门控乘法和 Transformer 的 FFN 很像，是让模型可以在逐元素层面控制信息流动。

Δ 为什么是”门”：三个参数各自的角色

把选择性 SSM 和门控 RNN 放在一起看，$\Delta$ 的作用最清楚。对单个通道，离散递推是

当 $a<0$（Mamba 用 $A=-\exp(A_{\text{log}})$ 保证这一点）时，$\exp(\Delta_t a)\in(0,1)$。于是 $\Delta_t$ 同时控制两件事：保留多少旧状态和写入多少新输入。$\Delta_t \to 0$ 时，$\exp(\Delta_t a) o 1$、第二项也趋近 0，状态原封不动，模型相当于忽略当前 token；$\Delta_t$ 很大时，旧状态被迅速衰减，新输入几乎覆盖状态。这正是 LSTM 遗忘门与输入门的合并版，但门值不是由另一组参数独立算出来的，而是由同一套连续时间动力学中的步长给出。

这也解释了参数化上的几个细节。Mamba 不直接学 $A$，而是学 $A_{\text{log}}$ 再取负指数：这样 $A$ 永远在负实轴上，状态不会指数爆炸。$\Delta$ 先通过一个低秩投影（论文默认秩为 1，即每个头一个标量，再广播到通道），再经过 softplus 保证为正，最后加上一个可学习的偏置——偏置的初始化决定了训练初期模型偏向”记住”还是”忘记”。这些细节看起来琐碎，却直接决定了长序列训练能否稳定。

$B_t$ 与 $C_t$ 则分别负责”写什么”和”读什么”。它们由输入投影得到，因此每个 token 都有一组专属的写入、读出权重。把三者放在一起，选择性 SSM 的一层其实在做：根据当前内容算出写入门、读出向量和遗忘速度，然后在一个固定大小的记忆上做一次读-改-写。这个”记忆”就是后来推理系统要管理的状态。

并行扫描：把随输入变化的递推重新变成可并行运算

选择性的引入，把上面那个”卷积核”方法废掉了：$\bar{A}_t, \bar{B}_t$ 都随 $t$ 变化，卷积核不再固定。Mamba 的做法是回到线性递推本身，但用并行扫描（parallel scan）算法在 GPU 上算。

线性递推 $h_t = a_t h_{t-1} + b_t$ 有一个结合律结构：两个递推段可以合并为一段。把每段表示为三元组 $(a, b)$，合并规则是

并行扫描在 $O(\log T)$ 层里把整个序列的递推”折叠”成一次前缀运算。每一层的运算都是逐元素乘加，可以映射到 GPU 的 warp shuffle 和共享内存。

这里有一个容易被忽略的账：扫描的运算量随序列长度线性增长，但它的算术强度很低——每个元素只做一次乘加，却要从 HBM（高带宽显存）里读入、写回，再读回来做下一轮。GPU 的算力远大于带宽，于是这类逐元素内核几乎一定跑在带宽上限附近。Mamba 论文的处理方式是把整个 selective scan 融进一个 GPU 内核，并在反向传播时用重计算代替存储中间状态，使这个融合层的显存需求与使用了 FlashAttention 的 Transformer 相当。论文报告，相比标准的 PyTorch 逐点实现，融进内核后的扫描训练时快约 40 倍（图 8 左侧）——这个倍数主要是省下中间状态读写带来的，而不是算得更少。

对推理来说，这个细节的意义在于：预填充阶段如果仍然用扫描，瓶颈会是带宽而不是算力。Mamba-2 换成矩阵乘的动机正来自此。

解码时只需要算一步，扫描退化为单步递推，简单到几乎可以忽略开销。这一点正是推理路线成立的基础。

解码时的计算量

设模型为单层、通道数 $D$、状态维 $N$、扩张因子 $E=2$，则内部维度 $d_\text{inner} = ED$。每个 token 解码时的主要运算：

• 输入投影：$x \to (X', B_t, C_t, \Delta_t, Z)$，共 $2D d_\text{inner} + 2N d_\text{inner} + 2N d_\text{inner} + D d_\text{inner}$ 量级的乘加（$\Delta$ 是一个标量投影）；
• conv1d：核长 $K$ 的卷积，$K d_\text{inner}$；
• 扫描：$\bar{A}_t$ 的逐元素计算与状态更新 $h_t = \bar{A}_t h_{t-1} + \bar{B}_t x'_t$，约 $2 d_\text{inner} N$；
• 读出：$y_t = C_t h_t$，$d_\text{inner} N$；
• 门控乘法与输出投影：$d_\text{inner} D$。

这里的量级说明两件事。第一，主体是矩阵-向量乘法（逐个 token 的投影）和逐元素运算，没有与序列长度相关的项。第二，状态更新部分的运算量与 $N$ 成正比，而 $N$ 通常取 16 左右：Mamba 论文的默认配置是 $d_\text{model}=1024$、$E=2$、$N=16$，每层状态大小约为 $2d_\text{inner}N = 2 \times 2048 \times 16 \times 2$ 字节（假设 FP16 存储 $\bar{A}$ 与 $h$ 各占 $d_\text{inner}N$），每层约 128 KB。论文 350M 档配置是 24 层、$d_\text{model}=1024$、$d_\text{inner}=2048$、$N=16$，于是单序列全部层的状态合计约 $24 \times 128\text{ KB} \approx 3.1\text{ MB}$。对比 KV Cache：同规模 Transformer 每层每 token 存 K、V 各 $d_\text{model}$ 个元素，FP16 即每层 4 KB、每 token 全部 24 层约 98 KB；8K 上下文约 800 MB，128K 上下文约 12.6 GB。Mamba 的 3.1 MB 与序列长度无关。

这张账里最值得记住的一点是：Mamba 的状态大小与上下文长度无关。128K 和 8K 的请求占显存完全一样；而 Transformer 的 KV 项会随长度线性膨胀，1024 个并发请求各带 128K 上下文时，KV 的体量已经不是单卡能装下的量级。

为什么 Transformer 的吞吐在长上下文下会掉，而 Mamba 不会

Transformer 解码的每个 token 要读全部 KV：设层数 $L$、每 token KV 字节数为 $c$、上下文长为 $t$，则每步读取量为 $L \cdot t \cdot c$。显存带宽固定时，阅读时间与 $t$ 成正比，于是解码速度与上下文长度成反比。这是”解码是内存带宽受限”的直接体现。

Mamba 解码每步只读状态（大小恒定）、权重（大小恒定）和当前输入。总读取量与 $t$ 无关，所以解码速度基本不随上下文增长变化。论文的效率测试显示，Mamba 的解码吞吐能达到同规模 Transformer 的 4 到 5 倍，原因不是单个请求算得更快，而是没有 KV Cache 之后，同样的显存可以装下大得多的批。论文甚至给出一个有些反直觉的例子：未训练的 6.9B Mamba，推理吞吐高于比它小 5 倍的 1.3B Transformer（Mamba 论文图 8 右侧与 E.5 节）。

这里的数字需要谨慎引用：它是同一规模下的对比，受权重、精度、批量、序列长度等条件影响，不能外推为”所有情况下都快 4 到 5 倍”。但方向的结论是清晰的：Mamba 把解码从”读历史”变成”读状态”，读取量与长度解耦。

训练时为什么能并行

Mamba 在训练时处理整段序列，依赖两条：

1. 并行扫描把层层递推折叠为 $\log T$ 轮逐元素运算；
2. 由于扫描可以在芯片上按块进行，$O(T \log T)$ 的总运算量里大部分是逐元素操作，在 GPU 上是”足够快”的。

但服务系统的预填充和训练不同：提示长度差异极大，批量里每个请求长度不同；显存和片上存储有限，分块调度要随时切换；而扫描对每个请求都要在序列维上做归约，这会与动态批处理、chunked prefill、跨请求拼接产生复杂交互。

Mamba-2 的 SSD：把 SSM 写成矩阵运算

Mamba-2 的思路是换一个视角。前面说过 SSM 可以写成与一个”半可分矩阵”相乘：

当 $A$ 是对角矩阵时，乘积 $\prod \bar{A}_j$ 可以分解为每个维度的标量连乘，于是 $M$ 有一个特殊结构：半可分（semiseparable）。矩阵的任意 $t,s$ 位置的值可以写成低秩分解形式，整块子矩阵的秩不超过状态维 $N$。这正是它与注意力的桥梁：注意力矩阵也是低秩的（秩不超过头维 $d$），只不过 $M$ 是带因果掩码的、并且”掩码”本身是被 $A$ 参数化的指数衰减。

Mamba-2 的核心结论是：当 $A$ 为标量时，SSM 的序列变换矩阵 $M$ 可以表示为

其中 $Q, K$ 由 $C, B$ 与 $\bar{A}$ 的幂次构造，$L$ 是一个由 $a_t$ 决定的掩码矩阵（当 $a_t=1$ 时退化为标准的因果掩码）。也就是说，SSM 可以看作一种”结构化掩码的注意力”：数据结构与注意力相同，只是掩码由状态衰减决定，且不使用 softmax。论文把这个视角命名为结构化状态空间对偶（Structured State Space Duality, SSD）。

这个视角的价值在于：同一层的前向可以有两种算法。序列很长时用线性递推（状态视角，每步 $O(N)$ 每通道）；序列较短或在预填充的分块内时，用分块矩阵乘法（二次视角中的对角块），交给 Tensor Core。

这个对偶不是数学上的巧合，它带来三个可操作的结论。

第一，SSM 可以继承注意力的工程生态。既然 $M=QK^\top\circ L$ 与注意力矩阵同构，那么张量并行、序列并行、FlashAttention 式的分块与重计算、甚至用于注意力的编译器优化，都可以平移过来。Mamba-2 论文明确把”利用 Transformer 系统优化的积累”列为 SSD 的动机之一。

第二，状态维可以被放大。Mamba-1 的 $N$ 只有 16，原因之一是纯扫描实现里状态扩大会线性拖慢速度；SSD 用矩阵乘处理块内，$N$ 从 16 提到 64 甚至 128，速度只轻微下降。论文的效率测试给出了量级：对 $N=64$ 的大状态，SSD 比 Mamba 的融合扫描快 2 到 8 倍；在序列长度 2K 附近与 FlashAttention-2 持平，16K 时约快 6 倍。这组数字说明 Mamba-2 的主要收益不是”模型更聪明”，而是”同样的状态规模终于能跑得动了”。

第三，“注意力”与”SSM”的边界变成了掩码的选择。标准注意力使用全 1 的下三角掩码；SSM 使用由 $a_t$ 生成的衰减掩码。两者都可以写成”低秩矩阵乘 + 掩码”，区别只在于掩码是否随输入变化、是否需要 softmax 归一化。这个统一视角让混合层不再是”两套完全不同的算子拼接”，而是同一算子在掩码与归一化上的不同参数化——这也是后面混合模型能够共用部分基础设施的前提。

分块算法怎么算

Mamba-2 的分块算法把序列切成大小为 $Q$ 的块。对每个块 $j$，做四步：

1. 状态读取：用块内所有位置的 $C$ 与块起始状态 $h_{j-1}$ 相乘，得到”从历史带来的部分”；
2. 块内对角：块内各位置之间的影响，用块内矩阵乘法 $Q_j K_j^\top \circ L$ 计算——这就是”掩码注意力”，可以走 Tensor Core 的 GEMM；
3. 块内累加：第 1、2 步的结果相加得到该块输出；
4. 状态写出：用块内所有位置的 $B$ 与输入，更新出块结束状态 $h_j$，供下一块使用。

块内是矩阵-矩阵乘（高吞吐），跨块是”块与状态”的矩阵乘（代价随块数线性增长）。因此，整个预填充的运算量约为 $TQ(P+N) + TNP$：第一项来自块内矩阵乘，块数 $T/Q$ 乘每块 $Q^2(P+N)$；第二项来自跨块状态传递。对固定的块大小 $Q$，两项都对 $T$ 线性；同时块内运算落在 GEMM 上，利用率比纯逐元素的扫描高很多。上面这张 Mamba-2 论文图 5 就是这个分块结构。

值得注意的是：分块算法里”块起始状态”的概念，与推理系统的 chunked prefill 天然一致。一个长提示可以分成若干块递推执行，块与块之间只需传递 $h$；这恰好是服务系统调度分块的前向单位。

一：状态不是 KV，但同样需要被分配、复用与抢占

一个纯 Mamba 层的状态大小是每层 $d_\text{inner}N$ 个元素，与序列长度无关，但与序列数成正比。显存分配的单位从”页”变成”槽”（slot），每个请求占一个固定大小的槽。问题随之而来：

• 调度器如何估算一个请求需要多少显存？纯注意力模型里，预算来自最大上下文长度；纯 Mamba 模型里，预算来自批大小，与长度无关。混合模型里两者都要。
• 请求被抢占时，状态不能像 KV 页那样留在池子里等它回来——因为状态已经被后续 token 覆盖了。这等于说，纯 Mamba 没有”重放历史”的能力，除非框架额外保存了状态快照或者 KV 副本。

这在工程上是一个硬约束：一个被抢占的 Mamba 请求，如果状态没有保存，恢复时只能从头重新预填充它的全部历史。推理框架为 Transformer 设计的抢占机制——换出 KV 块、换入后直接继续——在纯 Mamba 上不成立。

二：混合结构的页大小对齐

vLLM 文档的核心问题是：一个分页缓存池只能有一种页大小，而不同层类型对”一页该存多少东西”的要求不同。文档把页大小定义为

全注意力层要为每个 token 保留 K、V；滑窗层只为最近一个窗口保留；Mamba 层则要为每层保留一份状态。三者的”每 token 开销”不同，如果各自定义页大小，就无法共享同一个物理池。

vLLM 的处理分两步。第一步，把层按注意力类型分组，使组内层类型相同、所需块数相同，从而共享同一组块号；组与组之间按固定比例分配，例如 Gemma-2 的 1:1、Llama 4 的 3:1。第二步，解决混合 Mamba 模型里”状态比 KV 大”的问题。文档明确指出：对 Bamba、Jamba、Minimax 这类交错结构，每个 Mamba 层的状态大小可能远大于注意力层的 kv_hidden_size。在单一页大小的约束下，只能把注意力层的 $\text{block size}$ 调大，直到

再把 Mamba 状态按这个乘积补齐（padding）。文档也提醒，这会让注意力 block size 涨到 400 以上，代价过高；另一种按注意力层总页大小对齐的策略仍在完善中。这一段工程细节很能说明问题：Mamba 节省的是”随长度增长”的显存，不是”每请求”的固定开销。在混合模型里，状态甚至可能成为页大小设计的主导项。

前缀命中规则也要按组分别执行。文档的缓存键是 (block_hash, group_id)，同一个 token 前缀在不同组里独立缓存、独立淘汰；新请求到来时，对每个组分别求最长命中前缀，再取交集作为整个请求的命中长度。对全注意力层，命中要求该前缀的所有 token 都在池中；对滑窗层，只要求最近一个窗口；对 Mamba 层，则要求对应的状态块还在池中。这就避免了”一个组命中、另一个组没命中”时读到不一致的状态。

三：批内不同长度的请求如何拼接

注意力层可以把多个请求的 token 拼成一个变长序列，用块表把每段映射到各自的分页块。Mamba 层做不到：每个请求有自己的状态，状态不能拼。于是推理框架必须做两件事：

• 把状态放在按请求槽位索引的连续缓冲区里，而不是按 token 位置索引的分页表。vLLM 的混合模型实现里，Mamba 层需要单独的”状态存储”和张量，而不是复用 KV 的页表抽象。
• 在批内同步推进每个请求的状态。这一层的每个 token 只更新自己请求的状态，因此内核要按请求槽位寻址，而不是按序列位置寻址。这会让内核实现从”一维 token 索引”变成”请求 + token 两维索引”。

四：预填充与解码的状态语义不同

预填充阶段需要给一个新请求算初始状态。如果该请求命中了前缀缓存，那么它应当从缓存中的状态开始，而不是从头扫描。但前缀缓存里存的是什么？对纯 Mamba 层来说，只有”某个 token 位置结束时的状态”才有意义；对注意力层来说，是那个位置之前的所有 KV。两者对”位置”的粒度要求不同：

• 纯注意力模型的前缀缓存可以命中任意长度前缀，只要能找到对应的块；
• 纯 Mamba 模型只可能命中事先保存过状态快照的那些位置，其他位置要重算。

所以，混合模型的前缀缓存设计通常选择把”状态快照 + KV 块”成对存储，命中粒度以快照位置为准。这也是为什么在工业实践中，Mamba 前缀缓存经常不是逐 token 的，而是按块或按预先切好的段落做快照。

五：投机解码的草稿与回滚

投机解码（可参考 EAGLE 与 MTP 的拆解）在推理侧依赖两件事：用较小的代价先生成若干候选 token，再一次性验证这些候选。对注意力模型，验证失败时只需要丢弃未确认的 KV 页；对 Mamba 层，验证过程已经把这些候选 token 写进了状态，回滚意味着恢复验证前的状态。

这带来一个直接的实现约束：草稿阶段开始前必须保存一份状态副本，验证通过 $k$ 个 token 后丢弃副本、保留覆盖后的状态；若只通过 $m<k$ 个，则从副本出发重新推进 $m$ 步。副本的大小与状态相同，因此投机解码对 Mamba 层的额外显存开销与”最大草稿长度无关、与草稿次数正相关”。相比注意力模型只需按页回退，Mamba 的回滚成本更高，但也不是不可承受——它换来的好处是：草稿阶段本身不产生随长度增长的 KV 读取。

此外，验证阶段需要对所有候选位置并行计算 $B_t, C_t, x'_t$，然后逐位置推进状态——这又是一次”块级短序列内的扫描”。若草稿长度取 8 或 16，块内运算量很小，框架往往会选择直接用分块矩阵乘（把整段草稿当一个块处理），再一次性写回状态。

六：并行策略的选择面变了

注意力模型沿头维切张量并行非常自然：每个 rank 拿一部分头和对应的 KV。Mamba 层的张量并行要切的东西更多——输入投影、$B/C/\Delta$ 的投影、状态本身。一旦状态也要切，状态更新就变成了部分和部分之间的逐元素关系，通信次数与状态维相关。Mamba-2 论文在系统章节讨论了这一点：SSD 可以在序列维、头维和状态维上做出不同的并行选择；状态维若要切分，状态更新就需要额外的跨 rank 同步，因此常见实现倾向于让每个 rank 持有完整状态、只切特征。

这也解释了为什么工业上的混合模型常把 Mamba 层放在专家并行或数据并行下，而把少数注意力层单独做张量并行：两者的切分自由度不同，强行共用一套并行拓扑会让其中一方吃亏。

selective scan 的状态放在哪

Mamba-1 的融合扫描内核，标准做法是把 $\bar{A}, \bar{B}, C, \Delta$ 等参数先加载到片上 SRAM，再在片上按块做扫描，最后只把输出 $y$ 写回 HBM。这样做的前提是：每个块内的状态和参数总量能装进 SRAM。以 $N=16$、通道数 2048 为例，一个块（比如 2048 个 token）的中间量会远超 SRAM，所以实现上必须沿通道维和时间维同时分块，块大小是”通道块 × 时间块”的二维 tile。

对解码，情况简单得多：batch 里的每个请求只有一个 token，状态更新是 $h \leftarrow \bar{a}\odot h + \bar{b}\odot x$ 的逐元素操作，整个 batch 的状态可以常驻显存，按请求槽位索引读取。瓶颈变成了”读状态、读权重、算投影”三件事，其中权重读取是大头——每步都要把整个模型的权重扫一遍，所以小 batch 时解码是权重带宽受限，而不是状态带宽受限。这与 Transformer 解码的结论相同：batch 小的时候，两者都被权重读取支配；batch 大、上下文长的时候，KV Cache 的读取才会成为 Transformer 的额外负担，而 Mamba 没有这一项。

状态用 FP32 还是 FP16

状态是逐 token 累积的，误差会随步数累积。工程上常见的折中是：状态的存储可以是 FP16/BF16，但状态更新与 $\bar{a}$ 的计算用 FP32 累加，必要时把状态本身也放在 FP32 的缓冲区里，只在输出时降精度。Mamba-2 的官方实现允许单独指定 state dtype；推理框架是否把它暴露成配置，取决于该框架的实现成熟度。在长上下文、高并发场景下，FP32 状态会让状态显存翻倍，但显著降低数值漂移。这与 KV Cache 的精度选择逻辑类似，却更敏感：KV 的每一项是独立写入的，误差不会跨 token 传播；状态是递推的，一次舍入误差会沿着整条链累积。

分块大小与头维

Mamba-2 的分块算法中，块大小 $Q$ 决定了块内 GEMM 的规模和跨块状态传递的频率。$Q$ 太小，块内矩阵乘的形状不友好，Tensor Core 利用率低；$Q$ 太大，块内 $Q^2$ 项会重新变成主导，同时中间矩阵占用更多 SRAM。官方实现默认用 256，并允许按头维 $P$ 与 SRAM 预算调整。头维本身一般取 64 或 128，与 Transformer 的 head_dim 对齐，这样同一个 GEMM 内核可以复用。

卷积与门控的融合

Mamba 块的顺序是：输入投影 → 因果 conv1d → SiLU → SSM → 与门控相乘 → 输出投影。推理实现里，conv1d 与 SSM 的输入投影可以融合为一次”读权重、算两路”的操作；门控乘法与输出投影也可以融合。融合的收益不是减少 FLOPs，而是减少中间张量的 HBM 往返。对解码这种每步只有一两个 token 的场景，中间张量的读写往往比计算更贵，因此融合程度直接决定小 batch 解码的延迟。