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

• URL: http://arxiv.org/abs/2603.09555v1
• Date: Tue, 10 Mar 2026 12:03:00 GMT
• Title: Compiler-First State Space Duality and Portable $O(1)$ Autoregressive Caching for Inference
• Abstract summary: Mamba-2's state space duality algorithm maps cleanly onto what XLA's fusion and tiling passes actually optimise.<n>We implement the full inference path (prefill, cached autoregressive decoding) as shaped standard primitives under XLA.<n>The implementation runs unmodified on CPU, NVIDIA GPU, and Google Cloud TPU from a single JAX source.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: State-space model releases are typically coupled to fused CUDA and Triton kernels, inheriting a hard dependency on NVIDIA hardware. We show that Mamba-2's state space duality algorithm -- diagonal state structure, chunkable recurrence, and einsum-dominated compute with static control flow -- maps cleanly onto what XLA's fusion and tiling passes actually optimise, making custom kernels optional rather than required. We implement the full inference path (prefill, cached autoregressive decoding) as shaped standard primitives under XLA, without hand-written kernels, and realise the architecture's theoretical $O(1)$ state management as a compiled on-device cache requiring no host synchronisation during generation. The implementation runs unmodified on CPU, NVIDIA GPU, and Google Cloud TPU from a single JAX source. On TPU v6e across five model scales (130M--2.7B parameters), XLA-generated code reaches approximately 140 TFLOPS on single-stream prefill ($15%$ MFU) and up to $64%$ bandwidth utilisation on decode. Greedy decoding matches the PyTorch/CUDA reference token-for-token across 64 steps, with hidden-state agreement within float32 rounding tolerance. The pattern transfers to any SSM recurrence satisfying the same structural conditions, on any platform with a mature XLA backend. The implementation is publicly available at https://github.com/CosmoNaught/mamba2-jax and merged into the Bonsai JAX model library.

Related papers

• Physically Partitioned KVCache Format for CPU--GPU Load Balancing in MoE Inference [10.117126683546028]
  Single-GPU long-context inference with Mixture-of-Experts (MoE) models requires spilling the key-value cache ( KVCache) to CPU memory.<n>Existing KVCache abstractions offer only storage semantics over a monolithic object of a single physical state.<n>We propose InplaceKVCache, the first KVCache abstraction whose format fixes each byte's physical residency at write time.
  arXiv  Detail & Related papers  (2026-09-13T13:33:24Z)
• Pre-Compiled Pipeline Shards for Distributed LLM Inference on Intel AI PC Fleets [0.0]
  We show that a handful of AIPCs can serve models beyond the capability of any single one.<n>A four-node deployment of Lunar Lake AI PCs on Intel Tiber Cloud serves a single user at interactive speed.
  arXiv  Detail & Related papers  (2026-08-19T17:33:28Z)
• QSimAdv: A Late-Bound, Vendor-Agnostic Architecture for High-Performance Quantum-Circuit Simulation [1.435000056979246]
  We present QSimAdv, which makes late binding, rather than a common kernel.<n>We realize this design on NVIDIA GH200 and AMD MI250X/EPYC systems.
  arXiv  Detail & Related papers  (2026-08-11T23:17:29Z)
• Claw-SWE-Bench: A Benchmark for Evaluating OpenClaw-style Agent Harnesses on Coding Tasks [75.92297551160692]
  General-purpose agents such as OpenClaw are increasingly used as autonomous tool users.<n>We introduce Claw-SWE-Bench, a multilingual SWE-bench-style benchmark and adapter protocol.<n>Claw-SWE-Bench treats harness and cost accounting as first-class axes of SWE-style coding-agent evaluation.
  arXiv  Detail & Related papers  (2026-06-10T17:16:23Z)
• Scaling Parallel Sequence Models to Foundation-Scale Vision Encoders [123.58723804218151]
  Vision foundation models are bottlenecked by the quadratic cost of self-attention.<n>We present C-GSPN based on 2D spatial propagation propagation.
  arXiv  Detail & Related papers  (2026-05-30T14:29:43Z)
• Scaling DoRA: High-Rank Adaptation via Factored Norms and Fused Kernels [83.99688944263843]
  Weight-De Low-Rank Adaptation (DoRA) extends LoRA by decoupling weight magnitude from direction.<n>At d_in = 8192 and rank r = 384, a single module's norm requires about 512 MB of transient working memory in bf16.<n>A factored norm decomposes the squared norm into base, cross, and Gram terms computable through O(d_out r + r2) intermediates, eliminating the dense product.
  arXiv  Detail & Related papers  (2026-03-23T17:57:24Z)
• FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling [20.849774181776414]
  FlashAttention-4 achieves up to 1.3$times$ speedup over cuDNN 9.13 and 2.7$times$ over Triton on B200 GPUs with BF16.<n>We implement FlashAttention-4 entirely in CuTe- embedded in Python, achieving 20-30$times$ faster compile times compared to traditional C++ template-based approaches.
  arXiv  Detail & Related papers  (2026-03-05T18:24:49Z)
• Scaling State-Space Models on Multiple GPUs with Tensor Parallelism [0.24148976266903474]
  Selective state space models (SSMs) have rapidly become a compelling backbone for large language models.<n>But in deployment, their inference performance is often bounded by the memory capacity, bandwidth, and latency limits of a single GPU.<n>This paper presents a communication-efficient TP design for selective SSM inference that addresses three practical engineering challenges.
  arXiv  Detail & Related papers  (2026-02-24T17:47:54Z)
• Memory-Efficient Acceleration of Block Low-Rank Foundation Models on Resource Constrained GPUs [11.45717904490388]
  Recent advances in transformer-based foundation models have made them the default choice for many tasks.<n>Their rapidly growing size makes fitting a full model on a single GPU increasingly difficult and their computational cost prohibitive.<n>Block low-rank (BLR) compression techniques address this challenge by learning compact representations of weight matrices.
  arXiv  Detail & Related papers  (2025-12-24T00:41:13Z)
• dInfer: An Efficient Inference Framework for Diffusion Language Models [54.80918957287927]
  Diffusion-based large language models (dLLMs) have emerged as a promising alternative to autoregressive (AR) LLMs.<n>We present dInfer, an efficient and framework for dLLM inference.
  arXiv  Detail & Related papers  (2025-10-09T16:19:42Z)
• dParallel: Learnable Parallel Decoding for dLLMs [77.24184219948337]
  Diffusion large language models (dLLMs) offer parallel token prediction and lower inference latency.<n>Existing open-source models still require nearly token-length decoding steps to ensure performance.<n>We introduce dParallel, a simple and effective method that unlocks the inherent parallelism of dLLMs for fast sampling.
  arXiv  Detail & Related papers  (2025-09-30T16:32:52Z)
• Two-stream Beats One-stream: Asymmetric Siamese Network for Efficient Visual Tracking [54.124445709376154]
  We propose a novel asymmetric Siamese tracker named textbfAsymTrack for efficient tracking.<n>Building on this architecture, we devise an efficient template modulation mechanism to inject crucial cues into the search features.<n>Experiments demonstrate that AsymTrack offers superior speed-precision trade-offs across different platforms.
  arXiv  Detail & Related papers  (2025-03-01T14:44:54Z)
• vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving [53.972175896814505]
  Large Language Models (LLMs) are widely used across various domains, processing millions of daily requests. Large Language Models (LLMs) are widely used across various domains, processing millions of daily requests.
  arXiv  Detail & Related papers  (2024-07-22T14:37:58Z)
• Hardware-Aware Parallel Prompt Decoding for Memory-Efficient Acceleration of LLM Inference [23.633481089469836]
  Auto-regressive decoding of Large Language Models (LLMs) results in significant overheads in their hardware performance.<n>We propose a novel parallel prompt decoding that requires only $0.0002$% trainable parameters, enabling efficient training on a single A100-40GB GPU in just 16 hours.<n>Our approach demonstrates up to 2.49$times$ speedup and maintains a minimal memory overhead of just $0.0004$%.
  arXiv  Detail & Related papers  (2024-05-28T22:19:30Z)
• FlexLLM: Token-Level Co-Serving of LLM Inference and Finetuning with SLO Guarantees [19.58773369944074]
  Finetuning large language models (LLMs) is essential for task adaptation, yet today's serving stacks isolate inference and finetuning on separate GPU clusters.<n>We introduce FlexLLM, the first system to co-serve LLM inference and PEFT-based finetuning on shared GPUs by fusing at the token level.<n>At runtime, a novel token-level finetuning mechanism paired with a hybrid token scheduler dynamically interleaves inference and training tokens within each co-serving iteration.
  arXiv  Detail & Related papers  (2024-02-29T01:33:08Z)
• QMoE: Practical Sub-1-Bit Compression of Trillion-Parameter Models [64.34635279436054]
  Mixture-of-Experts (MoE) architectures offer a general solution to the high inference costs of large language models (LLMs) via sparse routing. We present a solution to this memory problem, in form of a new compression and execution framework called QMoE.
  arXiv  Detail & Related papers  (2023-10-25T17:24:53Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.