编译器优先的状态空间对偶与可移植的 O(1) 自回归缓存推理 Compiler-First State Space Duality and Portable O(1) Autoregressive Caching for Inference
纯JAX/XLA实现Mamba-2，免自定义内核且O(1)缓存
前置知识
状态空间模型（SSM）与 Mamba-2
状态空间模型是一类通过隐状态 $h(t)$ 和输入 $x(t)$ 之间的线性动力学 $h'(t) = Ah(t) + Bx(t)$、$y(t) = Ch(t) + Dx(t)$ 建模序列的神经网络。与注意力机制的 $O(L^2)$ 复杂度相比，SSM 在序列长度 $L$ 上具有 $O(L)$ 的线性复杂度，适合长序列建模。Mamba-2（Dao and Gu, 2024）将传统 SSM 的循环结构重新组织为状态空间对偶（State Space Duality, SSD）算法，并限制状态矩阵 $A$ 为每头对角标量，使其在保持表达力的同时具备并行化潜力。该论文实验覆盖 130M 到 2.7B 五个 Mamba-2 官方检查点。
论文的核心对象就是 Mamba-2 的 SSD 算法及其推理路径。理解 SSD 的对角化、分块、einsum 主导四个结构性质是理解为什么它能被现代编译器高效处理的前提。
状态空间对偶（SSD）的四个结构条件
论文将 Mamba-2 的 SSD 算法归纳为四个对编译器友好的结构条件：(i) 对角状态矩阵 $A$ 是每头的标量，使矩阵指数退化为标量指数；(ii) 分块循环（chunked recurrence），序列被划分为固定 $L=256$ 的块，块内并行、块间串行；(iii) einsum 主导计算，块内输出是跨 (batch, chunk, sequence, head, state) 轴的批量 einsum 收缩；(iv) 静态控制流，块数、核大小、头结构在编译期已知，下三角因果掩码是与 token 内容无关的常量。这四个条件共同决定 XLA 是否能生成融合且高效的目标代码。
这四个条件是论文的核心技术贡献。如果不熟悉 SSD 的代数结构以及它与传统循环 SSM 的差异，就无法理解为什么纯 JAX 实现能在不写自定义 CUDA 内核的情况下达到接近 Triton 内核的性能。
XLA 与 JAX 编译路径
XLA（Accelerated Linear Algebra）是 JAX 的底层编译器，将 Python 级别的 jax.numpy 操作 lowers 为 HLO（High Level Optimizer）中间表示，再通过 fusion、tiling、scheduling 等 passes 针对 TPU/GPU/CPU 生成目标代码。JAX 通过 jit 装饰器对函数进行追踪和编译，并通过 fori_loop 等原语表达设备端循环，避免 Python 主机与设备之间的每步往返。论文在 JAX 0.9.0、Flax NNX 0.12.4 下实现，使用 jax.lax.fori_loop 让整个解码循环留在编译后的 XLA 程序中执行。
论文的方法建立在 JAX/XLA 编译器栈之上。理解 jit 的工作方式、PyTree 在追踪中的作用、以及为什么 Python 端循环会成为瓶颈，是读懂本文优化策略的关键。
JAX PyTree 与 O(1) 自回归缓存
PyTree 是 JAX 中由数组叶子组成的嵌套容器结构（如 dict、list、NamedTuple），jax.tree_util.register_pytree_node 可以让任意 dataclass 参与 jit 追踪与变换。在自回归解码中，每层的 SSM 状态 $h \in \mathbb{R}^{H \times P \times N}$ 和卷积状态（$D_{conv} \times (k-1)$）共同构成前缀的固定大小摘要。把这些状态组织成 PyTree 并作为 jit 函数输入/输出，可以让 XLA 把它们作为编译期图的一部分持续在片上更新，无需主机同步。
这是论文 O(1) 缓存机制的核心数据结构。理解 PyTree 的注册机制以及为什么把状态作为显式输入能避免 host-device 往返，是掌握性能优化的关键。
MFU / HBU 与 Roofline 模型
模型 FLOP 利用率（Model FLOP Utilisation, MFU）定义为 $MFU = F_{XLA}/t_{wall} / (\text{peak TFLOPS} \times 10^{12})$，衡量实际计算吞吐占硬件峰值的比例；硬件带宽利用率（Hardware Bandwidth Utilisation, HBU）定义为 $HBU = B_{XLA}/t_{wall} / (\text{peak BW} \times 10^{9})$，衡量内存子系统效率。Roofline 模型根据工作强度（FLOPs/byte）判断工作负载是计算受限还是带宽受限。论文中 TPU v6e 的算力强度阈值约为 574 FLOPs/byte，batch size 1 的单流 prefill 低于此阈值，因此观察到 ~15% MFU 对应 roofline ceiling。
论文几乎所有性能数字都用 MFU/HBU 来解读。理解这两个指标的含义以及 roofline 模型如何判断瓶颈所在，才能正确评估 '15% MFU' 是好是坏，以及 64% HBU 意味着什么。
Triton / mamba_ssm 参考实现
Triton 是 NVIDIA 主导的 Python 级 GPU 内核编程语言，通过分块（tiling）和共享内存管理实现接近手写 CUDA 的性能。mamba_ssm v2.2.2 是 Mamba 官方发布的 PyTorch 实现，其核心 SSD 路径依赖 Triton（甚至 CUDA）手写内核，并在 NVIDIA 设备上达到 LLM 级别的吞吐。论文以 mamba_ssm v2.2.2 为基准对比 WikiText-103 困惑度与数值一致性，并证明纯 JAX 路径在严格 float32 + TF32 disabled 条件下能与参考实现在 ±0.0005 困惑度内一致。
理解 Triton 内核与 XLA 编译器路径在性能与可移植性之间的权衡，才能体会论文想要论证的核心观点：编译器优先的 SSD 实现能在多大程度上替代手写内核。
研究动机
Mamba-2 这类状态空间模型（SSM）在长序列建模上具有线性时间复杂度的理论优势，但其推理性能目前严重依赖手写的 CUDA 与 Triton 内核。这些内核与具体硬件栈深度耦合：官方 Mamba-2 路径主要在 NVIDIA CUDA 上表现最佳，向 AMD ROCm（如 LightOn AI, 2024）和 Apple Metal Performance Shaders（如 Purohit, 2026）的移植则需要社区维护独立的内核代码路径。结果是，同一个预训练检查点能否在可服务或可评估的硬件上运行，取决于是否有匹配的内核支持。对于没有匹配 CUDA 栈的用户来说，他们要么依赖单独的移植版本，要么必须重写内核，模型代码本身绕过了通常能够提供硬件可移植性的编译器栈。即使是 JAX 生态内的早期工作（radarFudan, 2024; McKinney, 2024; repytmargorp, 2024），要么只覆盖 Mamba-1，要么只提供最小的 Mamba-2 前向传播而缺少自回归缓存和性能评估；社区报告（Obeid, 2023）甚至指出，纯 JAX SSM 实现不依赖自定义内核时训练速度过慢，无实用价值。
本文的目标是本文的目标是在不写任何自定义 CUDA 或 Triton 内核的前提下，用标准 JAX 原语表达 Mamba-2 的完整推理路径（包括分块并行 prefill 与缓存式自回归解码），并使其从单一源代码在 Google Cloud TPU v6e 与 NVIDIA L40S 上都能达到接近参考实现的 roofline-consistent 利用率。作者希望证明 SSD 的代数结构本身足以让 XLA 编译器生成高效的融合代码，从而把硬件可移植性作为编译器优先路径的副产物。具体来说，论文设定了四项可验证的技术目标：(1) 在 TPU v6e 单流 prefill 达到 ~15% MFU（batch-1 roofline ceiling），decode 达到 ~64% HBU；(2) 在 NVIDIA L40S 上同一份源码重现序列长度无关的缓存式 decode；(3) WikiText-103 验证困惑度在五个模型规模（130M–2.7B）上与 Triton 参考（mamba_ssm v2.2.2）一致到 ±0.0005；(4) 隐状态在 float32 舍入容差内一致。
与已有工作不同的是，本文的独特切入角度是把 SSD 视为对编译器友好的算法，而不是为某个具体硬件写一个手工内核。作者将 Mamba-2 的 SSD 抽象为四个结构性条件（对角状态矩阵、分块循环、einsum 主导计算、静态控制流），并显式论证这四个条件足以让 XLA 的 fusion 与 tiling passes 产出有竞争力的代码。在 JAX 前端，本文进一步提出三项必须保留这些条件的实现纪律：精确 einsum shaping、静态而非动态掩码、设备端编译循环（jax.lax.fori_loop）以及精度管理（残差、衰减参数、归一化层均采用 float32）。与 Bonsai（JAX-ML Team, 2025）等已有 JAX SSM 库相比，本文新增了 SSD-to-XLA 结构性条件分析、TPU v6e 与 L40S 的跨硬件 roofline 表征、实现选择的消融实验，以及与 Triton 参考的严格数值与困惑度对比。换言之，前人证明纯 JAX SSM 路径不实用，本文则通过识别并显式保留 SSD 的结构性质证明这是可能的。
核心方法
本文的整体思路是把 SSD 算法视为一组对编译器友好的代数条件，然后在 JAX 前端用标准原语表达这些条件，使 XLA 的 fusion 与 tiling passes 能自动生成高效目标代码，避免任何手写 CUDA/Triton 内核。直觉上，Mamba-2 的核心计算可以分解为四件事：(a) 用 $B$、$C$、$\Delta$、输入 $X$ 的批量 einsum 计算块内输出（intra-chunk matmul）；(b) 用分段前缀和计算块内衰减矩阵 $\mathcal{L}$；(c) 块间的轻量级扫描传播 summary state；(d) O(1) 大小的 SSM/卷积缓存用于自回归解码。本文把这四件事全部写成 einsum 和静态索引的组合，让 XLA 看到的是一个干净的批量收缩图（block, head, chunk, sequence, state）+ 静态掩码 + 静态形状，然后由编译器自动融合。围绕这个图，作者施加四条 JAX 前端纪律：精确 einsum shaping（让矩阵维度在物理上连续）、静态下三角掩码（在常量矩阵上施加，让 XLA 折叠到融合链里）、设备端编译循环（jax.lax.fori_loop 把整个解码过程留在 HLO 里）、精度管理（残差、衰减指数、归一化方差都用 float32，关闭 TF32）。整个 SSD 核心函数约 60 行 Python (Santoni, 2025)，所有收缩都在 (batch, chunk, sequence, head, state) 轴上批量进行。
本文的核心创新是把'硬件可移植性'问题转化为'SSD 的结构条件能否通过 JAX 前端传递到 XLA'问题。作者识别出 SSD 的四个结构性条件是 XLA 高效代码生成的充分条件，然后给出三条 JAX 实现纪律使这些条件在原始图层面得以保留：(i) einsum shaping——所有重计算写成批量 einsum，且输入 reshape 后让 batch、head、chunk、sequence-within-chunk 维度在物理上连续，从而直接映射到目标硬件的矩阵单元；(ii) 静态掩码——下三角因果掩码作为 $L \times L$ 的常量矩阵在编译期存在，让 XLA 能把它与前缀和、减法、mask、指数化一起融合进同一区域（在 1.3B 模型上消融显示，从静态掩码切换到动态逐行掩码造成 82.8% 的吞吐损失，因为融合图在循环边界处断裂）；(iii) 设备端编译循环——自回归解码整体包装在 jax.lax.fori_loop 里，并通过把 per-layer 的 SSM 状态与卷积状态注册为 JAX PyTree，让 XLA 把这些固定大小缓存作为编译期图的一部分持续更新，无需主机同步（130M 规模上，主机驱动循环与编译循环的吞吐比为 662 vs 1588 tokens/s，前者慢 2.4×，随着模型增大到 780M 以上二者收敛）。这些纪律的本质区别是：与 Bonsai 等已有 JAX SSM 库相比，本文显式地以 XLA codegen 的视角审视每个原语选择，并把结构条件的代价量化为可观测的吞吐损失。
方法步骤详情
完整方法包括预填充（prefill）与缓存式自回归解码（cached decode）两条路径，二者共享 SSD 核心函数。预填充阶段（Algorithm 1）：输入为 prompt token id 序列 $x_{1:T}$ 与模型参数 $\theta$，块大小 $L=256$。步骤 1 将 token 嵌入并 reshape 为 $X \in \mathbb{R}^{B \times N_c \times L \times D}$，其中 $N_c = T/L$ 为块数。步骤 2 对每一层 $\ell$：(a) 通过 $\text{InputProj}^{(\ell)}$ 投影得到 $B_t$、$C_t$、$\Delta$、$X_{proj}$；(b) 离散化得到 $\bar{A} = \exp(\text{softplus}(A^{(\ell)}_{\log}) \cdot \Delta)$，强制 float32 上转；(c) 衰减矩阵 $\mathcal{L} = \text{tril}(\exp(\text{segsum}(\log \bar{A})))$，使用静态 jnp.tril 在常量矩阵上生成；(d) 块内输出 $Y_{diag} = (\mathcal{L} \odot C_t B_t^\top) X_{proj}$，这是批量 einsum；(e) 状态累积 $S = \text{einsum}(B_t, \bar{A}, X_{proj}) \to (B, N_c, H, P, N)$；(f) 块间扫描 $S' = \text{einsum}(\bar{A}_{chunk}, S)$，对 chunk 维度串行；(g) 合并 $Y = Y_{diag} + Y_{cross}(S', C_t)$；(h) 把每层最终块状态 $h^{(\ell)} = S'[:, -1] \in \mathbb{R}^{B \times H \times P \times N}$ 存为缓存。最终 logits $\hat{y}_{1:T} = \text{LMHead}(Y)$。缓存式解码阶段（Algorithm 2）：预填充返回每层 SSM 状态 $C.ssm^{(\ell)} \in \mathbb{R}^{B \times H \times P \times N}$ 与卷积状态 $C.conv^{(\ell)} \in \mathbb{R}^{B \times D_{conv} \times (k-1)}$，二者作为 PyTree 字段被 jax.lax.fori_loop 持续追踪。对每个解码步 $t \in [P+1, P+G-1]$：(a) 嵌入上一 token $\hat{x}_t$ 为 $u_t$；(b) 对每一层 $\ell$：$u^{(\ell)}_t = \text{InProj}^{(\ell)}(u_t)$，然后通过 roll_and_insert 操作把 $u^{(\ell)}_t$ 插入卷积窗口，执行深度可分离卷积得到 $z_t$，投影出 $B_t$、$C_t$、$\Delta_t$，离散化得到 $\bar{A}$ 和 $\bar{B} = \Delta_t \cdot B_t$，最后做一次 $O(1)$ 的 SSM 更新 $C.ssm^{(\ell)} \leftarrow \bar{A} C.ssm^{(\ell)} + \bar{B} z_t$；(c) 用 LMHead 得到 logits $y_t$，在设备上执行 argmax 得到 $\hat{x}_{t+1}$。整个 fori_loop 在一次 XLA launch 内完成，主机在生成完成前保持挂起。除主体路径外，作者还施加四条精度规则：(1) 残差连接保持 float32 防止层栈累积漂移；(2) 衰减参数以 log-space float32 持有，指数化在计算时进行（消融显示 bfloat16 指数化造成 0.013 最大绝对 logit 误差，足以偏移输出分布）；(3) 归一化层对输入 cast 到 float32 计算方差，再 cast 回去；(4) 默认 jax_default_matmul_precision 设为 highest 以关闭 NVIDIA 上的 TF32 舍入。
技术新颖性
本文的技术新颖性体现在三个维度。第一是结构化分析维度：作者首次系统地把 Mamba-2 的 SSD 算法抽象为四条对编译器友好的结构条件（对角 $A$、分块循环 $L$、einsum 主导计算、静态控制流），并显式论证这四条条件足以让 XLA 的 fusion 与 tiling passes 生成有竞争力的代码，从而把硬件可移植性问题从'为每个硬件写一份内核'降维为'在 JAX 前端保留这四条性质'。这一分析与 XLA 文献中已有的'编译器优先'Transformer 推理（Bradbury et al., 2018; Pope et al., 2023）一脉相承，但首次系统化地应用到了 SSM 领域。第二是实现纪律维度：作者把抽象的'结构条件'翻译成三条可量化的 JAX 前端约束（精确 einsum shaping、静态掩码、设备端编译循环）和四条精度规则，并通过消融实验给出了违反每条纪律的具体代价——静态掩码切换为动态逐行掩码造成 82.8% 吞吐损失，主机驱动循环取代编译循环在 130M 上慢 2.4×，bfloat16 衰减指数化引入 0.013 最大绝对 logit 误差。这些代价被量化为可观测数字，使论文从'原则上可行'升级到'有可验证约束的工程方案'。第三是 PyTree-as-cache 维度：把 per-layer SSM 状态与卷积状态组织为注册过的 JAX PyTree，从而让 jax.lax.fori_loop 能把它们作为编译期图的一部分持续更新，避免每步主机同步。这一数据结构上的小细节是实现 O(1) decode 与硬件可移植性同时成立的关键，与 Bonsai（JAX-ML Team, 2025）的 caching PR (#131, d8f8d11) 在设计上一致，但本文给出了完整的跨硬件评估、消融与数值一致性证明。
实验结果
论文从六个维度报告实验结果，每一项都有具体数字支撑。**(1) TPU v6e 单流 decode 吞吐**：cached scan 在所有五个模型规模上几乎与序列长度无关——130M 在 128/1024/4096 token 处分别达到 1588/1635/1641 tokens/s，2.7B 同样稳定在 94/95/95 tokens/s（Table 1）；相比之下 non-cached 在 4096 token 处分别跌至 56、18、9、7、3 tokens/s，即 130M 时缓存路径相对非缓存路径取得 29× 加速（1641 vs 56），2.7B 时取得 32× 加速（95 vs 3）。cached host（主机驱动循环）在 130M 上只有 662 tokens/s，约为 compiled scan 的 41.7%，对应 2.4× 加速缺口；该缺口在 780M 以上因每步计算主导而消失（2.7B 上 cached host 97 vs scan 94，在测量方差内）。**(2) 峰值内存**：cached decode 在所有序列长度下保持常数（130M: 545.6 MB, 2.7B: 10861.8 MB），non-cached 随序列线性增长，在 4096 token 时 2.7B 消耗 16392 MB 而 cached 仅 10861.8 MB（节省约 34%）（Table 11）。**(3) 硬件利用率**：prefill MFU（Table 2）在 prompt=4096 时从 130M 的 8.23% 上升到 2.7B 的 12.96%，prompt=1024 时 2.7B 达到 15.23%，对应论文标题强调的 ~15% roofline ceiling；decode HBU（Table 3）在 130M 到 2.7B 之间从 51.62% 上升到 63.43%，且每个模型在 128 到 4096 token 之间变化不超过 1.7 个百分点；absolut 64% HBU 在 2.7B@4096 处达到，论文将其解读为 TPU v6e 的硬件天花板附近。**(4) L40S 单流 decode**：与 TPU v6e 趋势一致，cached scan 在所有规模上序列长度无关（130M: 240.2/267.1/314.2 tokens/s，2.7B: 35.4/36.3/36.1 tokens/s），同源代码无需修改即可在 NVIDIA 上运行（Table 4）。**(5) WikiText-103 困惑度**：五个模型规模上 JAX 路径与 Triton 参考（mamba_ssm v2.2.2）的最大绝对差异为 0.0005（130M: 18.7019 vs 18.7023；370M: 13.1244 vs 13.1247；780M: 10.8886 vs 10.8892；1.3B: 9.5704 vs 9.5708；2.7B: 8.3250 vs 8.3252）（Table 5）；隐状态在 1e-5/1e-4（相对/绝对）容差内一致，logits（首 256 维）在 1e-5/2e-4 容差内一致（Table 6）。**(6) 消融实验**：在 1.3B 模型 TPU v6e BF16 prompt=1024 下，静态掩码（jnp.tril）prefill 吞吐 42631 tokens/s，动态逐行掩码（fori_loop + dynamic_slice）只有 7330 tokens/s（−82.8%），输出仍 bitwise 相同（Table 7）；bfloat16 衰减指数化在 24 层累积下产生 0.013 最大绝对 logit 误差，float32 baseline 为 0（Table 8）。**(7) 训练对比**（Table 13）：在 L40S 上小模型短序列场景下，JAX 路径比 Triton 参考快——130M@512 快 64.8%（25.9 vs 73.7 ms）、370M@512 快 57.3%（62.8 vs 147.0 ms）；但优势随规模与长度迅速反转，130M@2048 时慢 27.6%，780M@1024 慢 132.1%，780M@2048 慢 287.1%。**(8) JIT 编译开销**（Table 12）：从 130M 的 5.5/5.6/2.5 秒（prefill 1024/decode 128/decode 4096）增长到 2.7B 的 15.8/19.5/43.0 秒，2.7B@4096 的 43 秒对迭代研究构成可观开销，但作为一次性成本在大量推理调用下被摊销。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| TPU v6e 单流 decode 吞吐 (130M) | Tokens/second @ seq=4096 | 1641 (cached scan) | 56 (non-cached) | 29.3× 加速；缓存常数吞吐 |
| TPU v6e 单流 decode 吞吐 (2.7B) | Tokens/second @ seq=4096 | 95 (cached scan) | 3 (non-cached) | 31.7× 加速；缓存常数吞吐 |
| TPU v6e prefill 计算利用率 | MFU (%) @ 2.7B prompt=1024 | 15.23% (~140 TFLOPS) | TPU v6e peak 918 TFLOPS BF16 | 达到 batch-1 roofline ceiling |
| TPU v6e decode 带宽利用率 | HBU (%) @ 2.7B seq=4096 | 64.08% | TPU v6e peak 1600 GB/s | 硬件带宽天花板的 64% |
| L40S 单流 decode 吞吐 (130M) | Tokens/second @ seq=4096 | 314.2 (cached scan) | 20.3 (non-cached) | 15.5× 加速；同源代码跨平台 |
| WikiText-103 困惑度一致性 | PPL |Δ| vs Triton 参考 | max 0.0005 across 130M–2.7B | mamba_ssm v2.2.2 (Triton 参考) | ±0.0005 内一致 (130M: 18.7019 vs 18.7023) |
| 数值一致性 (隐状态) | Relative/Absolute tolerance | 1e-5 / 1e-4 (float32 舍入容差) | PyTorch mamba_ssm reference | 比特级一致 |
| 静态 vs 动态掩码消融 | Prefill tokens/second @ 1.3B | 42631 (jnp.tril 静态) | 7330 (fori_loop 动态逐行) | +481% (静态掩码避免融合图断裂) |
| 训练步速度 (130M@512, L40S) | Forward+backward time (ms) | 25.9 | 73.7 (Triton) | −64.8% (JAX 更快) |
| 峰值内存常数 (2.7B) | Peak memory (MB) across seq lengths | 10861.8 (constant for 128–4096) | 11035→16392 (non-cached, linear) | 固定大小缓存，与序列无关 |
局限与改进
作者在论文 §6 'Limitations' 中明确承认了几个限制：第一，profiling 仅覆盖 TPU v6e 和 NVIDIA L40S 两个硬件后端，其他 XLA 后端（TPU v4/v5e、CPU、AMD via OpenXLA）有不同的融合与调度策略，绝对利用率会有差异；第二，所有实验使用固定块大小 $L=256$（Dao and Gu, 2024 的默认值），块大小作为调优变量与编译器优先路径正交，其与硬件 tiling 的交互未测量；第三，推理仅在固定 batch size 下报告，连续批处理与动态内存分页（Kwon et al., 2023, vLLM）等调度技术未实现，但缓存原语与之兼容；第四，训练对比对 L40S 上小模型短序列场景有利，forward+backward 在 130M@512 快 64.8%，但随规模与长度反转——130M@2048 慢 27.6%，780M@2048 慢 287.1%，1.3B（seq≥4096）和 2.7B（所有长度）在 JAX 路径下超过 L40S 内存，匹配 Triton 在 LLM 规模上需要 Pallas/Triton 自定义反向融合内核，属于编译器优先路径之外的工作；第五，JIT 编译开销在 2.7B@4096 处达 43 秒，对交互式研究构成可观开销。从独立观察补充几点：第六，WikiText-103 困惑度一致性的实现使用了 `jax_default_matmul_precision='highest'` 关闭 TF32，并通过 stride 512 评估，这种 controlled setting 并不意味着生产环境的高精度模式；第七，'非缓存基线'在同一 JAX 函数上仅禁用 SSM 缓存，并不等价于一个公平的对手——它实际上代表了'不用任何优化时 Mamba-2 的下界'，缓存路径相对于真实 Triton 实现（如 mamba_ssm v2.2.2 的 decode）在 L40S 上的绝对吞吐差距未直接给出（Table 4 给出 cached scan vs cached host vs non-cached 三个内部对比，但未与 Triton baseline 直接比较 decode 吞吐）；第八，MFU 与 HBU 计算依赖于 XLA 成本分析报告的 $F_{XLA}$ 与 $B_{XLA}$，作者承认 $B_{XLA}$ 是未融合的字节数，因此报告的 HBU 是真实带宽效率的上界。
独立分析的弱点
论文虽然论证了'编译器优先 SSD 在推理上可行'，但仍存在几个可改进的弱点。**(1) 训练反向路径仍依赖手写内核**：Table 13 清楚显示 JAX 路径在 780M@1024 起被 Triton 反超 132.1%，在 2048 token 处被反超 287.1%，1.3B 和 2.7B 完全超内存。这说明前向的 XLA 友好性不能直接迁移到反向——反向需要保存中间激活、计算梯度，XLA 的自动微分与 Pallas/Triton 内核的融合反向之间存在结构性差距。改进方向：作者已经指出'匹配 Triton 在 LLM 规模上需要 Pallas/Triton 自定义融合反向'，这是后续工作的明确路径。**(2) JIT 编译开销对交互式研究不友好**：2.7B@4096 需要 43 秒编译（Table 12），且若序列长度或 batch size 改变，需重新编译。改进方向：可以探索缓存 XLA 编译产物（XLA 已经支持 AOT compilation），或为常用 (batch, seq) 组合预编译。**(3) WikiText-103 验证的混淆因素**：评估使用 stride 512、batch 1、float32、TF32 关闭，论文承认'差异接近 float32 舍入'，但没有报告多个随机种子的方差或与 mamba_ssm 之外的独立参考实现的对比。改进方向：在 Pythia、RWKV 等其他 SSM/Transformer 基线上做交叉验证，或使用 lm-evaluation-harness 等标准化 harness。**(4) 块大小硬编码为 256**：所有实验使用 $L=256$，没有给出 $L$ 敏感性扫描。改进方向：在 64、128、256、512 上做 ablation，观察 MFU 与困惑度的权衡。**(5) 注意力 / 多模态架构的支持缺失**：论文只验证了纯 Mamba-2 推理，但实际部署常常是 SSM 与 attention 的混合架构（如 Jamba、Zamba）。改进方向：把 SSD-as-XLA 的四条结构条件分析推广到 hybrid 架构，验证 SSD 的 einsum 主导块是否能与 attention 的 quadratic attention block 共存于同一编译图。**(6) 数值容差与生产精度模式的差距**：Table 5–6 的对齐性测试在关闭 TF32 的最高精度模式下取得，论文承认生产部署通常使用 default matmul precision，性能数字可能会改善但需要重新验证一致性。改进方向：给出 default precision 下的误差上界或自适应精度管理策略。
未来方向
作者在论文 §6 明确提出的未来工作包括：(a) 把分析扩展到其他 XLA 后端（TPU v4/v5e、CPU、AMD via OpenXLA），评估不同融合策略对绝对利用率的影响；(b) 探索块大小作为调优变量的影响，研究其与硬件 tiling 的交互；(c) 将缓存原语与连续批处理、动态内存分页等 vLLM 式调度技术集成；(d) 解决训练阶段的反向融合——开发 Pallas 或 Triton 内核以在 LLM 规模匹配 mamba_ssm。基于论文成果，其他有价值的延伸方向包括：(e) 把 SSD-as-XLA 的结构条件分析推广到 Mamba-3、RWKV-7、Hyena 等新一代线性/次二次序列模型，验证它们是否同样满足'四条编译器友好条件'；(f) 探索不同精度（FP8、INT8）下的 SSD 实现与误差累积特征——论文已经显示 bfloat16 衰减指数化产生 0.013 logit 误差，FP8 可能更严重，需要误差补偿；(g) 把 PyTree-as-cache 范式推广到其他需要 O(1) 状态维护的推理场景，如 KV-cache compression、tree attention、speculative decoding 的 draft model 缓存；(h) 在长上下文（>32K token）评估中验证论文的 'cached decode sequence-length independent' 结论是否仍然成立，因为 SSM 的 $O(L)$ prefill 优势理论上应进一步放大；(i) 与 ONNX、TVM 等其他编译器栈的对比研究，验证 XLA 的 fusion/tiling 优势是 XLA 特有还是编译器优先范式的一般属性；(j) 工业级部署研究——评估该路径在 Kubernetes、TPU slices、多 host 分片等真实生产环境下的稳定性与吞吐。
复现评估
论文在复现性方面表现优秀，几乎所有必要细节都公开。**代码与权重**：作者在 §1 和 §4.1 提供了 GitHub 仓库链接 https://github.com/CosmoNaught/mamba2-jax，包含 SSD 核心函数（~60 行 Python）、benchmark 脚本、配置文件和复现指令。五个模型检查点（state-spaces/mamba2-130m/370m/780m/1.3b/2.7b）从 HuggingFace 直接加载，与 mamba_ssm v2.2.2 参考实现使用同一组权重。**软件版本**（Table 9）：JAX 0.9.0、jaxlib 0.9.0.1（含 XLA/libtpu TPU runtime）、Python 3.12、Flax NNX 0.12.4、PyTorch 2.10.0、mamba_ssm v2.2.2；Bonsai 依赖项固定到具体 commit（core PR #103 @ a907b75, caching PR #131 @ d8f8d11），这避免了 Bonsai 后续 commit 引入的不可预期修改。**配置标志**：jax_default_matmul_precision='highest'（正确性）或 default（吞吐），torch.backends.cuda.matmul.allow_tf32=False（用于 golden outputs 和困惑度），BF16 compute dtype 用于所有吞吐与消融。**硬件**：Cloud TPU v6e 单芯片（918 TFLOPS BF16 peak、1600 GB/s HBM）与单张 NVIDIA L40S（362 TFLOPS BF16 peak、864 GB/s GDDR6），Google TPU Research Cloud (TRC) 提供部分 TPU 时长。**测量协议**：所有计时取 5 次运行均值，标准差低于 0.3%；训练步对比取 10 次 timed steps 在 10 次 warm-up 之后；显式同步屏障确保计时准确。**复现难度**：低-中。需要 TPU 访问或 L40S GPU、JAX 0.9.0 兼容性、约 32 GB GPU 内存（用于 2.7B 模型），其中 1.3B/2.7B 在 L40S 上需要非默认内存配置或 CPU offload。主要的'隐藏'复杂度在于 XLA 编译依赖和 TPU 访问——TPU v6e 通常需要 Google Cloud 项目与配额，这是外部研究者可能面临的最大障碍。
论文图表