Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels
Authors: Maximilian Beck
Korbinian Pöppel
Phillip Lippe
Sepp Hochreiter
Venue: NeurIPS 2025 poster ()
Keywords: xLSTM
FlashLinearAttention
FlashAttention
Linear Attention
Linear RNN
RNN
mLSTM
Attention
Self-Attention
Transformers
LLM
Kernels
Systems
Efficiency
Primary area: deep_learning
TL;DR: We introduce TiledFlashLinearAttention a faster kernel algorithm for Linear RNNs and mLSTMs by improved Sequence Parallelism.

Abstract: Linear RNNs with gating recently demonstrated competitive performance compared to Transformers in language modeling. Although their linear compute scaling in sequence length offers theoretical runtime advantages over Transformers, realizing these benefits in practice requires optimized custom kernels, as Transformers rely on the highly efficient Flash Attention kernels (Dao, 2024). Leveraging the chunkwise-parallel formulation of linear RNNs, Flash Linear Attention (FLA) (Yang & Zhang, 2024) shows that linear RNN kernels are faster than Flash Attention, by parallelizing over chunks of the input sequence. However, since the chunk size of FLA is limited, many intermediate states must be materialized in GPU memory. This leads to low arithmetic intensity and causes high memory consumption and IO cost, especially for long-context pre-training. In this work, we present Tiled Flash Linear Attention (TFLA), a novel kernel algorithm for linear RNNs, that enables arbitrary large chunk sizes and high arithmetic intensity by introducing an additional level of sequence parallelization within each chunk. First, we apply TFLA to the xLSTM with matrix memory, the mLSTM (Beck et al., 2024). Second, we propose an mLSTM variant with sigmoid input gate and reduced computation for even faster kernel runtimes at equal language modeling performance. In our speed benchmarks, we show that our new mLSTM kernels based on TFLA outperform highly optimized Flash Attention, Linear Attention and Mamba kernels, setting a new state of the art for efficient long-context sequence modeling primitives. Our code is available at: https://github.com/NX-AI/mlstm_kernels

Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels
Maximilian Beck!.2 Korbinian Pöppel Phillip Lippe? Sepp Hochreiter ELLIS Unit, LIT Al Lab, Institute for Machine Learning, JKU Linz, Austria
2 NXAI GmbH, Linz, Austria
Abstract
Linear RNNs with gating recently demonstrated competitive performance com- pared Transformers language modeling. Although their linear compute scaling in sequence length offers theoretical runtime advantages over Transform- ers, realizing these benefits in practice requires optimized custom kernels, Transformers rely on the highly efficient Flash Attention kernels (Dao, 2024). Leveraging the chunkwise-parallel formulation linear RNNs, Flash Linear At- tention (FLA) (Yang & Zhang. 2024) shows that linear RNN kernels are faster than Flash Attention, by parallelizing over chunks of the input sequence. However, since the chunk size of FLA is limited, many intermediate states must be materialized in GPU memory. This leads low arithmetic intensity and causes high memory consumption and IO cost, especially for long-contes, pre- teeitrriiing.... In this work, we present Tiled Flash Linear Attention (TFLA), a novel kernel algorithm for linear RNNs, that enables arbitrary large chunk sizes and high arithmetic intensity by introducing an additional level of sequence parallelization within each chunk. First, we apply TFLA to the xLSTM with matrix memory, the mLSTM (Beck et al., 2024). Second, we propose an mLSTM variant with sigmoid input gate and reduced com- putation for even faster kernel runtimes at equal language modeling performance. In our speed benchmarks, we show that our new mLSTM kernels based on TFLA outperform highly optimized Flash Attention, Linear Attention and Mamba kernels, setting a new state of the art for efficient long-context sequence modeling primitives. Our code is available at: https: //gt thub. com/NX Boan/N-AAI/miin
Introduction
With the trend of training models of ever increasing size with large datasets on thousands GPUs, becomes increasingly important to optimize the model architecture as well as its low-level implemen- tations for modern hardware. Transformers (Vaswani et al., 2017), which are the core architecture of nowadays state-of-the-art models are highly optimized, but the computational requirements of self-attention scale quadratically with sequence length. This creates significant challenges for both training and inference on long context.
Recently, recurrent alternatives with linear scaling sequence length (Beck et al., 2024: Sun et al., 2023; Dao & Gu, 2024; Yang et al., 2024b) promise efficiency gains, especially on long sequences and during inference while providing competitive performance. The success of these emerging recurrent architectures is based on two main pillars: (1) A parallel or chunkwise-pare al., 2023; Hua al., 2022), which, like Attention, calculates all outputs parallel during training, and (2) kernel implementations that are close or exceed training speeds of FlashAttention (Dao, 2024).
formulation (Sun
Yang et al. (2024b) show that their custom Flash Linear Attention (FLA) kernels, based on the chunkwise-parallee formulation of linear RNNs, achieve faster runtimes than FlashAttention. Now at Google Deepmind.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

They accomplish this by dividing the sequence into chunks and recurrently materializing only the initial RNN state of each chunk in GPU memory. Subsequently, in the parallel part they employ one level of sequence parallelism and compute the oulputs for each chunk in parallel. For small chunk size and long sequences, this leads to a large amount of intermediate states to be stored and loaded from GPU memory, which increases memory consumptionand decreases arithmetic intensity. Since modern GPUs see faster increase in computation throughput than memory bandwidth (Gholami al., 2024), it essential to minimize large memory 10 and increase arithmetic intensity. simple approach would be to increase the chunk size. However, the chunk size of FLA limited by the physical SRAM available on the GPU.
QKV
QK
QK QKvω
inputs
Co
C₁
C2
recurrent
kernel
(1)
parallel
kernel
inter chunk
intra chunk
H
H())
H0)
outputs
Figure 1: Tiled Flash Linear Attention (TFLA) consists of recurrent kernel and a parallel ker- nel, which process the input sequence in chunks QKV(k) (1 level sequence parallelism). The recurrent kernel materializes the memory state Ck-1 for each chunk. The parallel kernel com- putes the output states H for all chunks. TFLA uses tiling for the 3 matrix- patrii--uullpliiationns in the parallel kernel (2nd level of sequence parallelism) fully utilize the hardware and prevent materi- alization of many memory states.
To solve this problem, introduce Tiled Flash Linear Attention (TFLA) which enables unlim- ited chunk sizes by introducing second level of sequence parallelism via tiling of the ma- trix computations in sequence dimension within each chunk. This increases the arithmetic inten- sity of the kernels and allows us to efficiently balance memory consumption and IO vs. com- putation.
In this paper, we implement our Tiled Flash Linear Attention algorithm for the xLSTM with matrix memory the mLSTM Beck et al. (2024). The mLSTM is a linear RNN that uses exponential gating with scalar gates per head, along with an additional normalizer state for output normalization, This gating mechanism has demonstrated competitive performance compared to Transformers and Mamba on language modeling tasks at moderate scales. However, for comparisons at even larger scales, efficient kernels that leverage the truukwwse--paralee In our speed benchmarks, we show that our new mLSTM kernels based on TFLA outperform highly optimized Attention, Linear Attention and Mamba kernels.
formulation for the mLSTM were still missing.
After optimizing our kernels for the existing mLSTM computation, we seek ways to reduce kernel run- time by targeted modifications to the mLSTM. Towards this end, we propose mLSTMsig, an mLSTM with sigmoid input gate and reduced computation, that enables even faster kernel implementations at no performance drops on language modeling up to 1 1.4B parameter scale.
Finally, motivated by the equal performance of both mLSTM variants, we perform an empirical study inspired by transfer function analysis from control theory (Ogata, 2010) to understand their differences and characteristics. We find that both mL.STM variants exhibit the same transfer behavior and, moreover, our analysis suggests that the input gate biases should be initialized at larger negative values. extensive experiments on language modeling, we confirm that this initialization improves training stability as well as the overall performance of mLSTM models.
To summarize, in this work, we make the following contributions: (1) We introduce Tiled Flash Linear Attention, new chunkwise chunkwise-parallel kernel algorithm for Linear RNNs with two levels of sequence parallelism, that enables arbitrary large chunk sizes and apply to the mI.STM (Beck et al., 2024). (2) We introduce MsS. a faster mLSTM variant with sigmoid input gate with no performance losses up to 1.4B parameter scales. (3) We improve the training stability and performance of the mLSTM through careful gate initialization guided by our empirical transfer behavior analysis.
2 mLSTM Formulations
The mL.STM cell is the fully parallelizable part of the xI.STM (Beck et al., 2024). has matrix memory and exponential gating.

---

2.1 Recurrent Formulation
In its recurrent formulation, the mLSTM cell processes the series of input vectors R2 E R for time steps & {I, .I,..., , } mapping a state (ht-1, 1 given an imput z,, Here, h, E R denotes the hidden state, C, E & denotes the cell state responsible for long-term memory, n& E RR denotes the normalizer state, and my E IK denotes the max state. Together normalizer and max state control the magnitude of the exponential input gate and ensure stability (see Appendix D.1). The recurrent SLM formulation is given by the following state update equations:
to a successor state (he, Ct Rey me)
my = max (logo(t) + m ie]
C. fe Ca-1 + it kt Vi
(2)
nt = f, n + is kt
(3)
C. (4t/Vdop ddgh)
h; max [|] (g: dgt)) exppp exp(-maa)
(4)
h = O2 3 NORM(h.)
(5)
The scalar forget and input gates i,, fr E IR are computed as fç = exp (logo(f)) + m4-1 my and i, ex me) with the pre-activations {ie, [e] w (i.13 + birrg.. respectively. The vector output gate o, E Rh is given by o2 = a (6,) with the are-activations oN = W. I; + b. and the sigmoid function a. The norm layer NORM in (5) can be either RMS norm (Zhang & Sennrich, 2019) or LayerNorm (Ba et al., 2016). Typically, multiple of these cells operate simultaneously as parallel heads, similar to Transformers (Vaswani et al., 2017).
2.2 Chunkwise-Parallee Formulation
The chunkwise-paralde formulation is trade-off between the parallel and the fully recurrent formu- lation. has a recurrent part and a (quadratic) parallel part, with an overall sub-quadratic scaling in sequence length. Similar to the fully parallel formulation (see Appendix B.1), we assume that all inputs are available at once. We then split the sequence of length T into N₂ = [T/L] chunks of length L and use & E {1, , Ne) for the chunk index. We rearrange the input and forget gates, as well as the queries, keys, and values into chunkwise matrices, where the chunk index becomes the first dimension. For example, the forget gate pre-activations E RT are rearranged into a matrix
F F=(F")) f(22 ...., ((N)) E N.N.,,, where each row ((t) = (f f(k +)). fuN.)) e RR contains the pre-activations of the chunk k. The input gate pre-activations follow analogously. Similarly, the queries, keys and values are rearranged into chunkwise tensors Q, K E N XLxd,e
and V E R xxxd1. Here, the query matrix ( (q(x (q(:-1)N_+1;..... qNN.)) R Lxdak contains
the query vectors of chunk kt Keys, and values follow analogously. For notational simplicity we drop the leading Ne dimension and omit normalization layer and the output gate, i.e. consider h, as hidden state outputs.
Chunkwise Gates. Given the logarithmic forget gates p(k) log (F()) E IR- and input gates {(k) log exp(i ((()
RRL.
t:
2
3
4
5
6
7
8
we can compute the logarithmic chunkwise gates gk R. and
Co
C
Co
bR, a& E RRL as gk sum ((A)). b cumsum
(R) and
((K) (()) We refer to Appendix B.2 for more details on the chunkwise gates. In Figure 2, we show the summed forget gates EA contain the forget gate contribution of all forget gates within a chunk. The cumulative forget gate vectors by contain the forget gate contributions from the beginning of the chunk up to the current time step within the current chunk. The cumulative input gate vectors a contain the input gates for every timestep as well as the forget gate contributions from the current time step the end of the chunk.
ag = rev_evc mew_cummum
QKV10
QKV
Figure 2: Illustration of
the chunkwise gates aμ, by and g& with chunk size L = 4. Each arrow denotes an element in the gate vec- tors. See Figure 9 in Ap- pendix B.2 for more details.
3

---

Inter-chunk Recurrent Contribution.
The inter-chunk recurrence is given by
C CB.Ck- + a O x0)
((K)
(6)
T 1, a K(
72 8k 1
where Ek and as are the stabilized chunkwise gates. This recurrent part resembles the fully recurrent formulation in Section 2.1, but instead of computing the intermediate states for every timestep t, compute them directly for every L time steps without materializing the states in between.
Intra-chunk Parallel Contribution. The recurrent part is followed by the intra-chunk parallel contribution:
(k)
-oc
for i <
D
bj +I(A)
(())
(8)
by
for i j
((K)
(k) K (A)
D ((k) о (())
ddgh
where (R) E TRLXL is the stabilized gate matrix. Compared to the fully parallel part from Ap- pendix B.I, the quadratic cost of the matrices ( g(k) E LRL is greatly reduced, since the chunk size typically small compared to the sequence length
Output Computation. Finally, the contributions from the intra-chunk parallel part H H are combined with the inter-chunk recurrent part H (8) obtain the hidden states H R Lxdxr for each chunk & (see Figure 1):
Q (k) Vagt
H inter
D
O
Ck - Q Ck-1
(k) Intra = (())
v(),
(10)
H (k)
H (k) inter
H (k)) intra
/ /
(R) denom
(11)
where h iemom E R is normalization factor. Appendix B.2 and B.3 provide detailed description of chunkwise-paraller forward and backward pass. Appendix provides the FLOP and memory operation counts for all formulations.
3 Tiled Flash Linear Attention
Flash Linear Attention (Yang al., 2024b) introduces fast kernel algorithm for the chunkwise formulation for Linear Attention (cf. Section 2.2 without gates) and shows that their implementation faster than optimized FlashAttention (Dao, 2024). This speedup is achieved by single level sequence parallelism, where the states C₁ are first materialized in GPU memory and then the outputs H (k) computed in parallel. However, since in Flash Linear Attention the chunk size parameter determines the tile sizes in SRAM, the maximum chunk size is limited (typically I, 64) by the physical SRAM size of the GPU. Therefore, we have to materialize many states in HBM, where the number of states No []. This leads low arithmetic intensity and high GPU memory consumption, which poses challenges especially for long-context pre-training.
We begin with brief review of fundamentals GPUs for writing efficient kernels in Appendix C.I.
More Efficient Kernels via Two Level Sequence Parallelism. To address the issue of limited chunk sizes, Tiled Flash Linear Attention (TFLA) introduces two levels of sequence parallelism, which enables fast kernels and a trade-off between memory consumption and computational efficiency (see Figure 6). The first level is the parallelization over the chunks of the sequence, which requires to compute and materialize intermediate states C₁ in GPU High Bandwidth memory (HBM). For this we use recurrent kernel similar to previous work (Yang et al., 2024b). The second level the parallelization within each chunk, which is achieved by tiling the intra chunk attention matrix along chunk dimension. This second level of parallelism enables large chunk sizes and hence reduces the memory consumption for the intermediate states we have to store and load N. = [T/L] intermediate states in HBM on each kernel call, where T is the sequence length and L. is the chunk size. In addition to the two levels of sequence parallelism and the naive parallelization over the batch

---

and head dimensions, TFLA also parallelizes over the embedding dimension. This enables arbitrary large head dimensions and results massive parallelization over five dimensions, which is crucial for achieving high performance modern GPUs. We analyze the theoretical runtime of our TFLA kernels in Appendix G.
We review the matrix multiplication operations of the intra-chunk parallel part of the mLSTM in order to show how we efficiently parallelize these operations. For simplicity we omit the the gate computations and normalization, these do not influence the work partitioning. We also omit the leading batch, head and chunk dimension, over which we can parallelize naively as they do not interact with the matrix multiplication (see Table 1). In simplified form, the intra-chunk parallel forward pass of the mLSTM STM (and other linear RNNs) for a chunk l can be written as three matrix multiplications, which we fuse into single kernel:
Forward Pass.
H (Lig xdnx)
((k)
K(A)T (dag x Lxc)
((kk) (Loo xduv)
+
Ck-1 (Log)) (deu xdwu)
Q(k)
(dgg))de))
(12)
H
In Appendix A.2, we show that TFLA can be applied to any linear RNN that either follows or can be reformulated into this form. In order to parallelize the computation in (12), we introduce the block sizes BLhg. BLks, Bagk and Bahe for the attention matrix, query, key, value and hidden state dimensions Lig. Lkve dgg and daws along which we either parallelize or accumulate by using a loop inside the kernel.
Liu
dgt
Lupe
Bisad dip
Bast
Loc
gth
Lin
Q'
Ki
Blse
Bag
Biss
des
Broe
Bra
8:4
Tike
V/3i
Bihe
Bruu
Bara
dee
dgt
dus
In Figure we show our TFLA tiling strategy for the forward pass H kernel. We parallelize across the outer sequence dimension Lig with Ning / Br.ng programs, and across the outer embedding dimension die with Naho- = / Bahe programs. We loop over the inner dimensions Live and dgk, which are tiled by the block sizes BL and Bags respectively.
Bug
Bite det
Bagt
Ling
C²-1
Bahe
Bigh
Bat
Figure 3: TFLA Intra-Chunkl Tiling. We loop over BL and Bagl (indicated by arrows) and paral- lelize over B and Bake (indicated by dashed lines) blocks. + denotes block-wise accumulation.
Tiled Computation. For the mLSTM we cannot simply accumulate the results of the matrix multiplications (*) along the Liv dimension and ({A) due to the stabilization of the exponential inter input gate with the max state m.- The max state tracks the maximum of the forget and input gates over time and used to stabilize the exponential input gate similar to the safe softmax computation (Milakov & Gimelshein, 2018). Since we compute the hidden state output (R) in blocks along the chunk size (i.e. time) dimension Lkv. need to rescale during accumulation of the block results for ((kk)
intra and the overall results into (k) in the same way as Fltettnt (Dao, 2024). We provide details on the rescaling in Section B.2. For the backward pass there is no rescaling
necessary as we store the max states in the forward pass and reuse them in the backward pass. The pseudocode for the forward pass of TFLA for the mLSTM is listed in Algorithm 1.
Backward Pass. The parallelization strategy for the backward pass of TFLA is more complex than for the forward pass, since we need to compute three output tensors the gradients for the queries, keys and values, of which each has an intra-chunk and inter-chunk part. However, in Section C.4 we show that the individual gradients can be mapped to three matrix multiplications similar to the forward pass. In TFLA, we then implement a separate kernel for each gradient and use the same work partitioning as in the forward pass but swap the loop and parallelization dimensions, accordingly. Table summarizes the work partitioning of our TFLA kernels.
Table 1:
TFLA kernel parallelization and loop dimensions. Parallelization dimensions are indicated by P and loop dimensions by L. The last column shows the first two dimen- sions of the 3D kernel launch grid. The last di- mension of all kernels is Nehunk Nhead N Match-
Kernel H(M)
Lng
Lee
dgg
dnu
Thread Block Grid
P
L
I.
P
Dave Heaks
duu Lag.
SQ
P
L
P
L
Bajs' DEAS"
Fans
(K))))
P
L
del
May
8V(6)
I.
I.
P
Day Bike
duc Lac.
5

---

Gain before Norm
Gain
before
Gain after
7.5
Pruu Cata Praictivation
5.0
16
5.0
1.6
14
1.4
0.5
12
1.2
-25
-2.
1.0
0.6
-7.
044
-75
0.4
-00.
0.
-100
0.2
-1002500 2/5
-50-2500 25
-10-3.500
-1.0-..0.
0.0
Foreet
Preactivation
Forast Presctivation
Forget
Preactivation
MLSTMexp (Exponential Input Gate)
(b) mLSTMsig (Sigmoid Input Gate)
Transfer behavior of the mLSTM before and after the RMS-norm layer (c =1e-6) for different input and forget gate values. The color shows the gain of the mL.TM defined (16). After the norm layer mLSTMexp and mLSTMMig exhibit the same transfer behavior.
Figure 4:
4 Faster mLSTM with Sigmoid Input Gate
The MLSTM with exponential gating (i.e. exponential input gate) introduced by Beck et al. (2024) requires compute and keep track of two additional states, the normalizer state nu and max state my, as we show in Appendix D.I. Both will increase kernel runtime: The normalizer must be computed through summations, and tracking the max state throughout the tiled computation in TFLA (see Section 3 and C.2) prevents efficient fusing of loops within the kernel (see Appendix C.3).
Additionally, our analysis in Section 4.2 suggests initialize the input gate biases larger negative values (e.g. -10), such that the input gate pre-activations can grow slowly during training. We observe that most of these values stay below O during training (see Figure 15 in Appendix E). Therefore, we seek an alternative activation function which is similar to the exponential function in the negative range, but bounded in the positive range. This suggests
the sigmoid function [x which converges to exp(x) for x + oo and for x + oo.
(xx) [+exp(-x)
exp(x)
4.1 mLSTM with Sigmoid Input Gate
The sigmoid function can be computed in two ways as given above. Deoee.Deending on the sign of x it can be ensured that the argument of exp is always smaller than 0 avoid numerical overflow. Therefore, we do not need to control the magnitude x externally with a max state and as a consequence also drop the normalizer state (see Appendix D.1). This yields the mLSTM with sigmoid input gate (henceforth referred to as mLSTMsigg in its recurrent formulation as
C. (Ei) C--1 + oli)) ke v. hns c, (gt/ (g//vdg))
(13)
(14)
h a (ō.) O NORM (hr.
(15)
where the query, key, and value vectors qr. kt Ut and the gate preactivations i,, f.. o, remain the same for the mLSTM with exponential input gate (from now referred to mLSTMexp) in Section 2.1. We confirm that our TFLA mLSTMsig forward kernel is over 30% faster than the mLSTMexp forward (see Section 5.2), and show that MLSTMsig performs equally well compared to mL.STMexp in our language modeling experiments up to ..4B parameters (see Section 5.
4.2
Normalization of mLSTM and Linear RNNs
Motivated by the performance of mLSTMsig, seek to understand the differences between mL- STMsig and mLSTMexp empirically. To approach this, we draw inspiration from the concept of frequency response and transfer function analysis for control systems design, where typically the amplitude ratio or gain of output and input signals for different frequencies is considered (Ogata, 2010, Ch. 7). In our case, we analyze the transfer behavior of mLSTMsig and mLSTMexp for random inputs q:, k; and v, and different input gate and forget gate preactivations i, and f..
We will see that the normalization layer y = NORM(x), will play a crucial role in our analysis. The default norm layer in language modeling, the RMS norm (Zhang & Sennrich, 2019) with input vector input vector æ E Rdd and output vector y E Rd is defined as y RMS(2)
iz @ 7 where
RMS(x)) V xf + e, with with E IRd being learnable scale parameter. The epsilon parameter R is small constant typically set to 1e-6 avoid division by zero.

---

Torch FlashAttn
FlashAtth 3 GLA (FLA)
Simple GLA (FLA)
Mamba
MLSTMesp (TFLA XL chunk)
CuDNN
Mamba
5L55Hexp (limit chunk)
MLSTMsig (TRLA chuni)
15.0
60
50
Time [ms]]
10.0
7.5
30
5.0-
20
2.5
0.0
512
1024
2048
4096
8192
16384
32768
65536
512
1024
2048
4006
16384
32768 65536
Sequence Longth
Sequence Longth
(b) Training (Forward and Backward pass) Figure 5: TFLA Kernel Runtime Benchmark for embedding dimension 4096 and 65,536 tokens on NVIDIA H100 GPUs. In training, our TFLA kernels are faster than FlashAttention 3 for longer sequences and over 2x faster than Mamba 2 kernels for all sequence lengths.
(a) Inference (Forward pass)
Transfer Behavior of the mLSTM. We analyze the transfer behavior by computing the gain of the mLSTM cells from random inputs sampled from (0, 1) to hidden states before and after the norm layer for varying input and forget gate values. More specifically, we compute the gains Gbefore and
Gafter as
| 11 mmux
Gatter (N:RM(h:)ll
Gbefore
and
(16)
Imax
where ||=||| max == [21x(1211...... (xal) and we average over the time dimension. For more details see App. D.2. In Figure 4 we observe that the transfer behavior of mLSTMsig without normalizer is identical to mLSTMexp with normalizer and max state. Both exhibit a transition from suppressing (G = 0) to passing (G = 1) the signal at larger negative input gate preactivation values, which could partly explain the matching performance in our language modeling experiments.
Normalization Layers in other Gated L Linear RNNs Interestingly, almost all other gated linear RNN variants also place a normalization layer after the RNN cell (Sun et al., 2023; Dao & Gu, 2024; Qin et al., 2024b; Yang et al., 2024b). Often this is justified with improved training stability, but a more thorough discussion is missing (Lieber et al., 2024). Qin et al. (2022) analyze the effect of the norm layer after a non- non-gated, kernel-based linear attention layer (Katharopoulos et al., 2020) and show that this effectively prevents unbounded gradients. We also confirm that the norm layer has significant impact on training stability and the gradient norm during training. In Section 5.1 we show that initializing the input gate bias at larger negative values, as suggested by our transfer behavior analysis in Figure 4, prevents large gradient norm variance and spikes during training. Relatedly, the general effect of layer normalization in the Transformer architecture has been investigated in several studies (Xiong et al., 2020; Zhu et al., 2025).
Effect of Normalization on Gating in Linear RNNs. We hypothesize that at this point the normalization layer does not only have a stabilizing effect by controlling the magnitude of the layer activations through rescaling. but also actively participates mechanism of the linear RNN. For example, if the squared norm of C, q, which is controlled by input and forget gates through C,,, is smaller than the epsilon, the denominator in the NORM(x) layer is dominated by c and the output moves towards zero (indicated by the purple area in Fig. 4). Hence, by moving through the x-y plane in Fig. 4, the gates could learn suppress amplify any input in the sequence. In Section D.2 we show additional experiments on the effect of varying the normalization layer epsilons and different modifications of the normalizers for the mLSTM.
the information routing or gating
5 Experiments
In this section, we examine the performance of the two mLSTM variants mLSTMexp (mLSTM with exponential input gate) and mLSTMsig (mLSTM with sigmoid input gate). We compare two kernel algorithms: (1) limit_chunkk chunk: A kernel that is limited in chunk size L. (2) xl_chunk: Our Tiled Flash Linear Attention (TFLA) kernels with unlimited chunk size. For details
Section We assess the performance of ml STMMss compared to SS in Section 5.1 and benchmark the runtime of our kernels against other baselines in Section 5.2. In App. E. we verify the numerical correctness of our kernels.
7

---

5.1 Language Modeling with mLSTM
We train three different model sizes (160M, 400M, 1.4B parameters) with context lengths 4096 and 8192 on the DCLM dataset (Li et al., 2024). We include Llama2 style Transformer models (Touvron et al., 2023b) as reference in our comparison and describe our experiment setup, model architecture and training recipe in Appendix E.2.
Performance in Language Modeling. We compare mLSTMsig and mLSTMexp models on next-token prediction with different num- ber of heads or head dimensions. Table 2 and Table 6 show the results for context length 4096 and 8192, respectively. We find that our limit t_iit_chunk and xl_chunk kernels yield the same loss (up small numerical deviations) for almost all head dimensions. For some head dimensions, we observe gradient norm or loss spikes for the xl_chunk kernels, which affect the final loss. As main result we find that ml STgg performs equally well compared to SS.xx....
Table 2: Validation Perplexity at context length 4096. EXP and SIG denote mLSTMexp and mLSTMsig. LIMIT and XL correspond to
and xl_chunk kernels.
Size
Tokens
Heads
Llama
EXP
EXP
SIG
LIMIT
21.03
21.18
21.03
160M
19B
20.89
21.03
21.06
21.05
16:66
16.66 66
16.67
400M
24B
8
16.55
16.80
16.67
16.85
16.60
16.61
16.61
13.31
13.35
13.34
1.4B
33B
13.20
13.22
13.21
16
13.64
13.20
13.87
13.22
Effect of Input Gate Bias Initialization. We analyze the effect of the input gate bias initialization on training stability and performance of our mLSTM models in Appendix E.2. We observe in Figure 12 and 13, that initializing the input gate biases to -10 effectively mitigates large gradient norm spikes and variance during training for both mLSTMexp and mLSTMsig. We therefore conclude that the additional input gate not only improves performance (see Table 7), but also improves training stability, if initialized correctly.
Effect of Norm Layer Epsilon. In Appendix E.2, we investigate the effect of the norm layer epsilon on language modeling performance for mLSTMexp. Our transfer behavior analysis in Figure 4 suggests, that there exists an interplay between norm layer epsilon and input gate bias initialization. We confirm this in our grid search in Figure 14 and find that the best performing configuration is the default epsilon = with input gate biases initialized to -10.
5.2 Kernel Benchmark
We compare the runtime of our mLSTM limit_chunk and TFLA xl_chunk kernels with kernel implementations of the state-of-the-art sequence modeling primitives FlashAttention (Dao, 2024; Shah et al., 2024), Mamba (Gu & Dao, 2024; Dao & Gu, 2024) and GLA Yang et al. (2024b). In Appendix E.3 compare with other kernels from the Flash Linear Attention library (Yang & Zhang, 2024). We run our benchmarks on NVIDIA H100 GPUs.
Runtime Benchmark. We use the standard embedding dimension of 4096 for 7B parameter models and adapt the head dimensions per kernel accordingly. For example for FlashAttention we use 32 heads with head dim 128 and for the mLSTM we use 16 heads with head dim 256. Following the practice of Shah et al. (2024), we keep the number of tokens constant at 65,536 and vary sequence length and batch size accordingly. For further details see Appendix E.3. Figure shows the runtime benchmark results for inference, i.e. forward pass only, (left) and for training, i.c. pass (right). Our mLSTMexp TFLA xl_chunk kernels with two level sequence parallelism is about 25% faster than our limi t t_chknk kernels. Through targeted modifications of the input gate of the ML.TMM we save computation and enable more efficient kernel implementations for the forward pass of mLSTMsig (see Sec. 4). This yields another speedup of over 30% for the forward pass of the mLSTMsig TFLA kernel over the mL.SMexpp TFLA kernel.
In training, our TFLA kemelss are faster than FlashAttention 3 for longer sequences and more than 2x faster than Mamba kernels for all sequence lengths. We perform additional runtime benchmarks for varying head dimensions and a more in-depth comparison to the FLA (Yang et al., 2024b) and LightinngAttention22 (Qin al., 2024a) kernels in Appendix E.3.
8

---

Time [ms]
GPU Memory (GB)
Runtime vs. Memory Trade-off.
80
The chunk size parameter L bal- ances the computation between the two levels of sequence parallelism (see Sec. 3). Smaller chunk sizes increase memory consumption, be- cause more chunks are materialized in memory, but they reduce the quadratic compute FLOPs in the parallel part. Larger chunk sizes have the opposite effect. They decrease memory con- sumption, but increase quadratic com- pute FLOPs. In Figure 6, sure this trade-off for our mLSTMsig TFLA x1 chunk kernels.
12.5
GPU Memooy [GB]
10.0
Time [ms]
40
7.5
5.0
20
2.5
128
256
512
1024
2048
4096
Chunk Size
Figure 6: Memory vs. Runtime Trade-off of TFLA Forward- Backward Pass. We show the mLSTMsig for embedding dimension 4096 (8 heads with dim 512), sequence length 8 8192 and batch size 8. By varying the chunk size parameter, TFLA kernels effectively balance memory vs. runtime.
5.3 Theoretical Runtime and Arithmetic Intensity
In Figure we empirically observe that there exists an optimal chunk size (between 128 and 256) at which the runtime is minimized. In Appendix G, we compute the theoretical runtime optimal chunksize and the arithmetic intensity of TFLA depending on the chunk size by taking the FLOPs, memory operations and hardware accelerator specification into account. We find that the theoretical runtime optimum exceeds the empirically observed one (see Figure 21 in Appendix G.1), suggesting that our current kernel implementations may not yet fully exploit the available performance potential. We analyze the effect of the chunk size parameter on runtime, FL.OPs, and arithmetic intensity in detail in Appendix F and G summarize our findings as follows:
(I) The chunk size mediates trade-off' between runtime and GPU memory usage [Figure 6]. (2) L determines the total compute in FLOPs: L = 1 matches the recurrent formulation, while L = T matches the parallel one [Figure 19J. (3) There exists an optimal chunk size L [1. 7'] that minimized the total FLOP count [Equation (103), Figure 19, Figure 20J. (4) Increasing L raises the arithmetic intensity of TFLA kernels [Equation (109), Figure 22]. (5) The chunk size determines whether the kernel is memory- bound or compute- compute-bound on a given hardware [Figure 23, Figure 22]. (6) FLOPs/s alone can be misleading; the optimal chunk size should be chosen based on total runtime [Figure 24, Figure 21]. (7) The runtime-optimal chunk size scales proportionally with the square root of the head dimension and the accelerator's 's computational intensity [Figure 25, Figure 26). (8) Newer hardware generations require larger chunk sizes to approach peak performance. [Figure 26, Figure 21].
Related Work
Tiled Flash Linear Attention (TFLA) integrates the concept of tiling along one sequence dimension of the attention matrix for improved work partitioning (Dao, 2024) with the strategy of dividing the sequence into chunks (Yang et al., 2024b), yielding two levels of sequence parallelism (see Figure 7).
Flash Attention. Flash Attention (Dao al., 2022) is an 1O-aware implementation of softmax attention introduced by (Vaswani et al., 2017). It uses the idea of tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. In this way the quadratic attention matrix QK is never materialized in HBM, which reduces the memory requirement from quadratic with sequence length to linear, and significantly speeds up the kernel due to reduced memory I0 cost. However, the computation still remains quadratic with sequence length. Flash Attention 2 (Dao, 2024) improves the work partitioning by parallelizing the attention computation over the sequence dimension in addition to the naive parallelization over batch and head dimension. Flash Attention 3 (Shah et al., 2024) leverages new hardware features of recent GPU generations (e.g. NVIDIA Hopper GPUs) such as FP8 precision exploiting asynchrony of Tensor cores and Tensor Memory Accelerators (TMA) to speed up Flash Attention.
TFLA is 10-aware and parallelizes over one sequence dimension of the intra-chunk QKT matrix as the second level of sequence parallelism. New hardware features will also speed up future TFLA implementations.

---

Tiled Flash Linear Attention
Flash Linear Attention
Flash Attention
+
Intra chunk
QK
full sequence in tlles -iiited sed dim
QK'
inter chunk intra chunk turatti chaniic
inter chunk
intra chunk churk NLEw
re-weighting
Figure 7: Tiled Flash Linear Attention (TFLA) combines Flash Linear Attention and Flash Attention 2. Flash Attention (middle) employs tiling across the head dimensions (green & light blue), while Flash Attention (right) uses tiling in the sequence dimension (orange & purple). Tiled Flash Linear Attention (left) combines both tiling strategies.
Flash Linear Attention. Flash Linear Attention (FLA) (Yang et al., 2024b; Yang Zhang, 2024) makes use of the fact that linear attention can be interpreted as linear RNN (Katharopoulos et al., 2020). It then leverages the chunkwise- parialll formulation of linear RNNs (Hua et al., 2022; Sun et al., 2023) for efficient kernel implementations, that process the sequence in chunks. More specifically, Yang et al. (2024b) propose two FLA variants: version that materializes intermediate states in HBM and non-materialization version. The materialization version consists of two kernels: The first is a recurrent kernel that materializes the first intermediate states of every chunk. The second kernel then processes all chunks in parallel and computes the outputs within the chunks. The non-materialization version was proposed concurrently by Qin et al. (2024a) and does not employ parallelism the sequence dimension, but processes the inputs sequentially in chunks.
TFLA uses the idea of chunking of the sequence for the first level of sequence parallelism.
Application of TFLA to other Linear RNNs. While TFLA is designed for efficient mLSTM kernels (Beck al., 2024), its formulation also extends naturally other linear RNNs such as RetNet, Mamba 2, and DeltaNet (Sun et al., 2023, 2024; Dao & Gu, 2024; Yang al., 2024a). A more detailed discussion of these extensions is provided Appendix
Conclusion and Future Work
With Tiled Flash Linear Attention (TFLA) we introduce algorithm for linear RNN and mLSTM kernels with two levels of sequence parallelism. Our TFLA kernels for the mLSTM with exponential input gate (mLSTMexp) achieve state-of-the art kernel execution speeds, while remaining flexible to trade off GPU memory consumption and runtime. To further improve kernel runtimes, we propose mLSTMsig, mLSTM variant with sigmoid input gate, that reduces computation and increases speed. Our experiments show that both mLSTM variants perform equally well on language modeling.
Although we enhance training stability through careful gate initialization informed by our empirical transfer behavior analysis, future work could explore instabilities arising from numerical errors in kernel implementations in greater depth. Finally, the programming techniques and hardware features used to optimize Flash Attention (Shah et al., 2024) could also be applied to our TFLA algorithm to approach peak performance next-generation hardware, as suggested by theoretical runtime analysis. This establishes TFLA as a strong candidate for a foundational primitive in future long-context language models.
Acknowledgements
We thank Sebastian Böck, Richard Kurle, Patrick Blies, Antonio Orvieto, Maximilian Stadler, Thomas Schmied, Kajetan Schweighofer and Sebastian Lehner for helpful discussions and feedback.
References
Alharthi, M. and Mahmood, A. XISTMTime: Long-term time series forecasting with xLSTM. AI, 5 488-1482- 2024. ISSN 2673-2688. doi: 10.3390/ai5030011. URL. https: http://wwppooo..o com/ 2673-2688/5/3/71,
10

---

Alkin, B., Beck, M., Pöppel, K., Hochreiter, S., and Brandstetter, J. Vision-LSTM: xLSTM as generic vision backbone. International Conference on Learning Representations, 2025. URL https://openreviee met/forum?id-SSitDDDi
Austin, J., Douglas, S., Frostig, R., Levskaya, A., Chen, C., Vikram, S., Lebron, F., Choy, P.. Ramasesh, V., Webson, A., and Pope, R. How to scale your model. 2025. Retrieved from
Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer Normalization. arXiv, 1607.06450, 2016. URL https:/ /apps://axxivv org/abs//tt/oooooooooooogo abs/ 011.116677.06800 06450.
Beck, M., Pöppel, K., Spanring, M., Auer, A., Prudnikova, O., Kopp, M., Klambauer, G., Brandstetter, J., and Hochreiter, S. xI.STM: Extended long short-term memory. In Thirty-eighuh Conference on Neural Information Processing Systems, 2024. URL /iitts://arkt/e org/ hat/bb//2400 240524255..0 04517.
Beck, M., Pöppel, K., Lippe, P., Kurle, R., Blies, P. M., Klambauer, G., Böck, S.. and Hochreiter, XLSTM 7b: A recurrent LLM for fast and efficient inference. In Forty-second International Con- ference on Machine Learning, 2025. URL https::/opeereeeeee wwwpeerrviieeeeetforriiiiiiiioee d3DDddeees....
Black, S., Biderman, S., Hallahan, E., Anthony, Q., Gao, L.,, Golding, I.., He, H., Leahy, C., McDonell, K., Phang. J., Pieler, M., Prashanth, U. S., Purohit, S., Reynolds, L.., Tow, J., Wang. B., and Weinbach, S. GPT-NeoX-200 An open-source autoregressive language model. In Proceedings of the ACL Workshop on Challenges & Perspectives Creating Large Language Models, 2022. URL https: /tappii//axxv...///// org/ abss 2204 06745.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http: //gi thub. jomm/jaxmm/jax
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P. Sastry, G., Askell, A., et al. Language models are Few-Shot Learners. Advances neural information processing systems, 33:1877-1901, 2020.
Dao, T. FlashAttention-2 22: Faster Attention with Better Parallelism and Work Partitioning. In The Twelfth International Conference Learning Representations, 2024. URL https://openreew net/forum?id 2nn.frrnmmm.in2222xxh......
Dao, T. and Gu, A. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality, Forty-first International Conference on Machine Learning, 2024. URL https: //t::://ooo /pwp/reveeee.... www.nnt/oorrmmiidzzeeeem
td.
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, FlashAttention: Fast and memory-efficient exact attention with 1O-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Fu, D. Y., Kumbong, H., Nguyen, E., and Ré, C. FlashFFTConv: Efficient convolutions for long sequences with tensor cores. In International Conference on Learning Representations, 2024.
Gholami, A., Yao, Z., Kim, S., Hooper, C., Mahoney, M. W., and Keutzer, K. Al and Memory Wall. URL
IEEE Micro, 44(03))33-39, May 2024. ISSN 1937-4143. 433. doi: 10.1009/MM.20443300 pttps://doi.ieee ieeecomputersociety org////.orr/eeee
1 0 1109/MM. 2024. 3373763
Grazzi, R., Siems, J., Franke, J. K., Zela, A., Hutter, F., and Pontil, M. Unlocking state-tracking lin- ear RNNs through negative eigenvalues. In International Conference on Learning Representations, 2025. URL
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. In International Conference on Learning Representations, 2024. URL. httpttt///oonneeeeeeee 1 fq0507H.
Heek, J., Levskaya, A., Oliver, A., Ritter, M., Rondepierre, B., Steiner, A., and van Zee, M. Flax: neural network library and ecosystem for JAX, 2024. URL http: //gt thub. cww/cooo//o//o/oo comm/gogle//fla..

---

Hochreiter, S. Untersuchungen zu dynamischen neuronalen Netzen. Master's thesis, Technische Universität München, 1991.
Hochreiter, S. and Schmidhuber, Long short-term memory. Neural Computation, 98)):1335-1780 1997.
Hua, W., Dai, Z.,, Liu, H., and Le, Q. Transformer quality in linear time. In Chaudhuri, K., Jegelka, S., Song. L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learming Research, pp. 9099-9117. PMLR, 17-23 Jul 2022. URL https: / http:////occceeess.e/eee/.../see mlr. www.ssmww..........eeeeeeee html.
Jia, X., Donat, A., Huang, X., Zhao, X., Blessing, D., Zhou, H., Wang, H. A., Zhang, H., Wang, Q.. Lioutikov, R., and Neumann, G. X-IL: Exploring the design space of imitation learning policies. arXiv, 12330, 2025. URL httts::/aarke org/ aaaaabbs/22..// 12330.
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Proceedings of the International Conference on Machine Learning, 2020.
Katsch, T. GoteLoop: Fully data-controlled linear recurrence for sequence modeling. ArXiv, 2311.01927, 2023.
Kong, Y., Wang, Z. Nie, Y., Zhou, T., Zohren, S., Liang, Y., Sun, P., and Wen, Q. Unlocking the power of Istm for long term time series forecasting. arXiv, 2408.10006 2025. URL https: / //arxiv.org orrlv..org/aao
10006.
Kraus, M., Divo, F. Dhami, D. S., and Kersting, K. xlstm-mixer: Multivariate time series forecasting by mixing via scalar memories. arXiv, 2410.16928, 2024. URL httpp:://axii.org/aaassss 16928.
Li Fang, A., Smyrnis, G., Ivgi, Jordan, M., Gadre, S., Bansal, H., Guha, E., Keh, S., Arora, K., Garg, S., Xin, R., Muennighoff, N., Heckel, R., Mercat, J., Chen, M., Gururangan, S., Wortsman, M., Albalak, A., Bitton, Y., Nezhurina, M., Abbas, A., Hsieh, C.-Y,, Ghosh, D., Gardner, J., Kilian, M., Zhang, H., Shao, R., Pratt, S., Sanyal, S., Ilharco, G., Daras, G., Marathe, K., Gokaslan, A., Zhang. J.. Chandu, K., Nguyen, T. Vasiljevic. I., Kakade, S., Song. S., Sanghavi, S., Faghri, F., Oh, S.. Zettlemoyer, L., Lo, K., El-Nouby, A., Pouransari, H., Toshev, A., Wang, S., Groeneveld, D., Soldaini, L., Koh, P. W.. Jitsev, J., Kollar, T., Dimakis, A. G., Carmon, Y., Dave, A., Schmidt, L., and Shankar, V. Datacomp-lit In search of the next generation of training sets for language models. arXiv, 1476.11794, 2024. URL ht /ttp///axvvvv//e org/ abs 2406 11794.
Lieber. O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedigos, I., Safahi, E., Meirom, S., Belinkov, Y., Shalev- Shwartz, S., Abend, O., Alon, R., Asida, T., Bergman, A., Glozman, R., Gokhman, M., Manevich, A., Ratner, N., Rozen, N., Shwartz, E.. Zusman, M., and Shoham, Y. Jamba: A hybrid transformer-mamba language model. arXiv, 2403. 1183..9888 1987, 2024. URL //atps://arxiv.... org/abs/ 19887.
Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. URL
Milakov, M. and Gimelshein, N. Online normalizer calculation for softmax. ArXiv, 1805.02867, 2018.
MiniMax, Li, A., Gong, B., Yang. B., Shan, B., Liu, C.,, Zhu, C., Zhang, C., Guo, C.,, Chen, D., Li. D., Jiao, E., Li, G., Zhang, G., Sun, H., Dong, H., Zhu, J.. Zhuang, J.. Song, J., Zhu, J., Han, J., J., Jie, Xu, J., Yan, J., Zhang, K., Xiao, K., Kang, K., Han, L., Wang, L., Yu, L., Feng, L.. Zheng, L., Chai, L., Xing, L., Ju, M., Chi, M., Zhang. M., Huang. P., Niu, P., Li, P., Zhao, P., Yang. Q. Xu, Q. Wang, Q. Wang. Q., Li, Q., Leng. R., Shi, S., Yu, S., Li, S., Zhu, S., Huang, T., Liang,
T., Sun, W. Sun, W., Cheng. W., Li, W., Song, X., Su, X., Han, X., Zhang, X., Hou, X., Min, X.,
Zou, X., Shen, X., Gong, Y., Zhu, Y., Zhou, Y., Zhong Y., Hu, Y., Fan, Y., Yu, Y., Yang, Y., Li,
Y., Huang, Y., Li, Y., Huang. Y. Xu, Y., Mao, Y., Li, Z., Li, Z., Tao, Z., Ying, Z.,, Cong, Z., Qin, Z., Fan, Z.,, Yu, Z., Jiang, Z., and Wu, Z. MiniMax-01, Scaling foundation models with lightning attention. Xrviv, 2501.08313, 2025. URL
3.10//b/2220..08111.. 08313.
12

---

Movahedi, S., Sarnthein, F., Cirone, N. M., and Orvieto, A. Fixed-point mns: From diagonal to dense in few iterations. arXiv, 2503.10799, 2025. URL http:://arxvvv.rr//aaaaaaaee
10799.
Ogata, K. Modern control engineering. Prentice- Hail electrical engineering series. Instrumentation and controls series. Prentice-Hall, Boston, 5th ed edition, 2010. ISBN 978-0-13-615673-44
Orvieto, A., Smith, S. I.., Gu, A., Fernando, A., Gulcehre, C., Pascanu, R., and De, S. Resurrecting recurrent neural networks for long sequences. arXiv, 2303.06349, 2023. URL https://axxiv. org/abs/2300 06349.
Paszke, A., Gross, S., Massa, F,, Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, I.., Desmaison, A., Köpf, A., Yang, E., Vitoo Z.., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. Pytorch: An imperative style, high-performance org.vvorr//aaaoe www.//////////o///////////ae abs 191//911222 922.29202 01703.
deep learning library. arXiv, 19 1912.01703. 2019. URL https: http///axvvvvrrreeee
Peng, B., Goldstein, D., Anthony, Q., Albalak, A., Alcaide, E., Biderman, S., Cheah, E., Du, X., Ferdinan, T., Hou, H., Kazienko, P., GV, K. K., Kocoń, Koptyra, B., Krishna, Jr., R. M., Lin, J., Muennighoff, N., Obeid, F. Saito, A., Song. G., Tu, H., Wirawan, C., Wožniak, S., Zhang, R., Zhao, B., Zhao, Q. Zhou, P., Zhu, J., and Zhu, R.-J. Eagle and finch: Rwky with matrix-valuet states and dynamic recurrence. arXiv, 2404.05892, 2024. URL https: /ttp///axivvvv//a////////eee/aaa .rg abs//404......./////// 05892.
Peng, B., Zhang, R., Goldstein, D., Alcaide, E., Hou, H., Lu, Merrill, W., Song, G., Tan, K.. Utpala, S., Wilce, N., Wind, J. S., Wu, T., Wuttke, D., and Zhou-Zheng, C. Rwkv-7 "goose" with expressive dynamic state evolution. arXiv, 2503.14456, 2025. URL wwtph:ttss////iii org/ abs//503.. 14456.
Pöppel, K., Beck, M., and Hochreiter, S. FlashRNN: 1/O-aware optimization of traditional RNNs modern hardware. In The Thirteenth International Conference on Learning Representations, 2025. net/forumRidd-02ZITPPP
URL https:://pppnt//
Qin, Z., Han, X., Sun, W., Li, D., Kong, L.,, Barnes, N., and Zhong. Y. The devil in linear transformer. In Goldberg, Y., Kozareva, Z., and Zhang. Y. (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 7025-7041, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10. 00.18333v1/2022. emnlp-main.4733 URL http::::///aaaaaholoogooooo org 2022. www.p-main..1777.. 473/.
Qin, Z., Sun, W., Li, D., Shen, X., Sun, W., and Zhong, Y. Lightning attention-2: A free lunch for handling unlimited sequence lengths in large language models. arXiv, 2401.04658, 2024a. URL org/ abss/4011 04658.
https://arxias
Qin, Z., Yang, S., Sun, W. Shen, X., Li, D., Sun, W., and Zhong, Y. HGRN2: Gated linear RNNs with state expansion. In First Conference on Language Modeling, 2024b. URL https: //openreview heteeeeww.eeeeeeeeeeeeeeeeeeeeee food/foorummdidySaa.ICCC..
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, guage models are unsupervised multitask learners. OpenAl, 2019. URL https: //cdn.openal www/beetteeeeeeeeeeeeee -anggua warpagge-mooelss.e htteemmmmdeesseeaeeeeeeeeeeee hwwpe//aaaaaaaeeellleeeoeee uww..eriiseeeeiieddde multitask_learners pdf.
Lan-
Schlag, I., Irie, K., and Schmidhuber, J. Linear transformers are secretly fast weight programmers. In International Conference Machine Learning. 2021. URL https://proceeddings mlr v1 v139/schlag21a a. html.
ptwpeii.n........e//////.. press/
Schmidinger, N., Schneckenreiter, L., Scidl, P., Schimunek, J., Hoedt, P.-J., Brandstetter, J., Mayr, A., Luukkonen, S., Hochreiter, S., and Klambauer, G. Bio-xLSTM: Generative modeling, representa- tion and in-context learning of biological and chemical sequences. In International Conference on Learning Representations, 2025. URL https
Schmied, T., Adler, T., Patil, V., Beck, M., Pöppel, K., Brandstetter, J., Klambauer, G., Pascanu, R., and Hochreiter, S. large recurrent action model: xLSTM enables fast inference for robotics tasks. arXiv,
2025, URL /ttps::/arxxv..rrg///e
22391.
Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. 2024. URL https://arxii..o wwtp:ewwwwwwi....../oo//ee/peeee
13

---

Shazeer, N. GLU variants improve transformer. arXiv, 2002.05202, 2020. URL https:///rxivy org/abs/2002. 002./20...........2.2
Siems, J., Carstensen, T., Zela, A.. Hutter, F., Pontil, M., and Grazzi, R. Deltaproduct: Increasing the expressivity of deltanet through products of houscholders. JCLR 2025 Workshop on Foundation Models in the Wild, 2025. URL https: http///ppeeeeeeeeeeeeeeeeeeeeeet htt.:eete//ooommidnnneeeeeee
Spector, B. F., Arora, S., Singhal, A., Fu, D. Y., and Ré, C. ThunderKittens: Simple, fast, and adorable ai kernels. arXiv, 2410.20399, 2024. URL https://arxxvv.rrg///e
Sun, Y., Dong, L.,, Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. Retentive network: A successor to transformer for large language models. ArXiv, 2307.08621, 2023.
Sun, Y., Dong, L., Zhu, Y., Huang, S., Wang, W. Ma, S., Zhang, Q. Wang. J., and Wei, F. You Only Cache Once: Decoder-dccoder architectures for language models. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL forum?1d=251oww576rr
net/
Team, G. Gemma Improving open language models at practical size. arXiv, 2408.001 18, 2024.
URL https://arxiv www.v.org//bbo. swrrg/bbs/4400 001110000111 00118.
Tillet, P. Triton, 2024. URL https:
Tillet, P., Kung, H. T., and Cox, D. Triton: 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, pp. 10-19, New York. NY, USA, 2019. Association for Computing Machinery. ISBN 978 77814503661999 196.. doi: 10.1145/331550. 3329973. URL https:///oi...... 0rg/10.11// 43/13355588 3329973.
Touvron, H., Lavril, T., Izacard, G.. Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E.. Azhar, F., Rodriguez, A., Joulin, A., Grave, E.. and Lample, G. Llama: Open and efficient foundation language models. arXiv, 2302.13971, 2023a. URL https::/arxiv.oor abs/2302. 3971.
Touvron, H., Martin, L.. Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I., Korenev, A., Koura, S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I., Nic, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, F. M., Subramanian, R., Tan, X. E., Tang, B., Taylor, R., Williams, A., Kuan, X., Xu, P., Yan, 7., Zarov, I., Zhang, Y., Fan, A., Kambadur, M., Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. I.lama 2: Open foundation and fine-tuned chat models. arXiv, 2307.09288, 2023b. URL https: /ttps://axxvvvvvv
abs 2307 23...0922222 09288.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.., Gomez, N., Kaiser, L. u., and Polosukhin, I. Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances Neural Information Processing Systems, volume 30, 2017. URL
org/ abs/ 7701.03762....
Vikram, S., Jones, C., and Lebedev, S. Jax-triton, 2022. URL jax-triton.
Waleffe, R., Byeon, W. Riach, D., Norick, B., Korthikanti, V., Dao, T., Gu, A., Hatamizadeh, A.,
Singh, S., Narayanan, D., Kuishreshtha, G., Singh, V., Casper, J., Kautz, J., Shoeybi, M., and Catanzaro, B. An empirical study of mamba-based language models. arXiv. 2406.07887, 2024. URL https://arxii www.oogg///b/aeeeeoooooeee 2406 07887.
Williams, S., Waterman, A., and Patterson, D. Roofline: an insightful visual performance model for multicore architectures. Commun. ACM, 5264))65-76, April 2009. ISSN 0001-0782. doi: 10. 1 4445/44974655... 1498785. URL
org/ 01145/1497665 111866.1499789 1498785.
14

---

Xiong, R., Yang, Y., He. D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., and Liu, On layer normalization in the transformer architecture. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 10524-10533. PMLR, 13-18 Jul 2020. URL https://proceed/s
mlr. www...rees/vviv. wwtsss/v11////// 19/xiong2xx........ html.
Yang. S. and Zhang, Y. FLA: A Triton-based library for hardware-efficien implementations of linear attention mechanism. January 2024. URL https://gi whwbbbttttth...... con/susscconglin/ flash-1inear-attentee
Yang. S., Kautz, and Hatamizadeh, Gated delta networks: Improving mamba2 with delta rule. arXiv, 2412.06464, 2024a. URL https://aaxe: www./oorggoooooo ht/g/////41..../.// 06464.
Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated linear attention transformers with Kardware-efficien training. In Forty-first International Conference on Machine Learning. 2024b. URL /ttps://openrevis/ net nett/forum2dd:s555xffti
Yang. S., Wang, B., Zhang, Y., Shen, Y., and Kim, Y. Parallelizing linear transformers with the delta rule over sequence length. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024c. URL https: http:///ppeeeeeeeeeeeeeeee eveeer met//orum?id=yy8RAA
Zhang, B. and Sennrich, R. Root mean square layer normalization. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-But -uch F., Fox, E., and Garnett, R. (eds.), Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https://aamee org/ /bs/ 007477.
Zhu, J., Chen, X., He. K., LeCun, Y., and Liu, Z. Transformers without normalization. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 14901-14911, 2025.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: provide extensive experiments the main paper and the appendix.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
We point out efficiency limitations of our Triton kernel implementations recent hardware generations in the theoretical runtime analysis and also report runs with numerical problems.
Justification:
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: There are no Theorems such in the paper.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide a thorough description of our experimental setup as well as Pseudocode for our algorithm in the Appendix. We will also release the software library for our kernel implementations.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: Our code will made publicly available.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All details provided in the Appendix, well as the accompanying code.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Our language modeling experiments were too expensive to run multiple seeds.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Answer: [Yes]
16

---

Justification: We state that our experiments were conducted on NVIDIA H100 GPUs.
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https:
Answer: [Yes]
Justification: We conform with the code of ethics.
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We provide impact statement in the appendix, just before the references.
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper does not pose such risks.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All software licenses are respected.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[Yes]
Justification: The code that will be released with the paper contains proper documentation.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing or research with human subjects 15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration not required.
Answer: [NA]
Justification: The core method development this research does not involve LLMs as any important, original, or non-standard components.
17

---

Appendix
TABLE OF CONTENTS
Appendix A Related Work
19
A.1. Other Related Work
19
A.2. Application of TFLA to other Linear RNNs
19
Appendix Extended mLSTM Formulations
20
B.1. Fully Parallel Formulation
20
B.2. Detailed Chunkwise-Paae Formulation
21
B.3. Chunkwise-Paralne Parallel Backward Pass
Appendix C Extended Tiled Flash Linear Attention
27
C.1. GPU Fundamentals
27
C.2. Tiled Computation
28
C.3. TFLA Forward Pass
28
C.4. TFLA Backward Pass
30
Appendix D Extenned mLSTM with Sigmoid Input Gate
31
D.1. Stabilization of the Exponential Input Gate
D.2. Empirical Transfer Behavior Analysis of the mLSTM
32
Appendix Extended Experiments
33
E. Numerical Validation of TFLA Kernels
33
E.2. Extended Language Modeling Experiments with mLSTM
35
E.3. Extended Kernel Benchmark
40
Appendix F- FLOP and Memory Operation Counts for the mLSTM
44
F.I. Exact FL.PPs and Memory Operations Count for the mLSTM
44
F.2. Simplified FLOOP Count Summary for the mLSTM
46
F.3. FLOP Comparison between mLSTMexp and mLSTMsig
46
F.4. FLOP Comparison between different mLSTM Formulations
47
F.5. FLOP-Optimal Chunk Size for mLSTM
48
Appendix G Theoretical Runtime Analysis for TFLA mLSTM Kernels
49
G.1. Theoretical Runtime
49
G.2. Arithmetic Intensity
G.3. Roofline Analysis
52
G.4. Runtime-Optimal Chunk Size
54
G.5. FLOP and Theoretical Runtime Analysis Summary
56
18

---

A Related Work
A.1 Other Related Work
Other Hardware-Aware Optimizations. Optimized, hardware-awarc implementations enable the exploration of new primitives or new model architectures. FlashRNN Pöppel et al. (2025) introduces framework of IO-aware optimized CUDA kernels in order to simplify research on traditional, non eonppaallelizable RNNs. Mamba (Gu & Dao, 2024) enables large scale language modeling ex- periments (Waleffe al., 2024) with an efficient parallel scan algorithm in their optimized CUDA kernels. FlashFFTConv (Fu et al., 2024) provides efficient implementations for FFT convolutions for modern hardware by reducing I0 and leveraging specialized matrix multiply units. DeltaNet Yang et al. (2024c,a) introduces an efficient algorithm for training linear Transformers with the delta rule (DeltaNet) (Schlag et al., 2021), which enables to scale up DeltaNet to standard language modeling settings.
Our TFLA kernel algorithm provides effective method to balance the runtime and memory for linear RNN kernels based on their chunkwise- thunkwsse-parallel formulation, paving the way to even larger model training setups, e.g. for multi-billion parameter xLSTM models (Beck et al., 2025).
Gating mechanisms for Linear RNNs. Many different gating techniques for linear RNNs have been explored (Sun al., 2023; Beck et al., 2024; Yang et al., 2024b; Gu & Dao, 2024; Dao & Gu, 2024; Sun et al., 2024; Qin et al., 2024b; Orvieto et al., 2023; Katsch, 2023; Peng et al., 2024). We propose mLSTMsig, a variant of mLSTM with a sigmoid input gate similar the original LSTM (Hochreiter, 1991; Hochreiter & Schmidhuber, 1997) and empirically analyze the transfer behavior.
mLSTM Applications. The mLSTM, which is the xLSTM with matrix memory, has already been adopted in several domains beyond language modeling. For example, Alkin et al. (2025) find that the mLSTM can serve generic backbone for computer vision architectures. In the field of robotics, the mLSTM architecture has been explored as a large recurring action model, which uses the efficient thear-time inference complexity of the mLSTM (Schmied et al., 2025), and the backbone for imitation learning policies (Jia et al., 2025). In the domain of biological and chemical sequences, the mLSTM has been studied for generative modeling, representation, and in-context learning (Schmidinger et al., 2025). Finally, ML.TM has also been examined in the domain time series forecasting (Alharthi & Mahmood, 2024), where sLSTM (Beck et al., 2024) has also been applied (Kraus et al., 2024; Kong et al., 2025).
These and other applications will benefit from our TFLA kernels, which provide efffccintt and scalable implementations of the mLSTM, that can be easily integrated into existing models. We believe that our TFLA mLSTM kernels will increase the adoption of mLSTM in other application areas well.
A.2 Application of TFLA other Linear RNNs
We have applied TFLA to linear RNNs with scalar headwise gates such as mLSTMexp and mL- STMsig. In this section, show how TFLA could be applied other Linear RNNs, but leave the implementation to future work.
Linear RNNs with Scalar Headwise Gates. With minimal modifications, TFLA can be applied to other scalar headwise gated linear RNNs, such as, for example, Retention (Sun et al., 2023), Gated Retention (Sun et al., 2024), Simple GLA (Yang et al., 2024b) or Mamba (Dao & Gu, 2024). In general, TFLA implementations of these linear RNNs can be obtained by modifying the forget and input gate parameterizations (of the TFLA mLSTMsig). Since neither of the aforementioned models has input gate in the mLSTM sense, fix the input gate of TFLA mLSTMsig to one (i.e. (ii)) 1). Then, for Simple GLA or Retention we keep the headwise scalar sigmoid forget gate (f.) or set it constant decay parameter o (f)) r.
To implement Gated Retention - which introduces a sigmoid forget gate modulated by a temperature parameter 7 we modify the activation function of the forget gate to incorporate this temperature term. Beyond coupling the forget gate parameter A with the keys k; (represented by the parameter B) through the step size parameter A, Mamba 2 also adopts the linear attention structure from
19

---

Equation (12). In this formulation, Mamba's 's notation C, B, and X corresponds to our queries Q. keys K, and values V, respectively.
Linear RNNs with Delta Rule. Recently, Linear RNNs with non-diagonal state transition matrices have become popular due their increased expressivity that enable them solve synthetic state tracking tasks (Grazzi et al., 2025; Siems et al., 2025; Peng et al., 2025; Movahedi et al., 2025). One method to implement such non-diagonal linear RNNs efficiently is (Gated) DeltaNet (Yang et al., 2024c,a), which introduces a hardware efficient algorithm for linear transformers with the delta rule. The core of DeltaNet's efficient implementation is the ceannwwsse-ppaallle formulation of the delta-rule, which is implemented using primitives from Flash Linear Attention.
In our notation, the chunkwise-paree parailllppaaalllll formulation of DeltaNet is given by the inter-chunk recurrence
Ck (due xdhoo
Ck-1 + (dut x ddww)
K(K)T (dox bic))
(ik) (
W(( (lio xdge) (dit xdico)
CK-1
(17)
V-(k)
and the combination between inter-chunk and intra-chunk contribution
H (Lho xdne)
(8)
Ck-1
+ (Uing (1 xdgi) (dat x bxwo)
Q(k)
K(R)"
U (I xdno)
W()) (Ieo xdxx) (dge xdhn)
Ck-1
(Ting xdgi) (dge xdhe)
V-(k)
(18)
where omit the mask M applied to QK" of DeltaNet for clarity. We refer Yang et al. (2024c) for the definition of the matrices (k) and W(k)
By defining the new values V*(k) Uk) Cx-1
recover the core formula of TFLA
for the intra-chunk parallel forward pass (see Equation (12)) to which we can apply our TFLA tiling strategy from Section 3. For computing the matrices U1), W()) and the new values V*(k) one
could follow the same strategy as in DeltaNet, where (k) and (k(k are computed in a separate kernel before the recurrent kernel. Then, the new values V*(k) are computed and stored HBM in the recurrent kernel together with the memory states Cs.
We conclude that TFLA can be applied any linear RNN that either follows can be reformulated into the simplified, chunkwise-paralldd form of Equation (12).
B Extended mLSTM Formulations
B.1 Fully Parallel Formulation
For the parallel formulation it is assumed that all inputs are available at once. Then, the queries. keys and values qt. ke, U, can be stacked into the matrices Q, K E T'xxs V E RTxxa, in order to compute all hidden states H e Txd parallel using the following equations:
D log F +
(19)
m- = max Dij,
(20)
D exp m)
(21)
= 1
QKT
(22)
v dyk
S @ D
(23)
n = max ( 1/ exp( (es)
(24)
H = (5 o (n (n-1)) V.
(25)
20

---

where 1 E R² is a vector of ones. The logarithmic forget gate activation matrix log F E RTXT is computed by
00
for i <
log Fij
0
for =
(26)
log
[l==++1
( (1))=
=
logo
fk
for >
Similarly, the input gate pre-activation matrix RTxT is given by
00 for i < j I,
(27)
i for i >
Note that in contrast the recurrent formulation, in the parallel formulation the states C are not materialized, i.e. computed explicitly. This comes at the cost of computing the quadratic matrices D,S RTxT, with an overall quadratic scaling sequence length
B.2 Detailed Chunkwise-Paralle Formulation
In this section, we provide more detailed formulas for the thunkwise-paralle formulation of the mLSTM from Section 2.2.
Chunkwise Gates. Given the logarithmic forget gates ((()
log o logo (b) E R and input gates
(k) log exp(i (8) E R&,, compute the logarithmic chunkwise gates
L E R,
sum (())) =
Σ
(k)
(28)
1=1
by cumsum (p(h)
L, with bky Σ ((K) for j 1,2,..., L
(29)
1
Ag rev rev_cumsum (F:)) + 7(0) RR,, with akj Σ
((())
(() + 1)
for j = 1,2,..... L.
(30)
where [1:] denotes (in numpy notation) that the first index is excluded as it is done in the sum notation of Equation (30). Additionally, in Figure we illustrate the chunkwise gate computation and show a Py Porch code snippet for computing the chunkwise gates in Figure 8
21

---

def
vecl: torch. Tensor, (B, NH, S) vecF: torch. Tensor, # (B, NH, 5) chunk_size: int,
):
B, NH, vecI. shape assert x chunk_size 8, f"S=(S) not divisible by chunk_size=( chunk_size)"
_ddvice = vecI device
NC = S // chunk_size
chunk_size
11
12.
# compute vecB
vecF _logil logsigmoid (vecF to(dtype=torch.Floes
vecF llgglogggcuunnnm chunked rearrange (vecF_logsig, "b nh (nc nh
1", nc=NC, 1=L)
15
vecB = vecf sec__logsig_chunker cumsum (dim==1))
compute vecA vecI_chunked rearrange (vecI, "b nh (nc 1) b nh 1", nc=NC, 1=L
# unstable vecA computation:
# vecA (vecB[.... -1, None] vecB) vecI # (B, NH, NC, L)
stable vecA computation:
22
vecA
23
torch.cate [
24
1:3. (-1) c-msumm-1))) (-)) torch zeros ((B, NH, NC, 1), Revice=_devices dtype=torch. float32),
26
]. dim=-1,
29
)
30
vecI_chunker ) (B, NH, NC, L) return vecB, vecA
31
Figure 8: PyTorch function to compute the logarithmic chunkwise gates for mLSTMexp.
Inter-chunk Recurrent Contribution.
The inter-chunk recurrence given by
m
(inter)
max
+ m max au:
(inter)
(31)
C₁ cxp
[k + m
(inter)
m
(inter)
Ck-1 +
exp
aa
m
(inter)
(k)
(32)
(inter)
(inter)
(inter)
3 K(()) T
n4 exp
ge + m k-
m
nk-1 1 + exp
[a
mi
(33)
simplified form we can write the inter-chunk recurrence as
C₁ = Ea C + a4 @ k T (())
- C+- + K (() ((k)
(34)
n₂ = E n + ag c k(()]
= g + (k) (k)
(35)
with the running max state integrated into the gates.
22

---

chunkwise recurrent:
fully recurrent:
C
C
ag a k(te
Vily
C Cc R,n,
h C. 9:
Ch-1
chunk size L=4
b
b
by
t
1
2
3
5
6
7
8
9
10
11
12
Co
C1
C2
C3
a1
a3
gi
82
g3 QKV
QKV
K (2)
G
f Is + fg + +
+ f2 f + f2 fs f +
fo + fio to + ho + fo + fun + far + f12
log bj
log by
log by
Es
fa + fo + fr
fg
12 + 1
fe + fr + fu
to fio + fiz
loga;
fo + fa
+
12
log a2
fg + fg
16
log ag
fu + fix f12
is
111
$12
loggi
G + la + Lg +
lo = Ls
[g + b fg
log = [g + Go + fn + f12
Figure 9: Illustration of the chunkwise gate computation.
Intra-chunk Parallel Contribution. contribution given by
The recurrent part is followed by the intra-chunk parallel
(k)
-00
for i < j
D
(36)
by
bj i
for > j
(intra) m %
D (k)
max
(37)
,
D (k) cxp(DD (A)
(intra)
m
(38)
8(k)
Q(h) K (R)
(39)
vdgh
((() = ((k) o (K)
(40)
where exp is acting component-wise.
Output computation. with the inter-chunk recurrent part H (k)
The contributions from the intra-chunk parallel part (())) are combined
inter to obtain the hidden states H (8) for each chunk & (see
23

---

Figure 1):
m (combine)
max
by
mk-1
(inter) (intra)
m,
(41)
Q(k)
H [k]
exp
b&
m m.
(inter)
(combine)
Ck-1
(42)
ddz
Q(A)
Ck-1
(43)
Ck-1
(44)
H intra
[k)
S
(k)
v(k)
(45)
[bg ② /QQ()//O/ /dbb))))) Ck-1
(10) V(a)
(k)
(46)
max
(b& O (Q(K)/ / /db))))) n +5 (k)
11, exp
mm
(combine)
CR-1+ + 5 (H) ((k)
(47)
T + 5(h) 1/, exp I
(combine)
max
(k) Ca-1 S (k) v(())/ (o
denom
(b)
(48)
B.3 Chunkwise-Parallel Backward Pass
In this section we provide a detailed description of the backward pass of the chunkwise- parallel mLSTM.
Gradients Through Normalizer States. Following Sun al. (2023), do not compute the gradients through the normalizer states n. The gradients cancel out due to the Layer- or RMS-Norm on the mLSTM cell hidden states H, since the normalizer state is constant over the embedding or feature dimension, which is the normalization dimension.
Inter-chunk Recurrent Backward Pass. Given the incoming memory cell state gradients from the next chunk 8C2 and the hidden state output gradients ( for chunk k, we can compute the inter-chunk recurrent backward pass. The query, key and value gradients _Q(() (8) 5 (k) and ő (V) of the inter-chunk recurrent part computed by:
inter
inter
inter
(k)
H
8H
(49)
(k)
denom
(A))) inter
K
(k)
8Ck
(50)
8K
(k)
(k)
SC
(51)
KK inter
(k)
SK (N)
@ Ch 1T
(52)
(k)
8Q
8H
CL-
(53)
(M)
8Q inter
sQ
b
1
(54)
ddg
The memory cell state gradients Sc have incoming contributions from the next timestep (ree)
and output / They are given as
8C = (rc)) k-1 + +5C/ (aut)
(Tc)
(55)
k-1
= =g O 8C2 + Q (k)
5H (k)
(56)
24

---

Finally, we can compute the cumulative gate gradients 8g., dag and ob, for chunk as
9g, 1 (Ck-1 8C
(57)
бg* = òg, O EK
(58)
adg = ( K (M)
@ K (k))
(59)
σαg sαk a
(60)
8b = (oQ (k)
(61)
dak
8b obg bk-
(62)
Intra-chunk Parallel Backward Pass.
Given the mLSTM hidden state output gradients 8H (k) the
intra chunk query, key and value gradients sQ (k)
5K({) and VV*(k) gradients computed by
8H (M)
(R)
8H
(63)
(k)
denom
S(R)
Q (1)
(() T
(64)
Vdgh
5 (K))
O (()
(65)
8V (k)
S (4)
jH (N)
(66)
intra
(S(A) HH (R)
(67)
8S (K) SS
(M)
(68)
8Q
(4)
(((k)
K(()
(69)
intm
Vdgk
k intra
(K(R) T
(k)
(70)
dyk
In order to compute the cumulative intra gate gradients, we compute the gradients through the gate matrix D (R), which is computed from the cumulative forget gates
b cumsum ((()) E R
(71)
6/
(Aww cumsum
(())
L
(72)
kv
where we use the logarithmic forget gates log (f). We denote the dimensions as I.g and L for the query and key-value dimensions, respectively. Omitting the masking operation, we compute the gate matrix as
D(k) b(()) 1 1, b (ke)
=(R) T
1. ke
(73)
where 1g E Rhe and 1& Lhk are vectors of ones used to indicate broadcast operations, and ((R)
kg E IR liv are the logarithmic input gates for chunk k.
The gradients are computed as
5D (k) is
S())
(74)
5b
8D
1
(75)
6b (kv)
2D (4:) 1g
(76)
8D (R)
dik
(77)
25

---

Combined input and gate gradients. The intra and inter chunk gradients are combined by sum-
ming up the contributions. This yields for the query, key and value gradients
(Q 8Q ttto + 8Q (A)
(78)
8K (h) = áv
inter + ò K (k)
(79)
intra
dv VV()) ò v
(80)
inter
intra
The input and forget gate gradients (k) and F(() can be computed from the cumulative gate gradients
og., ôbg and dag with the following equalities
(F()) 8g*
(81)
op(h) = rev rev_ rev_eunsuum(bbi
(82)
68 (A)
rev_ "vv_eumsum(faa)
(83)
61 (k) Sak
(84)
26

---

C Extended Tiled Flash Linear Attention
C.1 GPU Fundamentals
We review the GPU fundamentals for writing efficient kemels. Since we perform our experiments on NVIDIA GPUs, our review is targeted towards NVIDIA's terminology, though the principles also apply to other hardware. For a more extensive overview we refer to (Spector et al., 2024).
GPU Overview. GPU (Graphies Processing Unit) is specialized processor designed to efficiently handle large-scale parallel computation tasks, such as matrix multiplications in neural networks. These tasks are divided into small programs called kernels, that are executed on GPUs. A kernel loads data from high bandwidth memory (HBM), performs work on it, and writes the results back to HBM. For writing efficient kernels, it important to understand the software hierarchy of the GPU, which closely follows its physical hardware hierarchy.
GPU Hierarchy. At the lowest level the GPU runs multiple Threads, operating on small but fast register memory in parallel. On the software side usually multiple (e.g. 32) Threads are grouped together into Warps. Again, multiple Warps grouped into Thread blocks which together execute kernel on a physical core, called streaming multiprocessor (SM). Warps or Threads within the same Thread block can communicate data through special on-chip shared memory (SRAM). When executing a kernel, grid (with typically 3 dimensions) of Thread blocks that run in parallel is launched on the GPU. All Thread blocks have access to the large but slow off-chip high-bandwidre memory (HBM), which has both the largest latency and least bandwidth of all GPU memories. For efficient kernels is important to minimize memory read and writes from and to HBM.
Specialized Compute Units. Modern GPUs have specialized compute units called tensor cores
that accelerate matrix multiplications on GPUs. Tensor cores have most of the GPU compute and are accessed at the warp or block level. For efficient kernels it is important to maximize lensor core utilization.
Triton Language. Triton is a GPU kernel programming language with an associated compiler, that provides bythen- environment for GPU programming. The user can load data from HBM via a t1. load instruction and store data to HBM via store. t1. dot is an instruction, that leverages tensor cores for matrix multiplications. While this Triton interface of increases productivity in writing very fast custom kernels, peak performance can be achieved sometimes only with CUDA kernels. We write our kernels in Triton and leave CUDA implementation for future work. In contrast to NVIDIAs programming model CUDA, which provides access to all levels of the GPU hierarchy, Triton programs operate the Thread block level and hide register and thread management from the user. Therefore, we describe TFLA on the more abstract Thread block or program level in the following section.
27

---

C.2 Tiled Computation
For the tiled computation of the intra-chunk hidden state contribution Haura within chunk, we consider blocks of the matrix S = ( s(2) and the gate matrix D (D (
with ( (4) R x BLKe. Here, the superscript denotes the block index along the
Liv dimension (and not the chunk index). Similarly, consider blocks of the value matrix
viu with v(i) I Bea Basv, We then accumulate the unnormalized hidden state blocks va
intra,num E BBAs (.LAuxBaug Bans and the corresponding normalizer (()) E BLku as
m max D
(85)
(()) (S
c ex (1)
m )))
(86)
(1) (S())
c exp((D
(1) -m ))) (
(87)
intra.mum
m (2) max
m max ij
(2)
(88)
(2) = exp(m (1) m (3)) (0) + (S
exp D (2)
? 1
(89)
H (2)
exp(m (1) m (2) intra.uum
(1)
+ SS (2)
exp (2)
m () (
(90)
intra.mum
After computing this intra-chunk part, we need to do one more rescaling step to combine the intra- chunk and inter-chunk parts of the hidden state output H**) since H
( and H (R) were computed
with different max states. Therefore, we compute the final hidden state output H ( as
m
(combinc)
max
b2 + mK-1 m
(inter) (2)
(91)
(k) Ck-1 + exp
m (2) - m k
(combine)
5
(()
(b)
(92)
max Q
n2 1 exp
m (2)
m
(combine)
[((21
exp
(-m
m
(combine)
where we assume that m (2) the block maximum and ( the normalizer after the last BLkU block of the intra-chunk computation for chunk k.
C.3 TFLA Forward Pass
For notational simplicity we drop the & index for the query, key and value matrices Q Xdokk K E R Xdge and V E L xd., respectively. We make use of reweighting (as discussed in Appendix C.2) in order keep track of the maximum value over the gate matrix tiles, similar to (Dao et al., 2022).
The forward pass algorithm of TFLA for one thread block is described in Algorithm 1.
Note that the loop in line 27 of Algorithm is the same as the loop in line In both loops we load the same blocks of the matrix Q. Fusing these loops would avoid loading this data twice. Unfortunately. fusing these loops efficiently is problematic due the online computation of the maximum mold
(combine) and b on the final mnew (see and mew in the loop in line 4 and the dependence of m Appendix D.I and C.2).
We address this issue in Section 4 by modifying the input gate of the MLSTM.
28

---

Algorithm TFLA Intra-Chunk Forward Pass for mLSTMMxp (H Kernel)
Require:
Matrices E Rlina K R Lix Xdak, V E R Blauxdnee
States CR-1 E R xd,, 1 E R.ok..
Input- and cumulative forget gate vectors ig, by RLaag Block sizes Bagk. Baho. Brng and BL where Br.nq d BLKU- Block index Vig and Block HV index idha-
1: Initialize mold; m.ew E L to -o in SRAM.
& Compute intra-chunk contribution
2: Initialize accumulators H E IR R2Ay Bao and n (intra) E Renu SRAM.
3: Load by² E IR BLAu from HBM to SRAM.
4: for i = 1 to (1-BBi do
BLAD
5: Initialize accumulator ( IR RgBBagxBua in SRAM. dgt
6: for j= 1 to Bäok
do
7: Load Q() E
BENg Back and K() E BRke x Baak from HBM to SRAM.
8: Accumulate S(i) Q g()
9:
end for
Load b( ERBoNo and (kv)
l/
E BERU from HBM to SRAM.
11:
Compute (())
(1)
b/
+ (kv) E BRaiggBBum
if BLKU 2 itg Blkg then
Apply causal mask to D
13:
14:
end
if
Compute Menw maximum (Mold, rowmax D
(9) }.
15:
Compute ( C (8)
16:
-mnw).
Compute S
(
17:
digh
Load V() E RRB x
Bans for Block dhe from HBM to SRAM.
19:
Accumulate Hinm exp(Motd Moew)) Hintm S V.
20:
Accumulate n (intra) exp(m (Moid mnew) (intra) +5 1.
(1)
Update mold m
22:
end for
Compute inter-chunk contribution (inter) R from HBM to SRAM.
23: Load m
24: Compute m
(sombiine)
maximum
{ (g)
(b(])
+ m mnew J-
(inter)
25: Compute D exp (b(a)
(inter)
(combine)
m m
26: Initialize accumulators Hinter Brag x Bas for Block iabo and (inter) Brag in SRAM. Note: This is the same loop as the inner one above. They cannot be merged because of the max state computation.
27:
for to
d,y
do
Bagk
28:
Load Q(1)
R
BLag Bank and
(
1-1
Bagi Barv for Block iabo from HBM to SRAM.
Compute Q
Q
(
b
dit
30:
Accumulate Hineer
k-1-
(
31: Load n RR Bayk
(
32:
Accumulate n (inter) Compute H (comb)
(1)
n
(
33: end for
Combine inter- and intra-chunk contributions
H + exp( Mmn
combine
H Inter-
35: Compute n (sumb) n (intra)
exppmnew
combine
(inter)
36: Compute (k)
H (omb)
max In (epp( ? (ommin))
37: Store H (k) n (comb) and m
(combinc) to HBM.
29

---

C.4 TFLA Backward Pass
For the TFLA backward pass, we need compute the gradients of the queries, keys and values (Q(h), (k) and (V()), Omitting the gate computations and normalization, write simplified
version of these gradients as
(Q (k) (Log (dok)
(H (Fruy ddi)) (dho Lae))
V
K(k) (Fov Xdel))
+
CK-1 (Te xdiv) (dyo xdge)
oH (k)
(93)
6Q
8Q
(k(k) (Trer xdux)
V(k)
(H(
Q(A) (kno xdgz)
V(k)
òCk
(94)
(Tike xdua)) (dhv x Enne)
(Toke xdiv) (dno xdgl)
(V (Jee xdhoo
K(k)
Q(A)
(H( (Ting Xdno)
8C (Tec xdgn) (dys xdao)
K(k)
(95)
(Jov xdgut) (dug x Lino)
8V
(R)
100
see that each of the query, key and value gradients has a similar structure as the forward pass in Equation (12). They be computed with the same work partitioning scheme, where parallelize over the outer chunk size and outer embedding dimension the matrix multiplications and loop over the inner dimensions, respectively. For example, for the key gradients (k) we parallelize over the outer chunk size Live and the outer embedding dimension dgk and loop over the inner dimensions
Lig and dac. Table summarizes the TFLA work partitioning scheme for the forward and backward pass kernels.
30

---

D Extended mLSTM with Sigmoid Input Gate
D.1 Stabilization of the Exponential Input Gate
In this section we show how the exponential input gate is stabilized with the max state me (Beck et al., 2024). The stabilization is based on the idea Safe Softmax (Milakov & Gimelshein, 2018). We will see that the max state stabilization ensures that the argument of the exponential input gate activation always smaller than 1. We will also see that the normalizer state guarantees cancellation of the max state, so that the overall outputs of the mLSTM remain unaffected by the max state.
Without stabilization mLSTM hidden state output is computed as
C q: max [|n grl 1
h or ?
(96)
where we omit the scaling factor / dgk for To simplify we also omit the lower bound and the absolute value on the dot product in the denominator. We obtain
qt
h4 = a (öt)
O n, q:
(97)
Inserting the update formulas for the memory cell state C, and the normalizer state n, gives
(o(fi) CL-1 + exp(i) ko U. (o(ft) 7 + exp(i.) k; q:
qi
h; a (ō.)
(98)
We now show that from this unstabilized version of the mLSTM three steps. At first we use the identity oii) = exp(log(a(fi))) extend the fraction in Equation (98) by exp(-m;) and select m, max(log(o(f))) in} to be the maximum of the two arguments of the exponential function. This gives
derive the stabilized form in
(exp(lo((a(T))) m.) C.-1 exp(ie mi) Ko V.
C. q: exp(-m:) h; = o (ō:)
q
(ör) Ø n, q: exp(-mi)
(cxp(log(o (f:) mm) n4-1 + exp(ie m,) ki
(99)
In this way, we ensure that the arguments of the exponential function are always smaller than 1, such that numerical overflow due large values can never occur.
As next step reparameterize C₁ and n, to C₁ and 7o.
Ö C. exp( exp(-m) + Č C.-1
+ C.-1 = Č,-1 exp(mi-1) 74-1 n exp(-m4-1) $ 74-1 7د cxp(m=-1)
(100)
n. n4 cxp(-m:)
Finally, we replace C and n, with the stabilized states Ö2 and nh in the recurrence. We arrive at
(xxp(log(o(E))) + m mm) Č,-1 + exp(ie - m.) Ke v.." T
h = a (ō,) o
exp((lg((((.)) m4-1 mi) ñ4-1 explie mz) k.) q: (101)
C qt
= a (oz
n4 q:
Now we choose the max state as m, max (ax(((g(a(f))) + m4-1, id and arrive the stabilized mLSTM formulas by changing the denominator to max {| qul, exp(m)-1)). We have add exp(m)-1) also to the right side of the maximum, so that it cancels out.
To summarize, we see that the normalizer necessary for the max state to cancel out and the exponential input gate argument is bounded through the max state.
31

---

D.2 Empirical Transfer Behavior Analysis of the mLSTM
We provide details on the transfer behavior analysis of mLSTMMxp and mLSTMsig in Section 4.2.
Experiment Setup. We analyze the transfer behavior of the mLSTM for single head and a single input sequence of length T =512. The inputs are for the queries, keys and values qr. k2 and 8, are sampled from the standard normal distribution (O, 1). We set the head dimensions to dg-- =128 and d,. =128. As norm layer (a)M(()) we use the RMS-nome Changing the norm to layernorm does not alter the results, as for this experiment we set the mean of the inputs to zero. For every plot we measure the gains Ghefore and Gatter (as defined in (16)) for imput and forget gate preactivation values in the ranges [-12, 8] and [-5, 12], respectively.
Effect of Normalization Layer Epsilon on Transfer Behavior. Based on our analysis on the normalization layer after the gated linear RNN operation in Section 4.2, we hypothesize that the normalization layer and especially the norm epsilon this experiment, probe the effect of the epsilon value the transfer behavior of the mLSTM. Figure 10a and Figure I la show the transfer behavior of mLSTMexp and mLSTMsig for e =Ile-2, le-6, 1e-8), respectively.
integral the gating mechanism. In
We observe that the epsilon acts in the same way for mLSTMexp and mLSTMsig. Increasing e causes an offset of the gain in positive y-direction, increasing in negative y-direction. We set our default value e 1e-6, which yields the best performance in our experiments (see Sec. 5.1).
Normalizers of mLSTMexp and mLSTMsig, In this experiment, we test the effect of different normalizers n in Equation 24 for mLSTMexp and mLSTMsig. The parallel formulation in Section B. is presented for the mLSTM with exponential input gate, but applies similarly to the mLSTM with sigmoid input gate. For the default ml.STMsgg, we set the normalizer to n = and modify the calculation of the gate matrix D for sigmoid input gates.
In Figure 10, show the results of different normalizers for the mLSTM with exponential input gate. Only the default mLSTMexp with correct normalizer and max state (in Fig. 10a) shows a transfer behavior that depends on the input gate.
In contrast, in Figure 1 la and 1 llb we observe that incorporating a normalizer similar to mLSTMMxpp (excluding the max state) into mLSTMsig does not alter its transfer behavior.
The other two normalizer variants for mLSTMsig in Figure 1e and lld show clearly different transfer behavior and do not train successfully. Similarly, the variants in Figure 10b and 10c also fail to train successfully.
In summary, we find that if the mLSTM exhibits the characteristic gate dependent transfer behavior it trains successfully and shows good performance in our language modeling experiments. In order to achieve this behavior for the mLSTMexp we need to normalize correctly as derived in Section D. Adding normalizer the mLSTMsig does not change performance and transfer behavior, if the normalizer incorporates lower bound on the dot- dot-product n, qi- However, our default mLSTMsig omits the normalizer in order to reduce computational cost and runtime.
32

---

Gain
Gain EPS-1e:02
Norm,
EPS-1e.08
Input Gate Preacivation
50
18
25
0.0
-25
1.0
-3.0
0.6
-7.5
0.0
5.0
-5.0-2.5 5.0
-5.0-2.000 2.5
10.0
Forget Presctivation
Forget Preactivation
Forpe: Preactivation
Forget Gate Preactivation
(a) mLSTMexp - n max
1/ (exp((mm) (default)
Gain before Norm
Gain after Norm, EP5-1e-02
Gain after Norm, EPS- 1e-06
Gain after Norm, EPS=1e-08
7.5
Input Gate Preactivation
5.-
1.6
2.5
0.0
-2.
1.0
-3.0
04
-00.0
-5.0-2.5 0.0 25 5.0 7.5 10.0
-3.0-2.3 0.0 2.5 5.0 7.5 10.0
-5.0-2.5 0.0 2.5 5.0 7.5 10.0
-5.0-2.5 0.0. 25 5.0 7.5 10.0
Forget Gate Preactivation
Forget Gate Preactivation
Forge: Gate Preactivation
Forget Gate Preactivation
(b) mLSTMexp n max
1|, 1) (wrong m state)
Gain befora Norm
Gain after Norm, EPS-1e.0 1a-02
Gain after Norm, EPS-1e-06
Gain after Narm, EPS=1e-08
7.5
Input Cate Preactvattoo
5.0-
14
-2.5
0.8
7.5
-10.0
0.2
-5.0-2.5 0.0 2.5 5.0 10.0
-3.0-2.5
2.5
10.0
-5.0-2.5 0.0 5.0 10.0
-5.0-2.5 0.0 2.5
10.0
Forget Gate Preactivation
Forget Gate Preactivation
Forget Gate Preactvation
Forget Gate Preactivation
(c) mLSTMexp--1 - n (no normalizer)
Figure Transfer behavior of the mLSTM with exponential input gate for different normalization layer epsilons (EPS) and different normalizer variants. Only the default normalization shows the input gate dependent transfer behavior. Varying the normalization layer epsilon causes a shift of the gain curve in y-direction.
E Extended Experiments
In this section, we provide additional experiments and details Section 5.
E.1 Numerical Validation of TFLA Kernels
Before we begin experiments langauge modeling, we first verify that our kernels yield the same result as a reference implementation in pure JAX based on the fully parallel formulation (see Appendix B.1).
Validation Perplexity Match (Table We compare the validation perplexity at the end of training for 160M parameter mLSTMexp and mLSTMsig models trained 19B tokens. We context length 4096 since the parallel JAX implementation go oft-of-memory for longer contexts. Model architecture and training recipe follows general setup described in Appendix E.2.
In Table we confirm that our kernels yield the same results as our reference implementation JAX.
33

---

Gain before
Gain EPS-1e-02
after 1e-06
EPS=1e-08
Gate Preactivatio
2.5
0.0
-2.5
-5.0
0.6
-7.5
0.0
5.0
-5.0
10.0
2.5 0.0
2.5
5.0
10.0
-5.0-2.5 0.0 2.5 5.0 7.5 10.0
Preactivation
Preactivation
(a) mLg
n2 1 (default)
Gein betore Norm
Gein after Norm, 12
Gein after Norm, EPS-1e-00
Gain after Norm, EPS=1e
7.5-
Cate Peectivvtion
14
-2.
-5.0
0.5
-75
-00.0
-5.0-2.5
Forget Gate Preactivation
Forget Gate Preactivation
Forget Gate Preactivation
Forget Preactivation
(b) mLSTMsig
max (|S 1)
(same normalizer mL.TMMxpp)
Preacivvtion
5.5
25
0.0-
3.6
-23
-5.0
18
-73
-5.0 0.0
5.0
5.0-2.5 10.0
-5.0-2.5 2.5 5.0 -5.0-2.5 2.5
10.0
Forget Preactivation
Preactivation
Preactivation
(c) mLSTMsig - n |S 1| (no max different colorbar)
Gein before Norm
Gein after Norm, EPS=1e-02
after Norm, EPS=1e-06
Gain after Norm, EPS=1e-00
Input Cate Preactvation
5.0
4.8
25
0.0
-25
3.0
-5.0
12
10.0
-5.0-2.5 0.0 25 5.0 7.5 10.0
-5.0-2.5 0.0 255 5.0 75 10.0
-5.0-2.5 0.0 25 5. 75 10.0
-5.0-2.5 0.0 2.5 5. 7.5 10.0
Forgat Gate Preactivation
Forqet Gate Preactivation
Force: Gate Preactivatio
Forget Gate Preactivation
(d) mLSTMsig - n S (no | different colorbar)
Figure 11: Transfer behavior of the mLSTM with sigmoid input gate for different normalization layer epsilons (EPS) and different normalizer variants. Removing the normalizer from mLSTMsig (which is default setting in (a)) has effect the transfer behavior. If the normalizer is added, should be bounded by (see (b)). Varying the normalization layer epsilon causes a shift of the gain curve in y-direction.
Table 3: Validation Perplexity for 160M parameter models at context length 4096 trained on 19B tokens.
EXP
SIG
Heads
JAX PARALLEL
LIMIT chunk
XL chunk
JAX PARALLEL
XL chunk
6
21.02
21.03 21.03
21.18 21.07
21.01 21.02
21.05 21.06
12
21.01
34

---

E.2 Extended Language Modeling Experiments with mLSTM
In this section we provide details on our experiment setup, model architecture and training recipe and add additional performance results on context length 8192 as well as analyze the effect of the epsilon parameter in the norm layer.
Software and Hardware Setup. We run our language modeling experiments in JAX 0.4.34 (Brad- bury et al., 2018) and use FLAX 0.9.0 (Heek et al., 2024) implement our models. We implement our kernels in Triton 3.1.0 (Tillet et al.. 2019; Tillet, 2024) and use JAX-Triton 0.2.0 (Vikram et al., 2022) integrate the kernels into JAX. Our kernel benchmark experiments are run in PyTorch 2.5.1 (Paszke et al., 2019), because most kernel baselines are available in PyTorch. All experiments are run on NVIDIA H100 80GB GPUs.
Model Architecture. The model architecture for MLSTeexpp and MLSMssg follows the design of most dense Transformer decoder only large language models (Radford et al., 2019; Brown et al., 2020; Touvron et al., 2023a,b).
An embedding layer, is followed by stack of blocks and language model head that produces the output logits (i.e. the values before softmax), which typically consists of a normalization layer and linear (unembedding) layer. We apply logit soft-capping (Team, 2024), such that the value of the logits stay between -c and e for a specific cap value c. We choose c =30. The logits are capped with the following function:
softeap(x) = c- c-tanh(a/cc)
(102)
We use the GPT-NeoX tokenizer (Black et al., 2022) with vocabulary size 50257 and do not tie the weights for the embedding layers and and the last (unembedding) layer.
Each block consists of two layers, where each layer has skip connection and a normalization layer before the layer input (i.e. we use the pre-norm block architecture). As normalization layer we use the R.MS-nor (Zhang & Sennrich, 2019) with epsilon e = =le-6.
The first layer is a sequence-mix layer, that mixes the tokens along the sequence or time dimension. For standard Transformers this is the Attention operation (Vaswani et al., 2017). In our case, we replace Attention by the mLSTM operation with exponential or sigmoid input gate, Similar to Attention, mLSTM processes each token in multiple parallel heads. The second layer in the block feedforward linear layer that mixes the tokens per timestep channelwise. We use the SwiGLU feedforward linear layers (Shazeer, 2020; Touvron et al., 2023a).
For the mLSTM we set the head dimension for the queries and keys to be half of the values, i.c. dyp = 0.5 dn.. We use Layernorm (Ba et al., 2016) as NORM(x)) operation with epsilon =1e- in our experiments. 2 We apply soft-capping from equation (102) on the input and forget gate preactivations, as we found that this improves training stability. For the gate preactivations we set =15
We provide the remaining model parameters in Table 4.
Training Recipe. We train our models with the AdamW optimizer (Loshchilov & Hutter, 2019) with 31 009, 32 -095 and c =1e-8. We use learning rates and batch sizes as specified Table 4. We apply a weight decay of 0.1 all linear layers (including the last linear layer or unembedding) and exclude biases and the token embeddings from weight decay. We clip the gradient norm at 0.5. We use cosine learning rate scheduler with a linear warmup for the first 750 steps and decay to 0.1 of the peak learning rate, followed by a linear cooldown to 0 for the last 1000 steps. We list the number of training steps for every model size in Table 4. During pre-training we ensure that no information is leaked across document borders by resetting the memory states at the beginning of each new document. We implement this by manually setting the forget gate preactivations to large negative values at the beginning of each new document.
Additional Performance Results (Table 6). In Table we show the validation perplexity for MLSMMexp and mLSTM for context length 8192 (the results for context length 4096 are shown in
We confirmed empirically that the type of normalization layer does not affect the performance as well as our qualitative results on transfer behavior and gradient norm variance. Therefore, we generally prefer RMS-norm faster.
35

---

Table Training and Model Architecture Hyperparameters for our model sizes 160M, 400M and
1.4B.
Model Size
Emheeding Dim
Head Dim
Batch Size
Trkens 4k ctx
Tokens
Blocks
Heads
LR
Steps
8k ctx
6
128
160M
12
768
3c-3
128
36k
19B
38B
12
64
8
128
400M
24
1024
le-3
128
46k
24B
48B
16
64
4
512
1.4B
24
2048
8
256
8e-4
256
31k
33B
65B
16
128
Table 2). For some head dimension configurations we observed irrecoverable gradient norm spikes during training (indicated by -).
Table 5: Validation Perplexity at context length 8192. EXP and SIG denote mLSTMexp and mLSTMsig. LIMIT and XL correspond limi t_ chunk and x1 ch_uk kernels. indicates that the run experienced irrecoverable loss spikes during training.
EXP LIMIT
EXP
SIG
Size
Tokens
Heads
Llama
XL
XL
20.29
20.43
20.46
160M
38B
19.99
20.31
20.42
20.52
15.91
16.01
16.08
400M
48B
16
16.05
15.95
16.01
4
12.69
12.71
12.91
1.4B
65B
12.62
12.65
12.67
16
12.97
12.59
12.75
Table 6: Training Step Time in seconds at context length 8192. EXP and SIG denote MLSMMexp and mLSTMsig. LIMIT and XL correspond limi c__chnk and x1_ chunk kernels. All runs the same training setup with global batch size of 256 sharded across 32 GPUs corresponding to local batch size 8.
Batch Size
EXP
EXP
SIG
Size
Heads
Llama
Global Local
LIMIT
XL
XL
4
1.77
1.73
1.4B
256
1.66
16
2,39
1.60
1.61
1.62
Effect of Trainable Input Gate (Table 7). We investigate the effect of the input gate the performance. Table 7 shows that having the input gate learnable consistently improves performance for both mLSTMexp and mLSTMsig.
Effect of Input Gate Bias Initialization (Figure 12 and 13), In our transfer behavior analysis in Section 4.2 we find that there is transition from suppressing the signal to passing the signal at weights of the gates Wii.1 to 0, the biases of the input and forget gates determine the actual position in the x-y plane in the beginning of training. Initially, with input gate biases initialized to O, we observe a high gradient norm variance, which was more pronounced for mLSTMsig (see Figure 12a and 13a).
negative input gate values around -8 (see Figure 4). Since we initialize
36

---

Table Validation Perplexity for 160M mLSTMs at context length 4096 with learnable and fixed input gate (bias initialized at -00).
EXP LIMIT
SIG XL
Inpur Gate
Fixed
21.23
21.24
Learnable
20.95
21.04
Therefore, we test to initialize the input gate biases at larger negative values. The forget gate biases initialized equally spaced in the range [3.6]. As the weights w(i.f) grow during training, so do the gate preactivations and the model could learn to gradually move into the dynamical region of Figure 4, where the input signal is passed.
Indeed, as we observe in Figure 12 and 13 initializing the input gate biases to 100 effectively mitigates gradient norm spikes and reduces high gradient norm variance during training for both mLSTMexp and Msig.. We therefore conclude that the additional input gate not only improves performance (see Table 7), but also improves training stability, if initialized correctly.
We use the limit chunk kernel for mLSTMexp and our xl_chunk kernel for mLSTMsig and confirm that we obtain the same behavior with the __ccknk kernel for ml m..TTMep.
gnd tgal Mesiee te Bocd 57556 1910 sudteal
beee Foom berd 1SP wrreed
30
25
53
35
3.4
3.2
3.1
254
30%
154
Grain/ grad, momm ST88.
Urait/ www.d
E-s gen lite Eags sirad Nic endin arcend saiignn Dww.i traetiine
Sestiie- 20-3 gen Saate-2 Tibel LoMkeviin food
102
MLSTMexp with Trainable Input Gate
(b) ml STMxxp with Fixed Input Gate
Figure 12: Trainable and fixed exponential input gate for bias initializations [0, -2, -5, -10] and norm epsilon =le-6.
Effect of Normalization Layer Epsilon on Performance (Figure 14). In our empirical transfer behavior analysis of the mLSTM Section 4.2 and D.2 we find that the transfer behavior depends on the input and forget gate preactivations, as well as the normalization layer epsilon (see Fig- ure 10a and la). Therefore, we perform a grid search over different normalization layer epsilons and input gate bias initializations for the mLSTM with exponential input gate with 160M parameters and 6 heads at context length 4096. We show the results in Figure 14.
We observe that there is diagonal region from norm layer epsilon and input gate bias (e, by (1)=(1 -10) to (le-4, -5) with improved performance. This indicate that if we increase the norm layer epsilon can or should also increase the input gate bias initialization, as the shift of the gain curve in positive
37

---

Eoskcn dwddd Jusa
Stilloegg gra Stdd srivn tanniigng Fasd
30
14
12
244
33
34
11
15c
204
35%
120 330 22.
traini-grad
133
(a) mLSTMsig with Trainable Input Gate
(b) mLSTMsig with Fixed Input Gate
Figure 13: Trainable and fixed sigmoid input gate for bias initializations [0, -2, -5, -10] and norm epsilon =le-6.
-2 21.25
21.32
21.33
21.22
23.30
21.35
Input Gate Bias Init
21.30
Aaliddat m Perplexity
21.29
21.28
21.15
23.22
21.25
21.20
-7
21.29
21.23
21.16
21.02
22.73
21.15
21.10
-10-
21.21
21.04
21.14
21.31
22.13
21.05
21.00
1e: 08
1e: 06
1e:05
1e:04
1e:02
Norm EP5
Figure 14: Validation Perplexity of mLSTMexp with 160M parameters with heads. Grid search over norm layer epsilon and input gate bias initialization. The diagonal region of improved performance indicates, that there exists an interplay between the norm layer epsilon and input gate bias initialization. This supports the hypothesis that the norm layer is important for the gating mechanism.
y-direction for larger epsilons in Figure 10a suggests. This supports our hypothesis in Section 4.2, that the norm layer important for the gating mechanism.
We use (c, (1)=(1-6,, -10) as our default configuration.
Input Gate Activations over Training (Figure 15). We show the maximum input gate pre- activations (maximum over batch, sequence and head dimension) over training for mLSTMexp and SLSTMsig with 160M parameters in Figure 15. Both models have the input gate bias initialized to -10.
38

---

Layer
Laper1
Layer2
TECO STOEE
393M miSTiin
10N ST.TTAA
TEEM WELTTUee
1220 STLTTM
35
15
455
30k
33%
20C
224
354
-.5
-15
Zrep
Step
Layer
Layer?
Layer
IGCM 1LS1
302M miSTMeq
10N HESTRMR
1000 15
1220 mLSSMMa Mai
12
11k
36k
10k
20c
10k
30k
Sep
Step
Dee
Figure 15: Maximum input gate pre-activation values īn over training for mL.SMexp and ml STMsig with 160M parameters. Maximum taken over batch, sequence and head dimension. Both models have the input gate bias initialized to -10. most cases the input gate pre-activations remain below zero.
39

---

E.3 Extended Kernel Benchmark
In this section, we provide details on our benchmark setup and add additional benchmark results.
Details on GPU Memory Measurement. by the kernels. For this, we use the PyTorch torch.cuda.ma. aaaddamax_mmmmoo memooryalllocaeed peak memory allocated during one kernel iteration. We make sure that the memory statistics are reset after each iteration and that the PyTorch caches are cleared before the start of each benchmark.
Figure and 16 we measure the GPU memory used alllocated API to measure the
Details on the Runtime Benchmark (Figure 5). In our TFLA kernel runtime benchmark in Section 5.2, Figure 5 we report the median runtime of 30 iterations, after 10 warmup iterations in milliseconds. We run all kernels in 16 precision.
We use the standard embedding dimension of 4096 for 7B Transformer models for our benchmark. Since different models and kernels have different default input sizes at this embedding dimension, we adapt the head dimension, number of heads and remaining input dimensions for each kernel accordingly. Following the practice of Shah et al. (2024) we keep the number of tokens constant at 65,536 and vary the sequence length (i.e. = [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536)) and batch size accordingly (i.e. Nbuch 65536/T).
We benchmark the following mLSTM kernels:
mLSTMexp (FLA limit samit_con t_miikk::hh::: Our own baseline kernel for the mLSTM with exponen- tial input gate with limited chunk size based on FLA. Similar to FLA this kernel employs only single level sequence parallelism across chunks. We report the best performing chunk size of 64. The chunk size of 128 would still fit in SRAM, but is considerably slower.
MLSTMexp (TFLA x1__chunk] TFLA kernel for the ml STM with exponential input gate with two levels of sequence parallelism. We set the chunk size to the best performing chunk size of 128.
mLSTMsig (TFLA xl_chunk): TFLA kernel for the mLSTM with sigmoid input gate. We set the chunk size to 128, but find chunk size 256 to perform equally well in terms of runtime (see Fig. 16 and 6).
For all mLSTM kernels for the values. Similar to GLA (Yang et al., 2024b), we set the query and key head dimension to dgk= dn./2, i.. dgk= 128.
16 heads, which results in head dimension dhe 4096/ /16 256
We compare our mLSTM kernels with the following baselines:
Torch FlashAttention:
PyTorch 2.5.1 implementation of FlashAttention 2.
Accessed via SDPBackend. FLASHLATTENTION
cuDNN Attentio::o NVIDIA cuDNN implementation of FlashAttention 2 integrated
in Py Tyrchh222 2.5.
Accessed via SDPBackend. BUTTIICCCNNNNAANNTTIIII
FlashAttention 3: FlashAttention 3 implementation", which has been optimized for NVIDIA H100 GPUs (Shah et al., 2024).
GLA (FLA): Gated Linear Attention Triton kernel based on the Flash Linear Attention algorithm with one level of sequence parallelism (Yang et al., 2024b). Implementation from the official FLA repository, version 0.1
Simple GLA (FLA): A simple version of GLA with scalar forget gates per head. This primitive is not published implementation for kernels for RetNet (Sun et al., 2023) or Mamba 2 (Dao & Gu, 2024) in the FLA library Yang & Zhang (2024). Moreover, Simple GLA is similar to mLSTMsig, but has no input gate. Therefore, we find interesting to add as baseline. We use the implementation from the official FLA repository, version 0.
new sequence modeling primitive but serves as reference
See
'See https:://githbb htwpth/..w.aaaalllaaaaaaaaaattte
Seee https://gi wwwggh.t...... wwtp:///ww.oooooggaaa---aaaaaaae
40

---

Mamba: Mamba CUDA kernel Gu & Dao (2024). Implementation from the official Mamba repository, version 2.2.4.
Mamba 2: Mamba 2 Triton kernels Dao & Gu (2024). Implementation from the official Mamba repository, version 2.2.4.
For all FlashAttention baselines we use 32 heads with head dimension 128 for queries, keys and values. For the Flash Linear Attention (FLA) kernels GLA and Simple GLA, we use the identical head configuration as for our TFLA mLSTM kernels (i.e. 16 heads, Avv = 256, dgk = 128). For Mamba,
embedding dimension of 4096 and set the state dimension 16 similar to Gu & Dao (2024). For Mamba 2, we use their default head dimension of 64 and set the number of heads to 4096/64 = 64. Note that smaller head dimension can yield faster runtimes (see Figure 18).
We show the results of this benchmark for varying sequence length and constant number of tokens in Figure 5. When comparing the forward pass runtime only, we find that Mamba2 and Simple GLA kernels are slightly faster than our mLSTMsig kernels. However, this difference within ms. In training. when forward and backward pass runtime is measured, FlashAttention 3 for longer sequence lengths and more than two times faster than Mamba kernels for all sequence lengths. Only Simple GLA (FLA) can keep up in training speed with our TFLA mLSTM kernels. Therefore, we compare the runtime and memory usage for a larger head dimension in Figure 16 and find that this comes at the cost of almost 2 times the GPU memory usage compared to our TFLA mLSTM kernels. These memory savings are achieved by leveraging a larger chunk size, enabled through the two levels of sequence parallelism outlined in Section 3.
TFLA kernels are faster than
Runtime and Memory Comparison with FLA Kernels (Figure 16). In this experiment we compare the runtime and memory consumption of our TFLA mLSTM kernels with prominent kernels from the Flash Linear Attention library. We use similar setup to our previous benchmark, but perform this comparison with 8 heads at larger head dimension of 512 for the values and 256 for the queries and keys, since both Beck et al. (2024) and Yang et al. (2024b) report better language modeling performance for larger head dimensions.
In addition to GL.A (chunk) and Simple GLA (chunk), we also compare with GLA (fused) which is the non-materializationn version of Gated Linear Attention (GLA) (Yang et al., 2024b).
The non-materializzation version of GLA has been also proposed by Qin et al. (2024a) as Lightning Attention-2 (see also Section A). For the forward pass it fuses the inter- and intra-chunk part of the mhunkwise-paralle Linear Attention formulation (see Section 2.2) and therefore does not materialize the hidden states in GPU memory.
Interestingly, our experiments we find that even though the non-materialization version uses the least GPU memory of all FLA kernels, it is neither faster nor more memory efficient in training than our TFLA mL.TMM kernels (see Figure 16). While Simple GLAA is slightly faster (within 3 ms or 15%), it uses almost twice the GPU memory compared to TFLA mLSTM kernels. The speed of Simple GLA can be partly explained to the fact that it computes less FLOPs (no input-gate) and the fact that in the forward pass the memory cell states are materialized in bf bfloot16 16, while TFLA materializes states in float32 for improved stability, which causes twice the memory 10 per state.
Runtime and Memory Comparison with Lightning trantminggAteention2 Attention2 Kernels (Figure 17). Similar to the previous experiment, we compare the runtime and memory consumption of our TFLA mLSTM kernels with LightningAttention hrbiid large language model MiniMax-01, which combines lightning attention (a linear attention variant with data independent decay) with softmax attention (MiniMax al., 2025). MiniMax-01 is proposed as very efficient long-con mong-contxtt language model, which makes the comparison between LightningAttteniin and our TFLA mLSTM kernels interesting.
(Qin et al., 2024a). LightniggAttennionn2 is the core of the recent
LightningAAtentionn also uses the chunkwise-paralle formulation for linear RNNs (see Section 2.2). However, in contrast to Simple GLA and TFLA it does not split the computation in recurrent and parallel part, but instead processes all chunks fully recurrent (see Section A for more details).
We find that LightningAAttention supports only identical head dimensions for queries, keys and values up to 128. For this reason, we discuss this comparison separately from the other experiments.
See tps::///ttthee com//stat-sppceeeecccc com/sstate-spaces/mambb.
41

---

GLA (chunk)
GLA (fused)
ML5TMe
Me at1024]
mLSTMsia[256]
Simple GLA (chunk)
mLS TMxxp (limit chunk)
mL5TNsin[2048]
mLST MMsi [51g[
mLS MLSTMsin[128]
16
70
14
60
GPU Memory (G8)
Time [ms]
12
50
20
512
1024
2048
4096
8192
16384 32768 65536
512
1024
2048
4096
8192
16384
32768 65536
Sequence Length
Sequence Length
Figure 16: Runtime and Memory Comparison with FLA Kemels. Left: Runtime (Forward Backward Pass). Right: GPU Memory Usage.
We use 8 heads and head dimension of 512 for values, and 256 for queries and keys. Simple GLAA (the fastest FLA kernel in our experiments) is slightly faster than our TFLA mL.STMsig kernels but uses almost twice much GPU memory.
NightnAttnNH664
mLS Mexp (limit chunk) ML557M)gNHH64((56]
MLSIMMiggN:32[4095]
MLSTMsigNH32(102 mLS5MMiggN:32(512
mLSTMssigNH322256 MLSTMSigNH33/112)
LightWttnNH32
mLSTMsigNH33/22048
70
GPU Memory [GB]
8.5
Time [ms]
50
30
7.0
1024
2048
8192 16384
65536
1024 2048
4096
32768
Sequence Length
Sequence Length
Figure 17: Runtime and Memory Comparison with LightningAtttenion2.. Left: Runtime (Forward Backward). Right: GPU Memory. We use 32 and 64 heads with head dimension 128 and 64 for queries, keys and values. Lightning At- tention has the least memory usage of all kernels, but more than times slower than our TFLA mLSTM at the larger head dimension of 128.
We compare our TFLA mLSTM kernels with LightningAttennoon for 32 and 64 heads, corresponding to head dimension 128 and 64. We keep the number of tokens fixed to 65536 and vary sequence length and batch size the same way as above.
We show the results in Figure 17. Since LightningAAttnntnn does not materialize intermediate states, it has the least GPU memory usage with 6.2 GB. However, this GPU memory efficiency comes the cost of a more than 3 times longer runtime compared to our TFLA mLSTMsig kemel with chunk size 256, which uses about 7.3 GB of GPU memory. This highlights that there exists trade-off between GPU memory usage and runtime for linear RNN kernels based on the chunkwise-parallel formulation. Our experiments demonstrate that our TFLA kernel al algorithm provides effective method to balance this trade-off via the chunk size parameter (see Figure 6).
Runtime Benchmark for Varying Head Dimensions (Figure 18). It has been reported several other works that larger head dimensions (compared to common Self-Atrention head dimensions) lead to improved language modeling performance for linear RNNs (Sun al., 2023; Beck al., 2024; Yang et al., 2024b). Consequently, is desirable for linear RNN kernels to be fast and efficient across
42

---

GLA (chunk)
GLA (lused)
MLSTMexp (TPLA XL chunk)
Simple GLA (chunk)
MLSTMexp (limit chunki
mLSTMsig (TFLA XL chunk)
17.5
80
15.0
70
60
Time [ms]
[ms]
10.0
7.5
30
5.0
20
10
32
64
128
256
512
1024
32
128
256
512
1024
Head Dimension
Head Dimension
Figure 18: Head Dimension Benchmark for FLA and TFLA mLSTM kernels. Left: Forward Pass.
Right: Forward and Backward Pass.
We measure the runtime for sequence length 8192 and batch size for different head dimensions. We the same head dimension for queries, keys and values. Our TFLA mLSTM kernels show fast runtimes even for very large head dimensions.
wide range head dimensions. In this experiment, we evaluate whether our new TFLA kernels exhibit this property.
We vary the head dimension from 32 1024 and adapt the number of heads for total embedding dimension of 4096 and measure the runtime for inputs of sequence length 8192 and batch size 4. We the same head dimension for queries, keys and values.
For the FL.A kernels the head dimensions 32 and 64 did not run, due Triton compiler errors. As the FLA library is still being developed at the time of writing this paper, expect this to be fixed soon.
We observe that for small head dimensions (i.e. and 64) mLSTM limit chunk kernel fast as our TFLA mLSTM kernels in training.
In summary, results in Figure 18 confirm that TFLA kernels achieve fast runtimes across wide range of head dimensions.
43

---

FLOP and Memory Operation Counts for the mLSTM
We count the number of floating point operations (FLOPs) and the memory operations (load and stores in bytes) in a forward pass (with batch size 1) of the mLSTM with exponential and sigmoid input gate, We use factor of 2 to describe the multiply accumulate cost for FLOPs.
We do not count FLOPs that belong to recomputation, that happens within kernels. For example, when we parallelize across the embedding dimension the forward kernel H each of the dne/ Bake blocks recomputes the matrix S. Similarly, we do not count the additional memory-loadinn operations that are necessary for the recomputations. During training. we typically have fixed context lengths. Therefore, we do count loading the initial state and storing the final state.
We use factors denoted as Fop to describe the number of FLOPs for operation OP (e.g. Feep for the exponential function). By default, we set all of these factors to 1. We do not neglect these factors, as the impact depends on the selected chunk size in some terms, which can be chosen freely in TFLA. Moreover, might want to do an more fine-gran tineigriieed FLOP analysis, where account for the differences in compute cost of some operations.
We use the factors bytesx to denote the size of each element in the tensor (e.g. byte8 for the query, key and value tensors). Typically during training the queries, keys and values are stored in bfloat16 (i.e. bytes,tv 2), while the memory cell states are kept float (i.e. byteSçam 4). We summarize the notation used in this section in Table 8.
In the remainder this section, we count the FL.OP and memory operation counts for the chunkwise- parallel, fully-parallel and recurrent STT formulations (Section F.1 and F.2), analyze the difference in FLOPs counts between mLSTMexp and mLSTMsig (Section F.3) as well as between the different formulations (Section F.4), and finally compute the FLOP-optimal chunk size for the chunkwise- parallel formulation (Section F.5).
Table 8: Notation for FLOP and Memory Operation Counts.
Symbol
Description
Nbsth
Batch size
Nhead
Number of heads
Nehunk
Number of chunks
T
Sequence length
Chunk size
dno
Head dimension for values and hidden states
dgb
Head dimension for queries and keys
FLOPPs for the operation OP (e.g. exp) Feausal
Fop
Factor that accounts for causality, typically 0.5
bytesx
Number of bytes used for each element in tensor X
F1 Exact FLOPs and Memory Operations Count for the mLSTM
Chunkwise-Parran Parallel Formulation (Table 9, 10). We count the FLOPs (Table 9) and memory oper- ations (Table 10) of the hhunkwise-paraalel mLSTM formulation (see Section 2.2 and Appendix B.2). All counts are for single head and a single chunk. To obtain the total counts for full sequence, multiply these counts by the number heads Nhedd and chunks Nehunk = T/L.
Fully Parallel Formulation (Table 11, 12). We count the FLOPPs (Table 11) and memory operations (Table 12) of the fully parallel mLSTM formulation (see Appendix B. All counts are for a single head and full sequence of length T.. To obtain the total counts, multiply by the number of heads Nhead-
Recurrent Formulation (Table 13, 14). We count the FLOPs (Table 13) and memory operations (Table 14) the recurrent mLSTM formulation (see Section 2.1 and 4.1). For the memory operations we assume that the states are materialized after every timestep, which is the setting during text generation. All counts are for a single head and a single time step. To obtain the total counts for full sequence, multiply by the sequence length T and the number of heads Nacad-
44

---

Table 9: FLOP counts for the chunkwise-paralle mLSTM formulation for MexT and mLg All terms denote the FLOP count per head and chunk.
FLOPs
MLSTMexp
mLSTMsig
Recurrent computation of the inter chunk states
+ 1
21 L LFes Fesp
Gates:
+L11 + Fia) + Fum + Fas
+2L( Fiie)
Numerator:
2dgadho + 21d,, h Ldux
2d tno + 0 2Ldg + Ldue
Denominator:
2dgt + 2Ldgt
Parallel computation of the intra chunk ourputs
Cumulative Forget Gates:
LLL 1) + F(Fig Fug) Feausa x (1' Feg Fmnal L(1 Fam))
1IL 2L(FFin Fug) P Pensel X (1. + Fess))
Gate Matrix:
Intra Outputs:
Feaual x (21 (dge + dnu) + 31.2)
Fenusal x (21 + d22)) 31.)
Parallel computation of the inter chunk outputs
Inter Outputs:
2Ld 3Ldgr
2dd + Ldga
Combination of inter and intra chunk outputs
Output Combination:
2Ldny L(1 Fmxx
Fats
Fesp)
Ldhe
Table 10: Memory operation counts for the chunkwise-parallee mLSTM formulation for mLST-
Mexp and mLSTMsig. All terms denote the memory operation count per head and chunk.
Bytes
mLSTMexp
nLSTssggg
Inter-chunk Recurrent Kernel
Load: L(dgk + dhw) x bytes gke + 2L x bytes: (dyrche + duk + 1) x bytes
L(dg + dnw) x bytes gkv + 2L x bytes;f dgicho bytes
Store:
Intra-chunk Parallel Kernel
L(2dg; + dnu) x byteSgkv 2L x bytes;
L(2dge dno) x bytes skn + bytes;
Load:
+dqx + det + 1) x byteScum
+ +dgadav x bytes
Store:
Ldno x bytes,kr + 2L x bytes
Ldhe x bytes
4L x bytes,f +3 (dho + dgp) x bytes, gkr
4L x bytes,
Total:
+3L (dho + dgp)) x bytes,
+2 (L + dandge dgp + x bytescmn
+ 2dhnolag x bytescmn
Table 11: FLOP counts for the fully parallel mLSTM formulation for mLSTMexp and mLSTMsig. All terms denote the FLOP count for a full sequence per head.
FLOPs
mLSTMexp
mLSTMsig
Cumulative Forget Gates:
Fin)
TTT + 1) 2T + Feig) TC Fes + Faux F Fenos Fenusa x d 21
Gate Matrix:
1 Fep Faux + Fmnsk) Femual (21 +2 Femust (3+ Pan.) TT(Fosp Fonx))
Attention Logits:
Normalization:
Outputs:
Feausll x 2T
Feausal X 2T
Table 12:
Memory operation counts for the fully parallel mLSTM formulation for mLSTMexp and mLSTMsig. All terms denote the memory operation count for full sequence per head.
Bytes
MLSTMexp
mLSTMsig
Load:
T(2dge dne) bytes gke + 27' x bytes, Thho x bytes 27 x bytes
T(2d,t + dne) x bytes 21' bytes; Tdho x bytes
Store:
Total:
27((bytes, + (day + dyx) bytes + bytescma)
(bytes, + (dev + dyx) bytes
Table 13: FLOP counts for the recurrent mLSTM formulation for mLSTMexp and mLSTMsig. All terms denote the FLOP count for single timestep per head.
FLOPs
mLSTMexp
mLSTMsig
Gates:
4 + 2F cxp + Flog + Psig + Fmax 4dghdhv 6dgk + dnw + 1 Fabss + Fmax 2dhu + dikk
2F 4d
Memory Cell Update:
Denominator & Scale:
Output:
2d + duk
45

---

Table 14: Memory operation counts for the recurrent mLSTM formulation for mLSTMcxp and mLSTMsig. All terms denote the memory operation count for single timestep per head. We assume the states are materialized at every timestep.
Bytes
ml mnMeeeexp
ml .T..MMiig
(2dgk + dho) x bytes akv + 2 x bytes:j dhv + d,t + 1) x byteScmn
(2dgt + dav) x bytes, okv + 2 bytes,y +dgidho x byteScmn
Load:
Store:
dhe x byle8 + (dysdho + dgk 1) x byleSCmp
dhe x byles,kn + dekche byle8Cmm
2 x bytes, + 2(dnr dut)) x bytesgku
2 x bytes, 2(dn. + dyk) x bytes,te +2 +2(nhedgt dok 1) bytesemn
Total:
+ +2dnudgh bytescmm
E.2
Simplified FLOP Count Summary for the mLSTM
In this section we simplify the FLOP count for the MLSTM by setting all factors For to 1. We leave the causal factor Feausal unspecified, but typically set it 0.5 or slightly larger. Since the attention logit matrix (i.e. the quadratic matrix S) always computed in blocks due to the blockwise nature of tensor core operation, usually some parts of the upper triangular matrix computed and then masked out. To account for this, the factor Feansal can be set to a value larger than 0.5 (e.g. 0.66). In Figure 19 show the impact of the causal factor on the overall FLOP count for the mLSTMsig.
Tables 15, 16 and 17 summarize the simplified FLOP counts for the chunkwise-paralley fully parallel and recurrent mLSTM formulation.
Total Flop Count Summary (Table 18). In Table 18 we summarize the total FLOP counts for all formulations of the mLSTM with exponential and sigmoid input gate for single head and full sequence of length (i.e. batch size Neatch = 1). To obtain the total FLOP counts for one sequence, we multiply the chunkwise- laaallss-parrllel FLOP counts per chunk by the number of chunks Nehunk T/L and the recurrent FLOPs per step by the sequence length T.
Table 15:
Simplified FLOP counts for the chunkwise-parallee mLSTM formulation for mLSTM- exp and mLSTMsig. All terms denote the FL.OP count per head and chunk. We set all factors FOr to 1.
FLOPs
mLSTMexp
mLSTMsig
Recurrent computation of the inter chunk states
Gates:
0.51. + 6.5L +5 2dgt dn + Ldy + Ldgg 2dgt + 2Ld
0.5L² + 7.5L +1 2dg 2Ldgtdh Ldgx
Numerator:
Denominator:
Parallel computation of the intra chunk outputs Cumulative Forget Gates:
0.5L² + 225LL
0.5L² + 2.5L
Gate Matrix:
Feausal (5L. 2L)
Feausal 312 Feausal x (21 + dno) + 31.2)
Intra Outputs:
Feausal X (dd. + dro) + 312)
Parallel computation of the inter chunk outputs Inter Outputs:
2 2ddg + 3Ldgt
2Ldgkdne Ldgt
Combination of inter and intra chunk outputs 2L 4L
Output Combination:
Lda.
F.3 FLOP Comparison between MLSTMexp and mLSTMsig
The MLSTM with sigmoid input gate does not have normalizer and max state. Therefore, it has fewer FLOPs and memory operations compared to mL.TM with exponential input gate. In this section we quantify this difference in FLOP counts between mLSTMexp and mLSTMsig.
FLOP Count Difference between mLSTMexp and mLSTMsig (Table 18). We compute the FLOP difference of ml STMexp and ml ML.TMssg in the last column of Table 18. We observe that in the leading terms, there is no difference and conclude that the FLOP difference between ml. STMexp and MLSTMsig is small. For example for head dimension dyfk
dnu 64, we find that mLSTMex,
46

---

Table 16: Simplified FLOP counts for the fully parallel mLSTM formulation for mLSTMexp and mLSTMsig. All terms denote the FLOP count for full sequence per head. We set all factors For to 1.
FLOPS
ml.STMexp
mLSTMsig
Cumulative Forget Gates:
0.51 2.5T
0.57 + 4.57
Gate Matrix:
6T
6T
Attention Logits:
P Femusal (2T + 27
Peausl 22 + 27
Normalization:
Feausal x (4T + +2))
Outputs:
F Faausal X 27
Feausal 2I
Table 17: Simplified FLOP counts for the recurrent mLSTM formulation for mLSTMexp and mLSTMsig. All terms denote the FLOP count for a single timestep per head. We set all factors For to 1.
FLOPs
mLSTMexp
mLSTMsig
Gates:
9 4dgtdhe 6dgs + dnv + 2dh + dgk
2 Ady
Memory Cell Update:
Denominator & Scale:
Output:
2dhodgt + dyk
has less than 2% more FLOPS, while for dgp = dhv = 512 the mLSTMexp has only about 0.2% more FLOPs for all formulations.
mLSTMsig faster because has fewer FLOPs? We find that the mLSTMexp has only slightly more FLOPS than ml STMsg.. Therefore, the speed difference between ml MTxxpp and STMsss cannot be explained by the FLOP count difference alone. However, even though the absolute and relative FLOP count difference is small, the FLOPPs that differ are the more "expensive" FLOPs, i.e. pointwise operations and vector operations, which are more expensive than matrix multiplications as they are performed on tensor cores. So to seems negligible, is expected that the actural relative runtime difference is larger than the relative FLOP difference, since the FLOPs that differ "slower non tensor core" FLOPs. For the forward pass this indicates that the main reason for the 30% speedup of mLSTMsig over mLSTMexp is the efficient fusion of loops in the mLSTMsig kernel (see Appendix C.3).
the question, even though the FLOP difference
E.4 FLOP Comparison between different mLSTM Formulations
The main advantage of TFLA over Flash Linear Attention is the freely configurable chunksize, wich allows to effectively trade off between memory consumption and runtime (see Figure 6) and we will see in Figure 19 also between the total number FLOPs.
We compare the FLOP counts of the chunkwise-paralle formulation for different chunk sizes with the fully parallel and the recurrent formulation. We use the simplified FLOP counts for the mLSTMsig from Table 18 for this analysis.
Table 18: Total Simplified FLOP counts for the chunkwise shunkwise-parallel, fully parallel and recurrent formulation of MLSTMexp and mLSTMsig. All terms denote the FLOP count per head for full sequence of length? T. We set all factors For to 1.
FLOPs
MLSTMexp
mLSTMsig
Difference
F.Fe (2ldg + dua) 8) + TL Funal
TIFend (21dg + dua)) + +TL
2TLF 2T (Ad 3dm 424
chunkwise-paralich
(4dgdna 6d,e + (2dgddo 20h +-))
4doa
(4 2d,a dne
+7 (2daadh.
T (2idge dno) 6)
7 Fans (2ldgt deo) +6.57" 4.5T
fully parallel:
+ PT t 6.5T + 25T
+27 Foon 2.5T
recurrent:
(6d dne 7dg + dao 12)
(Bdgydhu dug
T (6dge dao + 10)
47

---

dx56 d,,=512
1g d,,
dx= db,
0.5
D.66
Fausam1.0
FLOPS
Recurrent
Parallel
Chunkwise
1 2 4 1632 128 512 20488192 Chunk Size L (logscale)
2 1632 128 512 20488192 Chunk Size (logscale)
2 8633 128 512 20488192 Chunk Size (logscale)
Figure 19: FLOP counts for the recurrent, fully parallel and chunkwise-parallel mLSTM formulation of mLSTMsig. Left: dg =26, d1.=512. Middle: dg==128, =256 Right: d,x=64, d1.=128. We count the number of FLOPs for a one head and one sequence of length T=8192 for different head dimensions and vary the chunk size L. By varying the chunk size L the chunkwise-paran formulation FLOP counts transition between the recurrent and fully parallel FLOP counts. Smaller head dimensions decrease the overall FLOP count.
Chunkwise-paralle FLOPs vary between Recurrent and Fully Parallel FLOPs (Figure 19). We plot the FLOP counts of the recurrent, fully parallel and maankwise-parallle mLSTM formulation of mLSTMsig for different chunk sizes in Figure 19. We observe, that the chunkwise-paralll FLOP counts transition between the recurrent and fully parallel FLOP counts when varying the chunk size from (recurrent) to T (fully parallel). Smaller head dimensions decrease the overall FLOP count. We also show the impact of the causal factor Feausal on the overall FLOP count. The causal factor Feausal accounts for the causality of the mLSTM and can vary between 0.5 and 1.0. Small values of 0.5 indicate that only past values are computed, while values of 1.0 indicate that all values are computed and then masked out. Efficient implementations achieve values close to 0.5 (see also Appendix F.2).
F55 FLOP-Optimal Chunk Size
The number of FLOPs of the thunkwise-parallee formulation of the mLSTM depend on the chunk size L. In this section we compute the FLOP- chunk size, i.e. the chunk size LopuFLOP that minimizes the FLOP count. We use the simplified FLOP counts for mLSTMsig from Table 18 for this analysis. We denote this FLOP count as FLOPs STMsig.ww(()) (L))
To compute TP substitute d,p Pht and then set the derivative of St..S.sssMMsi...pppp. (T.) with respect to I. to zero and solve for L. This gives
2d hopk +
Lop.FLOP
(103)
V 2. Feausall (dro (1 + Pyk) 3) +
The FLOP-optimal chunk size depends on the head dimension dhu and the projection factor Pyk and grows proportional to the square root of dav (i.e. LOP.FLOP ()))))
FLOP-Optimal Chunk Size grows with Head Dimension (Figure 20). We plot the FLOP-optimal chunk size over head dimension dnu for different projection factors Pyk, that determine the query-key head dimension dgt. = PykdAr in Figure 20. The FLOP-optimal chunk size grows proportional to the square root of due (i.e. (((dhw)). but remains small for typical head dimensions. The projection factor Pyk bends the curve, but does not change the overall trend.
In order to minimize the FLOPs to compute, LoM.FLOP indicates we should use rather small chunk sizes (e.g. 1=16 for d1,=512). However, the chunk size L does not only affect the FLOPs but also the memory 10 (for e.g. for loading and storing the memory cell states). Therefore, in order to find the optimal chunksize that minimizes the runtime, need consider the memory 10 well, which we do in the next section (see Appendix G).
48

---

dg 0.5dhv (pau=0.5)
dox 1.0dhr (px-1100
35
0
30
F..6
Chunk Size Lon.foo
25
F.0
15
10
512
256
512
768
1024
Head Dimension dhr
Head Dimension
Figure 20: FLOP-Opting Chunk Size LoOt for mLSTMsig. Left: P.5 Right: P,t=1.0. We plot the FLOP-optimal chunk size over head dimension d,u for different projection factors Puke that determine the query-key head dimension dgt Pokdhe. The FLOP-optimal chunk size grows proportional to the square root of die (i.c. O(ydhe)), but remains small for typical head dimensions.
G Theoretical Runtime Analysis for TFLA mLSTM Kernels
In Section 5.2 and E.3 we measure the runtime and memory consumption, experimentally. We see that TFLA needs substantially less GPU memory than other baselines and that there exists optimal chunksize at which the runtime minimized.
This runtime minimum is limited and defined by the physical constraints of our hardware (in our case NVIDIA H100 GPUs). Typically, these constraints are how fast the GPU can compute floating point operations (FLOPs) measured in FLOPs per second (FLOPs/s), how fast the GPU can load and store data from and to high-bandwidth memory (HBM) measured as memory bandwidth in bytes per second (B/s), and how much total (HBM) memory is available to store the data in bytes (B) (Austin et al., 2025, Part 1). If the kernel runtime limited by the maximum FLOPs/s, we say the kernel is compute-bound and if it is limited by the memory bandwidth, say the kernel is memory bbund.
In this section, our aim the example of mLSTMsgg belongs and what the optimal chunk size would be given the physical constraints of our hardware. We will see that taking hardware constraints into account the optimal chunk size will be much larger than the FLOP optimal chunk size found in the previous section (see Appendix F.5).
theoretically understand which region our TFLA kernel algorithm for
We begin with modeling the theoretical runtime of our TFLA kernels in Section G.I before we compute the arithmetic intensity of TFLA in Section G.2 and use the arithmetic intensity estimate upper bounds on the peak performance on modern hardware in Section G.3. Finally, in Section G.4 we compute the theoretical runtime optimal chunk size and conclude with summary of the statements from our analysis in Appendix G.5.
G.1 Theoretical Runtime
The theoretical runtime of a kernel consists the time to compute the FLOPs T TELOPs and the time to load and store the inputs, outputs and intermediates from and to the GPU memory Tyyees
Given the number of floating point operations FLOPs algo memory operations in bytes load and store Bytes algo for specific algorithm, the accelerator speed flops in FLOPS/s and the accelerator memory bandwidth, we can compute the runtimes in seconds as
FLOPS algo
Bytes algo
T FLOPsalgo
and T Bytesalgo
(104)
a FLOPS
Bytes
49

---

For the accelerator speed a FLOPs and the accelerator memory bandwidth 3 Bytes, we use the hardware specifications of NVIDIA V100", A100", H100" and B200 GPUs, which we summarize in Table 19. If there is overlap between the computation and the memory operations in other words if the data is not loaded asynchronously to the computation, the total runtime is the sum of the two, i.e.
Talgoupper alg + Rytes,algo
(105)
If the computation and memory operations can be completely overlapped, the total runtime is the maximum of the two, i.e.
Talgo.lower max (T FLoPs.agg 7 Bytes.algo)
(106)
This means the runtime is lower bounded by the maximum of the two and upper bounded by their sum (Austin et al., 2025, Part 1). We use these formulas to compute the theoretical runtime of the TFLA mLSTMsig kernel.
Table 19:
Hardware Accelerator Specification for NVIDIA GPUs used in this analysis. Values without sparsity. If only the value with sparsity is known, we divide by
bfloat16 [FLOP//s]
Memory Bandwidth [Byte/s]
Arithmetic Intensity
GPU
Year
[FLOP/bytel
VI00 SXM2
2017
120e12
0,9e12
133
A100 SXM
2020
312e12
1,935e12
161
H100 SXM
2022
989e12
3.35e12
295
B200 HGX
2025
2250e12
7.7e12
292
Theoretical Runtime of TFLA mLSTMsig Forward Pass. To compute the theoretical runtime of the TFLA mLSTMsig forward pass mLSSMssg we use the FLOP and memory operation counts for the mounkwise-paralle formulation from Table 18 and Table 10 in Appendix F. We denote the FLOP count as FLOPSmLSTHE STSMMsgg and the memory operation count in bytes as myteSILSTMsii'
We assume that memory operations are not overlapped with computation, because (1) in our current implementation of TFLA, first materialize all states in the recurrent kernel before launch the parallel kernel (see Figure 1) and (2) we do not use advanced hardware features of NVIDIA GPUs (yet), such as asynchronous memory loading, which would allow to overlap memory operations with computation. Therefore, the total theoretical runtime is the sum of the FLOP and memory operation runtimes, i.e.
FLOPS
Bytes
T STMsig
(107)
O'
Bytes
By inserting the expressions from above and multiplying by the number of heads Nhead and batch size Nhuch, we obtain the total runtime of the mLSTMsig forward pass as
T 6LdRn (1 + Pak) + 8L + 2bytes ChePas
(6Ldnn
Nhatch Nhead E
TmI.STMsig
B
1.2 Feausal (2d + 2dhr + +6) + 12? + T. (Adh..Ppg + 2dhxPgh + dhu + 11) + 2ihh. + 5
FLOPs
(108)
where assume that the queries, keys, values and input and forget gate stored in bfloat16, i.e. =bytes,
The theoretical runtime depends on the model architecture (e.g. head dimension dnu or query-key projection factor P,k), kernel parameters (e.g. chunk size L, Feausal or bytes and the hardware accelerator specifications (e.g. a FLOPs or ( Byns).
hhttps://ww.nnw http://wwww..aaaaacocoaaaaaaaaaa
50

---

Theoretical Runtime over Chunk Size (Figure 21). We show the theoretical runtime of our mLSTMsig kernel for 7B model size (i.e. d1. p05= N for 8 7=8192 and bytes =4 on different NVIDIA A100, H100 and B200 GPUs (according to Table 19) in Figure 21. We observe that newer GPUs (e-g. H100 or B200) are faster and have a higher memory bandwidth, which results in a lower runtime. Moreover, for newer GPUs, the runtime becomes less sensitive to the chunk size L, as the curve becomes flatter. Moreover, there exists an optimal chunk size that minimizes the runtime, which is determined by the physical constraints of the hardware. This optimal chunk size increases on more recent GPUs (e.g. B200), We explore the runtime-optimal chunk size more in depth in Appendix G.4.
0.5
14
0066
12
Runtime [ms]]
10-
A100 Theoretical Runtime
H100 Theoretical Runtime
6-
B200 Theoretical Runtime
4
Theeretical Chunk Size Optimal Runtime
2-
mLSTMrig TFLA Forward Measured
32
64
128 256 512 1024 2048 4096 8192 Chunk Size L (logscale)
Figure 21: Theoretical Runtime of TFLA mLSTMsig Forward Pass with varying Chunk Size L on different accelerators. We plot the theoretical runtime of mLSTMsig with 7B size (i.e. d52 P9t=0.5, 88) for = 7=8192 and bytescmm =4 on NVIDIA A A100, H100 and B200 GPUs (according Tab. 19). We also plot the measured runtime of the mLSTMsig kernel on NVIDIA H100. Newer GPUs (e.g. H100, B200) are faster and have higher memory bandwidth, which results lower theoretical runtime.
Discrepancy between Measured and Theoretical Runtime (Figure 21) In Figure 21, we addi- tionally compare the measured runtime of mLSTMsig kemels on NVIDIA H100 with the theoretical runtime. While the measured runtime is higher than the theoretical runtime, the qualitative trend of the runtime over the chunk size is similar.
There several reasons for the discrepancy between the measured and theoretical runtime. First, the model of the runtime of our MLSTMsig kernels in Equation (108) has approximation errors. While we account for the uncertainty the causal factor F'cusal. when counting the FLOPs and memory operations, we still make approximations and simplifications. For example, we do not account for the recomputation and reloading of data within the kernels (see Appendix F). Addditionally, we do not model the fact that our computation consists of multiple kernels (e.g. the recurrent and parallel kernel (see Figure 1)). Hence, we do not include the delay when multiple kernels are launched in sequence.
Second, in addition to model errors our kernels are not optimized to reach peak performance. For example, do not use advanced features of the H100 GPUs, such as asynchronous memory loading, which would allow to overlap memory operations with computation.
Finally, the hardware specifications (c.g. FLOPs) specify the peak performance of the accelerator for matrix multiplications. Our computation for the mLSTM (see Section 2.2) includes several pointwise and vector operations, which are not performed on tensor cores and are therefore slower.
Therefore, we do not expect our theoretical runtime model to perfectly match the measured runtime, but it provides good approximation of the overall qualitative runtime behavior of our mLSTM kernels, that could guide further optimization efforts.
G.2 Arithmetic Intensity
For a specific kernel we can compute the time to compute the FLOPs Telops and the time to load and store from and to the GPU memory TBytss We can then distinguish two cases: (1) TFOPP Byyss (2) TFLOPs <
The runtime is dominated by the computation time. We call this being compute- Bytes: The runtime is dominated by memory loading. We call this being memory-bound.
Instead of comparing the times, we can also compare the arithmetic intesity or operational intensity of our algorithm and our hardware to determine if the kernel is compute- bound or memory-bonnd.
51

---

The arithmetic intensity directly relates the number of FLOPs to the number of bytes of GPU memory traffic (Williams al., 2009). We can compute the arithmetic intensity for our accelerator Iace or our algorithm Zaloo by computing the ratio between the accelerator speed and the memory bandwidth or the number of FLOPs and the number bytes loaded and stored, i.e.
a FLOPs
FLOPs algo
I acc
and
algo Bytes algo
(109)
3 Bytes
Then, If > I ace, the kernel is likely to be compute-bound and (2) if algo < ac. the kernel likely wemory-bound. The accelcrator arithmetic intensity wce is the minimum arithmetic intensity required to achieve maximum performance of the accelerator.
Arithmetic Intensity of TFLA mSSTMsig Forward Pass. We compute the arithmetic intensity of the TFLA MsSS forward pass ILLSTMsig by computing the ratio between the total FLOP count FLOPSmLSTMES and the memory operation count in bytes mytesmisTMaig This gives
L Feaual (2dhuPak + 2dh 6) + L + (4df. 2dhuPak + dno + 11) 2d +
I mLSTMsig
6Ld (1 + Pqk) + 8L + Cmn
(110)
The arithmetic intensity depends on the model architecture (e.g. head dimension dno or query-key projection factor Pyk) and the kernel parameters (e.g. chunk size L, Feausal or bytesCmmn)
Arithmetic Intensity over Chunk Size (Figure 22). We plot the arithmetic intensity of the TFLA mLSTMssi forward pass over the chunk size T. in Figure 22 and vary the head dimension d,s and the precision of the memory cell states bytes Cmn Additionally, we indicate the arithmetic intensity Icc for different NVIDIA GPUs (according to Table 19). Values above the accelerator arithmetic intensity indicate that the kernel is likely to be compute-bound for the corresponding chunk sizes, while values below indicate that kernel is likely memory-bound
We observe that the arithmetic intensity increases for larger chunk sizes and that the kernel arithmetic intensity curve crosses the accelerator arithmetic intensity larger chunk sizes. This means, we can move from the memory- bound regime to the compute-bound regime by increasing the chunk size parameter T..
de 0.5d (05=0.55) States:
dee 0.5dm (0.)= bfioat16
Anthmetic Intensiy [FLOP/byte]
Fravsy 0.5
Funs
F 11.0
A100 - 160
da, 512
d,, = 128
16
128 256 512 1024 2048 4096
16
64
512 1024 2048 4096 Chunk Size (logscale)
128
256
Chunk Size (logscalee
Figure 22: Arithmetic Intensity of TFLA mLSTMsig Forward Pass with varying Chunk Size L. Left: States in float32, i.c bytes =4. Right: States in bfloat16, i.c. bytes =2. The arithmetic intensity increases for larger chunk sizes.
G.3 Roofline Analysis
The roofline model a performance model that combines the arithmetic intensity and the accelerator specifications to determine the upper bound on the peak performance of a kernel (Williams et al., 2009). To do so, the roofline model plots the arithmetic intensity of the algorithm on the x-axis and the attainable performance in FLOPs per second (FLOPs/s) on the y-axis. For each hardware accelerator we plot the roofline, which the maximum performance of the accelerator for given arithmetic intensity
Roofline(a FLOPs, Bytes) min (3 Bytes I algo: FLOPs)
(111)
52

---

If we then plot the arithmetic intensity of kernel as column that hits the roof, either it hits the flat part of the roof, meaning performnace is compute-bound or performance is ultimately memory bound Williams et al. (2009).
Roofline Model for TFLA mLSTMsig Forward Pass (Figure 23). We perform roofline analysis for our TFLA mLSTMsig forward kernels in Figure 23 for different chunk sizes L and plot the rooflines for NVIDIA V100, A A100, and H100 GPUs (according Table 19).
We observe that smaller chunk sizes are memory-bound, while larger chunk sizes are compute-bound (similar to Figure 22). Moreover, we find that our TFLA Triton kernels for the ml.STMsig, which benchmark on NVIDIA H100, are still far from the attainable peak performance (intersection with H100 roofline). This highlights the potential for further optimization.
1e12
Accelerator Roofline
1000
H100
Memory vs. Compute
Kernels
Attainable FLOPS/s
800
Chunk Sizes (Measured):
700
600
500
128
400
256
300
A100
200
1024
100
V100
50
100
150
200
250
300
350 400 450
500
Arithmetic intensity [FLOP/byte]
Figure 23: Roofline Model for TFLA mLSTMsig Forward Pass with varying Chunk Size L. We measure the performance of mLSTMsig with d5., P,s=0.5, 88, and batch size at sequence length 8192 and compare the performance with the roofline of NVIDIA V100, A A100 and H100 GPUs (according Table 19). Smaller chunk sizes are memory-bound, while larger chunk sizes are compute-bounde Our TFLA Triton kernels for the mLSTMsig are still far from the attainable peak performance (intersection with H100 roofline).
FLOPs/s is not the optimal Performance Metric for TFLA (Figure 24). In Figure 23 we observe and the roofline model also suggests that we can increase the performance of our TFLA kernels by increasing the chunk size However, while this might increase the performance in FLOPs/s, might not decrease the overall runtime, number of FLOPs increase with the chunk size L due to the increased quadratic term in the FLOP count (see Figure 19).
We confirm this in Figure 24, where we plot the FLOPs/s and the total FLOPs over the chunk size and compare the values with the actual runtime. We observe that the FLOPs/s continuously increase with the chunk size L. but the total FLOPs also increase. This means that while we could reach peak performance in FLOPs/s by increasing the chunk size the actual runtime is not necessarily minimized by doing so. Therefore, for our TFLA kernels we should use the actual runtime as final performance metric and determine the runtime-optimal chunk size, which do next.
53

---

G.4 Runtime-Optimal Chunk Size
The main advantage of TFLA is that FLOPs and memory IO and minimize the runtime (see Figure and 24). The reason for this tradcoff that the chunk size L affects the number of FLOPs and the memory I0 (e.g. for loading and storing the memory cell states) (see Appendix F). In this section, we use this insight and our theoretical runtime model from Appendix G.1 to determine the theoretical runtime-optimal chunk size for the TFLA mLSTMssg forward kernel.
choose arbitrary chunk sizes T. to trade off between
Runtime-Optimal Chunk Size for TFLA mLSTMsig Forward Pass. In order to compute the runtime-optimal chunk size LoptRuminme for the TFLA mLSTMsig forward pass, the theoretical runtime T MLSSMsig from Equation (108). We first differentiate the theoretical runtime with respect to the chunk size L and set the derivative to find the minimum runtime, i.e.
mLSTMsig
0.
(112)
OL
We then solve the equation for the chunk size I. to find the runtime-optimal chunk size Rapt.ummime This yields
2 dl Pyk + + 2 Iace Pyp bytes
LuptRimmimee
(113)
2.eausal (div (1 + Pyp) + 3)
where Tace is the accelerator arithmetic intensity from Equation (109). Compared to the FLOP-optimal chunk size LopFFLOP (see Equation (103)), the runtime-optimal chunk size Runtiim additionally depends on the arithmetic intensity Iace of our hardware and the precision of our states bytescmm The runtime-optimal chunk size grows proportional to the square root of the head dimension due (i.e. (lvdr))) and the accelerator arithmetic intensity Ince (i.e. ((vInc))). We visualize these trends in Figure 25 and 26.
Bayes
Runtime-Optimal Chunk Size depends on Model Architecture (Figure 25). We plot the runtime- optimal chunksize for the TFLA mLSTMsig forward pass Rupiine the head dimension dıe in Figure 25 for memory cell states in float32 and bfloat16 16 and NVIDIA A100 and H100 arithmetic intensities.
The runtime-optimal chunk size grows proportional to the square root of the head dimension dnv (i.e. O (d..)) and is much larger than the FI.OP-optimal chunk size.
Time [ms]
(F666
Time [ms]
1e12
1el2
2.5
20
100
2.0
Time [ms]
15
1.5
10-
40
10
1.0
5
0.5
64
128
256
512
1024
2048
4096
64
128
256
512
1024
2048 4096
Chunk Size
Chunk Size
(a) FLOPs/s Chunk Size L
(b) FLOPs over Chunk Size L
Figure 24: FLOPs/s and FLOPs over Chunk Size L for TFLA mLSTMsig Forward Pass. We compare the FL.PS// and the total FLOPss with the actual runtime at different chunk sizes. We measure the performance of mLSTMsig with d1.2 0055 Nn and batch size 8 at sequence length 8192. The FLOPs/s increase with the chunk size L, but the total FLOPPs also increase. The actual runtime is not necessarily minimized by increasing the chunk size L.
54

---

dpp 0.55dh, 50.5 States:
dyp 0.5diu (px=0.55 bfloat16
900
Fau.
Chunk Size Lasttuntme
800
Foua==0.66
700
600
Fua0==10
500
A100 Intensity = PLOP/
400
Intensily FLOP/byle
200
TRLA mLLSMMsi Optimal Chunk Size (dh,=512)
200
32 128
256
512
768
1024
32
128
256
512
768
1024
div
Runtime-Optimal Chunk Size Runtime head dimension diu for mLSTMsig. Left: Memory cell states in float32. Right: Memory cell states in bfloat16. We plot the runtime-optimal chunk size over head dimension dhe with dyk 0.5d (Pxx=0.5). Similar to the FLOP-optimal chunk size (Fig. 20), the runtime-optimal chunk size grows proportional to the square root of dıv (i.e. O( dno))) ), but much larger than the FLOP-optimal chunk size.
Figure 25:
Our measured runtime-optimal chunk size for mLSTMsig on NVIDIA H100 is around 256, which smaller than the theoretical runtime-optimal chunk size. This discrepancy is due to the approximations our theoretical runtime model and the fact that our kernels are not yet optimized to reach peak performance NVIDIA H100 GPUs (see Appendix G.I,, but already faster than almost all all other baseline kernels and than Flash Attention 3, which optimized for NVIDIA H100 GPUs (see Section 5.2).
Couttmm-Optinnee Chunk Size depends on Hardware Accelerator (Figure 26). We plot the uutimee-ppimms chunk size for the TFLA mLSTMsig forward pass LoptRunime over the hardware accelerator intensity in Figure 26 for different head dimensions dav- We highlight NVIDIA GPU accelerator intensities for common GPUs (e.g. V100, A100. H100).
The runtime-optimal chunk size grows proportional the square root the accelerator intensity Iace (i.e. O. More recent GPUs (like e.g. H100) have higher accelerator intensities, which results in larger runtime-optimal chunk size. If the trend increasing accelerator intensities continues, TFLA that enables arbitrary large chunk sizes will become increasingly important.
dgg
0.5d (0x= 0055) States: float32
700
V100
A390
0.5
hay.hunune
600
00.6
500
F
400
300
deu 128
TFLA mLSTMsig Optimal Chunk Size (dh,5112)
200
100
200
250
Accelerator Intensity [FLOP/byte]
Figure 26: Runtime-Optiman Chunk Size Lopt Runtime over Hardware Accelerator Intensity for mL- STMsig. We plot the runtime-optimal chunk size over the hardware accelerator intensity for different head dimensions da.. We highlight NVIDIA GPU accelerator intensities for commong GPUs (e.g. V100, A A100, H100). The runtime-optimal chunk size grows with the accelerator intensity (i.e. Ropttuntimme O (vCce))
55

---

G.5 FLOP and Theoretical Runtime Analysis Summary
The configurable chunk size is the core advantage of Tiled Flash Linear Attention. We summarize the statements about the chunk size:
1. The chunk size I. mediates trade-off between runtime and GPU memory usage. [Figure 6]
2. L determines the total compute in FLOPs: L matches the recurrent formulation, while
L T matches the parallel one. [Figure 19]
3. There exists an optimal chunk size [1. 2] that minimized the total FLOP count.
[Equation (103), Figure 19, Figure 20]
4. Increasing L raises the arithmetic intensity of TFLA kernels. [Equation (109), Figure 22]
5. The chunk size determines whether the kernel is memory-bound compute-bound given hardware.
[Figure 23, Figure 22]
6. FL.OPs/s alone can be misleading; the optimal chunk size should be chosen based on total runtime.
[Figure 24, Figure 21]
7. uuntime-opimmal chunk size scales proportionally with the square root of the head dimen- sion and the accelerator's 's computational intensity.
[Figure 25, Figure 26]
8. Newer hardware generations require larger chunk sizes to approach peak performance. [Figure 26, Figure 21]
56


Reviews:
[{"review_id": "p5pMPn9VCH", "reviewer": "Reviewer_Qinz", "summary": "This paper proposes a two-level tiling algorithm for efficiently training linear attention models. In particular, it focuses on optimizing the mLSTM model with new sigmoid-based gating. Results show that the training algorithm can greatly accelerate training while mLSTM achieves comparable performance with other linear attention models.", "questions": "Can you explain the intuition behind achieving high arithmetic intensity in the main paper? What are the computational trade-offs involved?\nProviding training speedup results at long context lengths (e.g., 65k tokens) would make the efficiency claims more convincing and practically relevant.", "limitations": "No", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\nA highly optimized mLSTM variant with strong empirical performance that demonstrates the effectiveness of the combined algorithmic and architectural improvements.\nThe ablation study on gate design is insightful and provides valuable understanding of the component contributions.\n\n\nWeaknesses:\nFrom a presentation standpoint, the paper would benefit from directly presenting mLSTM-sig + the new training algorithm as a unified contribution, which would be much clearer than the current flow. The current presentation suggests mLSTM and the algorithm are more tightly coupled than they actually are.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer Qinz for the helpful feedback. We appreciate that the reviewer highlights the strong empirical performance and finds our ablation study insightful. \n\nWeaknesses: \n\nWe did not present mLSTMsig + TFLA as a unified contribution as TFLA could be applied to other Linear RNN variants as we outline in Section A.3.  \nWe agree that mLSTM and TFLA are not tightly coupled and in our updated version of the paper we will emphasize that mLSTM is one possible application of TFLA.\n\nQuestions:\n\nUnfortunately, at this time we are not able to run new large scale experiments at extremely long contexts due to compute constraints. However, we extracted the training times for our largest scale experiments at 1.4B parameters for 8k context lengths (i.e. Table 5) from our experiment logs. We confirm that the speedups in kernel runtime from Figure 5 transfer to speedups in training step time. More specifically, we confirm that our TFLA kernels yield more than 30% end-to-end speed up over the Llama baseline.\nWe will add this table to the final version of the paper.\n\nTraining step times for the 1.4B models at 8192 context length with the same head dimensions as for our kernel benchmark in Figure 5. All runs use the same training setup with global batch size of 256 sharded across 32 GPUs, which corresponds to a local batch size of 8:\n| Model/Kernel                                 | Time (s) |\n|----------------------------------------------|----------|\n| Llama                                        | 2.39     |\n| mLSTMexp limit_chunk kernel (chunk size 64)  | 1.67     |\n| mLSTMexp xl_chunk kernel (chunk size 128)    | 1.66     |\n| mLSTMsig xl_chunk kernel (chunk size 128)    | 1.64     |\n\n**Arithmetic Intensity and computational trade-offs** \n\nTFLA trades-off memory vs. compute (See Figure 6). \nDue to the enhanced intra chunk tiling strategy (Figure 3), the chunk size can be chosen arbitrarily, overcoming the limitation by physical SRAM on the GPU. \n\nThe chunk size L determines the number of chunks and hence number of materialized memory states $N_C = T / L$ (see Sec. 2.2) necessary to compute the forward and backward pass. \nThe higher the chunk size, the fewer memory states (less memory) are materialized (i.e. stored to HBM). However, the number of FLOPs increases (see Fig. 19).\n\nThe arithmetic intensity of an algorithm is defined as its FLOP per byte loaded or stored ratio. The arithmetic intensity for GPUs or TPUs is defined similarly by their FLOPS/s per bytes/s ratio (see Sec. G.1 Eq. (109)).\n\n\nSince a larger chunk size leads to more FLOPs and less bytes stored and loaded, the arithmetic intensity increases with larger chunk size. We show this relation in Fig. 22. \nAt the point where the algorithm arithmetic intensity matches or exceeds the accelerators arithmetic intensity the workload can reach peak performance (see Fig. 23). We analyze this in Appendix G.\n\nWe will add the intuition and explanations to the final version of the main paper.\n\nWe thank the reviewer again for their helpful feedback, and hope they find their questions addressed."}, {"review_id": "MwjWXlZXDB", "reviewer": "Reviewer_i6Dq", "summary": "This paper introduces Tiled Flash Linear Attention (TFLA), a novel kernel for linear RNNs that improves existing Flash Linear Attention (FLA) kernels. The key innovation is the introduction of two levels of sequence parallelism: first over chunks of the input sequence, and second within each chunk through tiling of matrix computations. This modification is key, as it enables larger chunk sizes while maintaining high arithmetic intensity. \nThe authors apply TFLA to the mLSTM architecture and also propose mLSTMsig, a new variant that uses sigmoid input gates instead of exponential gates to accelerate computation. \nThe paper shows a comprehensive theoretical analysis including FLOP counts, memory operation analysis, and theoretical runtime modeling. \nExperimental results demonstrate that TFLA kernels outperform highly optimized baselines including FlashAttention 3 and Mamba 2 kernels on H100 GPUs, particularly for longer sequences during training.\n\n- Small note: Use \\citep instead of \\citet in Line 103.", "questions": "- Can you provide more detailed analysis of numerical stability implications when using larger chunk sizes?\n\n- What specific guidance can you offer practitioners for choosing between TFLA and existing methods based on problem characteristics?\n\n- How sensitive are the results to the chunk size selection strategy, and can this process be automated?\n\n- Can you expand the experimental validation to other linear RNN architectures beyond mLSTM to better support the generalizability claims?", "limitations": "The authors briefly discuss the limitations in the conclusion and theoretical analysis sections. They acknowledge that their Triton implementations may not achieve peak performance and that the theoretical model has approximation errors. However, the paper could better address the numerical stability issues observed in some experiments and provide clearer guidance on when these might occur.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strong Points**\n\n\n- The paper clearly identifies and addresses the constraint on chunk sizes due to GPU SRAM limitations, which leads to excessive memory consumption and IO costs for long sequences and is a real limitation of existing FLA implementations.\n\n- The two-level sequence parallelism design (chunking + intra-chunk tiling) is technically sound and follows established GPU optimization principles. The mathematical formulations are consistent and, as far as I could check, the implementation details seem correct to me\n\n- The paper provides both theoretical analysis (runtime complexity, arithmetic intensity, FLOP counts) and empirical validation through language modeling experiments and detailed kernel benchmarks against strong baselines\n\n- The paper shows concrete speedups over highly optimized baselines including FlashAttention-3 and Mamba-2 kernels, with particularly strong results for longer sequences.\n\n- Sufficient implementation details are provided, with clear algorithmic descriptions. \n\n- Sections 5.2 and 5.3 are great. Figure 6/21 (and the analysis there-in) is very insightful. Well done!\n\n- Overall, the paper is really well-written and the technical details are introduced in a clear manner. \n\n\n\n\n**Weak Points**\n\n- The core contribution essentially combines existing techniques (chunking from FLA + tiling from FlashAttention). While the specific combination is novel, the individual components are well-established, somewhat limiting the novelty of this work.\n\n- The primary application is to mLSTM, which is a recent architecture with limited adoption. While the paper claims broader applicability to other linear RNNs, this is mostly theoretical with limited empirical validation.\n\n- The comparison between mLSTMexp and mLSTMsig shows minimal performance differences, questioning the novelty value of the sigmoid variant. Also, the motivation behind seems a bit related to [Dynamic Tanh (DyT)](https://arxiv.org/abs/2503.10622), which I think deserves a citation here.\n\nIn a more concrete manner, there are some minor mathematical details that needs some rethinking to improve reading:\n\n- The jump from equations 1-5 to equation 6-7 are unclear. \n- The way $\\tilde{D}$ is defined is confusing in a first glance.\n\nBeing a bit picky, although I really liked Figure 3, a more descriptive caption would significantly improve the readability and understanding of the main idea of the paper: intra chunk tilling. \n\nFinally, the related work section in the main paper is very shallow. In contrast, the related work section in the appendix is comprehensive and elucidative. I suggest the author to pull it to the main paper in order to give a broader view of the field of Linear Attention for non-familiar readers.", "was_revised": true, "final_justification": "I have decided to keep my score unchanged. This is paper represents a solid work, and the authors have addressed my concerns convincingly during the rebuttal.", "rebuttal": "We thank reviewer i6Dq for the helpful feedback, which helps to improve our paper. \nWe appreciate that the reviewer finds our paper clearly identifies and addresses limitations in existing FLA implementations, that our paper is well-written with the technical details introduced in a clear manner.\n\nSuggestions:\n\nWe thank the reviewer for the helpful suggestions regarding the related work and figure captions, citations and some minor mathematical details. We will rework the particular sections  and update our final version of the paper. \n\nWeaknesses:\n\nWe would like to highlight that the main motivation of mLSTMsig is its reduced computation (no normalizer and max state) and hence faster runtimes at equal modeling performance, with the main differences visible in the faster runtimes of the kernels in Fig. 5.\n\nIn Section 2.2 we explain that the transition from equations 1-5 to 6-7 happens due to rearranging the sequence dimension into N_c chunks of chunk size L. Therefore, we change the time index from t to chunk index k. \nAs highlighted by the reviewer, this transition is not trivial and therefore we tried to make it more understandable by adding illustrative Figures 2 and Figure 9 (in the Appendix), which make the transition explicit for a small problem of chunk size L=4. \nWe hope that especially Figure 9 will help to follow the formulas in the main paper. \nFurthermore due to space limitations the equations in Section 2.2 are a compact form of the detailed chunk-wise formulation in Appendix B.2.\n\nQuestions:\n\nInvestigating numerical stability implications is a common problem in many optimized deep learning frameworks and training setups, especially with mixed precision settings.\nWe investigated the mean absolute errors for our mLSTMexp XL chunk kernels and did not find major sensitivity regarding the chunk size. \n\nWe show the mean absolute errors for mLSTMexp XL chunk kernels forward compared to the pytorch reference implementation for context length 8192 at the same head dimension as in our kernel benchmark in Figure 5 (d_qk=128, d_hv=256) for Gaussian normal distributed random inputs in float32 and bfloat16:\n| Chunk size | Float32              | Bfloat16              |\n|------------|----------------------|-----------------------|\n| 64         | 7.8761462e-04        | 2.9082941e-03          |\n| 128        | 7.8597685e-04        | 2.9058390e-03          |\n| 256        | 7.8570450e-04        | 2.9046188e-03          |\n| 512        | 7.8560099e-04        | 2.9047025e-03          |\n| 1024       | 7.8498963e-04        | 2.9049833e-03          |\n| 2048       | 7.8578707e-04        | 2.9037838e-03          |\n\n**Chunksize selection and TFLA in practice**\n\nThe runtime optimal chunk size depends on the underlying hardware accelerator and the model head dimensions (see Fig. 21, 25 and Eq. 113). This means that typically you have to tune the chunksize once for each head dimension configuration when switching to a specific accelerator.\nThis process can be automated and in our source code we provide a kernel benchmarking module which could help with this.\n\nHowever, in some scenarios you might want to reduce peak memory usage, e.g. for long contexts and large models, where you trade-off runtime for GPU memory (see Fig. 6).\nIn that case you want to maximize the batch size that fits in one GPU, which is typically limited by GPU memory size. \nIn this case the chunk size L would be another hyperparameter (in addition to gradient checkpointing or gradient accumulation for example).\n\n**Extending Empirical Validation to other Architectures**\n\nIn our paper we implemented TFLA for two different mLSTM variants with and without a normalizer state and different gate activation functions. \nWe would like to highlight that RetNet (fixed forget gate, no input gate) and Simple GLA (no input gate) are a subset of mLSTMsig and are hence already implemented. This already demonstrates its generalizability. Since reimplementing and applying TFLA to other advanced linear RNNs requires a significant amount of resources, we leave this for future work, but provide guidance on how it might be applied in Appendix A.3.\n\n\nWe thank the reviewer for their valuable feedback and hope that our clarifications help reinforce the merits of the submission."}, {"review_id": "AaYodLwqgm", "reviewer": "Reviewer_uJbu", "summary": "The paper addresses the problem of accelerating training and inference of LLM on GPUs. It proposes an implementation applicable to any linear RNNs types of LLMs, called Tiled Flash Linear Attention (TFLA). TFLA builds upon FlashLinearAttention (the version of Flash attention for linear RNN), which is based on a chunkwise-parallel implementation (recurrent inter-chunk, parallel intra-chunk). Inspired by FlashAttention2, TFLA adds another level of intra-chunk parallelism enabling arbitrary chunk size. This flexibility allows to chose the chunk size to balance the number of operations and the number of costly memory accesses to the GPU High Bandwith Memory, depending on model architecture and hardware characteristics, to achieve optimal runtime. The method is applied on the mLSTM architecture, on which some modifications are proposed to enable even faster inference. The resulting mLSTM architecture with the chunkwise-parallel acceleration yields faster training than efficient implementations of FlashAttention and linear RNN (Mamba 2), while the inference time is comparable to other linear RNN implementations.", "questions": "1)\tImprovement of the inference runtime. Why is there a bigger difference in inference and training time compared to other linear RNN methods (e.g. Mamba 2) ? How could inference speed be improved ? Indeed, training time is a major limitation to AI development, but inference time is crucial for reducing the costs of AI deployment.\n2)\tImprovement of the method for modern GPUs. As the FLOPs/s in hardware increases faster than the memory bandwidth, the authors observe that the optimal chunk size increases in modern GPUs. Moreover, the proposed method yields lower benefits in more recent GPUs than older ones. Knowing this trend, do the authors see limitations of their method as GPUs improve ? Could the authors comment on the benefit of their method on future hardware ?\n3)\tCompatibility with other linear RNNs. Knowing that Mamba 2 is compatible with such acceleration and is currently the fastest in inference runtime, it would be interesting to know how this acceleration would performed compared to the one originally implemented for Mamba2.\n4)\tExtension of the method to other hardware. The proposed method is intended for cloud-type GPUs. Would this method be relevant for embedded GPUs (e.g. Jetson), although they are targeted for inference loads rather than training loads ?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths :\n- Quality : The paper seems technically correct and complete, with analysis allowing to understand the challenges and the proposed method. Limitations and future works are properly aknowledge, consisting mainly in further optimizations to get closer to the theoretical peak performance.\n- Clarity : The paper provides sufficient details to reproduce the experiments. There are many analysis (theoretical and experimental) to understand the proposed method and how to use it.\n- Significance : The applicability of the proposed method is large. Indeed, the paper shows that the proposed acceleration is not only applicable to the mLSTM architecture but also to most of linear RNNs.\n- Originality : The paper provides an interesting analysis of the method, with theoretical and experimental results on key performance indicators (FLOPs/s, runtime...) and especially how they should be used.\n\nWeaknesses\n- Quality : The authors propose modifications on the mLSTM layer to be implemented more efficiently (for instance replacing an exponential activation function by a sigmoid, and removing the normalization state), therefore they need to evaluate the impact of the proposed modifications on the accuracy of the model. This evaluation is relatively limited (only 1 dataset, with only 1 run per experiment).\n- Signifiance : While the method yields state-of-the-art training time, it is not the case for inference time. Moreover, the authros observed that the proposed method yields lower benefits in more recent GPUs than older ones.", "was_revised": true, "final_justification": "After reading the other reviewers’ comments and the authors’ response, I want to maintain my acceptance recommendation. My concerns regarded the inference runtime compared to other efficient linear RNN (Mamba2) and the authors clarified that it depends on the head dimension that was chosen for each of the models. Another concern was about the relevance of the method for modern GPU. The authors clarified that, by enabling arbitrary large chunk sizes, the method is well suited to modern GPUs that benefits from larger chunk size. The authors did a very good work in analysing their method, facilitating its practical use. The remaining limitation concerns the limited evaluation of the proposed mLSTM variant. In particular, I suggest to evaluate accuracy in downstream tasks, rather than evaluating only perplexity. This may further motivate the adoption of the model. I believe this can be addressed in future work.", "rebuttal": "We thank the reviewer for the helpful feedback which helped to improve our paper. \nWe appreciate that the reviewer highlights the general applicability of our method and sees strengths in our extensive analysis in the paper.\n\nWe would like to address the reviewers questions as follows:\n\n1.)\nIn Figure 5a (Forward pass only / Inference runtime measurements) the runtime of Mamba2 forward pass is lower than for the mLSTM TFLA kernels. This can be explained by the fact that Mamba2 uses smaller head dimensions (or the respective equivalent) of 64 by default. However, for mLSTM and other linear RNNs the head dimension is typically larger (which we confirm in our experiments). Therefore, we kept larger head dimensions for the mLSTM in Fig. 5 and added an ablation on the effect of the head dimension at fixed batch size in Fig. 18. We find that TFLA runtimes are faster for smaller head dimensions. (A side note: The broad range of possible head dimensions for TFLA is another advantage compared to Mamba2.) \n\nIn the current form we optimized the kernels for training stability and kept the precision dtype of the C states in float32 in forward and backward. Future optimizations for maximal inference performance could explore lower precisions (e.g. bfloat16 for the memory states) for example. \n\n2.) \nIn our theoretical runtime analysis in Fig. 21 and 22 we show that newer hardware generations tend to have a higher arithmetic intensity and that the (theoretical) runtime optimal chunk size increases with newer hardware generations due to the higher arithmetic intensity. \n\nSince our TFLA algorithm enables arbitrary large chunk sizes it is well suited for this trend, as our analysis shows that large chunk sizes will be necessary for optimal performance on future hardware generations.\n\nFinally, as the region of optimal runtime in Fig. 21 becomes wider for newer hardware generations, we suspect that we could increase the chunk size even further to reduce GPU memory consumption and potentially use larger batch sizes to increase efficiency.\n\n3.)\nPlease see 1.)\n\n4.) \nWe believe that the general principles (arithmetic intensity, memory vs. compute trade-off) from our theoretical runtime analysis of TFLA apply also to other GPU models, such as NVIDIA Jetson GPUs.\nHowever, at the time of writing the paper we had only H100 GPUs available. Therefore, we could not test TFLA on other GPUs, such as Jetson GPUs for embedded systems. \nIn the future we plan to test TFLA also on other hardware, which might also require rewriting the kernels in other frameworks than Triton.\n\n\nWe sincerely thank the reviewer for their thoughtful feedback, and hope our responses have addressed all remaining questions."}, {"review_id": "vcyWum0L85", "reviewer": "Reviewer_5h29", "summary": "This paper introduces Tiled Flash Linear Attention (TFLA), a new kernel algorithm for linear RNNs designed for high efficiency on modern GPUs. TFLA extends prior work like Flash Linear Attention by adding a second level of sequence parallelism via tiling within each chunk. This technique enables arbitrarily large chunk sizes, improving arithmetic intensity and reducing memory I/O. The authors apply TFLA to xLSTM (mLSTM) and propose a simplified, faster variant, mLSTMsig. Extensive benchmarks on H100 GPUs demonstrate that TFLA kernels outperform highly optimized baselines, including FlashAttention 3 and Mamba 2 in training speed.", "questions": "Please see the weakness.", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The paper proposes a well-motivated and clever kernel optimization technique, TFLA. By combining the chunking parallelism of linear RNNs with the intra-chunk tiling from attention mechanisms, it effectively addresses the SRAM size limitation of previous methods. This enables a flexible and powerful trade-off between computation and memory usage.\n\n2. The experimental results are extensive and compelling. The proposed TFLA kernels demonstrate significant speedups over state-of-the-art and highly optimized baselines, including FlashAttention 3 and Mamba 2, on modern hardware. This establishes a new state of the art for efficient long context sequence modeling primitives.\n\n3. Beyond kernel optimization, the paper provides valuable insights into the model architecture itself. The introduction and analysis of mLSTMsig, along with the empirical study of gate initialization and its effect on training stability, adds significant depth and practical value to the work.\n\nWeaknesses:\n\n1. The core idea of TFLA, while effective, is a combination of two existing techniques: chunk-wise parallelism from Flash Linear Attention and intra-chunk tiling from FlashAttention. The novelty lies in the synthesis and implementation of these ideas rather than a fundamentally new algorithmic paradigm.\n\n2. The benchmark comparison in Figure 5 uses different head configurations for the proposed model versus key baselines. For example, TFLA uses 16 heads while FlashAttention uses 32 heads. While the total embedding dimension is constant, this difference can influence hardware performance and make a direct, apples-to-apples speed comparison less clear.\n\n3. The paper reports training instabilities like gradient spikes for certain configurations, particularly for the xl_chunk kernels shown in Table 2 and Table 5. While the authors identify a mitigation strategy through bias initialization, the root cause of the instability is not fully diagnosed. It is unclear if it is an inherent model property or a numerical precision issue in the custom kernel.", "was_revised": true, "final_justification": "This paper proposes a well-motivated and clever kernel optimization technique.", "rebuttal": "We thank the reviewer for the thorough and constructive feedback and appreciate the recognition of the TFLA kernel as a well-motivated and effective solution to SRAM limitations, along with the acknowledgement of the strong experimental results and architectural contributions—particularly the analysis of mLSTMsig and gate initialization.\n\nWe would like to address the reviewers concerns in the following:\n\nWe explain the benchmark setup of Figure 5 in detail in Section 5.2 and Appendix A.3. The main motivation for the benchmark setup is to assess the kernel runtime in a setting, which would be used for training a model of size ~7B parameters with embedding dimension typically 4096. \n\nWe now aimed to select also typical head dimensions that would be used for Llama or Mamba2 models. Specifically, we used configs from huggingface, where available. For Llama for example we used the Meta-Llama-8B, which has 32 heads and for Mamba2 we used Codestral-Mamba2 which has a head dim of 64 as reference. \n\nFor our model, we preferred slightly larger head dimensions (16 heads, d_hv = 256, d_qk = 128) as this leads to a larger memory state and similar perplexity. \nTo analyse the effect of the head dimension on the runtime we perform a benchmark with varying head dimensions in Figure 18 in the Appendix. We observe that smaller head dimensions (i.e. more heads) lead to even faster runtimes. This means that the runtimes of mLSTM are expected to be lower if we choose the same (small) head dim as Mamba2, but in practice one would choose larger head dimensions for mLSTM, which is why we made this decision.\n\nRegarding the numerical issues, we verified correctness of the kernel implementations in float32 and confirmed that training with different implementations JAX vs. our Triton kernels yield the same results (see Table 3).\n\nWe sincerely thank the reviewer for their thoughtful feedback, and hope our responses have addressed the reviewers questions."}]

Author rebuttal:

Meta-review: The paper proposes Tiled Flash Linear Attention (TFLA), a new GPU kernel for Linear RNNs and mLSTMs. Builds on Flash Linear Attention (FLA) but removes chunk-size limitations by adding two-level sequence parallelism: 1) Chunkwise parallelism (as in FLA), 2) Intra-chunk tiling (new contribution, inspired by FlashAttention2). Benchmarks show state-of-the-art training efficiency over FlashAttention 3, Linear Attention, and Mamba 2, especially for long-context training. 
It is not a wholly new paradigm, but a good synthesis of existing GPU optimization techniques with a practical implementation. The paper demonstrates good efficiency gains for long-context training, accompanied by a strong theoretical and empirical analysis, as well as practical improvements to mLSTM.
The weaknesses of this paper include limited novelty, a narrow evaluation scope (mostly mLSTM), weaker inference improvements, and some clarity issues.
Overall, reviewers generally lean towards accept this paper.
Decision: Accept (poster)