AI 推理系列（十四）：Mamba 的 Selective State Space——不用 Attention，模型如何“记住”长序列？

一个持续读取生产日志的模型可能已经处理了数十万行事件，但业务只希望它保持对当前服务状态、最近错误和关键配置变化的判断。Transformer 通常把历史 Token 转换为逐层 K/V 并继续保存；上下文越长，缓存容量和 Decode 读取量越高。Mamba 采用不同的状态模型：历史不会以可直接回看的 Token 缓存长期存在，而是随着输入不断压缩到每层固定维度的内部状态中。

这种差异使 Mamba 在长序列与持续流式处理场景中很有吸引力。它的序列计算量随长度线性增长，自回归阶段也不需要让请求状态随已处理 Token 数持续扩张。但“固定状态”同时意味着不可逆的信息压缩。理解 Mamba 时，不能只看到它没有标准 Attention，还要看它如何选择写入信息、如何在 GPU 上执行状态递推，以及哪些任务会暴露有限状态容量的边界。

Transformer 与 Mamba 保存的是两种不同历史

标准自回归 Transformer 在 Prefill 阶段为历史 Token 生成各层的 Key 和 Value。Decode 新 Token 时，新 Query 可以再次访问这些缓存，因此历史仍以相对细粒度的形式存在。即使使用 GQA、MLA 或量化降低缓存规模，缓存通常仍随序列长度增加。

Mamba 的自回归路径更接近状态机。每个新 Token 到来后，模型使用当前输入和旧状态计算新状态，再从新状态产生输出。下一步只需要保留更新后的状态以及局部卷积所需的少量历史，而不需要保存此前每个 Token 的完整 K/V。

维度	Transformer Attention	Mamba 选择性 SSM
历史表示	每个历史位置对应的 K/V	固定维度的递归状态
Decode 状态规模	通常随上下文长度增长	与已处理长度基本无关
历史访问	可按 Attention 权重直接访问不同位置	只能使用已压缩进状态的信息
长序列计算	标准 Full Attention 成本增长更快	序列方向为线性扫描
主要风险	KV Cache、Attention 计算和带宽压力	有限状态容量与信息压缩损失
典型优势	精确检索、复制和跨位置组合能力强	持续流、长序列和固定状态推理

这里的“不用 Attention”描述的是 Mamba 主体块的核心计算路线，不代表所有实际模型只能在纯 Attention 与纯 SSM 之间二选一。工程系统可以构建 Attention、SSM、卷积或其他模块的混合架构，用少量显式检索能力补充递归状态的不足。

SSM 如何把历史折叠进状态

State Space Model 可以先用一组示意状态方程理解。设当前输入为 x_t，上一步状态为 h_{t-1}，当前状态和输出可以写成：

h_t = A_t h_{t-1} + B_t x_t
y_t = C_t h_t

这不是完整实现代码，而是状态流的抽象。A_t 控制旧状态如何保留和演化，B_t 决定当前输入怎样写入状态，C_t 决定从状态中读出哪些信息。传统线性时不变 SSM 的这些变换对不同位置基本遵循固定规律，因此很适合具有稳定动力学的连续信号。

语言和代码却包含大量离散切换。一条日志中的请求 ID、错误码和配置值可能需要保留很久，标点、模板化前缀或重复时间戳则可能只具有局部作用。若所有 Token 都按近似相同的规则更新状态，模型很难根据内容决定“这次应该覆盖旧信息，还是让当前输入迅速衰减”。这也是早期 SSM 面对离散、信息密集序列时的重要限制。

Selective State Space 如何根据输入控制记忆

Mamba 的关键变化是让部分状态参数依赖当前输入。原论文中的选择机制使离散化步长以及与写入、读出相关的参数能够由 Token 表示动态产生。工程上可以把它理解为三个内容相关的动作：

• 保留：当前输入是否应让已有状态继续维持较长时间。
• 写入：当前 Token 的哪些特征应该进入状态，以及写入强度多大。
• 读出：当前输出应该从已有状态中暴露哪些信息。

例如日志流中出现新的错误码时，输入相关参数可以让相关状态分量发生明显更新；连续出现大量普通访问日志时，模型可以抑制这些特征对长期状态的影响。这里的“选择”不是一个人工编写的布尔规则，也不是显式保存某个 Token，而是训练后形成的连续门控和状态动力学。

flowchart LR
    A[当前 Token] --> B[生成选择参数]
    C[旧 SSM 状态] --> D[状态更新]
    B --> D
    A --> D
    D --> E[新 SSM 状态]
    E --> F[读出当前特征]
    F --> G[输出投影]
    E --> H[传给下一 Token]

选择性机制给 Mamba 增加了类似内容寻址的能力，但它与 Attention 的路径仍然不同。Attention 在计算当前 Token 时回到历史位置中比较 Query 与 Key；Mamba 在每个位置到来时决定历史如何被修改。一个强调“需要时检索”，另一个强调“进入时压缩”。因此，选择性状态更新能减少无关信息污染，却不能保证被压缩掉的细节之后还可以精确恢复。

为什么选择性让普通卷积捷径失效

线性时不变 SSM 的一个优势是，它既能写成递归状态更新，也能转换为卷积形式。训练完整序列时，可以提前构造卷积核并对所有位置并行计算，从而绕开逐 Token 的串行依赖。

Mamba 让参数依赖输入后，每个位置的状态转移不再完全相同。模型不能简单使用一条固定卷积核处理整段序列，因为位置 t 的更新规则会被 x_t 改变。选择性提高了表达能力，却破坏了原先最直接的并行计算方式。

如果退回朴素递归：

Token 1 -> State 1 -> Token 2 -> State 2 -> Token 3 -> State 3

GPU 就必须等待前一个状态完成，难以像 Transformer 的大型矩阵乘法那样一次处理大量位置。Mamba 要同时获得内容相关选择和训练吞吐，必须为这类输入相关递推设计新的并行执行算法。

Selective Scan 如何兼顾训练与 Decode

Selective Scan 利用状态更新的关联结构，把长序列的递推组织成并行 Scan。它不消除位置之间的数学依赖，而是重新分组和合并状态变换，使 GPU 可以并行处理多个区间，再合成与顺序递推一致的结果。

训练与推理因此采用两种等价视角：

阶段	主要执行形式	保存的数据	主要优化目标
训练 / 长序列 Prefill	并行 Selective Scan	中间激活与分块状态	提高序列并行度、减少 HBM 往返
自回归 Decode	单步递归更新	每层 SSM State 与 Conv State	固定请求状态、低单步开销
批量流式推理	多请求状态批处理	每个请求一组固定状态	提高小状态更新的 GPU 利用率

原始 Mamba 工作同时强调 hardware-aware 实现。某些中间参数可以在更快的片上存储中生成和消费，而不必全部物化到 HBM。状态扫描的理论复杂度只是起点；真正的性能取决于数据布局、融合程度、分块大小、重计算策略以及是否减少了高成本内存流量。

对于持续日志 Agent，Prefill 可以用 Scan 快速建立初始状态；之后每收到一条新日志，只更新该请求的 Conv State 和 SSM State。即使累计输入继续增长，请求侧状态也不会像 KV Cache 一样按 Token 数增加。但模型对很久以前日志的了解只存在于状态压缩结果中，不能通过增大运行时间自动获得无限精确记忆。

线性复杂度为什么不保证端到端一定更快

O(N) 比标准 Full Attention 的序列增长形式更有优势，但 Big-O 不能直接预测 GPU 延迟。Transformer 的矩阵乘法形状规整，能够充分使用 Tensor Core；递归状态更新的张量通常更小，若 kernel 启动频繁、融合不足或 batch 太小，硬件利用率可能偏低。

Mamba 的推理优势还会随工作负载变化。长输入、持续生成和状态缓存受限场景更容易体现固定状态的价值；短 Prompt、短输出时，状态更新与投影层的常数成本可能掩盖复杂度差异。高并发服务需要把多个请求的状态组合成足够大的 batch，否则单请求状态虽然小，GPU 却可能因工作粒度不足而空闲。

生产评估应至少拆分以下指标：

• 不同序列长度下的 Prefill tokens/s 和显存峰值。
• 单请求与不同 Batch Size 下的 Decode latency、TPOT 和吞吐。
• 每请求 Conv State、SSM State 的实际字节数。
• Scan kernel、线性投影和卷积在 profiler 中的耗时占比。
• 长序列质量随距离增加的退化，而不只是能否完成 Forward。
• 框架是否使用专用 CUDA kernel，还是回退到通用算子组合。

如果模型理论状态固定，但服务端为了兼容接口仍保留大量不必要的历史张量，或者底层没有命中优化 kernel，架构优势不会自动转化为端到端收益。

Mamba-2 如何把 SSM 重新组织成矩阵计算

Mamba-2 的 State Space Duality（SSD）展示了一类结构化 Attention 与一类 SSM 之间的对应关系。它不是简单地把 Attention 加回 Mamba，而是利用统一结构把状态计算重新分块，使更多工作可以落到 GPU 擅长的矩阵乘法上。

这种双重视角很重要：同一个序列算子可以在训练时采用适合大矩阵的块算法，在 Decode 时继续采用固定状态递推。模型设计不再只追求更低算术复杂度，而是同时考虑计算是否能映射到现有硬件的高吞吐原语。

Mamba-2 论文报告核心层在特定实验配置下相对第一代 Selective Scan 获得明显速度提升。这个结果说明执行结构可以成为模型架构的一部分，但不能直接推导任意模型端到端都会获得相同比例加速。嵌入层、输出投影、通信、数据加载和其他非 SSM 组件仍会限制整体收益。

Mamba-3 为什么继续增加状态表达能力

固定大小状态带来稳定内存成本，也设置了信息容量上限。模型读取越来越长的序列时，新的信息必须与已有信息共享同一组状态维度。若任务要求精确跟踪多次覆盖的变量、复制很久以前的字符串或保留大量独立实体，状态可能出现干扰和遗忘。

用户提供的 Mamba-3 资料将改进重点放在状态表达能力和硬件利用率上，包括更丰富的离散化、复数状态以及 Multi-Input Multi-Output（MIMO）设计。MIMO 允许一次状态更新处理更多输入和输出通道，让硬件执行更多有效矩阵工作。论文报告在其配置中可以增加 Decode FLOPs 而保持相近 Wall-clock 延迟，这反映了 GPU 性能工程中的常见现象：当原算子未充分占用硬件时，增加可并行计算不一定同比增加时间。

这些结果仍然属于特定模型和硬件实验。部署时应分别验证质量、状态大小、kernel 时间和端到端延迟，不能把“更多 FLOPs 近似免费”当作所有 GPU、精度和 Batch Size 都成立的规则。

Mamba 与长上下文 Transformer 不是同一种能力

一个 Transformer 可以在上下文范围内直接给某个远端 Token 分配高 Attention 权重。Mamba 没有对应的原始历史数组，远端信息必须已经被编码进当前状态。因此，两者对“长序列”的支持含义不同。

任务特征	Mamba 更有吸引力的原因	可能更偏向 Attention 的原因
持续日志、传感器、音频流	状态固定，适合无限累计输入	需要精确回看具体历史事件时仍需外部存储
长文本语言建模	线性扫描降低序列执行压力	远距离复制与精确检索可能受状态容量限制
在线逐 Token 生成	不维护随长度增长的 KV Cache	Transformer Serving 生态更成熟
代码和结构化状态追踪	可学习递归更新规则	多变量精确回溯可能需要显式位置访问
RAG 与文档问答	可高效编码长流	文档证据通常需要可定位、可引用的检索路径

Mamba 的固定状态也不等于业务长期 Memory。长期 Agent 若需要记住用户约束、任务决策和数小时前的具体事件，仍应把关键信息写入数据库、事件日志、摘要或检索系统。SSM State 更适合作为模型运行时的隐式短期状态，不能替代可审计、可更新和可精确查询的外部记忆。

推理框架需要管理完全不同的请求状态

Transformer 推理引擎围绕 KV Block、分页分配、Prefix Cache 和 Attention kernel 建立了成熟的数据结构。Mamba 请求通常维护卷积窗口状态和 SSM State，调度器需要完成不同的生命周期管理：请求创建时分配固定状态，Prefill 后写入最终状态，Decode 时原地更新，结束时整体释放。

这种状态结构减少了随上下文增长的碎片，但会带来新的实现要求：

• 不同模型层和 Mamba 版本可能具有不同状态布局。
• Continuous Batching 需要在请求加入和退出时紧凑重排状态。
• Beam Search 或候选分叉需要复制状态，而不是共享只读 KV 前缀。
• Pipeline Parallelism、Tensor Parallelism 与状态切分方式需要配套设计。
• 通用 Transformer kernel 不能自动高效执行 Selective Scan 或单步 SSM 更新。

官方 Mamba 实现和 TensorRT-LLM 中的相关模型路径说明这类架构已经进入实际软件栈，但“框架存在代码路径”不代表任意模型配置、GPU 和精度都具备同等成熟度。上线前需要确认支持矩阵、实际 kernel 路径、数值一致性以及并发调度行为。

什么时候应该选择 Mamba 路线

Mamba 更适合从工作负载约束出发评估，而不是作为 Transformer 的通用替代品。若单条序列持续增长、KV Cache 已成为容量瓶颈、任务允许把历史压缩为状态，并且运行环境具有成熟的 Scan/Decode kernel，它具有明确的工程价值。

若业务核心是从几十万 Token 中精确定位一段证据、复制远端字符串、跨多个历史位置进行组合推理，或者现有系统高度依赖 Prefix Caching、PagedAttention 和成熟 Transformer Serving，迁移成本和质量风险可能高于理论复杂度收益。此时可以考虑混合模型，或继续通过 GQA、MLA、KV Cache 量化、稀疏 Attention 和外部检索控制 Transformer 成本。

最终需要同时回答两个问题：固定状态是否足以承载任务所需信息，以及目标硬件能否把 SSM 的线性计算真正执行得足够快。Mamba 的价值不只是删除 Attention，而是把“保存全部历史”改写成“持续维护可选择更新的状态”。这条路线能否胜出，取决于压缩后的状态是否仍保留业务真正需要的记忆，以及软件栈能否把这种状态更新转化为稳定的吞吐和延迟收益。