# Compiler-First State Space Duality and Portable O(1) Autoregressive Caching for Inference

arXiv: https://arxiv.org/abs/2603.09555 · HTML: https://mlsanity.com/paper/2603.09555/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 32 references that are in the map.
- First posted on arXiv 10 Mar 2026.

## Builds on (most cited first, 14 of 14)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Efficiently Scaling Transformer Inference](https://mlsanity.com/paper/2211.05102/) (2022)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [Is Mamba Effective for Time Series Forecasting?](https://mlsanity.com/paper/2403.11144/) (2024)

## Built on by (most cited first, 2 of 2)

- [LinearKV: One Cached State Suffices for Position-Independent Caching in Hybrid LLMs](https://mlsanity.com/paper/2608.11231/) (2026)
- [Accelerator Choice Is Not Enough: AlphaFold2 Inference on Cloud TPUs](https://mlsanity.com/paper/2609.34818/) (2026)

## Abstract (arXiv)

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.