mlsanityData to 7 Oct 2026

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

arXiv 2603.09555 ↗First posted 10 Mar 2026cs.LG cs.AI cs.DCcited by 2

Not on the map. Why?

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.

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

Builds on

14 of 32 references.

References · 14

  1. 2017Attention Is All You Needcited by 59,932
  2. 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,767
  3. 2022PaLM: Scaling Language Modeling with Pathwayscited by 5,304
  4. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
  5. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  6. 2016Pointer Sentinel Mixture Modelscited by 2,992
  7. 2021Efficiently Modeling Long Sequences with Structured State Spacescited by 2,424
  8. 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 1,979
  9. 2024Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Dualitycited by 1,283
  10. 2022Simplified State Space Layers for Sequence Modelingcited by 776
  11. 2020Long Range Arena: A Benchmark for Efficient Transformerscited by 696
  12. 2022Efficiently Scaling Transformer Inferencecited by 529
2 more
  1. 2023Gated Linear Attention Transformers with Hardware-Efficient Trainingcited by 490
  2. 2024Is Mamba Effective for Time Series Forecasting?cited by 110

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 2.

2 papers.

  1. 2026LinearKV: One Cached State Suffices for Position-Independent Caching in Hybrid LLMscited by 3
  2. 2026Accelerator Choice Is Not Enough: AlphaFold2 Inference on Cloud TPUscited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2026ECHO: Elastic Speculative Decoding with Sparse Gating for High-Concurrency Scenarios
  2. 2026NetKV: Network-Aware Decode Instance Selection for Disaggregated LLM Inference
  3. 2025SpecMamba: Accelerating Mamba Inference on FPGA with Speculative Decoding
  4. 2026SonicSampler: Unified Tile-Aware Kernels for LLM Sampling and Speculative Verification
  5. 2025Correctness Forensics for Batch Speculative Decoding: Diagnosing the Ragged Tensor Problem
  6. 2026MARS: Enabling Autoregressive Models Multi-Token Generation
  7. 2026DualPath: Breaking the Storage Bandwidth Bottleneck in Agentic LLM Inference
  8. 2026DualDecoder: Accelerate Long Context LLM Inference by Predictive Prefetch