用编译器友好方式实现Mamba-2推理，支持跨硬件零成本缓存加速。
Compiler-First State Space Duality and Portable $O(1)$ Autoregressive Caching for Inference
- 将状态空间双重性结构转为标准JAX算子，无需定制内核
- 单个TPU上解码速度比重算快27至36倍，能耗效率达理论极限
- 代码可直接移植到NVIDIA显卡，适合追求高效部署的研究者
高性能Mamba-2推理通常依赖融合的CUDA与Triton内核，限制了在不同加速器上的可移植性。我们发现状态空间双重性（SSD）递推具有编译器友好的结构：每头对角化动态、固定大小分块、以einsum为主导的计算以及静态控制流。通过标准JAX原语表达该结构，实现了无需自定义内核的单一源码推理路径，包含注册的JAX PyTree缓存和编译后的设备端自回归循环。在单个Google Cloud TPU v6e上，批量为1的prefill达到约140 TFLOPS，即15%模型浮点利用率（MFU），达到此场景下的理论上限；缓存解码最高实现64%硬件带宽利用率（HBU）。在4096令牌上下文下，缓存解码速度相较全前缀重算提升27至36倍，覆盖从130M到2.7B参数的五个Mamba-2检查点。相同源码在NVIDIA L40S上无需修改即可运行，解码性能与序列长度无关。WikiText-103验证困惑度与Triton参考实现mamba_ssm v2.2.2相差不超过±0.0005，隐藏状态符合float32舍入容差。代码已公开于https://github.com/CosmoNaught/mamba2-jax。
原文摘要 · Abstract (English)
High-throughput Mamba-2 inference is usually tied to fused CUDA and Triton kernels, limiting portability across accelerator backends. We show that the state space duality (SSD) recurrence has a compiler-friendly structure: diagonal per-head dynamics, fixed-size chunking, einsum-dominated compute, and static control flow. Expressing this structure in standard JAX primitives gives a single-source inference path with no custom kernels, a registered JAX PyTree cache, and a compiled on-device autoregressive loop. On a single Google Cloud TPU v6e, batch-1 prefill reaches approximately 140 TFLOPS, or 15% model FLOP utilisation (MFU), the roofline ceiling for this regime, and cached decode reaches up to 64% hardware bandwidth utilisation (HBU). At a 4096-token context, cached decode is 27x--36x faster than full-prefix recomputation across five Mamba-2 checkpoints from 130M to 2.7B parameters. The same source runs unmodified on NVIDIA L40S, where cached decode remains sequence-length independent across all model scales. WikiText-103 validation perplexity matches the Triton reference mamba_ssm v2.2.2 within +/-0.0005 points, and hidden states agree to float32 rounding tolerance. Code is available at https://github.com/CosmoNaught/mamba2-jax.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。