For years, deep sequence modeling operated under a structural divide. On one side stood the Transformer architecture, anchored by softmax attention. Softmax attention scales quadratically in compute and memory with sequence length (O(T^2)), but its operations map cleanly to dense General Matrix Multiplications (GEMMs), maximizing utilization of GPU Tensor Cores. On the other side stood Structured State Space Models (SSMs), such as S4 and Mamba-1. Structured SSMs scale linearly in time (O(T · N)) and maintain a constant state footprint (O(N)) during autoregressive decoding, but their training relies on associative parallel scans that are memory-bandwidth-bound on GPU vector ALUs.
In 2024, researchers Tri Dao and Albert Gu bridged this divide with the framework of Structured State Space Duality (SSD). Published alongside the Mamba-2 architecture, SSD demonstrates that linear attention mechanisms and structured state-space models are not competing paradigms. Mathematically, both are dual representations of the same class of structured transformations: 1-semiseparable matrices.
By formulating state space models as structured matrix multiplications, SSD allows SSM computations to run directly on GPU Tensor Cores via block decomposition, achieving 2x to 8x faster training throughput while expanding state capacity by up to 16x.
The Mathematical Bridge: 1-Semiseparable Matrices
A discrete-time state-space model maps a 1D input sequence x of length T to an output sequence y through an intermediate hidden state h_t of dimension N:
h_t = A_t h_{t-1} + B_t x_t y_t = C_t h_t
In the selective SSM formulation introduced in Mamba-1, A_t, B_t, and C_t are dynamic functions of the input x_t. While this input-dependent selection enables strong reasoning and associative recall, computing the recurrence across long sequences requires parallel associative scans. On modern accelerators like Nvidia A100 and H100 GPUs, elementwise associative scans run on streaming multiprocessor (SM) vector units, achieving roughly 20 to 30 TFLOPS, compared to the 300 to 1,000+ TFLOPS available on matrix Tensor Cores.
+-------------------------------------------------------------------------+
| STRUCTURED STATE SPACE DUALITY |
| |
| Recurrent View (SSM) Quadratic View (Attention) |
| h_t = A_t h_{t-1} + B_t x_t Y = ( (C B^T) ⊙ L ) X |
| y_t = C_t h_t |
| Linear time: O(T * N) Matrix multiply: GEMMs |
| Memory-bound on Vector ALUs Compute-bound on Tensor Cores|
| |
| UNIFIED THROUGH |
| 1-Semiseparable Transformation Matrix M |
+-------------------------------------------------------------------------+The SSD framework restricts the transition matrix A_t to a scalar-times-identity structure: A_t = a_t I, where a_t is a scalar decay factor per time step. Under this parameterization, the sequence-level transformation Y = MX can be written as an explicit lower-triangular matrix M, where each entry for row i and column j (i >= j) is defined as:
M_{ij} = C_i^T * (∏_{k=j+1}^i a_k) * B_j
This matrix M is a structured 1-semiseparable matrix. Crucially, it factorizes into the Hadamard (elementwise) product of a low-rank matrix and a semiseparable decay matrix:
M = (C B^T) ⊙ L
Where L_{ij} = ∏_{k=j+1}^i a_k for i >= j, and 0 for i < j.
The Duality Revealed
This factorization reveals the mathematical equivalence between SSMs and Attention:
- Standard Linear Attention: When a_k = 1 for all steps, the decay matrix L becomes the standard causal all-ones lower-triangular matrix. The transformation simplifies to Y = (C B^T ⊙ M_causal) X, which is the exact formulation of unnormalized Causal Linear Attention where C acts as Queries (Q), B acts as Keys (K), and X acts as Values (V).
- Decayed Linear Attention: When a_k < 1, the transformation applies exponential decay over time, weighting recent tokens higher than distant tokens.
- Structured SSM: The exact same computation can be executed sequentially in O(T) time via state accumulation (h_t = a_t h_{t-1} + B_t x_t).
Thus, structured state-space models are linear attention mechanisms equipped with structured exponential decay masks, and causal linear attention is a state-space model with identity state transitions.
The SSD Block Decomposition Algorithm
While establishing theoretical duality clarifies model mechanics, the practical power of SSD lies in its hybrid computation algorithm. By partitioning the semiseparable matrix M into blocks, the SSD algorithm combines the compute density of matrix multiplication with the linear memory scaling of recurrent scans.
The sequence of length T is divided into contiguous chunks of size Q (typically Q=64 or Q=128), yielding T/Q blocks:
1. Diagonal Blocks (Intra-Chunk Computation via GEMMs)
Each diagonal block represents the attention computation within an isolated chunk of Q tokens:
Y_diag^{(c)} = ((C_c B_c^T) ⊙ L_local) X_c
Because Q is small (64), this operation is computed using dense matrix multiplications on Tensor Cores. The intermediate matrices are small enough to fit entirely inside fast on-chip SRAM (shared memory), avoiding round-trip reads and writes to high-bandwidth memory (HBM).
2. Off-Diagonal Blocks (Inter-Chunk State Recurrence)
Because the matrix M is 1-semiseparable, every off-diagonal block factorizes into low-rank components. Interactions between past chunks and current chunks do not require materializing large pairwise attention grids. Instead, computation proceeds in three steps:
- Chunk-Level State Summarization: Each chunk c computes its cumulative output state vector h_c from local inputs: h_c = ∑_{j in chunk c} (∏_{k=j+1}^Q a_k) B_j X_j.
- Inter-Chunk Recurrent Scan: A lightweight associative scan passes states between chunk boundaries across the T/Q blocks. Because the sequence length is compressed by a factor of Q (for example, 4,096 tokens become 64 chunk states), the scan overhead is reduced by nearly two orders of magnitude.
- Chunk State Projection: The incoming hidden state h_{c-1} is projected across the tokens of chunk c using Tensor Core matrix multiplications: Y_off-diag^{(c)} = C_c * h_{c-1}.
+------------------------------------------------------------------------+
| SSD CHUNK EXECUTION FLOW |
| |
| Input X, B, C ----> [ Intra-Chunk GEMM ] (SRAM, Tensor Cores) |
| | | |
| v v |
| [ Chunk State ] ---> [ Inter-Chunk Scan ] ---> [ Output Projection ] |
| Summary (Length T / Q) | |
| v |
| Combined Output Y |
+------------------------------------------------------------------------+By decomposing the computation, over 85% to 90% of the total floating-point operations are performed via dense matrix multiplications on Tensor Cores, while preserving O(T) linear scaling for end-to-end processing.
Architectural Scaling and Mamba-2
The algorithmic efficiency of SSD directly informed the architecture of Mamba-2. In Mamba-1, hardware bottlenecks in the parallel scan kernel forced the designers to limit the hidden state dimension to N=16. Increasing N in Mamba-1 caused steep memory bandwidth penalties.
Under SSD, state transformations are expressed as matrix multiplications where expanding N increases compute intensity rather than memory traffic. Mamba-2 scales N from 16 to 64, 128, and 256, expanding the memory retention capacity of each layer by 4x to 16x without throughput degradation.
Key architectural shifts between generations:
- Core Formulation: Mamba-1 used selective recurrent scans; Mamba-2 unifies recurrence and attention via 1-semiseparable matrix transformations.
- Primary Execution Unit: Mamba-1 was bound to SM vector ALUs; Mamba-2 runs primary workloads on GPU matrix Tensor Cores.
- State Dimension (N): Mamba-1 used N = 16; Mamba-2 scales to N = 64, 128, and 256.
- Hardware Utilization: Mamba-1 achieved 20% to 30% Model FLOPs Utilization (MFU); Mamba-2 reaches 50% to 65% MFU, matching FlashAttention-2.
- Head Organization: Mamba-1 used single-state per channel; Mamba-2 introduces multi-head and grouped-query SSM patterns.
- Parallelization: Mamba-1 required custom scan sharding; Mamba-2 supports standard Tensor Parallelism (TP) over head dimensions.
Multi-Head SSM Patterns
Mamba-2 standardizes the SSM layout into multi-head tensor shapes: (Batch, Sequence, Heads, Head_Dimension). This design choice aligns SSM layers with modern multi-head attention (MHA) and grouped-query attention (GQA) architectures.
Projections for X, B, and C are bundled into a single input projection matrix at the start of the block, and the output is computed through a standard linear projection. This structural symmetry enables standard tensor parallelism (Megatron-LM style) across multiple GPUs by partitioning attention-like head dimensions directly.
Performance and Serving Economics
In empirical benchmarks across language modeling tasks, Mamba-2 matches or outperforms Transformer baselines while providing substantial latency and memory advantages during inference:
- Training Throughput: On 8k sequence lengths, Mamba-2 achieves training speeds comparable to FlashAttention-2 on Nvidia H100 GPUs, executing between 2x and 8x faster than Mamba-1 implementations.
- Inference Footprint: During autoregressive generation, Mamba-2 maintains a fixed-size recurrent state. Unlike standard Transformers whose key-value (KV) cache grows linearly with context length (O(T)), Mamba-2 generates tokens with O(1) constant memory overhead and constant time per step.
- State Capacity: On associative recall benchmarks (such as Multi-Query Associative Recall), scaling the state dimension N to 128 or 256 allows Mamba-2 to retain fine-grained in-context information across tens of thousands of tokens where smaller SSM states experienced degradation.
Implications for Hybrid Models
While State Space Duality resolves the training speed bottleneck of SSMs, empirical evaluations indicate that full-attention layers retain superior performance for exact, needle-in-a-haystack associative retrieval and complex in-context copying tasks.
As a result, leading open-weight architectures increasingly adopt hybrid designs:
- Jamba (AI21 Labs): Interleaves Transformer attention layers with Mamba layers in a 1:7 or 1:3 ratio, cutting KV cache footprint by up to 80% while retaining global retrieval fidelity.
- Samba and Zamba: Combine SSD/Mamba layers with shared global attention mechanisms to scale effective context windows past 256k tokens on standard hardware.
State Space Duality establishes that attention and recurrence are two ends of a continuous mathematical spectrum governed by semiseparable matrix structure. By shifting SSM execution onto Tensor Cores, SSD provides a foundation for sub-quadratic sequence architectures in frontier AI systems.
Sources
- Dao, T., & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. arXiv:2405.21060
- Tri Dao / Goomba Lab. State Space Duality (Mamba-2) Part III: The Algorithm. tridao.me
- State Spaces Research. Mamba Architecture and SSD Official Implementation. GitHub: state-spaces/mamba
- Katharopoulos, A., et al. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. arXiv:2006.16236
