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

Cosmo Santoni; Anmol Thapar

cs.LGMar 10, 2026

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

Authors: Cosmo Santoni, Anmol Thapar

Organizations: Imperial College London

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: Host–device boundary for compiled autoregressive decoding. The SSM cache is a JAX PyTree, so the compiled loop updates the $O(1)$ SSM and convolution state on device. The Python host launches the compiled XLA program once and remains inactive until generation completes.

Figure 2: Autoregressive generation on Cloud TPU v6e across five model scales and six sequence lengths. (a) Speedup from caching. (b) Generation latency. Cached (solid) grows linearly, non-cached (dashed) grows quadratically. (c) Per-step throughput. Cached throughput is flat regardless of sequence length, and non-cached throughput collapses.

		Throughput (Tokens/Second)
Model	Method	128	1024	4096
130M	Cached (scan)	1588	1635	1641
	Cached (host)	662	729	751
	Non-Cached	903	278	56
370M	Cached (scan)	626	641	641
	Cached (host)	392	391	390

Table 1: Decode strategy comparison on TPU v6e (batch size 1). Small inversions between cached scan and cached host at 780M (318 vs 325) and 2.7B (94 vs 97) are within measurement variance; both models are per-step-compute dominated, so the two paths have converged.

Figure 3: Peak memory during autoregressive generation on Cloud TPU v6e. Cached decoding (solid) is constant; non-cached (dashed) grows linearly. At sequence length 4096, the 2.7B non-cached path consumes over 16 GB versus a constant 10.9 GB cached.

Figure 4: Hardware utilisation on Cloud TPU v6e (batch size 1). (a) Prefill MFU versus model size at three prompt lengths. (b) Decode HBU versus sequence length. HBU varies by less than 1.7 percentage points across sequence lengths for every model.

	Prefill MFU (%) by Prompt Length
Model	1024	4096	8192
130M	6.22	8.23	7.68
370M	7.47	9.04	7.60
780M	10.62	11.33	8.20
1.3B	12.53	11.67	8.39
2.7B	15.23	12.96	9.71

Table 2: Prefill compute efficiency on TPU v6e (batch size 1, peak = 918 TFLOPS BF16).

	Decode HBU (%) by Sequence Length
Model	128	256	512	1024	2048	4096
130M	51.62	52.29	52.87	53.13	53.26	53.32
370M	57.88	58.65	59.07	59.29	59.32	59.32
780M	62.07	62.59	62.80	62.87	62.93	62.99
1.3B	61.22	61.55	61.69	61.77	61.86	61.87
2.7B	63.43	63.74	63.91	64.00	64.06	64.08

Table 3: Decode memory-bandwidth efficiency on TPU v6e (batch size 1, peak = 1600 GB/s).

		Throughput (Tokens/Second)
Model	Method	128	1024	4096
130M	Cached (scan)	240.2	267.1	314.2
	Cached (host)	178.4	141.9	188.5
	Non-Cached	203.3	115.8	20.3
370M	Cached (scan)	154.3	165.1	148.0
	Cached (host)	104.1	98.8	112.3

Table 4: Single-stream autoregressive decode throughput on NVIDIA L40S (batch size

1

). Sequence lengths mirror the TPU evaluation in Table 1 .

Model	Triton PPL	JAX PPL	\| $\Delta$ \|
130M	18.7023	18.7019	0.0004
370M	13.1247	13.1244	0.0003
780M	10.8892	10.8886	0.0005
1.3B	9.5708	9.5704	0.0004
2.7B	8.3252	8.3250	0.0002

Table 5: WikiText-103 validation perplexity (stride 512) for Triton reference ( mamba_ssm v2.2.2) and JAX implementation. |

\Delta

| is the absolute difference between the JAX implementation and the Triton reference; values within

\pm

0.0005 indicate functional equivalence.

Figure 5: WikiText-103 validation perplexity versus batch size for the $130$ M model (JAX implementation). Perplexity is invariant to batch size for this checkpoint.

Output	Relative Tolerance	Absolute Tolerance
Last hidden state	$1\times 10^{-5}$	$1\times 10^{-4}$
Logits (first 256)	$1\times 10^{-5}$	$2\times 10^{-4}$

Table 6: Numerical parity against the PyTorch reference ( mamba_ssm , 130M checkpoint, 512 tokens). Tolerances apply element-wise.

Masking Strategy	Prefill (tokens/s)	Output
Static mask ( jnp.tril )	42,631	Bitwise identical
Dynamic row-wise mask ( fori_loop )	7,330 ( $-82.8\%$ )	Bitwise identical

Table 7: Masking ablation on TPU v6e (1.3B model, BF16, prompt length 1024). Both produce bitwise identical output; the dynamic variant breaks XLA’s fusion chain.

Decay Dtype	Max Absolute Error (Logits)
float32 (baseline)	0.0
bfloat16	0.013

Table 8: Decay precision ablation on TPU v6e (130M checkpoint, 24 layers, BF16, prompt length 1024).

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Version / Value
JAX	0.9.0
jaxlib	0.9.0.1
XLA / libtpu (TPU runtime)	Bundled with jaxlib 0.9.0.1
Python	3.12
Flax (NNX)	0.12.4
PyTorch (golden outputs)	2.10.0

Appendix

Table 9: Software versions and configuration flags.

		Throughput (Tokens/Second) by Sequence Length
Model	Method	128	256	512	1024	2048	4096
130M	Cached	1588	1609	1627	1635	1639	1641
	Non-Cached	903	898	626	278	132	56
370M	Cached	626	634	639	641	641	641
	Non-Cached	495	348	205	124	39	18
780M	Cached	318	321	322	322	323	323

Appendix

Table 10: Autoregressive decoding throughput on TPU v6e (batch size 1, fori_loop scan path).

		Peak Memory (MB) by Sequence Length
Model	Method	128	256	512	1024	2048	4096
130M	Cached	545.6	545.6	545.6	545.6	545.6	545.6
	Non-Cached	565	585	624	702	857	1169
370M	Cached	1591.9	1591.9	1591.9	1591.9	1591.9	1591.9
	Non-Cached	1644	1696	1799	2007	2422	3251
780M	Cached	3322.8	3322.8	3322.8	3322.8	3322.8	3322.8

Appendix

Table 11: Peak memory during autoregressive generation on TPU v6e (batch size 1).

Figure 6: Fraction of hardware peak on Cloud TPU v6e (batch size $1$ ). Solid bars give best prefill MFU (% of $918$ TFLOPS); faded bars give mean decode HBU (% of $1600$ GB/s). Utilisation increases with model size in both regimes.

	JIT Compilation Time (seconds)
Model	Prefill (1024)	Decode (128)	Decode (4096)
130M	5.5	5.6	2.5
370M	10.2	13.0	6.4
780M	13.0	13.7	12.6
1.3B	10.2	14.9	21.4
2.7B	15.8	19.5	43.0

Appendix

Table 12: XLA JIT compilation time on TPU v6e. One-time costs; subsequent calls reuse the compiled program.

Model	Seq. length	JAX (ms)	Triton (ms)	$\Delta\%$
130M	512	25.9	73.7	$-64.8$
	1024	45.2	72.4	$-37.5$
	2048	86.7	68.0	$+27.6$
370M	512	62.8	147.0	$-57.3$
	1024	115.8	128.6	$-9.9$
	2048	229.6	151.4	$+51.7$

Appendix

Table 13: Reduced training-step comparison on a single NVIDIA L40S. Cells give mean forward

+

backward time (ms) over ten timed steps for the JAX compiler-first path and the Triton reference ( mamba_ssm v

2.2.2

), batch size

1

\Delta=(t_{\text{JAX}}-t_{\text{Triton}})/t_{\text{Triton}}

; negative means the JAX path is faster. The JAX path is faster for small models at short sequences and crosses over to several times slower as model size and sequence length grow. The

1.3

B model (sequence length

\geq 4096

) and the

2.7

B model (all sequence lengths) exceed L40S memory under the JAX path and are omitted.

Efficient Inference

Efficient inference reduces the compute, memory, latency, or energy needed to produce accurate predictions, posterior estimates, and decisions.

See all 107 papers

Mamba-3: Improved Sequence Modeling using State Space Principles
Mar 16, 2026Aakash Lahoti, Kevin Y. Li, Berlin Chen +5State Space Language Models Sequence Modeling
Carnegie Mellon University · Princeton University · Together AI +1
DART: Decoded Attention over Recurrent States for Efficient Long-Context Sequence Modeling
Aug 3, 2026Yixiao Qian, Song Chen, Pengkai Wang +3Memory-Augmented Language Models KV Caching
College of Control Science and Engineering, Zhejiang University · Department of Mathematics, National University of Singapore · Hong Kong Polytechnic University +1
TurboMamba: Faster Inference, Faster Training, and Better Prediction with Lexical Memory
Oct 14, 2026Hanxin ZhangHybrid Transformer-State Space Model Architectures Autoregressive Language Modeling
Department of Neurology, University of Pittsburgh

Mar 16, 2026J/K move · Enter open · S save

Mamba-3: Improved Sequence Modeling using State Space Principles

Aakash Lahoti, Kevin Y. Li, Berlin Chen, Caitlin Wang+4

Carnegie Mellon University · Princeton University · Together AI+1

State Space Language Models Sequence Modeling Efficient Language Model Inference State Space Models State Tracking

Open paper AskPDF arXiv

Scaling inference-time compute has emerged as an important driver of LLM performance, making inference efficiency a central focus of model design alongside model quality. While the current Transformer-based models deliver strong model quality, their quadratic compute and linear memory make inference expensive. This has spurred the development of sub-quadratic models with reduced linear compute and constant memory requirements. However, many recent linear models trade off model quality and capability for algorithmic efficiency, failing on tasks such as state tracking. Moreover, their theoretically linear inference remains hardware-inefficient in practice. Guided by an inference-first perspective, we introduce three core methodological improvements inspired by the state space model (SSM) viewpoint of linear models. We combine: (1) a more expressive recurrence derived from SSM discretization, (2) a complex-valued state update rule that enables richer state tracking, and (3) a multi-input, multi-output (MIMO) formulation for better model performance without increasing decode latency. Together with architectural refinements, our Mamba-3 model achieves significant gains across retrieval, state-tracking, and downstream language modeling tasks. At the 1.5B scale, Mamba-3 improves average downstream accuracy by 0.6 percentage points compared to the next best model (Gated DeltaNet), with Mamba-3's MIMO variant further improving accuracy by another 1.2 points for a total 1.8 point gain. Across state-size experiments, Mamba-3 achieves comparable perplexity to Mamba-2 despite using half of its predecessor's state size. Our evaluations demonstrate Mamba-3's ability to advance the performance-efficiency Pareto frontier.

Abstract

Efficient Inference

Explore similar work

Mamba-3: Improved Sequence Modeling using State Space Principles

DART: Decoded Attention over Recurrent States for Efficient Long-Context Sequence Modeling

TurboMamba: Faster Inference, Faster Training, and Better Prediction with Lexical Memory