Title: Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels
Venue: NeurIPS 2025
Accepted: yes
Keywords: xLSTM

FlashLinearAttention

FlashAttention

Linear Attention

Linear RNN

RNN

mLSTM

Attention

Self-Attention

Transformers

LLM

Kernels

Systems

Efficiency
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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\nPlease see the weakness.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths :\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.\n\nQUESTIONS:\n1)\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 ?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-01"}

comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths:

A highly optimized mLSTM variant with strong empirical performance that demonstrates the effectiveness of the combined algorithmic and architectural improvements.
The ablation study on gate design is insightful and provides valuable understanding of the component contributions.


Weaknesses:
From 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.

QUESTIONS:
Can you explain the intuition behind achieving high arithmetic intensity in the main paper? What are the computational trade-offs involved?
Providing training speedup results at long context lengths (e.g., 65k tokens) would make the efficiency claims more convincing and practically relevant.

LIMITATIONS:
No
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-07