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

Cosmo Santoni; Anmol Thapar

Source-linked AI summary

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

Cosmo Santoni, Anmol Thapar

arXiv:2603.09555v2Mar 10, 2026cs.LGcs.AIcs.DCcs.PF

TL;DR

Mamba-2 inference is commonly tied to custom CUDA and Triton kernels, creating a portability gap across accelerator backends. This paper expresses SSD’s compiler-friendly structure in standard JAX primitives with a registered cache and compiled on-device decoding, achieving roofline-level TPU utilisation, portable L40S decoding, and Triton-level numerical parity. The main scope boundaries are inference-focused evaluation, fixed batch policies, and substantial compilation or large-model limitations.

Problem
Custom CUDA and Triton kernels tie high-throughput Mamba-2 inference to specific accelerator backends, limiting a single portable serving path.
Method
The paper preserves SSD’s diagonal, chunked, einsum-dominated, and statically controlled structure in standard JAX, with a registered PyTree cache and compiled on-device autoregressive loop.
Results
The single-source implementation reaches 15% MFU prefill and 64% HBU decode on TPU v6e, sequence-length-independent cached decode on L40S, and WikiText-103 parity within ±0.0005 points across five scales.
Takeaways & Limitations
SSD’s structural constraints are sufficient for competitive Mamba-2 inference code generation across TPU and GPU targets without custom kernels.
Takeaways & Limitations
Evaluation targets inference with fixed batch policies; continuous batching and dynamic memory paging are not implemented, and 2.7B decode compilation at length 4096 takes 43 seconds.

Abstract

from arXiv · show

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.

1 Introduction

This work presents a compiler-first JAX implementation of Mamba-2 inference, using SSD’s structural regularities to avoid custom kernels and support portable cached decoding. It reports roofline-consistent TPU utilisation, sequence-length-independent L40S decoding, and numerical parity with Triton.

Motivation and approach: SSD’s diagonal dynamics, fixed-size chunks, einsum-dominated computation, and static control flow expose Mamba-2 to compiler optimisations in standard JAX primitives.These conditions map onto XLA fusion and tiling passes.
Contribution: The implementation provides chunked-parallel prefill, cached autoregressive decoding, and a registered JAX PyTree cache in a single inference path.The cache’s array leaves participate in JAX tracing.
Results: 140 TFLOPS and 15% MFU are reached for batch-1 prefill on TPU v6e, while cached decode reaches 64% HBU at the batch-1 roofline ceilings.These are reported as TPU v6e results.
Results: The same source reproduces sequence-length-independent cached decode across five model scales on NVIDIA L40S.This supports the claimed cross-hardware portability.
Validation: WikiText-103 perplexity matches the Triton reference within ±0.0005 points across all five checkpoints, with hidden states agreeing to float32 rounding tolerance.The reference is mamba_ssm v2.2.2.

2 Related Work

Prior Mamba implementations rely on specialised kernels or limited JAX ports, while this work extends compiler-first JAX inference with cache support and cross-hardware evaluation.

Kernelised SSM implementations: Reference Mamba and Mamba-2 implementations use fused CUDA or Triton kernels integral to their reported throughput.These kernels are tailored to selective scan or SSD’s chunked structure.
Kernelised SSM implementations: Community ports to AMD ROCm and Apple MPS are maintained as separate forks with distinct kernel codepaths.The ports therefore do not constitute one shared implementation source.
JAX implementations: Earlier JAX work targets Mamba-1 or supplies minimal Mamba-2 forward passes without an autoregressive cache or performance evaluation.S5 provides a separate JAX-native parallel-scan approach for structured state spaces.
Relation to Bonsai: This paper adds structural XLA analysis, cross-hardware roofline characterisation, implementation-choice ablations, and numerical validation against Triton.The evaluated module is part of Bonsai and uses its module structure and registered-PyTree cache.
Compiler-first inference: JAX delegates device-specific code generation to XLA, whereas torch.compile offers a comparable CUDA route but lacked a mature TPU backend at writing.This comparison situates the compiler-first approach among compilation routes.

3 Method

The method preserves SSD’s algebraic structure through the JAX front end, exposes it to XLA, and carries fixed-size recurrent state through a compiled on-device decoding loop.

SSD structure: SSD combines diagonal per-head dynamics, fixed-size chunking, einsum-dominated computation, and static control flow for compiler-friendly inference.The recurrence remains diagonal and chunked, while causal masking is static.
State-space duality: Mamba-2 makes B, C, and Δ input-dependent while restricting A to a diagonal scalar per head.This reduces the matrix exponential to scalar exponentiation in the discretised recurrence.
Chunked recurrence: Within each fixed-size chunk, SSD uses a parallel matrix computation; a lightweight scan propagates summary states between chunks.Chunk size L is 256 throughout and balances arithmetic intensity against sequential overhead.
Primitive-level implementation: Batched einsum contractions and precomputed lower-triangular masking expose fused computation to XLA.Runtime row-wise masking breaks fusion and reduces prefill throughput by 82.8%.
Autoregressive decoding: Compiled on-device loops keep the loop body, cache update, and argmax in one XLA program, avoiding per-token host synchronisation.A Python-driven loop is 2.4× slower at the 130M scale.
Autoregressive cache: The fixed-size hidden and convolution states make each next-token update O(1) in prefix length.The cache is stored as a registered JAX PyTree and traced into the compiled loop.

4 Evaluation

The evaluation measures compiled SSD inference on TPU v6e and NVIDIA L40S across five Mamba-2 checkpoints, comparing cached and non-cached decoding, memory, utilisation, and numerical agreement.

Setup: Five pretrained Mamba-2 checkpoints spanning 130M to 2.7B parameters were evaluated with batch size 1 on TPU v6e and NVIDIA L40S.The experiments use single-stream decoding and a fixed chunk size of L = 256.
Autoregressive throughput: Cached decoding provides sequence-length-independent per-step throughput, while non-cached throughput collapses as sequence length increases.On TPU v6e, cached latency grows linearly and non-cached latency grows quadratically across the evaluated sequence lengths.
Autoregressive throughput: 2.4× higher throughput is achieved by the compiled on-device loop than the host-driven cached loop at the 130M scale.The paths converge above 780M parameters because per-step computation dominates host–device round-trip overhead.
Memory: Cached decoding keeps peak memory constant, whereas non-cached decoding grows linearly with sequence length.The cache stores fixed-size per-layer SSM and convolution states, while the baseline materialises the full token buffer and intermediate activations at each step.
Hardware utilisation: 15% MFU prefill and 64% HBU decode are observed on TPU v6e at batch size 1, consistent with the stated roofline ceilings.Prefill is compute-bound, while cached decode is memory-bandwidth bound; decode HBU varies by less than 1.7 percentage points across sequence lengths.
Numerical validation: WikiText-103 validation perplexity differs from the Triton reference by at most 0.0005, with hidden-state differences remaining at float32 rounding scale.The comparison uses matched conditions, float32 arithmetic, and the same five checkpoints.

5 Discussion

Cached decode reaches 64% HBU because XLA operates on fixed-size SSM and convolution caches, while the registered PyTree carries O(1) state through the compiled loop without host round-trips.

64% HBU results from XLA tiling fixed-size SSM and convolution cache traffic into a stable memory pattern.Each decode step reads and writes fixed-size caches.
A host-driven decode loop is 2.4× slower than a compiled loop at 130M.
The registered JAX PyTree carries O(1) autoregressive state through compiled on-device control flow without host round-trips.

6 Limitations

The evaluation is bounded by specific hardware, chunking, batching, training, architectural, compiler-maturity, and compilation-cost conditions. These constraints limit how broadly the reported performance and portability results should be generalized.

Hardware scope: MFU and HBU are reported only on TPU v6e and NVIDIA L40S, while utilization on other XLA backends may vary.TPU v4, TPU v5e, CPU, and AMD GPUs via OpenXLA use different fusion strategies and scheduling.
Configuration scope: All experiments use chunk size L = 256, leaving its interaction with hardware tiling unmeasured.Chunk size is a tuning variable separate from the compiler-first pattern.
Inference scope: Inference results use fixed batch sizes and do not implement continuous batching or dynamic memory paging.The cache primitive is described as compatible with such schedulers, but their scheduling effects are not evaluated.
Training scope: Training is not the primary target, and the favorable reduced L40S comparison applies only to small models and short horizons.Forward+backward time is up to 2.8× lower at 130M parameters and 512 tokens, with the advantage shrinking as scale and sequence length grow.
Architectural scope: The analyzed structural conditions exclude architectures requiring runtime-index gather or scatter, warp synchronization, or data-dependent control flow.SSD does not require these compiler-hostile primitives.
Compiler scope: Early-stage XLA backends may not match mature TPU and GPU fusion or tiling quality, so reported utilization should not be extrapolated without measurement.Backend maturity can affect absolute performance.
Compilation cost: At 2.7B parameters and sequence length 4096, decode compilation takes 43 seconds and is amortized only across many inference calls.This one-time compilation cost can dominate iterative-research wall-clock.

7 Conclusion

The paper identifies four structural conditions that let XLA generate competitive Mamba-2 inference code without custom kernels. A single source reaches batch-1 TPU roofline ceilings for prefill and decode and supports sequence-length-independent cached decode on NVIDIA L40S.

Conclusion: Four conditions—diagonal state dynamics, chunkable recurrence, einsum-dominated compute, and static control flow—support competitive XLA code without custom kernels.These conditions characterize the compiler-friendly structure of SSD for Mamba-2 inference.
Conclusion: 15% model FLOP utilisation prefill and 64% hardware bandwidth utilisation decode are reached on TPU v6e at batch-1 roofline ceilings.The reported figures come from a single-source implementation.
Conclusion: Cached decode remains sequence-length independent on NVIDIA L40S across the evaluated model scales.The same implementation source runs across the two accelerator backends.

A.2 Benchmark Configurations

The benchmarks cover single-stream prefill and cached decoding on TPU v6e, cached decoding on L40S, reduced training-step comparisons, and WikiText-103 perplexity evaluation.

TPU v6e decode uses five models, six sequence lengths from 128 to 4096, three methods, and five timed runs after JIT warm-up.Prompt length is fixed at 16 tokens.
L40S decode follows the same single-stream protocol as the TPU v6e sweep.
TPU v6e prefill evaluates five models at prompt lengths of 1024, 4096, and 8192 tokens with five timed runs.XLA cost analysis is extracted for each configuration.
The L40S training-step sweep covers three checkpoints from 130M to 780M, batch size 1, and sequence lengths of 512, 1024, and 2048.The larger checkpoints and longer sequences exceed L40S memory under the JAX path and are omitted.
WikiText-103 perplexity is measured on the validation split with stride 512, batch size 1, float32 computation, and TF32 disabled for both implementations.

A.3 Reproducibility

The reproducibility materials provide software configurations and TPU v6e measurements for throughput, memory, utilisation, and compilation time across model sizes and sequence lengths.

The reproducibility record includes software versions, configuration flags, and dedicated TPU v6e throughput and memory tables.
Cached throughput is sequence-length independent, while the non-cached path slows sharply as sequence length increases.
Cached peak device memory remains constant, whereas non-cached peak memory grows with sequence length.
Prefill MFU and mean decode HBU both increase with model size across the measured TPU v6e regimes.
Figure 6 reports best prefill MFU against 918 TFLOPS and mean decode HBU against 1600 GB/s for batch size 1.Utilisation increases with model size in both regimes.
One-time XLA compilation time grows with model size and decode horizon, while subsequent calls reuse the compiled program.

B.5 Reduced Training-Step Comparison (NVIDIA L40S)

The reduced L40S comparison evaluates compiler-first JAX against the Triton reference using batch-1 forward-plus-backward timing on smaller checkpoints and shorter sequences.

At 130M and 512 tokens, the JAX path is 64.8% faster than the Triton reference.
The JAX advantage shrinks with model size and sequence length and becomes a slowdown by 2048 tokens for every measured checkpoint.
The comparison measures mean forward-plus-backward time over ten timed steps after ten warm-ups on a single NVIDIA L40S.The optimiser update is excluded because its JAX timing was dominated by a measurement artefact.
The SSD implementation uses einsum contractions for intra-chunk output, state accumulation, and inter-chunk recurrence updates.The equations expose batch, chunk, sequence, head, state, and projection axes in the contractions.