Mamba-3: Improved Sequence Modeling using State Space Principles (From a Hardware Perspective)

Abstract

With the growing importance of agentic workflows and test-time scaling, inference efficiency has become a central consideration in model design alongside model quality. In this work, we present Mamba-3, a new generation of state space models designed from an inference-first perspective to increase expressivity of prior Mamba architecture without compromising latency. Mamba-3 introduces three key changes. First, we develop an exponential-trapezoidal discretization scheme that yields a more expressive recurrence. Second, we incorporate complex-valued state transitions to enable state-tracking ability. Third, we propose a Multi-Input, Multi-Output (MIMO) formulation that increases arithmetic intensity during decoding without increasing latency. To optimize for speed and efficiency, we implement Mamba-3 through optimized GPU kernels in Triton, achieving the fastest prefill+decode latency at the 1.5B scale across varying sequence lengths compared to Mamba-2, Gated DeltaNet, and Llama-3.2 on vLLM. In addition, we are developing next-generation training kernels in CuTe DSL to further exploit low-level control over memory layout and warp specialization on modern NVIDIA GPUs. Taken together, Mamba-3 pushes the Pareto frontier by demonstrating stronger modeling capabilities without increasing state size.