Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels

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.

We would like to address the reviewers concerns in the following:

We 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.

We 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.

For 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. To 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.

Regarding 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).

We sincerely thank the reviewer for their thoughtful feedback, and hope our responses have addressed the reviewers questions.

We thank the reviewer for the helpful feedback which helped to improve our paper. We appreciate that the reviewer highlights the general applicability of our method and sees strengths in our extensive analysis in the paper.

We would like to address the reviewers questions as follows:

1.) In 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.)

In 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.

2.) In 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.

Since 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.

Finally, 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.

3.) Please see 1.)

4.) We 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. However, 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. In the future we plan to test TFLA also on other hardware, which might also require rewriting the kernels in other frameworks than Triton.

We sincerely thank the reviewer for their thoughtful feedback, and hope our responses have addressed all remaining questions.

We thank reviewer i6Dq for the helpful feedback, which helps to improve our paper. We 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.

Suggestions:

We 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.

Weaknesses:

We 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.

In 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. As 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. We hope that especially Figure 9 will help to follow the formulas in the main paper. Furthermore due to space limitations the equations in Section 2.2 are a compact form of the detailed chunk-wise formulation in Appendix B.2.

Questions:

Investigating numerical stability implications is a common problem in many optimized deep learning frameworks and training setups, especially with mixed precision settings. We investigated the mean absolute errors for our mLSTMexp XL chunk kernels and did not find major sensitivity regarding the chunk size.

We 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:

Chunksize selection and TFLA in practice

The 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. This process can be automated and in our source code we provide a kernel benchmarking module which could help with this.

However, 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). In that case you want to maximize the batch size that fits in one GPU, which is typically limited by GPU memory size. In this case the chunk size L would be another hyperparameter (in addition to gradient checkpointing or gradient accumulation for example).

Extending Empirical Validation to other Architectures

In our paper we implemented TFLA for two different mLSTM variants with and without a normalizer state and different gate activation functions. We 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.

We thank the reviewer for their valuable feedback and hope that our clarifications help reinforce the merits of the submission.

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.

Weaknesses:

We 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.
We 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.

Questions:

Unfortunately, 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. We will add this table to the final version of the paper.

Training 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:

Arithmetic Intensity and computational trade-offs

TFLA trades-off memory vs. compute (See Figure 6). Due 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.

The 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. The 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).

The 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)).

Since 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. At 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.

We will add the intuition and explanations to the final version of the main paper.

We thank the reviewer again for their helpful feedback, and hope they find their questions addressed.