Paper deep dive
M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling
Mayank Mishra, Shawn Tan, Ion Stoica, Joseph Gonzalez, Tri Dao
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%
Last extracted: 3/22/2026, 5:08:23 AM
Summary
M^2RNN is a novel non-linear Recurrent Neural Network architecture that utilizes matrix-valued hidden states and an outer product state expansion mechanism. It addresses the limitations of traditional vector-valued RNNs (poor language modeling and retrieval) and linear RNNs (limited state tracking expressivity) by providing a scalable, hardware-efficient design that leverages tensor cores and achieves superior performance on long-context and state-tracking benchmarks.
Entities (5)
Relation Signals (3)
M^2RNN → improvesperformanceon → LongBench
confidence 95% · Hybrid Gated DeltaNet models with a single M^2RNN layer also achieve superior long-context generalization... on LongBench.
M^2RNN → outperforms → Gated DeltaNet
confidence 95% · Hybrid M^2RNN outperforms equivalent Gated DeltaNet hybrids by 0.4–0.5 perplexity points
Transformers → limitedto → TC0
confidence 90% · Transformers are highly parallel but are limited to computations in the TC^0 complexity class
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Transformers are highly parallel but are limited to computations in the TC$^0$ complexity class, excluding tasks such as entity tracking and code execution that provably require greater expressive power. Motivated by this limitation, we revisit non-linear Recurrent Neural Networks (RNNs) for language modeling and introduce Matrix-to-Matrix RNN (M$^2$RNN): an architecture with matrix-valued hidden states and expressive non-linear state transitions. We demonstrate that the language modeling performance of non-linear RNNs is limited by their state size. We also demonstrate how the state size expansion mechanism enables efficient use of tensor cores. Empirically, M$^2$RNN achieves perfect state tracking generalization at sequence lengths not seen during training. These benefits also translate to large-scale language modeling. In hybrid settings that interleave recurrent layers with attention, Hybrid M$^2$RNN outperforms equivalent Gated DeltaNet hybrids by $0.4$-$0.5$ perplexity points on a 7B MoE model, while using $3\times$ smaller state sizes for the recurrent layers. Notably, replacing even a single recurrent layer with M$^2$RNN in an existing hybrid architecture yields accuracy gains comparable to Hybrid M$^2$RNN with minimal impact on training throughput. Further, the Hybrid Gated DeltaNet models with a single M$^2$RNN layer also achieve superior long-context generalization, outperforming state-of-the-art hybrid linear attention architectures by up to $8$ points on LongBench. Together, these results establish non-linear RNN layers as a compelling building block for efficient and scalable language models.
Tags
Links
- Source: https://arxiv.org/abs/2603.14360v1
- Canonical: https://arxiv.org/abs/2603.14360v1
Trouble viewing inline? Open PDF directly →
Full Text
103,976 characters extracted from source content.
Expand or collapse full text
M 2 RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling Mayank Mishra 12 , Shawn Tan 3 , Ion Stoica 1 , Joseph Gonzalez 1∗ , Tri Dao 45∗ 1 UC Berkeley, 2 Learning Machine, 3 MIT-IBM Watson Lab, 4 Princeton University, 5 Together AI ∗ Equal Mentoring Transformers are highly parallel but are limited to computations in the TC 0 complexity class, excluding tasks such as entity tracking and code execution that provably require greater expressive power. Motivated by this limitation, we revisit non-linear Recurrent Neural Networks (RNNs) for language modeling and introduceMatrix-to-Matrix RNN(M 2 RNN): an architecture with matrix-valued hidden states and expressive non-linear state transitions. We demonstrate that the language modeling performance of non-linear RNNs is limited by their state size. We also demon- strate how the state size expansion mechanism enables efficient use of tensor cores. Empirically, M 2 RNN achieves perfect state tracking generalization at sequence lengths not seen during training. These benefits also translate to large-scale language modeling. In hybrid settings that interleave recurrent layers with attention, Hybrid M 2 RNN outperforms equivalentGated DeltaNethybrids by 0.4–0.5 perplexity points on a 7B MoE model, while using 3×smaller state sizes for the recurrent layers. Notably, replacing even a single recurrent layer with M 2 RNN in an existing hybrid architecture yields accuracy gains comparable to Hybrid M 2 RNN with minimal impact on training throughput. Further, the HybridGated DeltaNetmodels with a single M 2 RNN layer also achieve superior long-context generalization, outperforming state-of-the-art hybrid linear attention architectures by up to 8 points on LongBench. Together, these results establish non-linear RNN layers as a compelling building block for efficient and scalable language models. Correspondence: Mayank Mishra @ mayank@learning-machine.ai § Training Code: https://github.com/open-lm-engine/lm-engine § Kernels: https://github.com/open-lm-engine/accelerated-model-architectures Models: https://huggingface.co/collections/open-lm-engine/m2rnn 1 Introduction Foundation models trained on massive datasets and adapted to downstream tasks have become the dominant paradigm in modern machine learning (Brown et al., 2020; Chowdhery et al., 2023; Grattafiori et al., 2024). At their core, these models rely on sequence-based architectures capable of processing diverse modalities including text, code, images, video, and speech (Brown et al., 2020; Chowdhery et al., 2023; Grattafiori et al., 2024; Mishra et al., 2024; Dosovitskiy, 2020; Lugosch et al., 2019). The prevailing architecture is the decoder-only Transformer, composed of alternating multi-layer perceptron (MLP) and causal attention layers (Vaswani et al., 2017). Much of attention’s success stems from its amenability to parallelization across the sequence length, the availability of hardware-efficient implementations (Dao et al., 2022; Dao, 2023; Shah et al., 2024; Zadouri et al., 2026), and its effective in-context retrieval capabilities. Attention-based Transformers have also demonstrated remarkable scalability, powering models ranging from hundreds of billions (Brown et al., 2020; Chowdhery et al., 2023; Workshop et al., 2022; Smith et al., 2022; Liu et al., 2024; Yang et al., 2025) to trillions of parameters (Team et al., 2025; Meta AI, 2025). 1 arXiv:2603.14360v1 [cs.LG] 15 Mar 2026 However, attention’s quadratic time complexity during training and linearly growing memory requirements during inference have motivated the development of more efficient alternatives (Tay et al., 2022). State Space Models (SSMs) and linear attention (Gu et al., 2021a,b; Gu and Dao, 2024; Dao and Gu, 2024; Yang et al., 2023, 2024b,a; Peng et al., 2023) have emerged as promising replacements. Katharopoulos et al. (2020) propose computing linear attention via a dot product of kernel features, Peng et al. (2021) approximate attention using random feature methods, Schlag et al. (2021) introduce the delta update rule, showing significant improvements in associative recall, and Gu and Dao (2024); Dao and Gu (2024) propose Mamba with an efficient training algorithm that scales linearly with sequence length. Linear attention admits both a recurrent and a parallel formulation. The constant-state size for linear attention enables memory-efficient and fast autoregressive inference (using the recurrent form) while the standard parallel form retains quadratic complexity during training. To bridge this gap, a chunkwise parallel formulation can be employed (Hua et al., 2022; Yang et al., 2023), preserving linear time complexity while batching computations within chunks to leverage matrix multiplication units (tensor cores on NVIDIA GPUs (NVIDIA Corporation)). SSMs like Mamba-1 (Gu and Dao, 2024) can alternatively be accelerated via the parallel scan algorithm (Hillis and Steele, 1986; Blelloch, 1990). Although linear RNNs have attracted considerable attention due to their efficient inference and training, they present several notable limitations: 1.Limited state tracking: Linear RNNs are provably less expressive than non-linear RNNs, particularly for hard state-tracking tasks such as code evaluation, entity tracking, and permutation composition (S 5 group permutation task). Consequently, they are no more expressive than Transformers (Merrill et al., 2024). 2.Poor in-context retrieval performance: Linear RNNs exhibit weak performance in in-context retrieval tasks (Yang et al., 2023, 2024b,a). Because their recurrent state is updated via a fixed-rank outer product, the state can be overwritten when the number of key-value associations exceeds its capacity, leading to degraded in-context retrieval performance. This manifests as poor performance on needle-in-a-haystack benchmarks and long-context recall tasks that require retrieving specific information from distant context. A common strategy is to interleave quadratic attention layers with sub-quadratic recurrent layers in order to achieve a balance between inference efficiency and downstream task performance in state-of-the-art production models. This alleviates the poor language modeling and long-context performance of linear RNNs while retaining their long-context training and inference efficiency advantages (Blakeman et al., 2025; Lieber et al., 2024; Ren et al., 2024). While non-linear RNNs have been shown to excel at hard state tracking problems (Merrill et al., 2024), they face several practical challenges: 1.Poor language modeling performance: Current non-linear RNNs significantly underperform linear RNNs in language modeling. We show that large state sizes are critical for strong language modeling performance, and that the inferior results of non-linear RNNs are primarily due to their smaller state sizes rather than their non-linearity. 2. Poor long-context retrieval performance: We further show that the smaller state sizes significantly hurt the retrieval and long-context performance. The vector-valued non-linear RNNs in our experiments consistently underperformMamba-2andGated DeltaNetby≈20 points in real-world retrieval tasks (Tables 3, 4). 3. Training inefficiency: Unlike linear RNNs, non-linear RNNs cannot be parallelized across the sequence length, making training significantly more expensive. This is compounded by poor hardware utilization: naïve recurrent GEMM execution prevents reuse of on-chip tiles across time steps, forcing costly HBM I/O and synchronization. Even block-wise variants (Pöppel et al., 2025) remain constrained by micro-batch tiling, leading to padding overhead and wasted FLOPs. We proposeMatrix-to-Matrix RNN(M 2 RNN), a non-linear RNN architecture that when combined with Hybrid linear RNNs results in a significant performance improvement in both language modeling and long- 2 context performance. Specifically, M 2 RNN layers address the following key limitations of linear RNNs and non-linear RNNs: 1.Improved state tracking:Matrix-to-Matrix RNN(M 2 RNN) layer achieves perfect state tracking, signifi- cantly improving state tracking performance over linear RNNs similar to non-linear RNNs and GRUs (Cho et al., 2014a). We further demonstrate that M 2 RNN layers can express all computations that can be performed by a non-linear RNN. 2.Improved language modeling, in-context retrieval and long-context performance: We demonstrate the importance of large state sizes for recurrent models and that it is a key driver behind the success of linear RNNs such asMamba-2(Dao and Gu, 2024) andGated DeltaNet(Yang et al., 2024a). We show that the outer product state expansion mechanism is directly applicable to non-linear RNNs. Similar to LSTMs (Hochreiter and Schmidhuber, 1997) and GRUs (Cho et al., 2014a), we use a forget gate to prevent gradient degradation across time steps (Chung et al., 2014), but unlike LSTMs and GRUs, our forget gate is independent of the recurrent state, enabling parallel computation. Hybrid M 2 RNN models outperform HybridMamba-2and HybridGated DeltaNeton language modeling benchmarks. The outer product state expansion mechanism also improves in-context retrieval and long-context performance which are sensitive to the state size. Further, we demonstrate that combining M 2 RNN layers with HybridMamba-2and HybridGated DeltaNetenhances long-context performance by up to 8 points at both 410M dense and 7B MoE scales (Tables 5, 6). 3.Improved hardware utilization: The outer product based state expansion mechanism enables efficient tensor core utilization without the wasted FLOPs of FlashRNN (Pöppel et al., 2025) due to padding along the batch dimension (Section 3.4). Since M 2 RNN layers are computationally expensive and non- parallelizable across time steps, we further explore using M 2 RNN layers sparingly. We demonstrate with experiments that using M 2 RNN layers sparingly (even just one layer) results in substantial improvements in all benchmarks. Beyond efficient single-device execution, training models at scale requires careful attention to systems design. We describe the forward and backward kernels for M 2 RNN layers implemented in Triton (Tillet et al., 2019), and present two strategies for applying tensor parallelism (TP) to M 2 RNN layers: a topology-aware approach that requires no additional communication, and a topology-independent approach that preserves parameter count across TP configurations at the cost of extra synchronization (Section 4.3). 2 Background We use bold letters (A) to represent matrices and small letters (a) to represent column vectors in the following sections. 2.1 State Tracking State tracking refers to determining the state of a system by observing a sequence of updates applied to it. The complexity of state tracking tasks is characterized by circuit complexity classes: TC 0 : TC 0 denotes the class of languages decidable by constant-depth threshold boolean circuits of polynomial size. Simple state tracking problems such as parity (equivalent to the permutation groupS 2 ) fall within this class. Merrill et al. (2022) show that Transformers (Vaswani et al., 2017) lie in TC 0 , making them capable of solving only these simpler tasks. NC 1 : NC 1 denotes the class of languages decidable by logarithmic-depth boolean circuits of polynomial size 1 . Many practical state tracking tasks such as tracking chess moves or evaluating programs reduce to problems at least as hard as theS 5 word problem, whereS k denotes the symmetric group onkelements (i.e., all permutations ofkobjects). SinceS 5 is NC 1 -complete, these tasks lie strictly outside the expressivity class of both Transformers (Merrill et al., 2022) and linear SSMs with input-independent or diagonal transition matrices (Merrill et al., 2024), motivating architectures with greater computational expressivity. 1 Depth is measured in the input size. 3 2.2 Linear RNNs Both linear attention mechanisms and SSMs admit recurrent and parallel formulations, leading the literature to refer to them collectively as recurrent neural networks (RNNs). However, Merrill et al. (2024) demonstrate a significant expressivity gap between these linear RNNs and non-linear RNNs. Non-linear RNNs, recurrent networks that apply a non-linear activation at each step of the recurrence, have traditionally been employed for sequence modeling (Elman, 1990), particularly for tasks requiring state tracking across time steps. While the inclusion of a non-linearity yields a much richer class of representable functions, it precludes the use of algorithms like parallel scan (Hillis and Steele, 1986; Blelloch, 1990) for efficient computation (Merrill et al., 2024). Nevertheless, the state-tracking capabilities enabled by non-linear RNNs can make their incurred cost a reasonable tradeoff. Many linear attention and SSM models take the following general form where Equation 1 is used to compute the recurrent state at the next time step and Equation 2 is used to compute the final output: H t = H t−1 A(x t ) + B(x t )(1) y t = H ⊤ t c(x t ) + d(x t )(2) whereH t ∈R K×V denotes the recurrent state at time stept,A(x t )∈R V×V is the transition matrix, and B(x t )∈R K×V is a function of the inputx t ∈R d .c(x t )∈R K×1 is the vector used to compute the output and d(x t )∈R V×1 is the residual stream.KandVrepresent the key and value head dimensions respectively. Note that the recurrent stateH t is a matrix rather than a vector. Katharopoulos et al. (2020) replace the similarity score in attention (exp(q ⊤ i k j )) (Vaswani et al., 2017) with a kernelκ(q i ,k j ) =⟨φ(q i ),φ(k j )⟩to compute the similarity scores between queries (q i ) and keys (k j ) whereφrepresents the feature map for the kernel function κ. This leads to the following simplified recurrence: H t = H t−1 + φ(k t )v t ⊤ (3) z t = z t−1 + φ(k t )(4) y t = H ⊤ t φ(q t ) z ⊤ t φ(q t ) (5) Prior works found that removing the normalization term in the denominator (z ⊤ t φ(q t )) in Equation 5 improves training stability (Sun et al., 2023) for linear attention and hence we omit this in going forward. The recurrence then becomes: H t = H t−1 + φ(k t )v t ⊤ (6) y t = H ⊤ t φ(q t )(7) A comprehensive summary of recurrent forms for various linear attention models can be found in Yang et al. (2024b) (Table 2). 2.3 Limitations of Linear RNNs 2.3.1 Limited State Tracking Merrill et al. (2024) demonstrate that hard state tracking tasks like the permutation groupS 5 cannot be solved when the transition matrix for a linear RNN is input-independent (A(x t ) =A) or diagonal (A(x t ) =diag(a(x t ))). The authors further suggest that input-dependent, non-diagonal transition matrices can overcome this limitation and can simulate deterministic finite-state automata (DFAs). Grazzi et al. (2024) extend this result, demonstrating that SSMs whose transition matricesA(x t ) have only positive eigenvalues (typical in Mamba andGated DeltaNet) cannot even solve simple state tracking problems like parity. Grazzi et al. (2024) evaluate Mamba-1 (Gu and Dao, 2024) andGated DeltaNet(Yang et al., 2024a) onS 3 and S 5 permutation groups. They demonstrate that Mamba-1 andGated DeltaNetcan solve the tasks when the eigenvalues ofAare allowed to be negative (specifically between [−1,1]). Siems et al. (2025) further demonstrate that allowing a product ofk−1 HouseHolder matrices further increases the expressivity and allows solving tasks like S k . Terzić et al. (2025) take an alternative approach via structured sparse matrices combined with straight-through gradient estimators, which preserves compatibility with parallel scan (Hillis and Steele, 1986; Blelloch, 1990) while improving expressivity. 4 2.3.2 Poor in-context retrieval performance Linear RNNs and SSMs significantly underperform parameter-matched Transformer (Vaswani et al., 2017) models on in-context retrieval tasks (Yang et al., 2023, 2024b,a). While recent linear RNNs such as DeltaNet (Yang et al., 2024b) andGated DeltaNet(Yang et al., 2024a) have closed this performance gap, there still remains a significant gap between Transformers (Vaswani et al., 2017) and linear RNNs. We experimentally demonstrate this further in Section 5.4.2 (Tables 3, 4). 2.4 Non-Linear RNNs The most basic non-linear RNN (Elman, 1990) takes the form: h t = tanh(Wh t−1 + x t )(8) y t = h t (9) whereh t ∈R d×1 is the hidden state at time stept,x t ∈R d×1 is the input, andW ∈R d×d is the (dense, input-independent) transition matrix. A key challenge for deep non-linear RNNs is vanishing and exploding gradients, which makes learning long-range dependencies difficult (Bengio et al., 1994). The Long Short-Term Memory (LSTM) (Hochreiter and Schmidhuber, 1997) and Gated Recurrent Unit (GRU) (Cho et al., 2014b) address this by introducing multiplicative gating mechanisms. Expressivity. Merrill (2019) analyze non-linear RNNs under finite precision and show that they are equivalent to finite-state automata, capable of recognizing all regular languages. Under infinite precision, saturated non-linear RNNs are strictly more expressive, capable of simulating arbitrary deterministic finite- state automata (DFAs) and thus solving hard state-tracking tasks such as permutation group composition that lie outside TC 0 (Merrill et al., 2024). This expressivity advantage over linear RNNs (which are confined to TC 0 under diagonal or input-independent transitions) motivates their use in architectures targeting tasks that require stateful computation beyond what Transformers can express. 2.5 Limitations of Non-Linear RNNs 2.5.1 Poor language modeling performance Despite their expressivity advantages, non-linear RNNs like LSTMs (Hochreiter and Schmidhuber, 1997) and GRUs (Cho et al., 2014a) significantly underperform both Transformers (Vaswani et al., 2017) and modern linear RNNs on language modeling benchmarks. A natural hypothesis is that non-linearity itself is the bottleneck, but we argue the gap is largely attributable to state size. Vector-valued non-linear RNNs maintain a hidden stateh t ∈R d , which is far smaller than the matrix-valued statesH t ∈R K×V used by linear attention models (Section 5.6) for the same number of model parameters. Naively increasingdto match this capacity is impractical: the transition matrixW∈R d×d grows quadratically in parameters, and for largedit can no longer fit in the shared memory of a single streaming multiprocessor (SM), eliminating the possibility of on-chip reuse. 2.5.2 Poor in-context retrieval performance Beyond language modeling, non-linear RNNs also struggle at in-context retrieval tasks. The fundamental bottleneck is again the state size: a vector-valued hidden stateh t ∈R d must compress the entire context into a fixed-size representation. As sequence length grows, the state becomes increasingly overloaded, making it difficult to retrieve specific key-value associations stored earlier in the context. This stands in contrast to Transformers, which retain all past tokens in an explicit key-value cache and can attend directly to any position. While linear RNNs with matrix-valued states fare better on retrieval due to their larger effective state capacity, they too lag behind Transformers on in-context retrieval and recall-intensive benchmarks (Yang et al., 2023, 2024b,a) (Section 5.4.2). Non-linear RNNs, constrained to vector-valued states, exhibit an even larger retrieval deficit. 5 2.5.3 Training Inefficiency Non-Parallelizability on Sequence Length. The element-wise non-linearity in non-linear RNNs breaks the associativity required by the parallel scan algorithm (Hillis and Steele, 1986; Blelloch, 1990), forcing sequential computation over the sequence length. Danieli et al. (2025) show that non-linear RNNs can be parallelized by casting the recurrence as a system ofL(Ldenotes the sequence length) non-linear equations and solving it via Newton’s method (Ortega and Rheinboldt, 2000). However, storing the Jacobian requires substantial memory and the recurrence becomes computationally expensive, rendering the approach impractical at scale. To make it feasible, Danieli et al. (2025) impose diagonal structure on the transition matrix, which eliminates cross-neuron mixing and substantially reduces expressivity. Moreover, each forward pass requires multiple Newton iterations, and the number of iterations needed for convergence is not known a priori for a given architecture. Consequently, executing the full sequential recurrence may ultimately be more efficient than iterative Newton updates. Expensive Computation due to Back-to-back GEMM for Non-Linear RNNs. Computing the GEMM (General Matrix Multiply)Wh t−1 on a GPU proceeds by tiling the output so that each output tile is assigned to a different streaming multiprocessor (SM). This works well for a single GEMM, as each SM can independently compute its tile without cross-SM communication. However, the tiling becomes problematic at the next time step when computingWh t : the output tiles from time stept−1 become the input tiles for time stept, but they reside on different SMs than where they are needed. Ideally, the SM that produced a given output tile would reuse it directly as input, but the standard tiling prevents this. As a result, intermediate results must be written to and reloaded from global memory (HBM), introducing costly synchronization and data movement. This issue is similar to what one would encounter when fusing two back-to-back GEMMs for MLP computation. To address this, FlashRNN (Pöppel et al., 2025) modify the RNN recurrence into a block-wise form, analogous to multi-head attention. This design is also adopted by xLSTM (Beck et al., 2024). By keeping the per-block state size small enough to reside in SM shared memory, they avoid HBM-level synchronization and achieve significant speedups. Poor Hardware Utilization due to Wasted FLOPs. FlashRNN (Pöppel et al., 2025) propose to parallelize a vector-valued non-linear recurrence over the batch dimension (B) and the number of heads (N). However, to enable the use of tensor cores on NVIDIA GPUs, Pöppel et al. (2025) propose to pad the batch dimensionBtoB padded (16 for using the WMMA instruction) (NVIDIA Corporation). This leads to a GEMM computation with shape M, N, K = B padded ,d,d at each time step. In practice, per-GPU batch sizes of 16 or more are rarely feasible when training large models due to GPU memory limitations (80GB on NVIDIA H100s), leading to substantial wasted FLOPs due to padding. For example, at a batch size ofB= 4, padding wastes 75% of the recurrence FLOPs. In Section 3.4, we describe how the M 2 RNN recurrence completely avoids the FLOPs wasted due to padding. 3 Matrix-to-Matrix RNN In this section, we describeMatrix-to-Matrix RNN(M 2 RNN), our proposed non-linear RNN layer which addresses the poor language modeling, poor retrieval performance and poor hardware utilization of vector- valued non-linear RNNs while also addressing the limited state tracking problem of linear RNNs. 3.1 M 2 RNN Layer To overcome the state size limitations of traditional RNNs, we developMatrix-to-Matrix RNN(M 2 RNN), which uses matrix-valued hidden states. We adopt an outer product state expansion strategy similar to that of linear attention (Katharopoulos et al., 2020) and SSMs (Gu and Dao, 2024; Dao and Gu, 2024), significantly increasing the state size without significantly increasing the number of parameters. We further incorporate a forget gate which is independent of the recurrent state and can be computed in parallel. We found that adding a residual connection around the M 2 RNN recurrence similar to SSMs helps to improve gradient flow 6 and mitigate the compounding effect of thetanhactivation across layers. The forward computation for M 2 RNN recurrence is: Z t = tanh(H t−1 W + k t v ⊤ t )(10) H t = f t H t−1 + (1− f t )Z t (11) y t = H ⊤ t q t + w r ⊙ v t (12) whereH t−1 ∈R K×V denotes the recurrent state at the previous time step. We initialize the initial state H 0 = 0 K×V (matrix of shapeK × Vfilled with 0s).Z t ∈R K×V denotes the state update andH t denotes the state at the current time step after the update.W∈R V×V is the transition matrix for the recurrence (independent of the input) while the query, key, value and forget gate (q t ,k t ∈R K×1 ,v t ∈R V×1 ,f t ∈[0,1]) are functions of the inputx t ∈R d×1 .y t ∈R V×1 represents the final output of the recurrence after the query readout and residual addition and w r represents the weight for the residual. 3.1.1 Forget Gate Following LSTMs (Hochreiter and Schmidhuber, 1997) and GRUs (Cho et al., 2014a), we augment M 2 RNN with a forget gate to mitigate the vanishing gradient problem. We use a scalar-valued (per-head) forget gate f t ∈[0,1]. While we experimented with a vector-valued forget gate of sizeKper head, it yielded only marginal improvements while substantially increasing the parameter count. The gated recurrent update is H t =f t H t−1 + (1−f t )Z t , wheref t = 1 fully preserves the current state andf t = 0 resets it entirely. Although the importance of forget gates is well established in the literature (Hochreiter and Schmidhuber, 1997; Cho et al., 2014a; Chung et al., 2014), we provide confirmatory ablations in Section 5.6. Note that the forget gate (f t ) is only a function of the inputx t and is independent of the previous stateH t−1 unlike LSTMs (Hochreiter and Schmidhuber, 1997) and GRUs (Cho et al., 2014a). The forget gate (f t,n ) 2 (forget gate value forn th head at time step t) for M 2 RNN layers is computed as: f t,n = ψ(x t ) = 1 (1 + e x t +β n ) α n (13) Figure 1 Forget gate (f t ) behavior as a function of the input (x t ) for different values ofα n andβ n .ndenotes then th head. whereα n andβ n are learnable parameters that are independent of the input. They are initialized to different values for different heads. A smaller value ofα n for then th head results in a slower decay rate. The parameter β n controls the location at which the forget gate saturates (see Figure 1). A similar gate is used in Mamba-1 (Gu and Dao, 2024),Mamba-2(Dao and Gu, 2024) andGated DeltaNet(Yang et al., 2024a). Each head is initialized to a different value ofα n andβ n which leads to different forgetting characteristics. We initialize α n ∼ Uniform(α min ,α max ) andβ n ∼ LogUniform(β min ,β max ). Note that whenα n = 1, the forget gate simplifies to ψ(x t ) = σ(−x t − β n ) which corresponds to a sigmoid with a reflected input and a bias shift. 2 We drop the notation for head f t,n in the rest of the paper and use f t for simplicity. 7 3.1.2 M 2 RNN Block The queries, keys, and values (q t ,k t ,v t ) are computed via linear projections followed by a depthwise causal short-convolution (kernel size 4) and SiLU activation (Hendrycks, 2016; Elfwing et al., 2018). The output gate g t uses SiLU without a convolution, and the forget gatef t uses the parameterized function from Section 3.1.1 to ensuref t ∈[0,1]. Crucially, all pre-recurrence projections depend only on the inputx t and can therefore be computed in parallel across the sequence. The full M 2 RNN layer (for a single head) is: q t = SiLU(conv1d(W q x t + b q ))∈R K×1 (14) k t = SiLU(conv1d(W k x t + b k ))∈R K×1 (15) v t = SiLU(conv1d(W v x t + b v ))∈R V×1 (16) f t = ψ(W f x t )∈ [0, 1](17) g t = SiLU(W g x t )∈R V×1 (18) Z t = tanh(H t−1 W + k t v ⊤ t )∈R K×V (19) H t = f t H t−1 + (1− f t )Z t ∈R K×V (20) y t = H ⊤ t q t + w r ⊙ v t ∈R V×1 (21) y gt = RMSNorm(y t ⊙ g t )∈R V×1 (22) o t = W o y gt ∈R d×1 (23) (24) where x t ∈R d . Following Mamba-1 (Gu and Dao, 2024),Mamba-2(Dao and Gu, 2024), andGated DeltaNet(Yang et al., 2024a), we adopt a hybrid architecture combining the H3 block (Fu et al., 2022) with a gated MLP, a design that has proven effective in production LLMs (Figure 2). Transition Matrix Initialization. We experimented with three initialization schemes for the transition matrixW: normal initialization (W∼N(0,σ 2 )), orthogonal initialization (W ⊤ =I), and identity initialization (W=I) (Le et al., 2015). Consistent with findings for vector-valued non-linear RNNs, orthogonal initialization outperforms normal initialization. Identity initialization performs on par with orthogonal initialization, so we adopt it for all models. Gradient Clipping. Training non-linear RNNs with backpropagation through time (BPTT) is susceptible to exploding gradients, as the gradient of the loss with respect to early hidden states involves repeated multiplication by the transition matrixWand the Jacobian of thetanhnonlinearity. Although the forget gate f t attenuates gradients flowing through the state, the transition matrixWcan still cause gradient magnitudes to grow unboundedly. To stabilize training, we apply per-step gradient clipping to the gradient of the recurrent state H t during BPTT. 3.2 Improved State Tracking Theorem 1. The M 2 RNN recurrence can represent all tasks representable by non-linear vector-valued RNNs and hence can represent regular languages. We give the proof for Theorem 1 in Appendix B. Siems et al. (2025) show that Gated DeltaProduct [−1,1] 3 parameterized as a product ofk−1 Householder matrices can, in principle, represent theS k task. In particular, this implies that theS 3 task should be representable using Gated DeltaProduct [−1,1] with product of two Householder matrices. To evaluate this claim empirically, we compare M 2 RNN, GRU (Cho et al., 2014a), Gated DeltaNet [−1, 1] 4 (Yang et al., 3 Gated DeltaProduct [−1, 1] denotes Gated DeltaProduct with the ability to have negative eigenvalues. 4 Gated DeltaNet [−1, 1] denotes Gated DeltaNet with the ability to have negative eigenvalues. 8 M 2 RNN Layer Norm Norm MLP/MoE + LinearLinearLinear Causal Convolution M 2 RNN recurrence x RMSNorm Linear v t k t g t q t ψ f t + Figure 2 Visualization of theMatrix-to-Matrix RNNlayer. This block replaces attention and is combined with MLP and RMSNorm modules as in the Transformer. 2024a), and Gated DeltaProduct [−1,1] (Siems et al., 2025) parameterized as a product of two Householder matrices on theS 3 task. All models are trained on sequences of length 128 and evaluated on context lengths up to 512. Figure 3 Accuracy on the permutation groupS 3 state-tracking task for M 2 RNN,Gated DeltaNet[−1,1], Gated DeltaProduct [−1,1] and GRU. The vertical line at 128 denotes the training context length for all models. For the Gated DeltaProduct, we use product of 2 HouseHolder matrices since it can solve theS 3 task in theory (Grazzi et al., 2024). However, as the evaluation context length is increased beyond the training context length, we observe accuracy degradation. As can be seen from Figure 3,Gated DeltaNet[−1,1] fails to learn the task altogether. Gated DeltaProd- uct [−1,1] achieves near-perfect accuracy up to context length 256, however, it fails to generalize to longer 9 sequence lengths not seen during training. These findings suggest that, despite its theoretical expressivity, the model exhibits limited length generalization even on the simple state-trackingS 3 task. However, both GRU (Cho et al., 2014a) and M 2 RNN generalize perfectly to unseen context lengths achieving≥99.5% accuracy up to 512 sequence length. This suggests that theoretical expressivity alone does not guarantee robust state tracking in practice, and that non-linear RNNs like M 2 RNN and GRU exhibit stronger length generalization on this task. Motivated by this, we investigate whether non-linear RNNs exhibit this length generalization advantage on language modeling when equipped with an expanded state size. 3.3 Improved Language Modeling, In-Context Retrieval and Long-Context Performance The outer-product formulation for state size expansion for RNNs has become standard in many SSMs and linear RNNs (Yang et al., 2023; Gu and Dao, 2024; Dao and Gu, 2024; Yang et al., 2024b,a). Qin et al. (2024) demonstrate that larger state sizes improve language modeling performance even when achieved through low-rank parameterizations. We also demonstrate this with empirical evidence in Section 5.6. Additionally,Mamba-2(Dao and Gu, 2024) increases the effective state size through a grouping strategy that shares parameters across heads using a multi-value attention formulation, allowing the model to scale the state dimension without increasing the overall parameter count. Building on this idea, we use an outer product with the multi-value attention formulation proposed by Dao and Gu (2024) to increase the achievable state size under a fixed parameter budget (see Table 8 in Appendix A). The expanded matrix-valued state directly addresses the two core weaknesses of vector-valued non-linear RNNs identified in Section 2.5: poor language modeling and poor retrieval performance. Language modeling. The low language modeling performance of non-linear RNNs such as LSTMs (Hochreiter and Schmidhuber, 1997) and GRUs (Cho et al., 2014a) relative to Transformers (Vaswani et al., 2017) and linear RNNs has traditionally been attributed to the non-linearity itself. We argue instead that the primary bottleneck is state capacity. By expanding the hidden state from a vectorh t ∈R d to a matrix H t ∈R K×V , M 2 RNN dramatically increases the information that can be stored at each recurrent step without a proportional increase in parameter count. We provide empirical evidence for this in Section 5.6, showing that larger state sizes yield consistent language modeling improvements under a fixed parameter budget. In-context retrieval and long-context performance. The outer product writek t v ⊤ t significantly expands the state size enabling efficient storage of key-value associations and structured in-context retrieval via the readoutH ⊤ t q t . This capacity is absent in vector-valued non-linear RNNs, which cannot store enough associations due to the limited state size. We demonstrate the resulting gains in Tables 3, 4. We further demonstrate the impact of the state size expansion in long-context tasks in Tables 5, 6. 3.4 Improved Hardware Utilization 3.4.1 Improved Tiling The outer product expansion enables M 2 RNN to be efficiently parallelized across both the batch dimension Band the number of headsN. In this formulation, each GPU streaming multiprocessor (SM) executes an independent recurrence over a matrix-valued state of sizeK × V. Each step of the recurrence involves a GEMM with dimensions M, N, K = K,V,V . Importantly, the matrix-valued recurrence in M 2 RNN requires no padding as long asKandVare multiples of 16 allowing utilizing the tensor cores. For example, whenK= 64 andV= 16, the GEMM computation H t−1 Wcan utilize the WGMMA instruction on NVIDIA Hopper GPUs (NVIDIA Corporation), ensuring high hardware efficiency. Crucially, the GEMM shape in M 2 RNN is independent of the batch sizeB. This contrasts with vector-valued non-linear RNNs, where GEMM dimensions depend onB, often necessitating batch-size padding to achieve tensor core compatibility. As a result, the outer product state expansion mechanism achieves improved parallelism without wasting FLOPs, avoiding the padding overhead described in Section 2.5.3. 10 3.4.2 Efficient Integration of M 2 RNN Layers in Hybrid Architectures Due to the expensive nature of M 2 RNN layers, we explore using them more sparingly. Our results show that replacing only a singleMamba-2orGated DeltaNetlayer with M 2 RNN in a hybrid architecture achieves the same accuracy improvements as the full hybrid M 2 RNN model, while keeping training throughput within 6% of the Hybrid Gated DeltaNet at both 4k and 16k context lengths. This provides a scalable strategy for incorporating M 2 RNN layers into large language models. 4 Distributed Training and Systems Optimizations In this section, we discuss the system considerations for M 2 RNN layers, including the forward and backward algorithms (Section 4.1, 4.2) and tensor parallelism (Section 4.3). Algorithm 1 Forward Computation for M 2 RNN layer with multi-value attention formulation. The blue colour indicates computation or tensors residing in on-chip memory. input : H 0 ∈R BNKV , Q, K∈R BTK , V∈R BTNV , W∈R NV V , F∈ [0, 1] BTN output : Y∈R BTNV for n = 0 to N − 1 in parallel do▷ Each head on a separate SM Load W← W[n, :, :] from HBM▷ Transition matrix to SRAM (reused across batch) for b = 0 to B− 1 in parallel do Load H← H 0 [b,n, :, :] from HBM▷ Initial hidden state (K× V ) to SRAM for t = 0 to T − 1 do▷ Sequential recurrence over time Load q,k,v,f ← Q[b,t, :], K[b,t, :], V[b,t,n, :], F[b,t,n] from HBM Compute Z← tanh(HW + kv ⊤ )▷ Nonlinear transition with outer product input Compute H← fH + (1− f)Z▷ Gated update of hidden state Compute y ← H ⊤ q▷ Query readout from hidden state Store Y[b,t,n, :]← y to HBM▷ Write output; no intermediate H cached end for end for end for 4.1 Forward Algorithm We implement the forward computation for the M 2 RNN layer using the Triton (Tillet et al., 2019) DSL. The kernel parallelizes across the batch dimensionBand the number of headsN, assigning each (b,n) pair to a separate GPU SM. We illustrate our algorithm in Algorithm 1 for the case of multi-value formulation (N q =N k = 1 andN v ,N f ,N w =N) which we use for all our models. We avoid caching intermediate activations during the backward pass for use in backward computation since storing theHmatrix is prohibitively expensive in both IO and the memory footprint, sinceH t ∈R BNKV is much larger than the per time step keys and values (k t ∈R N k K ,v t ∈R N v V ). 4.2 Backward Algorithm For the backward computation, we recompute the forward recurrence without the final query readout and cacheH t at every time step to HBM. The backward kernel then readsH t at every time step from HBM to compute the gradients for the inputs. The backward kernel is implemented in the Triton DSL (Tillet et al., 2019), which does not expose direct programmability of shared memory. As a result, the kernel is heavily memory-bandwidth bound. We are developing an optimized implementation in CUTLASS to address this bottleneck. The algorithm is illustrated in Algorithm 2. 4.3 Tensor Parallelism Shoeybi et al. (2019) introduced tensor parallelism (TP), a model parallelism technique that partitions attention and MLP layers across multiple accelerators. This approach has become standard for training large- 11 Algorithm 2 Backward Computation for M 2 RNN layer with multi-value attention formulation. The blue colour indicates computation or tensors residing in on-chip memory. input : H 0 ∈R BNKV , Q∈R BTK , K∈R BTK , V∈R BTNV , W∈R NV V , F∈ [0, 1] BTN , H full ∈R BTNKV output : dQ∈R BTK ,dK∈R BTK ,dV∈R BTNV ,dW∈R NV V ,dF∈ [0, 1] BTN for n = 0 to N − 1 in parallel do▷ Each head on a separate SM Load W← W[n, :, :] from HBM▷ Transition matrix to SRAM for b = 0 to B− 1 in parallel do Load H← H 0 [b,n, :, :] from HBM▷ Initial hidden state to SRAM for t = 0 to T − 1 do▷ Pass 1: forward recomputation, cache all H t Load k,v,f ← K[b,t, :], V[b,t,n, :], F[b,t,n] from HBM Store H full [b,t,n, :, :]← fH + (1− f) tanh(HW + kv ⊤ ) to HBM ▷ Cache hidden state for backward end for dW← 0 V×V ▷ Initialize dW to 0 for t = T − 1 to 0 do▷ Pass 2: reverse sweep to compute gradients if t == 0 then Load H prev ← H 0 [b,n, :, :] from HBM else Load H prev ← H full [b,t− 1,n, :, :] from HBM end if Load q,k,v,f,dy ← Q[b,t, :], K[b,t, :], V[b,t,n, :], F[b,t,n], dY[b,t,n, :] from HBM Compute Z← tanh(H prev W + k t v ⊤ t )▷ Recompute pre-gate activation Compute H← fH prev + (1− f)Z▷ Recompute hidden state at t Compute dq ← H· dy▷ Gradient w.r.t. query Atomic Add dQ[b,t, :]← dq to HBM▷ Accumulate across shared heads Compute r ← q· dy ⊤ ▷ Outer product for temporary variable Compute df ← 1 ⊤ K [r⊙ (H prev − Z)] 1 V ▷ Gradient w.r.t. forget gate (scalar) Store dF[b,t,n]← df to HBM Compute dH← r⊙ f▷ Gradient through the gated residual Compute dZ← r⊙ (1− f)▷ Gradient through the nonlinear branch Compute dX← dZ⊙ (1 K×V − Z⊙ Z)▷ Backprop through tanh Compute dH← dZ⊙ (dX· W ⊤ )▷ Gradient w.r.t. previous hidden state Compute dW← dW + H ⊤ prev dX▷ Accumulate gradient for transition matrix Compute dk ← dX· v▷ Gradient w.r.t. key Atomic Add dK[b,t, :]← dk to HBM▷ Accumulate across shared heads Compute dv ← dX ⊤ k▷ Gradient w.r.t. value Store dV[b,t,n, :]← dv to HBM▷ Write to HBM end for Atomic Add dW[n, :, :]← dW to HBM▷ Accumulate dW across batch end for end for scale models (Chowdhery et al., 2023; Smith et al., 2022; Grattafiori et al., 2024; Brown et al., 2020; Workshop et al., 2022). The core challenge in applying TP for M 2 RNN is that the multi-value formulation shares a single query and key head across all value heads (N q =N k = 1,N v =N), preventing a straightforward partition of heads across GPUs. We present two strategies for implementing TP in M 2 RNN layers: a grouped-value approach that avoids extra communication, and a multi-value approach that preserves the parameter count at the cost of additional communication during training. We represent the TP world size as N TP . 4.3.1 TP Topology-Aware M 2 RNN Layers A straightforward way to implement TP for M 2 RNN layers is to adopt a grouped-value formulation: the number of query and key heads is set equal to the TP world size (N q =N k =N TP ), while the number of 12 LinearLinearLinear Causal Convolution M 2 RNN recurrence x RMSNorm 1 Linear v t k t g t q t ψ f t ReduceScatter (FWD) AllGather (BWD) AllGather (FWD) ReduceScatter (BWD) LinearLinearLinear Causal Convolution M 2 RNN recurrence x RMSNorm 2 Linear v t k t g t q t ψ f t ReduceScatter (FWD) AllGather (BWD) AllGather (FWD) ReduceScatter (BWD) Figure 4 TP topology-aware M 2 RNN layer running with TP on 2 GPUs. Note that RMSNorm 1 and RMSNorm 2 have different weights for the RMSNorm module on both GPUs. value heads remainsN v =N. Each GPU is assigned one query head, one key head, and N v N TP value heads 5 , and applies RMSNorm (Zhang and Sennrich, 2019) with different weights on different GPUs in the TP group (refer to Figure 4). Since this is equivalent to running a multi-value M 2 RNN recurrence on every GPU, this requires no modifications to the kernels. Advantages. Each GPU operates independently on its assigned heads, so no communication is introduced beyond the standard TP communication required for an attention layer. Disadvantages. The query and key projection sizes grow proportionally withN TP , coupling the parameter count to the TP topology used during training 6 . Consequently, a model trained with TP world sizeN TP cannot be straightforwardly deployed at a smaller world size without carrying the inflated query, key projections and the larger activation footprint that comes with them (specifically, the input to the output linear projection). This is analogous to the approach proposed inMamba-2(Dao and Gu, 2024) and the Desync-Residual strategy 7 (Zhang et al., 2025). 13 LinearLinearLinear Causal Convolution M 2 RNN recurrence x RMSNormTP Linear ψ ReduceScatter (FWD) AllGather (BWD) AllGather (FWD) ReduceScatter (BWD) LinearLinearLinear Causal Convolution M 2 RNN recurrence x RMSNormTP Linear v t k t g t q t ψ f t ReduceScatter (FWD) AllGather (BWD) AllGather (FWD) ReduceScatter (BWD) No OP (FWD) AllReduce (BWD) v t k t g t q t f t No OP (FWD) AllReduce (BWD) Figure 5 TP topology-independent M 2 RNN layer. Note that RMSNormTP have weights sharded along the model width dimension for the RMSNorm module on both GPUs requiring a synchronization in both the forward and backward computation. 4.3.2 TP Topology-Independent M 2 RNN Layers Alternatively, TP can be applied directly to the multi-value formulation (N q =N k = 1,N v =N). Each GPU processes N N TP value heads but shares the same query and key projections. This requires replacing RMSNorm (Zhang and Sennrich, 2019) modules with RMSNormTP (refer to Appendix C) where the weight for RMSNorm is sharded and an extra communication is introduced in RMSNormTP for the forward and backward computation. The communication in the forward is required to synchronize the RMSNorm normalization term and the synchronization in the backward is required for computing the activation gradient. Note that this idea is applicable toMamba-2(Dao and Gu, 2024) andGated DeltaNet(Yang et al., 2024a) layers as well. If the weights for the causal convolution and linear projections for queries and keys are initialized to the same values on all GPUs participating in TP, the sharding method illustrated in Figure 5 ensures that the weights remain synchronized during training. Advantages. This TP approach preserves the parameter count regardless of the number of GPUs and makes the M 2 RNN layers independent of the TP topology used during training. 5 We assume N v is divisible by N TP . 6 The RMSNorm weight size also grows with N TP but is significantly smaller in comparison. 7 Refer to https://arxiv.org/abs/2501.06589v2 (v2) since later versions of the paper dropped Desync-Residual. 14 Disadvantages. This approach introduces additional AllReduce communications beyond standard TP. In the forward pass, one extra AllReduce is needed to synchronize the RMSNorm normalization term 8 . The backward pass requires three extra AllReduces: one for the RMSNorm gradient and two to aggregate the query and key activation gradients, since each GPU computes partial gradients for the shared query and key heads from its assigned subset of value heads. Weight gradients for the query and key projections remain naturally synchronized across GPUs throughout training, provided the corresponding weights are initialized identically, so no additional communication is needed for them. One way to reduce backward communication is to make only the RMSNorm module TP topology-dependent (as in Section 4.3.1); since RMSNorm has very few parameters, this trade-off may be acceptable in practice. 5 Experiments We train dense models with 410M parameters (24 layers) and Mixture-of-Experts (MoE) (Shazeer et al., 2017) models with 7B total (1.1B active) parameters (40 layers). All models are trained on 100B tokens sampled from the high quality subset of Nemotron-C-v2 (NVIDIA et al., 2025). 5.1 Models 5.1.1 Homogeneous Architectures We compare M 2 RNN primarily againstMamba-2(Dao and Gu, 2024),Gated DeltaNet(Yang et al., 2024a), RNN and GRU (Cho et al., 2014a). We focus onMamba-2andGated DeltaNetbecause they (and their variants) have demonstrated strong performance in state-of-the-art production models (Team et al., 2025; NVIDIA et al., 2025; Lieber et al., 2024; IBM Granite Team, 2025). 5.1.2 2-way Hybrid Models We additionally run experiments in hybrid settings that combine recurrent layers with attention layers. Following common practice in production LLMs (IBM Granite Team, 2025; Lieber et al., 2024), we use 1 attention layer for every 7 recurrent layers (i.e., 1 out of every 8 layers is attention). We compare Hybrid M 2 RNN against Transformer++ (Vaswani et al., 2017), Hybrid Mamba-2 and Hybrid Gated DeltaNet. 5.1.3 3-way Hybrid Models We also experiment with augmenting HybridGated DeltaNetmodels with M 2 RNN layers, yielding 3-way hybrids denoted “Hybrid GDN + M 2 RNN-n”, wherenis the number of M 2 RNN layers. Forn= 1, we replace the first recurrent layer with M 2 RNN; forn >1, we replace the layer immediately preceding each attention layer, producing “Hybrid GDN + M 2 RNN-3” and “Hybrid GDN + M 2 RNN-5” for the 410M and 7B models respectively. 5.2 Model Training For all the models, we use RMSNorm (Zhang and Sennrich, 2019) with pre-normalization (Xiong et al., 2020). Additionally, Transformer++ uses RoPE positional embeddings (Su et al., 2024) while all linear and hybrid architectures use NoPE (no positional embeddings) (Kazemnejad et al., 2023). For MLP and MoE layers (Shazeer et al., 2017), we use the SwiGLU activation (Shazeer, 2020). We keep the model width, intermediate MLP width, number of layers, and all other hyperparameters consistent across architectures, adjusting only the parameters in the sequence mixing block. All models are trained with the AdamW optimizer (Loshchilov and Hutter, 2019) using a cosine learning rate schedule. We linearly warm up the learning rate over the first 10B tokens to a peak value of 3×10 −4 , followed by decaying to 10% of the peak (3×10 −5 ). We use gradient clipping of 1 and weight decay of 0.1. We don’t apply weight decay to the normalization parameters. All models are trained with a batch size of 1024 samples 8 The RMSNorm normalization term depends only on the batch size (B) and sequence length (T), so this communication is lightweight. 15 at a context length of 4096 tokens. We ensure that all models see identical data in identical order to ensure fair comparison. The 410M dense models are trained on 32 H100s and the 7B MoE models on 64 H100s. We use the lm-engine codebase (Mishra, 2024) written in PyTorch (Paszke et al., 2019) for training all our models. We also use FlashAttention-3 (Shah et al., 2024) and SonicMoE (Guo et al., 2025) with PyTorch compile (Ansel et al., 2024) for training throughput optimization. We also use fast kernels for causal convolutions,Mamba-2,Gated DeltaNet, and M 2 RNN where applicable. The models use mixed precision (Micikevicius et al., 2017) in BF16 (Kalamkar et al., 2019), with gradient AllReduce and accumulation in FP32 for numerical stability. 5.3 Language Modeling We evaluate all trained models on a suite of downstream benchmarks using EleutherAI’s LM Eval Harness (Gao et al., 2024): LAMBADA (Paperno et al., 2016), HellaSwag (Zellers et al., 2019), PIQA (Bisk et al., 2019), ARC (Clark et al., 2018), Winogrande (Sakaguchi et al., 2019), BoolQ (Clark et al., 2019), OpenBookQA (Mihaylov et al., 2018), CoPA (Roemmele et al., 2011), and SciQ (Welbl et al., 2017). We additionally report language modeling perplexity on Wikitext (Merity et al., 2017) and LAMBADA (Paperno et al., 2016). ModelWiki PPL LMB PPL LAMBADA HellaSwag PIQA ARC-E ARC-C WinoGrande BoolQ OBQA COPA SciQ Avg Acc Mamba-222.9338.5632.7847.7070.29 62.0430.5553.2858.69 33.20 69.00 85.50 54.30 Gated DeltaNet22.7034.1734.2147.1671.00 61.2428.0751.0758.84 32.80 70.00 85.50 53.99 RNN33.74317.3813.0044.1269.70 62.5430.6349.7245.75 33.40 63.00 77.10 48.90 GRU25.8063.8326.1048.7071.49 64.9030.9749.6457.68 33.80 71.00 82.90 53.72 M 2 RNN22.9233.6333.9847.7671.60 62.2930.2953.4350.21 34.80 72.00 86.00 54.24 Transformer++23.3242.8733.3645.0268.88 59.7629.0150.4354.56 32.00 72.00 84.00 52.90 Hybrid Mamba-221.5935.8435.3248.1170.89 62.5429.3552.7251.41 35.40 65.00 86.80 53.75 Hybrid Gated DeltaNet21.8929.0235.7847.1970.67 61.9129.7852.1752.57 34.00 70.00 86.30 54.04 Hybrid M 2 RNN21.5337.3135.6748.7470.29 63.3830.2951.4650.09 33.80 71.00 86.80 54.15 Hybrid Mamba-221.5935.8435.3248.1170.89 62.5429.3552.7251.41 35.40 65.00 86.80 53.75 Hybrid Mamba-2 + M 2 RNN-121.4836.9834.4748.2370.46 64.1431.5754.4649.54 33.60 70.00 85.80 54.23 Hybrid Mamba-2 + M 2 RNN-321.3935.7735.7748.4770.57 62.3331.8350.7551.44 33.00 65.00 87.40 53.65 Hybrid Gated DeltaNet21.8929.0235.7847.1970.67 61.9129.7852.1752.57 34.00 70.00 86.30 54.04 Hybrid Gated DeltaNet + M 2 RNN-121.3933.7834.8348.1170.73 61.2429.9551.7048.41 33.20 72.00 87.30 53.75 Hybrid Gated DeltaNet + M 2 RNN-321.2628.9337.1848.8369.64 63.8529.4453.3559.14 33.80 71.00 86.10 55.23 Table 1 Language Modeling on the 410M parameter dense model across commonsense reasoning benchmarks. All models are evaluated in the 0-shot setting. ModelWiki PPL LMB PPL LAMBADA HellaSwag PIQA ARC-E ARC-C WinoGrande BoolQ OBQA COPA SciQ Avg Acc Mamba-213.7311.0350.0967.4376.82 74.3343.5261.6464.01 41.00 80.00 93.40 65.22 Gated DeltaNet13.8911.1048.9266.9477.26 73.8240.5360.1465.41 39.60 78.00 92.20 64.28 RNN17.6543.6830.2564.4976.61 72.6441.0455.6458.35 40.20 79.00 87.60 60.58 GRU14.8015.3343.9467.3676.82 75.6743.5258.9663.79 41.00 82.00 91.30 64.44 M 2 RNN13.8011.4849.1867.5776.71 75.5144.9758.1764.16 41.40 82.00 92.00 65.17 Transformer++14.9417.2643.9062.9975.63 71.9738.9157.5463.12 37.20 75.00 91.40 61.76 Hybrid Mamba-213.1011.0550.1868.6276.55 72.6443.0960.6263.61 43.00 79.00 93.00 65.03 Hybrid Gated DeltaNet13.5111.0650.2067.2377.09 73.0641.6459.1262.39 40.80 82.00 92.90 64.64 Hybrid M 2 RNN13.00 10.5851.7468.4476.55 76.8143.0961.0963.21 40.40 81.00 93.30 65.56 Hybrid Mamba-213.1011.0550.1868.6276.55 72.6443.0960.6263.61 43.00 79.00 93.00 65.03 Hybrid Mamba-2 + M 2 RNN-113.01 10.8450.9268.3577.86 74.6642.2460.5462.57 40.40 83.00 92.80 65.33 Hybrid Mamba-2 + M 2 RNN-513.0211.5949.8468.1977.04 75.1742.5859.0465.26 42.40 83.00 93.90 65.64 Hybrid Gated DeltaNet13.5111.0650.2067.2377.09 73.0641.6459.1262.39 40.80 82.00 92.90 64.64 Hybrid Gated DeltaNet + M 2 RNN-113.0710.3351.5468.6978.18 74.4541.8160.3065.08 42.00 82.00 93.90 65.80 Hybrid Gated DeltaNet + M 2 RNN-512.8510.2951.8068.9976.82 76.7343.6060.6963.30 40.20 81.00 93.50 65.66 Table 2 Language Modeling on the 7B (1B active) MoE model across commonsense reasoning benchmarks. All models are evaluated in the 0-shot setting. 5.3.1 Homogeneous Architectures At the 410M scale,Gated DeltaNetachieves the best Wikitext (Merity et al., 2017) perplexity among all baselines, with M 2 RNN matchingMamba-2(Dao and Gu, 2024) (within 0.01 points). On LAMBADA (Paperno et al., 2016), M 2 RNN outperforms bothMamba-2andGated DeltaNet(Yang et al., 2024a) by 4.93 and 0.54 perplexity points respectively (Table 1). We attribute M 2 RNN’s gap behindGated DeltaNeton Wikitext to its 3×smaller state size: while increasing the state size can close this gap, it also raises the per-step FLOP count of the recurrence, making it a less favorable tradeoff in practice. 16 At the 7B MoE scale,Mamba-2achieves the best Wikitext perplexity, followed by M 2 RNN, which outperforms Gated DeltaNetby 0.09 points (Table 2). On average downstream accuracy, M 2 RNN matchesMamba-2to within 0.06 points at both the 410M and 7B MoE scales while consistently outperformingGated DeltaNet(Ta- bles 1, 2). Vanilla RNN and GRU models significantly underperform M 2 RNN and other linear RNNs across all metrics. 5.3.2 2-way Hybrid Models Combining M 2 RNN with attention (Vaswani et al., 2017) (denoted Hybrid M 2 RNN) yields substantial gains on Wikitext (Merity et al., 2017). Hybrid M 2 RNN outperforms HybridMamba-2(Dao and Gu, 2024) by 0.06 and 0.1 perplexity points at 410M and 7B MoE scale, and HybridGated DeltaNet(Yang et al., 2024a) by 0.4 and 0.5 points at the 410M and 7B MoE scale. Average downstream accuracy also improves over all other hybrids at both scales (Tables 1, 2). Overall, 2-way hybrids improve upon both Transformer++ and homogeneous M 2 RNN models. We attribute these gains to M 2 RNN’s non-linear state transition: in the homogeneous setting, M 2 RNN’s advantage is constrained by its smaller state size relative toMamba-2andGated DeltaNet, however, in the hybrid setting, attention is responsible for in-context retrieval, allowing M 2 RNN’s non-linear recurrence to contribute expressivity that linear RNNs such asMamba-2andGated DeltaNetcannot provide. 5.3.3 3-way Hybrid Models We next explore 3-way hybrids that mix attention,Gated DeltaNet(orMamba-2), and M 2 RNN layers. We find that replacing even a single linear RNN layer with M 2 RNN (denoted HybridGated DeltaNet(or Mamba-2) + M 2 RNN-1) suffices to match the accuracy of the full Hybrid M 2 RNN model while preserving comparable training throughput (Section 5.7). This result indicates that M 2 RNN layers carry information that is substantially different from what Gated DeltaNet layers capture. HybridGated DeltaNet+ M 2 RNN-1 improves upon HybridGated DeltaNet, reducing Wikitext perplexity by 0.5 and 0.44 points at the 410M and 7B (1B active) scales, respectively. We also observe improvements on HybridMamba-2+ M 2 RNN-1 over HybridMamba-2of 0.1 perplexity points on both model scales. Increasing the number of M 2 RNN layers to 3 (5 for 7B MoE) yields additional gains of 0.13 and 0.22 perplexity points. Together, these results solidify the importance of non-linear recurrences in sequence mixing blocks. 5.4 In-Context Retrieval 5.4.1 In-Context Retrieval on RULER We evaluate the in-context retrieval ability of M 2 RNN relative to other linear RNNs on the RULER benchmark (Hsieh et al., 2024), focusing on the Single Needle In A Haystack (S-NIAH-1, 2, 3) (Nelson et al., 2024), Multi-Query (MQ), Multi-Key (MK), and Multi-Value (MV) variants: 1.S-NIAH: A key-value pair is embedded in context and the model is tasked to retrieve the value given the key. S-NIAH-1 uses a synthetic passage with word keys and numeric values, probing long-term retention. S-NIAH-2 and S-NIAH-3 use real-world essays with numeric and UUID values, respectively, testing efficient long-context memory management (Yang et al., 2024a). 2. MQ-NIAH: In the MQ-NIAH benchmark, multiple key–value pairs are embedded within the passage, and the task requires retrieving all needles associated with distinct keys. 3. MK-NIAH: The MK-NIAH benchmark inserts multiple key-value pairs in the passage. The task is to retrieve a single specified key while the other keys serve as distractors. 4.MV-NIAH: Multiple values share the same key and the model must retrieve all values associated with that key. Figures 6 and 7 report accuracy at the 410M scale forGated DeltaNetandMamba-2-based models, respectively; Figures 8 and 9 report the same for the 7B (1B active) MoE model. All models are trained with a 4,096 context length and evaluated up to 16,384. 17 HybridGated DeltaNet+ M 2 RNN: Replacing a fewGated DeltaNetlayers with M 2 RNN layers in the hybrid setting significantly improves retrieval at unseen context lengths on S-NIAH-2, and MQ/MK/MV-NIAH at the 410M parameter scale (Figure 6). At the 7B MoE scale, we observe that the 3-way consistently improves over the 2-way hybrid models except on S-NIAH-1 (Figure 8). Notably, augmenting with M 2 RNN layers also improves MQ/MK/MV-NIAH accuracy at context lengths seen during training (≤4,096) at both model scales. On S-NIAH-1, we find that both Hybrid M 2 RNN and HybridGated DeltaNetachieve 100% accuracy across all context lengths, however, adding M 2 RNN layers leads to accuracy degradation, likely due to the extremely synthetic passages used in the S-NIAH-1 task. HybridMamba-2+ M 2 RNN: We find that a single M 2 RNN layer with HybridMamba-2(Hybrid Mamba-2+ M 2 RNN-1) is insufficient to improve long-context retrieval at either model scale. However, replacing threeMamba-2layers with M 2 RNN layers at 410M and 5 layers at 7B MoE yields consistent improvement (Figures 7, 9). We attribute this to the weaker state transition ofMamba-2requiring more M 2 RNN layers to compensate. ModelSQuAD NQ DROP TriviaQA FDA SWDE Avg Mamba-231.215.019.550.515.632.427.4 Gated DeltaNet34.815.424.951.716.341.830.8 RNN9.65.914.629.70.05.010.8 GRU15.88.817.241.50.77.715.3 M 2 RNN34.013.022.952.310.628.526.9 Transformer++14.015.528.252.446.051.734.6 Hybrid Mamba-237.621.426.451.765.361.744.0 Hybrid Gated DeltaNet37.622.226.053.461.558.043.1 Hybrid M 2 RNN41.322.726.855.574.560.746.9 Hybrid Mamba-237.621.426.451.765.361.744.0 Hybrid Mamba-2 + M 2 RNN-138.623.826.355.761.360.244.3 Hybrid Mamba-2 + M 2 RNN-340.423.526.456.765.261.745.6 Hybrid Gated DeltaNet37.622.226.053.461.558.043.1 Hybrid Gated DeltaNet + M 2 RNN-139.823.525.654.667.262.545.5 Hybrid Gated DeltaNet + M 2 RNN-339.321.427.254.965.156.544.1 Table 3 In-context retrieval performance on real-world data at the 410M parameter scale. ModelSQuAD NQ DROP TriviaQA FDA SWDE Avg Mamba-241.726.530.065.033.054.441.8 Gated DeltaNet40.324.229.964.332.461.842.2 RNN20.012.220.453.10.89.819.4 GRU27.016.724.659.35.218.425.2 M 2 RNN37.921.128.262.527.649.937.9 Transformer++22.124.027.662.267.070.145.5 Hybrid Mamba-248.231.432.767.973.177.155.1 Hybrid Gated DeltaNet46.431.228.363.667.171.951.4 Hybrid M 2 RNN48.431.431.464.678.079.855.6 Hybrid Mamba-248.231.432.767.973.177.155.1 Hybrid Mamba-2 + M 2 RNN-146.630.633.667.178.475.055.2 Hybrid Mamba-2 + M 2 RNN-546.831.528.766.078.579.755.2 Hybrid Gated DeltaNet46.431.228.363.667.171.951.4 Hybrid Gated DeltaNet + M 2 RNN-146.429.632.166.473.177.454.2 Hybrid Gated DeltaNet + M 2 RNN-547.132.031.265.577.277.955.1 Table 4 In-context retrieval performance on real-world data at the 7B (1B active) MoE scale. 18 Figure 6 Zero-shot in-context retrieval performance on the RULER benchmark for M 2 RNN layers and Gated DeltaNet models on 410M parameter scale. The vertical lines at 4096 indicate the training context length. Figure 7 Zero-shot in-context retrieval performance on the RULER benchmark for M 2 RNN layers andMamba-2models on 410M parameter scale. The vertical lines at 4096 indicate the training context length. 5.4.2 In-Context Retrieval on Real-World Data Yang et al. (2024b) show that despite strong synthetic-task performance, DeltaNet still underperforms Mamba-2(Dao and Gu, 2024) on real-world retrieval. To assess whether this limitation persists for M 2 RNN, 19 Figure 8 Zero-shot in-context retrieval performance on the RULER benchmark for M 2 RNN layers and Gated DeltaNet models for the 7B MoE model. The vertical lines at 4096 indicate the training context length. Figure 9 Zero-shot in-context retrieval performance on the RULER benchmark for M 2 RNN layers andMamba-2models for the 7B MoE model. The vertical lines at 4096 indicate the training context length. we evaluate in-context retrieval on real-world data for both model scales (Tables 3, 4). As expected, Transformer++ (Vaswani et al., 2017) substantially outperforms all recurrent models due to its linearly growing key-value cache which helps with retrieval performance. Among purely recurrent architectures, 20 Gated DeltaNet(Yang et al., 2024a) outperformsMamba-2(Dao and Gu, 2024) and M 2 RNN which we attribute to the 3×larger state size ofGated DeltaNet. In the hybrid setting, all models significantly improve upon the Transformer++ baseline, with Hybrid M 2 RNN achieving the largest gains: 12.3 and 10.1 point gains across the 410M and 7B MoE model scales, outperforming both HybridMamba-2and HybridGated DeltaNet. Further, we find that adding M 2 RNN layers to HybridGated DeltaNet(HybridGated DeltaNet+ M 2 RNN) further improves upon HybridGated DeltaNet. We observe similar improvements for HybridMamba-2when combined with M 2 RNN (Tables 3, 4). 5.5 Long-Context Performance ModelGovReport QMSum MultiNews TREC TriviaQA SAMSum LCC RepoBench-P Avg Mamba-26.614.39.926.523.211.513.78.914.3 Gated DeltaNet7.618.212.019.034.121.314.911.817.3 RNN9.315.29.20.512.25.017.914.910.5 GRU7.016.510.81.021.48.512.59.010.8 M 2 RNN6.316.410.822.024.313.210.710.114.2 Transformer++5.65.76.211.09.74.612.09.98.1 Hybrid Mamba-210.04.810.216.09.49.210.512.010.3 Hybrid Gated DeltaNet8.216.816.119.528.715.98.79.215.4 Hybrid M 2 RNN10.014.07.217.535.514.913.410.315.4 Hybrid Mamba-210.04.810.216.09.49.210.512.010.3 Hybrid Mamba-2 + M 2 RNN-19.89.07.522.513.312.613.29.812.2 Hybrid Mamba-2 + M 2 RNN-310.515.813.229.031.712.110.713.317.1 Hybrid Gated DeltaNet8.216.816.119.528.715.98.79.215.4 Hybrid Gated DeltaNet + M 2 RNN-19.514.817.647.032.119.47.38.519.5 Hybrid Gated DeltaNet + M 2 RNN-316.018.114.043.041.524.712.713.923.0 Table 5 Performance comparison on the 410M model across long-context summarization, coding and few-shot learning benchmarks. ModelGovReport QMSum MultiNews TREC TriviaQA SAMSum LCC RepoBench-P Avg Mamba-25.116.013.442.059.722.013.48.722.5 Gated DeltaNet4.820.312.247.061.636.29.77.925.0 RNN9.316.710.54.027.513.017.115.814.2 GRU5.718.010.89.052.424.620.913.919.4 M 2 RNN8.119.111.826.551.330.720.419.323.4 Transformer++6.16.117.016.011.410.013.110.211.2 Hybrid Mamba-29.17.720.362.523.312.15.69.818.8 Hybrid Gated DeltaNet10.415.816.757.535.123.69.716.323.1 Hybrid M 2 RNN10.913.921.152.044.718.519.815.724.6 Hybrid Mamba-29.17.720.362.523.312.15.69.818.8 Hybrid Mamba-2 + M 2 RNN-18.63.921.225.515.011.57.511.613.1 Hybrid Mamba-2 + M 2 RNN-518.610.37.160.545.721.49.78.022.7 Hybrid Gated DeltaNet10.415.816.757.535.123.69.716.323.1 Hybrid Gated DeltaNet + M 2 RNN-117.216.56.069.564.934.120.115.130.4 Hybrid Gated DeltaNet + M 2 RNN-512.216.016.159.073.536.018.018.931.2 Table 6 Performance comparison on the 7B (1B active) MoE model across long-context summarization, coding and few-shot learning benchmarks. We evaluate long-context performance on LongBench (Bai et al., 2024) (Table 5, 6). In the 2-way hybrid setting, we find that Hybrid M 2 RNN matches HybridGated DeltaNetat 410M parameter scale and outperforms HybridGated DeltaNetby 1.5 points on the 7B MoE scale. We find that combining M 2 RNN layer with HybridGated DeltaNetresults in up to 8 points of average accuracy improvement across both model scales across long-context summarization, coding, and few-shot learning tasks. Similar improvements are observable with Hybrid Mamba-2 + M 2 RNN-3 (5 for 7B MoE) models on both model scales. 21 5.6 Ablations: Effect of State Size and Forget Gate We ablate the contributions of state size and the forget gate on the recurrent model performance (Table 7). First, we compare a traditional vector-valued RNN (406M parameters) against M 2 RNN (410M parameters). The matrix-valued M 2 RNN model outperforms the vector-valued RNN by more than 10 perplexity points on WikiText (Merity et al., 2017) and 280 points on LAMBADA (Paperno et al., 2016), underscoring the critical role of state size. ModelWiki PPL Lambada PPL state size RNN-406M33.74317.381,360 GRU-474M25.8063.831,360 M 2 RNN-410M22.9233.6386,016 Table 7 Perplexity comparison for RNN, GRU and M 2 RNN models with different state sizes. To disentangle the effect of gating from that of state size, we additionally train GRU (Cho et al., 2014a) models with the same state dimensions as the vector-valued RNN. Despite having 16% more parameters (474M) due to the additional reset gate, GRUs still underperform M 2 RNN by 3 perplexity points on WikiText and 30 points on LAMBADA. These results confirm that expanding the state size is the primary driver of performance improvements in recurrent architectures, with gating mechanisms potentially providing further gains. 5.7 Training Throughput We measure training throughput for M 2 RNN,Gated DeltaNet(Yang et al., 2024a),Mamba-2(Dao and Gu, 2024), and Transformer++ on 8×NVIDIA HGX H100 GPUs connected via NVLink, using our 7B (1B active) MoE configuration. As shown in Figure 10,Mamba-2consistently achieves the highest throughput across context lengths, while Transformer++ degrades significantly at longer contexts due to its quadratic complexity. Although M 2 RNN has linear time complexity, its higher constant factors make it more expensive than the linear baselines in the homogeneous setting. However, replacing a singleGated DeltaNetlayer with M 2 RNN layer to the HybridGated DeltaNetmodel (HybridGated DeltaNet+ M 2 RNN-1) achieves throughput comparable to HybridGated DeltaNetalone (within 6% at 16k context length) while consistently improving model quality across all benchmarks. We consider this a favorable tradeoff, making M 2 RNN layers practical when used sparingly. Our current implementation is based on the Triton DSL (Tillet et al., 2019) and is relatively unoptimized, leaving significant room for improvement that could further narrow the remaining throughput gap. 6 Conclusion We introduceMatrix-to-Matrix RNN(M 2 RNN), a non-linear RNN architecture with matrix-valued hidden states that addresses key limitations of existing linear recurrent models. By combining expressive non- linear state transitions with a forget gate and an outer product state expansion, M 2 RNN achieves strong state-tracking capabilities. Our experiments demonstrate that M 2 RNN consistently comes close toMamba-2(Dao and Gu, 2024) and Gated DeltaNet(Yang et al., 2024a) at both the 410M dense and 7B (1B active) MoE scales. More importantly, hybrid models combining M 2 RNN with attention layers significantly outperform equivalentMamba-2and Gated DeltaNethybrids in language modeling, in-context retrieval, and long-context generalization. We further show that even a single M 2 RNN layer yields substantial accuracy gains with only 6% throughput degradation, making it a practical enhancement for existing architectures. The matrix-valued recurrence in M 2 RNN eliminates the FLOPs wasted by padding in vector-valued recurrences when using the FlashRNN algorithm (Pöppel et al., 2025), enabling efficient tensor core utilization independent of batch size. Combined with our custom kernels, this makes M 2 RNN a viable component for training state- of-the-art production language models. 22 Figure 10 Training throughput measured in billion tokens per day on a machine with 8x NVIDIA HGX H100 GPUs for our 7B MoE configuration. We also present two complementary tensor parallelism strategies for M 2 RNN. The topology-aware approach adopts a grouped-value formulation that requires no additional communication beyond standard TP, while the topology-independent approach preserves the parameter count regardless of TP world size at the cost of a small number of extra AllReduce operations. Together, these strategies make M 2 RNN readily deployable across diverse multi-GPU configurations. Limitations and Future Work. The non-linear recurrence introduces additional computational overhead compared to linear alternatives. Future work could explore approximations or more efficient kernel implemen- tations to reduce this cost while preserving expressivity. Additionally, evaluating M 2 RNN at larger scales and with longer training contexts would further validate its potential as a foundational building block for efficient language models. Acknowledgement We would like to thank Songlin Yang, Bharat Runwal and Kevin Li for providing valuable feedback and discussion throughout the duration of this project. We gratefully acknowledge the support of the Schmidt Sciences AI2050 fellowship, the Google ML and Systems Junior Faculty Awards, and the Google Research Scholar program. References Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary DeVito, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, C. K. Luk, Bert Maher, Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Shunting Zhang, Michael Suo, Phil Tillet, Xu Zhao, Eikan Wang, Keren Zhou, Richard Zou, Xiaodong Wang, Ajit Mathews, William Wen, Gregory Chanan, Peng Wu, and Soumith Chintala. Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, ASPLOS ’24, page 929–947, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400703850. doi: 10.1145/3620665.3640366. https://doi.org/10.1145/3620665.3640366. Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 23 62nd annual meeting of the association for computational linguistics (volume 1: Long papers), pages 3119–3137, 2024. Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. Advances in Neural Information Processing Systems, 37:107547–107603, 2024. Y. Bengio, P. Simard, and P. Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE Transactions on Neural Networks, 5(2):157–166, 1994. doi: 10.1109/72.279181. Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019. https://arxiv.org/abs/1911.11641. Aaron Blakeman, Aaron Grattafiori, Aarti Basant, Abhibha Gupta, Abhinav Khattar, Adi Renduchintala, Aditya Vavre, Akanksha Shukla, Akhiad Bercovich, Aleksander Ficek, et al. Nvidia nemotron 3: Efficient and open intelligence. arXiv preprint arXiv:2512.20856, 2025. Guy E. Blelloch. Prefix sums and their applications. Technical Report CMU-CS-90-190, School of Computer Science, Carnegie Mellon University, November 1990. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020. Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint arXiv:1409.1259, 2014a. Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv:1406.1078, 2014b. Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1–113, 2023. Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555, 2014. Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions, 2019. https://arxiv.org/abs/1905.10044. Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018.https://arxiv.org/abs/1803.05457. Federico Danieli, Pau Rodriguez, Miguel Sarabia, Xavier Suau, and Luca Zappella. Pararnn: Unlocking parallel training of nonlinear rnns for large language models. arXiv preprint arXiv:2510.21450, 2025. Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023. Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060, 2024. Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344–16359, 2022. Alexey Dosovitskiy. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020. Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 107:3–11, 2018. Jeffrey L Elman. Finding structure in time. Cognitive science, 14(2):179–211, 1990. Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212.14052, 2022. Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria 24 Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024. https://zenodo.org/records/12608602. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, Danny Wyatt, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Govind Thattai, Graeme Nail, Gregoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Karthik Prasad, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Rantala-Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Ning Zhang, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng, Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohan Maheswari, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vítor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang, Xiaofang Wang, Xiaoqing Ellen Tan, Xide Xia, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Duc Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric-Tuan Le, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Filippos Kokkinos, Firat Ozgenel, Francesco Caggioni, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hakan Inan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Ilias Leontiadis, Irina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, Jean-Baptiste Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kiran Jagadeesh, Kun Huang, Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias 25 Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Rangaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shuqiang Zhang, Shuqiang Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Koehler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang, Xiaojian Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, Zhiwei Zhao, and Zhiyu Ma. The llama 3 herd of models, 2024. https://arxiv.org/abs/2407.21783. Riccardo Grazzi, Julien Siems, Arber Zela, Jörg KH Franke, Frank Hutter, and Massimiliano Pontil. Unlocking state-tracking in linear rnns through negative eigenvalues. arXiv preprint arXiv:2411.12537, 2024. Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First conference on language modeling, 2024. Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021a. Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems, 34:572–585, 2021b. Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, and Tri Dao. Sonicmoe: Accelerating moe with io and tile-aware optimizations. arXiv preprint arXiv:2512.14080, 2025. D Hendrycks. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415, 2016. W. Daniel Hillis and Guy L. Steele. Data parallel algorithms. Commun. ACM, 29(12):1170–1183, December 1986. ISSN 0001-0782. doi: 10.1145/7902.7903. https://doi.org/10.1145/7902.7903. Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9(8):1735–1780, 1997. doi: 10.1162/neco.1997.9.8.1735. Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What’s the real context size of your long-context language models? arXiv preprint arXiv:2404.06654, 2024. Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International conference on machine learning, pages 9099–9117. PMLR, 2022. IBM Granite Team.Granite 4.0 language models.https://huggingface.co/collections/ibm-granite/ granite-40-language-models, 2025. Hugging Face collection, accessed January 28, 2026. Dhiraj Kalamkar, Dheevatsa Mudigere, Naveen Mellempudi, Dipankar Das, Kunal Banerjee, Sasikanth Avancha, Dharma Teja Vooturi, Nataraj Jammalamadaka, Jianyu Huang, Hector Yuen, et al. A study of bfloat16 for deep learning training. arXiv preprint arXiv:1905.12322, 2019. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156–5165. PMLR, 2020. Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36:24892–24928, 2023. 26 Quoc V. Le, Navdeep Jaitly, and Geoffrey E. Hinton. A simple way to initialize recurrent networks of rectified linear units, 2015. https://arxiv.org/abs/1504.00941. Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887, 2024. Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024. Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. https://arxiv.org/abs/1711.05101. Loren Lugosch, Mirco Ravanelli, Patrick Ignoto, Vikrant Singh Tomar, and Yoshua Bengio. Speech model pre-training for end-to-end spoken language understanding. arXiv preprint arXiv:1904.03670, 2019. Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017. https://openreview.net/forum?id=Byj72udxe. William Merrill. Sequential neural networks as automata. arXiv preprint arXiv:1906.01615, 2019. William Merrill, Ashish Sabharwal, and Noah A. Smith. Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843–856, 2022. doi: 10.1162/tacl_a_00493. https://aclanthology.org/2022.tacl-1.49/. William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. arXiv preprint arXiv:2404.08819, 2024. Meta AI. The llama 4 herd: The beginning of a new era of natively multimodal ai innovation, 2025.https: //ai.meta.com/blog/llama-4-multimodal-intelligence/. Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint arXiv:1710.03740, 2017. Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering, 2018. https://arxiv.org/abs/1809.02789. Mayank Mishra. Lm engine: A hyper-optimized library for pretraining and finetuning, 2024.https://github.com/ open-lm-engine/lm-engine. Mayank Mishra, Matt Stallone, Gaoyuan Zhang, Yikang Shen, Aditya Prasad, Adriana Meza Soria, Michele Merler, Parameswaran Selvam, Saptha Surendran, Shivdeep Singh, et al. Granite code models: A family of open foundation models for code intelligence. arXiv preprint arXiv:2405.04324, 2024. Elliot Nelson, Georgios Kollias, Payel Das, Subhajit Chaudhury, and Soham Dan. Needle in the haystack for memory based large language models. arXiv preprint arXiv:2407.01437, 2024. NVIDIA, :, Aarti Basant, Abhijit Khairnar, Abhijit Paithankar, Abhinav Khattar, Adithya Renduchintala, Aditya Malte, Akhiad Bercovich, Akshay Hazare, Alejandra Rico, Aleksander Ficek, Alex Kondratenko, Alex Shaposhnikov, Alexander Bukharin, Ali Taghibakhshi, Amelia Barton, Ameya Sunil Mahabaleshwarkar, Amy Shen, Andrew Tao, Ann Guan, Anna Shors, Anubhav Mandarwal, Arham Mehta, Arun Venkatesan, Ashton Sharabiani, Ashwath Aithal, Ashwin Poojary, Ayush Dattagupta, Balaram Buddharaju, Banghua Zhu, Barnaby Simkin, Bilal Kartal, Bita Darvish Rouhani, Bobby Chen, Boris Ginsburg, Brandon Norick, Brian Yu, Bryan Catanzaro, Charles Wang, Charlie Truong, Chetan Mungekar, Chintan Patel, Chris Alexiuk, Christian Munley, Christopher Parisien, Dan Su, Daniel Afrimi, Daniel Korzekwa, Daniel Rohrer, Daria Gitman, David Mosallanezhad, Deepak Narayanan, Dima Rekesh, Dina Yared, Dmytro Pykhtar, Dong Ahn, Duncan Riach, Eileen Long, Elliott Ning, Eric Chung, Erick Galinkin, Evelina Bakhturina, Gargi Prasad, Gerald Shen, Haifeng Qian, Haim Elisha, Harsh Sharma, Hayley Ross, Helen Ngo, Herman Sahota, Hexin Wang, Hoo Chang Shin, Hua Huang, Iain Cunningham, Igor Gitman, Ivan Moshkov, Jaehun Jung, Jan Kautz, Jane Polak Scowcroft, Jared Casper, Jian Zhang, Jiaqi Zeng, Jimmy Zhang, Jinze Xue, Jocelyn Huang, Joey Conway, John Kamalu, Jonathan Cohen, Joseph Jennings, Julien Veron Vialard, Junkeun Yi, Jupinder Parmar, Kari Briski, Katherine Cheung, Katherine Luna, Keith Wyss, Keshav Santhanam, Kezhi Kong, Krzysztof Pawelec, Kumar Anik, Kunlun Li, Kushan Ahmadian, Lawrence McAfee, Laya Sleiman, Leon Derczynski, Luis Vega, Maer Rodrigues de Melo, Makesh Narsimhan Sreedhar, Marcin Chochowski, Mark Cai, Markus Kliegl, Marta Stepniewska-Dziubinska, Matvei Novikov, Mehrzad Samadi, Meredith Price, Meriem Boubdir, Michael Boone, Michael Evans, Michal Bien, Michal Zawalski, Miguel Martinez, Mike Chrzanowski, Mohammad Shoeybi, Mostofa Patwary, Namit Dhameja, Nave Assaf, Negar Habibi, Nidhi Bhatia, Nikki Pope, 27 Nima Tajbakhsh, Nirmal Kumar Juluru, Oleg Rybakov, Oleksii Hrinchuk, Oleksii Kuchaiev, Oluwatobi Olabiyi, Pablo Ribalta, Padmavathy Subramanian, Parth Chadha, Pavlo Molchanov, Peter Dykas, Peter Jin, Piotr Bialecki, Piotr Januszewski, Pradeep Thalasta, Prashant Gaikwad, Prasoon Varshney, Pritam Gundecha, Przemek Tredak, Rabeeh Karimi Mahabadi, Rajen Patel, Ran El-Yaniv, Ranjit Rajan, Ria Cheruvu, Rima Shahbazyan, Ritika Borkar, Ritu Gala, Roger Waleffe, Ruoxi Zhang, Russell J. Hewett, Ryan Prenger, Sahil Jain, Samuel Kriman, Sanjeev Satheesh, Saori Kaji, Sarah Yurick, Saurav Muralidharan, Sean Narenthiran, Seonmyeong Bak, Sepehr Sameni, Seungju Han, Shanmugam Ramasamy, Shaona Ghosh, Sharath Turuvekere Sreenivas, Shelby Thomas, Shizhe Diao, Shreya Gopal, Shrimai Prabhumoye, Shubham Toshniwal, Shuoyang Ding, Siddharth Singh, Siddhartha Jain, Somshubra Majumdar, Soumye Singhal, Stefania Alborghetti, Syeda Nahida Akter, Terry Kong, Tim Moon, Tomasz Hliwiak, Tomer Asida, Tony Wang, Tugrul Konuk, Twinkle Vashishth, Tyler Poon, Udi Karpas, Vahid Noroozi, Venkat Srinivasan, Vijay Korthikanti, Vikram Fugro, Vineeth Kalluru, Vitaly Kurin, Vitaly Lavrukhin, Wasi Uddin Ahmad, Wei Du, Wonmin Byeon, Ximing Lu, Xin Dong, Yashaswi Karnati, Yejin Choi, Yian Zhang, Ying Lin, Yonggan Fu, Yoshi Suhara, Zhen Dong, Zhiyu Li, Zhongbo Zhu, and Zijia Chen. Nvidia nemotron nano 2: An accurate and efficient hybrid mamba-transformer reasoning model, 2025. https://arxiv.org/abs/2508.14444. NVIDIA Corporation. Parallel Thread Execution (PTX ISA). NVIDIA.https://docs.nvidia.com/cuda/ parallel-thread-execution/. J. M. Ortega and W. C. Rheinboldt. Iterative Solution of Nonlinear Equations in Several Variables. Society for Industrial and Applied Mathematics, 2000. doi: 10.1137/1.9780898719468.https://epubs.siam.org/doi/abs/10. 1137/1.9780898719468. Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th annual meeting of the association for computational linguistics (volume 1: Long papers), pages 1525–1534, 2016. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019. Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, et al. Rwkv: Reinventing rnns for the transformer era. arXiv preprint arXiv:2305.13048, 2023. Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02143, 2021. Korbinian Pöppel, Maximilian Beck, and Sepp Hochreiter. Flashrnn: I/o-aware optimization of traditional rnns on modern hardware, 2025. https://arxiv.org/abs/2412.07752. Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. Hgrn2: Gated linear rnns with state expansion. arXiv preprint arXiv:2404.07904, 2024. Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.07522, 2024. Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S. Gordon. Choice of Plausible Alternatives: An Evaluation of Commonsense Causal Reasoning. In AAAI Spring Symposium on Logical Formalizations of Commonsense Reasoning, Stanford University, March 2011.http://ict.usc.edu/pubs/Choice%20of%20Plausible%20Alternatives-%20An% 20Evaluation%20of%20Commonsense%20Causal%20Reasoning.pdf. Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019. https://arxiv.org/abs/1907.10641. Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International conference on machine learning, pages 9355–9366. PMLR, 2021. Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems, 37: 68658–68685, 2024. Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020. Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. 28 Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538, 2017. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron- lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019. Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi. Deltaproduct: Improving state-tracking in linear rnns via householder products. arXiv preprint arXiv:2502.10297, 2025. Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint arXiv:2201.11990, 2022. Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024. Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621, 2023. Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, 2022.https://arxiv. org/abs/2009.06732. Kimi Team, Yifan Bai, Yiping Bao, Guanduo Chen, Jiahao Chen, Ningxin Chen, Ruijue Chen, Yanru Chen, Yuankun Chen, Yutian Chen, et al. Kimi k2: Open agentic intelligence. arXiv preprint arXiv:2507.20534, 2025. Aleksandar Terzić, Nicolas Menet, Michael Hersche, Thomas Hofmann, and Abbas Rahimi. Structured sparse transition matrices to enable state tracking in state-space models. arXiv preprint arXiv:2509.22284, 2025. Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, MAPL 2019, page 10–19, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450367196. doi: 10.1145/3315508.3329973. https://doi.org/10.1145/3315508.3329973. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions, 2017.https: //arxiv.org/abs/1707.06209. BigScience Workshop, Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilić, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, et al. Bloom: A 176b-parameter open-access multilingual language model. arXiv preprint arXiv:2211.05100, 2022. Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International conference on machine learning, pages 10524–10533. PMLR, 2020. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023. Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464, 2024a. Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. Advances in neural information processing systems, 37:115491–115522, 2024b. Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar, and Tri Dao. Flashattention-4: Algorithm and kernel pipelining co-design for asymmetric hardware scaling, 2026. https://arxiv.org/abs/2603.05451. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence?, 2019. https://arxiv.org/abs/1905.07830. Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in neural information processing systems, 32, 2019. 29 Muru Zhang, Mayank Mishra, Zhongzhu Zhou, William Brandon, Jue Wang, Yoon Kim, Jonathan Ragan-Kelley, Shuaiwen Leon Song, Ben Athiwaratkun, and Tri Dao. Ladder-residual: parallelism-aware architecture for accelerating large model inference with communication overlapping, 2025. https://arxiv.org/abs/2501.06589v2. A State Size Expansion Attention PatternW q W k W v W g W f W o Total (P)State Size (NKV ) Multi-HeadNKd m NKd m NV d m NV d m Nd m NV d m ≈ N(2K + 3V )d m ≈ PKV (2K+3V )d m Multi-QueryNKd m Kd m V d m NV d m Nd m NV d m ≈ N(K + 2V )d m ≈ PKV (K+2V )d m Multi-KeyKd m NKd m V d m NV d m Nd m NV d m ≈ N(K + 2V )d m ≈ PKV (K+2V )d m Multi-ValueKd m Kd m NV d m NV d m Nd m NV d m ≈ 3NV d m ≈ PK 3d m Table 8 Number of parameters in all linear projections.Ndenotes the number of heads andd m denotes the model’s embedding dimension. For all our models, we useK= 64 andV= 16 which results in the largest state size for the multi-value attention formulation for constant number of total parameters (P). B Representation Power of M 2 RNN Layers Theorem 1: The M 2 RNN recurrence can represent all tasks representable by non-linear vector-valued RNNs and hence can represent regular languages. Proof.We show that M 2 RNN can simulate a vector-valued nonlinear RNN, which is known to recognize all regular languages. Consider the parameter setting f t = 0, w r = 0 V , q t = k t = e 1 ∈R K , where e 1 = [1, 0,..., 0] ⊤ is the first standard basis vector. Under this choice, the matrix-valued recurrence reduces in its first row to Z t = tanh H t−1 W + v ⊤ t 0 ⊤ V . . . 0 ⊤ V ,(25) and the output is obtained from the first row, y t = (Z t ) 1,: . This is precisely a standard vector-valued nonlinear RNN withtanhactivation applied to the first row of the state matrix. Since vector-valued nonlinear RNNs can represent any regular language, it follows that M 2 RNN can represent any regular language. C RMSNormTP Module The forward computation for RMSNorm (Zhang and Sennrich, 2019) applied on the input activationsxis given by: s = 1 q 1 d P d j=1 x 2 j (26) y i = (w i ⊙ x i )s(27) 30 where i indicates the i th feature of the input or output, d is the number of features and w is some learnable weight for the RMSNorm module. We cachesduring the forward computation in HBM for backward computation. Similarly, the backward computation is given by: p i = w i ⊙∇ y i L(28) r = d X j=1 p j x j (29) ∇ x i L = sp i − rs 3 x i d (30) ∇ w i L = (∇ y i L⊙ x i )s(31) To use RMSNorm (Zhang and Sennrich, 2019) module with M 2 RNN layers with Tensor Parallelism (TP), we shard the weightwfor RMSNorm among the GPUs participating in TP. Since the incoming activationsxare also sharded along the feature dimension, to compute the inverse of the normalization terms, we require an AllReduce operation in the forward. We propose the RMSNormTP module which modifies the RMSNorm module to compute s in forward and r in backward as: s = 1 q 1 d P N TP n=1 P j∈D n x 2 j (32) r = N TP X n=1 X j∈D n p j x j (33) whereD n indicates the sharded features onn th GPU for TP. The summation in blue colour indicates an AllReduce operation across GPUs in Equations 32, 33. Note that for both AllReduce operations, we only need to communicate a single element. 31