Tiled flash linear attention for efficient linear RNN and xlstm kernels

Abstract

A method for processing an input sequence using a linear recurrent neural network (LRNN) is proposed. The input sequence is partitioned into chunks and materialized in the off-chip memory. Each chunk is then partitioned into tiles that fit within the on-chip memory. A parallel processing step is performed for the tiles, which includes loading the materialized recurrent memory state and computing output states in parallel using the on-chip memory. This enables efficient processing and training of LRNNs, making it suitable for applications with large input sequences. The method can be further optimized for specific hardware architectures.

Description

[0001] TILED FLASH LINEAR ATTENTION FOR EFFICIENT LINEAR RNN AND XLSTM KERNELS

[0002] TECHNICAL FIELD

[0003] The present disclosure generally relates to the field of machine learning models and neural networks, and more particularly to hardware-efficient implementations of such models and networks.

[0004] BACKGROUND

[0005] The present disclosure relates to the field of artificial neural networks used for sequence processing, particularly large language models implemented on high-performance computing hardware. These systems are increasingly utilized in applications ranging from natural language generation and translation to genomic sequence analysis and time-series forecasting. The efficient execution of these models relies heavily on specialized data processing apparatuses, such as graphics processing units or tensor processing units, which are characterized by massive parallelism and hierarchical memory structures comprising high-capacity off-chip memory and limited-capacity high-speed on-chip memory.

[0006] Transformers have established themselves as a dominant architecture in this field, relying on an attention mechanism to model dependencies between elements in a sequence. While effective, the standard self-attention mechanism exhibits quadratic computational complexity with respect to the sequence length. To mitigate this computational cost on hardware, optimized algorithms have been developed that reorganize memory access patterns. For example, techniques such as FlashAttention leverage tiling strategies to minimize data movement between memory levels, thereby accelerating the execution of Transformer models.

[0007] Recently, Linear Recurrent Neural Networks (RNNs) with gating mechanisms have demonstrated competitive performance compared to Transformers in language modeling tasks. These architectures utilize a recurrent state update rule that allows for linear compute scaling with respect to sequence length, offering theoretical runtime advantages over the quadratic scaling of Transformers. Realizing these theoretical benefits in practical scenarios requires optimized custom kernels that can exploit the specific hardware characteristics of the underlying processing apparatus. Leveraging a chunkwise-parallel formulation of linear RNNs, approaches such as Flash Linear Attention (FLA) have been proposed to accelerate processing by parallelizing computations over chunks of the input sequence.

[0008] However, the efficiency of such chunkwise-parallel formulations is constrained by the physical characteristics of the hardware memory hierarchy. Specifically, the size of the processing chunks in these approaches is typically coupled to the capacity of the fast on-chip memory. Because this on-chip memory is limited, the chunks must remain relatively small to avoid memory overflow. This coupling necessitates that the intermediate recurrent states — which carrythe historical context between chunks — must be materialized and written to the slower off-chip memory at the boundary of every small chunk.

[0009] This frequent materialization of intermediate states creates a significant bottleneck regarding arithmetic intensity. Arithmetic intensity, defined as the ratio of floating-point operations performed to the amount of data moved, determines whether a process is compute-bound or memory-bound. When the chunk size is restricted by on-chip memory limits, the system spends a disproportionate amount of time performing memory input / output (I / O) operations relative to actual computation. This high I / O cost prevents the hardware from reaching its peak computational throughput.

[0010] Furthermore, the requirement to materialize states for frequent, small chunks results in high memory consumption. In scenarios involving long-context pre-training, where sequence lengths can extend to tens or hundreds of thousands of tokens, storing an intermediate state for every small chunk accumulates a substantial memory footprint in the off-chip memory. This excessive memory usage can limit the maximum batch size or sequence length that can be processed, thereby constraining the scalability of the training process.

[0011] Another consideration involves the utilization of compute units when processing models with varying hyperparameter configurations, such as large head dimensions. Standard parallelization strategies that distribute work primarily across batch and head dimensions may encounter efficiency limits when the head dimension is large. In such cases, the number of available thread blocks may not align optimally with the available hardware resources, leading to low occupancy and idle compute cycles, further diminishing the overall efficiency of the kernel. Additionally, the challenges described above are exacerbated during the training phase of these models. The backward pass, required for gradient computation, typically necessitates access to the intermediate states generated during the forward pass. If the forward pass relies on frequent state materialization, the backward pass incurs similar or greater I / O penalties. Developing a backward pass strategy that maintains high arithmetic intensity without requiring excessive recomputation or storage remains a complex optimization task.

[0012] In view of these factors, a challenge remains in developing methods for executing linear recurrent models that effectively manage the trade-off between parallelization granularity, memory I / O frequency, and hardware utilization on systems with hierarchical memory architectures.

[0013] SUMMARY OF THE DISCLOSURE

[0014] It is an object of embodiments of the present disclosure to optimize the execution of linear Recurrent Neural Networks (RNNs) and Extended Long Short-Term Memory (xLSTM) architectures with matrix memory (mLSTM) on high-performance hardware, specifically by modifying the Flash Linear Attention (FLA) method to increase arithmetic intensity and reduce memory I / O bottlenecks when processing sequences, particularly where the chunk size wouldotherwise be limited by on-chip memory constraints. This objective is achieved by the subjectmatter defined in the independent claims. Advantageous modifications of embodiments of the present disclosure are defined in the dependent claims as well as in the description and the drawings.

[0015] Certain embodiments of the present disclosure relate to Tiled Flash Linear Attention (TFLA), a novel kernel algorithm for linear RNNs, that enables arbitrary large chunk sizes by introducing an additional level of sequence parallelization within each chunk. First, we apply TFLA to the xLSTM with matrix memory, the mLSTM. 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.

[0016] One aspect of the disclosed technology provides a method for processing an input sequence using a linear recurrent neural network (LRNN). The method may be computer-implemented. In embodiments, the LRNN may be implemented on a data processing apparatus. The data processing apparatus may have an off-chip memory, such as a High Bandwidth Memory (HBM) or a Graphics Processing Unit (GPU) global memory, and an on-chip memory, such as a Static Random Access Memory (SRAM) or a shared memory. This specific memory hierarchy configuration ensures that the method is optimized for current high-performance computing hardware where the bandwidth disparity between global and shared memory is a critical bottleneck.

[0017] In embodiments, the method may comprise partitioning the input sequence into a plurality of chunks, each chunk having a chunk size. The method may further comprise performing a recurrent processing step across the plurality of chunks to materialize a recurrent memory state for each chunk, preferably in the off-chip memory. By materializing states only at the chunk level, this approach reduces the frequency of writing to slower off-chip memory compared to prior art approaches that must materialize states more often due to smaller, SRAM-coupled chunk sizes.

[0018] In embodiments, the method may comprise partitioning each of the plurality of chunks into a plurality of tiles, each tile having a tile size smaller than the chunk size, such that a tile fits within the on-chip memory. The method may further comprise performing a parallel processing step for the plurality of tiles. The parallel processing step may comprise loading the materialized recurrent memory state of the corresponding chunk, or at least part thereof, from the off-chip memory into the on-chip memory. This may involve loading only the part from the memory state that belongs to the respective tile. The parallel processing step may further comprise computing output states for the tiles within said chunk in parallel using the on-chip memory. This intrachunk tiling effectively decouples the degree of parallelization (determined by the tile size) from the frequency of state materialization (determined by the chunk size), allowing for increasedarithmetic intensity by processing larger chunks without being constrained by the limited capacity of on-chip memory.

[0019] The disclosed technology addresses a fundamental physical constraint inherent to modern computing architectures: the bandwidth bottleneck commonly referred to as the "memory wall." Rather than merely implementing a mathematical operation on a general-purpose computer, the claimed method is inextricably tied to the specific physical attributes of the data processing apparatus, specifically the hierarchical arrangement of high-capacity off-chip memory and limited-capacity on-chip memory. By adapting the algorithmic structure, specifically the two-level partitioning into chunks and tiles, to the physical capacities of these distinct memory tiers, the disclosed technology produces a concrete technical effect: it minimizes the physical transfer of data across the memory bus. This reduction in data movement directly translates to lower energy consumption and reduced latency, constituting a specific technical modification to the internal functioning of the computer system that optimizes its resource usage beyond what is achieved by standard programming practices.

[0020] Furthermore, the disclosed technology provides a specific improvement to the functioning of the computer itself, resolving a problem rooted in computer technology. Conventional approaches are forced to stall computational units while waiting for memory I / O operations because the algorithmic chunk size is rigidly coupled to the small on-chip cache size. The present disclosure solves this technological problem by decoupling the parallelization granularity from the state materialization frequency. This is not an abstract idea or a disembodied mathematical formula; it is a practical application that reconfigures how the computer manages its memory hierarchy during execution. The specific steps of materializing states at a coarse granularity in off-chip memory while performing fine-grained parallel computations in on-chip memory integrate the processing logic with the hardware reality, resulting in a tangible increase in arithmetic intensity and enabling the data processing apparatus to perform complex sequence modeling tasks at speeds that were previously unattainable due to hardware limitations.

[0021] In embodiments, the recurrent processing step is executed by a recurrent compute kernel configured to process the chunks sequentially. This kernel efficiently handles the sequential dependencies of the linear RNN at the coarser chunk level, ensuring correct state propagation. In embodiments, the parallel processing step is executed by a parallel compute kernel configured to process the tiles in parallel to compute the output states for all chunks. This separation allows for massive parallelization of the intra-chunk computations, maximizing hardware utilization.

[0022] In embodiments, the parallel processing step further comprises parallelizing the computing of the output states over an embedding dimension of the linear recurrent neural network, such that computations for different embedding blocks are distributed across different compute units of the data processing apparatus. This prevents hardware under-utilization when processing models with large head dimensions, distributing the workload more evenly across available compute resources. In other words, the method may allow for arbitrary head dimensions, i.e.,embedding dimensions. In general, the term “embedding dimension” typically refers to the general token dimension of the model. In attention or linear RNNs, this embedding dimension is typically divided into heads, which are processed in parallel, and therefore the head dimension is the embedding dimension divided by the number of heads. This parallelization over the number of heads differs over existing Attention mechanisms in that we further divide the head dimension to parallelize across it and combine it with the other two-level parallelization dimensions across the sequence dimension. The implication of this further parallelizing over an embedding dimension (i.e. the head dimension) is that we can also freely configure the memory state size, because the memory state size is given as head dim x head dim. Since we can now “tune” the memory state size this might also allow for perfect adaptation to specific hardware (e.g. for mobile applications, robots, drones, etc.).

[0023] In embodiments, the method further comprises a backward pass for training the linear recurrent neural network, wherein the backward pass utilizes the partitioning into chunks and tiles, and wherein a parallelization dimension and / or a looping dimension of the backward pass are swapped relative to the parallel processing step of the forward pass. This ensures that the high arithmetic intensity achieved in the forward pass is maintained during gradient computation, preventing the backward pass from becoming a training bottleneck. In other words, the method may use the same method as for the main embodiment also in the backward pass, but with swapped parallelization dimension and looping dimension relative to the parallel processing step of the forward pass depending on the respective input gradient (i.e. the dQ, dK, dV).

[0024] In embodiments, the backward pass further comprises recomputing intermediate hidden states within the on-chip memory from the chunks, rather than loading stored intermediate hidden states from the off-chip memory. This recomputation strategy, often referred to as activation checkpointing or rematerialization, trades a small amount of extra computation for a significant reduction in memory bandwidth usage, further enhancing overall training efficiency.

[0025] In embodiments, the linear recurrent neural network is an xLSTM with matrix memory (mLSTM) and the recurrent memory state comprises a matrix memory state. This specific architecture leverages the efficiency of matrix operations within the recurrent framework.

[0026] In embodiments, the mLSTM utilizes an input gate activation function configured to behave similarly to an exponential function in a negative value range and to be bounded in a positive value range. This behavior ensures numerical stability while maintaining the desirable properties of exponential gating for memory updates.

[0027] In embodiments, the linear recurrent neural network is an mLSTM with sigmoid input gate (mLSTMsig), wherein one or more of an input gate and a forget gate of the network are determined using a sigmoid activation function. This variant simplifies the gating mechanism, facilitating hardware implementation.

[0028] In embodiments, the recurrent processing step and the parallel processing step are performed without maintaining a normalizer state or a stabilizer state. By eliminating these states, registerpressure is reduced, and the computational complexity per step is lowered, leading to faster execution.

[0029] In embodiments, the parallel processing step comprises executing a fused loop for matrix multiplications within the on-chip memory. Fusing loops minimizes kernel launch overhead and maximizes data locality within the fast on-chip memory.

[0030] In embodiments, computing the output states for the tiles comprises loading data from the on-chip memory into a register file of a processor, and performing matrix multiplication operations using said register file. This utilizes the fastest level of the memory hierarchy, the registers, to perform the core arithmetic operations, further optimizing performance.

[0031] In embodiments, the method further comprises initializing a bias of the input gate to a negative value, preferably a value less than or equal to -5. This specific initialization stabilizes training by ensuring that gate values start in a regime that prevents exploding gradients or unstable updates.

[0032] In embodiments, the chunk size is selected based on a computational intensity of the data processing apparatus, preferably such that the chunk size scales proportionally to the square root of said computational intensity. This hardware-aware configuration ensures that the algorithm adapts to the specific capabilities of the underlying hardware, maximizing arithmetic intensity for the given platform.

[0033] Another aspect of the disclosed technology provides a data processing apparatus, device or system, comprising means for carrying out (the steps of) the method(s) of any one of the embodiments of the disclosed technology. For example, one aspect provides a data processing apparatus comprising an off-chip memory, an on-chip memory, and one or more processors configured to carry out any of the disclosed methods.

[0034] Another aspect of the disclosed technology provides a computer program (product) comprising instructions which, when the program is executed by a computer, data processing apparatus, device or system, in particular the data processing apparatus mentioned above, cause the computer, data processing apparatus, device or system to carry out (the steps of) the method(s) of any one of the embodiments of the disclosed technology.

[0035] Another aspect of the disclosed technology provides a non-transitory computer-readable medium on which computer-executable instructions are stored to implement (the steps of) the method(s) of any one of the embodiments of the disclosed technology. For example, one aspect provides a computer-readable medium storing the computer program mentioned above.

[0036] Any embodiment of the disclosed technology may be combined with any other embodiment, not just merely with those expressly indicated.

[0037] It is an advantage of embodiments of the disclosed technology that the arithmetic intensity of sequence processing is significantly increased, thereby reducing memory I / O bottlenecks. This is achieved by the two-level sequence parallelization strategy, specifically the intra-chunk tiling,which decouples the degree of parallelization from the frequency of state materialization. By partitioning the input into large chunks for recurrence while processing small tiles in parallel within the on-chip memory, the system can scale the chunk size up to minimize expensive read / write operations to the off-chip memory, while ensuring that the working set for each calculation step remains small enough to fit within the fast on-chip memory.

[0038] It is an advantage of embodiments of the disclosed technology that the memory footprint required for storing intermediate recurrent states during long-context training is substantially reduced. By performing local state aggregation where the recurrent state is materialized to the off-chip memory only at the boundaries of large chunks rather than small tiles, the total storage requirement scales with the coarser chunk size rather than the finer tile size. This reduction allows for the training of linear RNNs on significantly longer sequences or with larger batch sizes on the same hardware hardware compared to prior art approaches where the materialization frequency is tightly coupled to the limited SRAM capacity.

[0039] It is an advantage of embodiments of the disclosed technology that hardware utilization is maximized even when processing models with large head dimensions. This is achieved by parallelizing the intra-chunk matrix multiplications over the embedding dimension, in addition to the batch and sequence dimensions. By distributing the computational workload of large embedding blocks across different compute units, the disclosed technology prevents the underutilization of thread blocks that typically occurs in prior art kernels when the number of heads is small or the head dimension is disproportionately large.

[0040] It is an advantage of embodiments of the disclosed technology that the high computational efficiency achieved in the forward pass is maintained during the training phase. This is achieved by a backward pass strategy that utilizes the same chunk-and-tile partitioning but swaps the parallelization and looping dimensions relative to the forward pass. This dimension swapping ensures that gradient computations also benefit from high arithmetic intensity, preventing the backward pass from becoming a bottleneck and facilitating the efficient training of large-scale models.

[0041] It is an advantage of embodiments of the disclosed technology that register pressure within the on-chip memory is reduced and computational complexity is lowered. This is achieved in specific embodiments, such as the mLSTMsig variant, by utilizing sigmoid activation functions for gating, which eliminates the mathematical necessity for maintaining normalizer ($n_t$) and stabilizer ($m_t$) states. The removal of these states frees up limited register space for other operations and allows for the execution of fused loops for matrix multiplications without the overhead of tracking running maximums, thereby increasing the overall kernel throughput. It is an advantage of embodiments of the disclosed technology that the training process is numerically stabilized, mitigating issues related to exploding gradients or unstable updates. This is achieved by initializing the bias of the input gate to a specific negative value range (e.g., less than or equal to -5). This initialization strategy ensures that the gate activations begin in aregime that moderates the magnitude of state updates during the early phases of training, facilitating smoother convergence for the linear recurrent neural network.

[0042] BRIEF DESCRIPTION OF THE DRAWINGS

[0043] The disclosed technology may be better understood by reference to the following drawings: FIG. 1 illustrates a schematic overview of a Tiled Flash Linear Attention (TFLA) framework processing an input sequence using a recurrent kernel and a parallel kernel, according to certain embodiments.

[0044] FIG. 2 illustrates a flowchart of a method for processing an input sequence utilizing a two-level partitioning strategy into chunks and tiles, according to certain embodiments.

[0045] FIG. 3 illustrates a schematic diagram of chunkwise gate computation for a linear recurrent neural network, according to certain embodiments.

[0046] FIG. 4 illustrates a diagram of a tiling strategy for a parallel compute kernel, specifically highlighting parallelization over an embedding dimension, according to certain embodiments. FIG. 5A illustrates a graph illustrating the transfer behavior and gain of an mLSTM with an exponential input gate (mLSTMexp), according to certain embodiments.

[0047] FIG. 5B illustrates a graph illustrating the transfer behavior and gain of an mLSTM with a sigmoid input gate (mLSTMsig), according to certain embodiments.

[0048] FIG. 6 illustrates a bar chart illustrating the trade-off between memory consumption and runtime for different chunk sizes, according to certain embodiments.

[0049] FIG. 7 illustrates a line graph comparing the training runtimes of the disclosed method against prior art methods across varying sequence lengths, according to certain embodiments.

[0050] FIG. 8 illustrates a schematic block diagram of computer hardware on which embodiments of the disclosed technology can be implemented.

[0051] DETAILED DESCRIPTION

[0052] In the following, representative embodiments illustrated in the accompanying drawings will be explained. It should be understood that the illustrated embodiments and the following descriptions refer to examples which are not intended to limit the embodiments to one preferred embodiment.Terms

[0053] The terms used herein should generally be construed as understood by the average person skilled in the art, unless explicitly indicated otherwise. The following explanations may guide the understanding:

[0054] A "linear recurrent neural network (LRNN)" may refer to a class of neural network architectures designed for processing sequential data, characterized by a recurrence relation that is linear with respect to the hidden state, often incorporating gating mechanisms. Examples include the xLSTM, specifically its matrix memory variant (mLSTM), as well as other architectures such as RetNet, Mamba, RWKV, and DeltaNet. These networks typically utilize matrix multiplications to update states and generate outputs, distinguishing them from non-linear RNNs like traditional LSTMs or GRUs, and enabling parallelization strategies such as the chunkwise-parallel formulation.

[0055] A "data processing apparatus" may refer to any computing hardware or system capable of executing the operations of the neural network, particularly those involving large-scale matrix computations. Examples include Graphics Processing Units (GPUs) like the NVIDIA H100, A100, or V100, as well as Tensor Processing Units (TPUs), Neural Processing Units (NPUs), or other specialized Al accelerators. In certain embodiments, the apparatus comprises a hierarchical memory structure including off-chip and on-chip memory.

[0056] An "off-chip memory" may refer to a high-capacity memory storage component located external to the processor's main compute die, characterized by higher latency and lower bandwidth compared to on-chip memory. Examples include High Bandwidth Memory (HBM), Graphics Double Data Rate (GDDR) memory, DRAM, or what is generally referred to as "global memory" or " GPU memory" in the context of GPU computing. This memory can be used to store the full model weights, the complete input sequence, and materialized intermediate states like the recurrent memory states for each chunk.

[0057] An "on-chip memory" may refer to a high-speed, lower-capacity memory located directly on the processor die, close to the compute units, facilitating fast data access and data sharing between threads or compute elements. Examples include Static Random Access Memory (SRAM), shared memory (in CUDA terminology), L1 or L2 caches, and registers. This memory can store the smaller "tiles" of data during the intra-chunk parallel processing to maximize arithmetic intensity.

[0058] A "chunk" may refer to a segment or partition of the input sequence, created by dividing the full sequence length into parts of a given length. Examples include partitions of size 128, 256, or larger, trading off memory usage and runtime. In the disclosed two-level parallelization strategy, chunks serve as the unit for the first level of parallelization where recurrent states are materialized in off-chip memory to handle inter-chunk dependencies.

[0059] A "tile" may refer to a further subdivision of a chunk, created by partitioning the chunk into smaller segments of a given size, specifically chosen to fit within the on-chip memory. Examplesinclude blocks of size 64 or smaller that allow for efficient matrix multiplication on tensor cores. Tiles serve as the unit for the second level of parallelization (intra-chunk tiling), enabling computations to proceed in parallel without materializing intermediate states for every step back to off-chip memory.

[0060] A "recurrent memory state" may refer to the hidden state or context information maintained by the LRNN that encapsulates the history of the sequence up to a certain point. Examples include the matrix memory state Ct(or Ck) in an mLSTM, which might be a dqkx dflvmatrix, and the normalizer state nt. This state can be materialized in off-chip memory at chunk boundaries to pass information between the chunkwise-parallel processing steps.

[0061] A "recurrent compute kernel" may refer to a software routine or function executed on the data processing apparatus that handles the sequential part of the computation. Examples include a GPU kernel responsible for iterating through the sequence of chunks, computing the necessary state updates (like the matrix Ck), and writing these materialized states to off-chip memory. This kernel ensures the correct propagation of information across the chunks.

[0062] A "parallel compute kernel" may refer to a software routine or function executed on the data processing apparatus that performs computations for multiple data elements simultaneously. Examples include a GPU kernel that processes the tiles within chunks in parallel to compute output states (H(k)), leveraging the massive parallelism of devices like GPUs. This kernel reads the materialized states and local input data to generate the final outputs for each position in the sequence.

[0063] An "embedding dimension" may refer to the size of the vector representation for each token or element in the input sequence, or the internal feature dimension of the model. Examples include dimensions denoted as dhv(value head dimension) or dqk(query / key head dimension). The disclosed invention specifically introduces parallelization over this dimension (e.g., splitting the computation for large dhvacross multiple thread blocks) to handle large models efficiently. An "input gate activation function" may refer to the mathematical function applied to the input gate pre-activations to determine how much new information flows into the memory state.

[0064] Examples include the exponential function (exp) used in standard xLSTM / mLSTMexp, and the sigmoid function (a) used in the proposed mLSTMsig variant. The term "bounded in a positive value range" specifically refers to functions like the sigmoid that do not grow to infinity for large positive inputs, contrasting with the exponential function.

[0065] An "mLSTMsig" may refer to a specific variant of the matrix LSTM architecture characterized by the use of a sigmoid function for the input gate activation. Examples described in the disclosure include an architecture which eliminates the need for a normalizer state (nt) and stabilizer state (mt) required by the exponential variant, thereby reducing register pressure and computational complexity.

[0066] A "fused loop" may refer to a programming optimization where multiple loops or operations are combined into a single kernel execution or a single loop structure to minimize overhead andmemory access. Examples include combining the matrix multiplications for the query, key, and value projections within the intra-chunk parallel kernel, or fusing the accumulation steps. In the context of mLSTMsig, the removal of the max state dependency allows for more efficient loop fusion compared to mLSTMexp.

[0067] " Computational intensity" (or arithmetic intensity) may refer to the ratio of floating-point operations (FLOPs) performed to the amount of data (bytes) transferred to and from memory during a computation. Examples include a metric used to determine if a kernel is computebound or memory-bound. The disclosed technology aims to increase this intensity by scaling the chunk size proportionally to the square root of the hardware's computational intensity, thereby performing more math per byte loaded from off-chip memory.

[0068] Tiled Flash Linear Attention Overview

[0069] FIG. 1 illustrates a schematic overview of a linear recurrent neural network 100 (LRNN) embodying the Tiled Flash Linear Attention (TFLA) framework according to certain embodiments. The linear recurrent neural network 100 processes an input sequence 102 by dividing it into chunks 104

[0070]

[0071] (e.g., QKV(1), QKV(2)). A recurrent kernel 110 manages the inter-chunk dependencies by materializing memory states 106 (C0, C1, C2). More precisely, the recurrent compute kernel 110 materializes the memory state Ck-1for each chunk. A parallel kernel 112 handles the intra-chunk computation, utilizing tiling strategies to compute output states 108 (H(1), H(2), H(3)). This architecture enables the decoupling of parallelization granularity from state materialization frequency. In certain embodiments, the linear recurrent neural network 100 uses tiling for the 3 matrix-multiplications in the parallel kernel to fully utilize the hardware and to prevent materialization of many memory states.

[0072] FIG. 2 illustrates a flowchart of a method 200 for processing an input sequence 102 using a linear recurrent neural network, such as the linear recurrent neural network 100 of FIG. 1.

[0073] In certain embodiments, the method 200 comprises a chunking step 202 in which the input sequence 102 is partitioned into a plurality of chunks 104, each chunk 104 having a chunk size. A "chunk" may refer to a segment of the input sequence length T divided into Ncparts of length L. For example, if the sequence length T is 4,096 and the chunk size L is 256, the sequence is partitioned into 16 chunks. This partitioning serves as the first level of sequence parallelization.

[0074] Optionally, the chunk size L may be a tunable parameter configured based on the memory capacity of the target hardware or the specific requirements of the model. In some implementations, the chunk size is selected to be significantly larger than what would typically fit in the fast on-chip memory of a graphics processing unit (GPU), as the method does not require the entire chunk to be processed atomically within the on-chip memory. Alternatively, the chunks may be variable in length, although fixed-length chunks are preferred for computational regularity.In certain embodiments, the method 200 further comprises performing a recurrent processing step 204 across the plurality of chunks 104 to materialize a recurrent memory state 106 for each chunk 104 in the off-chip memory. The recurrent memory state 106, such as a matrix Ck, encapsulates the history of the sequence up to the boundary of the chunk 104. Materializing this state "in the off-chip memory" implies writing the state data to a high-capacity storage area accessible by the compute units, such as High Bandwidth Memory (HBM).

[0075] Optionally, this recurrent processing step 204 is performed by a dedicated kernel that iterates through the chunks 104 sequentially, such as the recurrent compute kernel 110 of FIG. 1. In some implementations, only the final state of each chunk 104 is materialized to off-chip memory, while intermediate states within the chunk 104 are transient or not calculated at this stage. It is also possible that the recurrent memory state 106 includes additional state variables, such as a normalizer state nkor a stabilizer state mk, which can also materialized to the off-chip memory. In certain embodiments, the method 200 comprises a tiling step 206 in which each of the plurality of chunks 104 is partitioned into a plurality of tiles, each tile having a tile size smaller than the chunk size, such that a tile fits within the on-chip memory. A "tile" may refer to a subsegment of size B, where B « L. For instance, if the chunk size L is 256, the tile size B might be 64. This partitioning constitutes a second level of sequence parallelization. The constraint that the tile "fits within the on-chip memory" refers to the requirement that the input data (Queries, Keys, Values) and the intermediate accumulation buffers for that tile can be stored in the fast memory (e.g., SRAM) without eviction during the computation of that specific tile.

[0076] Optionally, the tile size is determined by the specific hardware characteristics, such as the size of the SRAM or L1 cache per streaming multiprocessor. In some implementations, the tile size corresponds to the block size used for matrix multiplication operations on tensor cores.

[0077] Alternatively, the tiling strategy may involve dynamic tile sizing based on the available shared memory resources at runtime.

[0078] In certain embodiments, the method 200 further comprises performing a parallel processing step 208 for the plurality of tiles, wherein the parallel processing step 208 comprises loading the materialized recurrent memory state 106 of the corresponding chunk 104 from the off-chip memory into the on-chip memory and computing output states 108 for the tiles within said chunk 104 in parallel using the on-chip memory. This involves fetching the initial condition for the chunk 104 (the materialized state from the previous chunk 104) once, and then performing computations for the tiles.

[0079] Optionally, computing the output states 108 includes adding the contribution of the recurrent state to the local intra-chunk computations. In some implementations, the intra-chunk computation is performed as a matrix multiplication involving the queries, keys, and values within the chunk 104, and this matrix multiplication is tiled according to the tile size. It is also possible that the parallel processing step utilizes a "local state aggregation" strategy, where a running state is maintained in registers or shared memory across the tiles of a chunk, initialized by the loaded materialized state.mLSTM Formulations

[0080] The mLSTM cell is the fully parallelizable part of the xLSTM. It has a matrix memory and exponential gating.

[0081] Recurrent Formulation

[0082] In its recurrent formulation, the mLSTM cell processes the series of input vectors xt∈ ℝdfor time steps t e {1,..., T} mapping a state

[0083]

[0084] to a successor

[0085] state (ht, Ct, nt, mt) given an input xt. Here, ht∈ ℝddenotes the hidden state, Cte denotes the cell state responsible for long-term memory, nt∈ ℝddenotes the normalizer state, and mt∈ ℝ denotes the max state. Together normalizer and max state control the magnitude of the exponential input gate and ensure stability. The recurrent mLSTM formulation is given by the following state update equations:

[0086] mt= max{log σ(f̃t) + mt-1, ĩt} (1 )

[0087] Ct= ftCt-1+ itktvtT(2)

[0088] nt= ftnt-1+ itkt(3)

[0089] h̃t= CtT(qt / √dqk) / max{|ntT(qt / √dqk)|, exp(-mt)}

[0090]

[0091] ht= ot⊙ NORM(h̃t) (5)

[0092] The scalar forget and input gates it, fte IK are computed as ft= exp(log σ(f̃t) + mt-1− mt) and it= exp(ĩt− mt) with the pre-activations

[0093]

[0094] + b{i,f}, respectively. The vector output gate ote JKdft” is given by ot= σ(õt) with the pre-activations õt= Woxt+ boand the sigmoid function o. The norm layer NORM in (5) can be either RMS norm or LayerNorm. Typically, multiple of these cells operate simultaneously as parallel heads, similar to Transformers.

[0095] Chunkwise-Parallel Formulation

[0096] The chunkwise-parallel formulation is a trade-off between the parallel and the fully recurrent formulation. It has a recurrent part and a (quadratic) parallel part, with an overall sub-quadratic scaling in sequence length. Similar to the fully parallel formulation, we assume that all inputs are available at once. We then split the sequence of length T into Nc= ⌈T / L⌉ chunks of length L and use k e {1,..., Nc} 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 f̃ ∈ ℝTare rearranged into a matrix f̃ =

[0097]

[0098] ..., f̃(N)) ∈ ℝN×L, where each row f̃(k)= (f(k-1)N+1, f(k-1)N+2, ..., fkN) ∈ ℝLcontains 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 iKWcXLxd‘ / fcand V eℝN×L×d. Here, the query matrix Q(k)= (

[0099]

[0100] q(k-1)N+1, ..., qkN) ∈ ℝL×dcontains the query vectors of chunk k. Keys, and values follow analogously. For notational simplicity we drop the leading Ncdimension and omit normalization layer and the output gate, i.e. consider h̃tas hidden state outputs.

[0101] Given the logarithmic forget gates f̃(k)= log σ(f̃(k)) ∈

[0102]

[0103] and input gates ĩ(k)= log exp(

[0104]

[0105] ĩ(k)) ∈ ℝL, we can compute the logarithmic chunkwise gates gk∈ ℝ, and bk, ak∈

[0106]

[0107] as gk= s

[0108]

[0109] um(f̃(k)), bk= cumsum(f̃(k)), and ak= rev_cumsum(f̃(k)) + ĩ(k). In FIG. 3, we show the summed forget gates gkcontain the forget gate contribution of all forget gates within a chunk. The cumulative forget gate vectors bkcontain 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 akcontain the input gates for every timestep as well as the forget gate contributions from the current time step to the end of the chunk.

[0110] Inter-chunk Recurrent Contribution.

[0111] The inter-chunk recurrence is given by

[0112] Ck= ḡkCk-1+ (āk⊙ K(k))TV(k)(6)

[0113] n

[0114]

[0115] k= gknk-±+ (dkO K(k))Tl, (7)

[0116] where gkand akare the stabilized chunkwise gates. This recurrent part resembles the fully recurrent formulation described above, but instead of computing the intermediate states for every timestep t, we compute them directly for every L time steps without materializing the states in between.

[0117] Intra-chunk Parallel Contribution.

[0118] The recurrent part is followed by the intra-chunk parallel contribution:

[0119] fori < j

[0120] nW = f

[0121] [(8

[0122] bk)

[0123] - bl + fori > j

[0124] S̃(k)= (1 / √dqkQ(k)K(k)T) ⊙ D(k)= S(k)⊙ D(k), (9)

[0125]

[0126] y / dqk

[0127] where D(k)∈ ℝL×Lis the stabilized gate matrix. Compared to the fully parallel part, the quadratic cost of the matrices

[0128]

[0129] ∈ ℝL×Lis greatly reduced, since the chunk size L is typically small compared to the sequence length T.

[0130] Output Computation.

[0131] Finally, the contributions from the intra-chunk parallel part H(k)intraare combined with the interchunk recurrent part H(k)interto obtain the hidden states H(k)∈ ℝL×dfor each chunk k (see FIG.

[0132] 1):n-S>r= (S* oya-t = Qmck.t, H^ = s< W‘>. (io)

[0133] V“qfc

[0134]

[0135] H(k)= (H(k)inter+ H(k)intra) / h(k)denom, (11)

[0136] where h(k)denom∈

[0137]

[0138] is a normalization factor.

[0139] Tiled Flash Linear Attention

[0140] Flash Linear Attention introduces a fast kernel algorithm for the chunkwise formulation for Linear Attention and shows that their implementation is faster than optimized FlashAttention. This speedup is achieved by single level sequence parallelism, where the states Ckare first materialized in GPU memory and then the outputs H(k)are 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 L = 64) by the physical SRAM size of the GPU.

[0141] Therefore, we have to materialize many states in HBM, where the number of states is Nc= ⌈T / L⌉. This leads to low arithmetic intensity and high GPU memory consumption, which poses challenges especially for long-context pre-training.

[0142] More Efficient Kernels via Two Level Sequence Parallelism.

[0143] 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. The first level is the parallelization over the chunks of the sequence, which requires to compute and materialize intermediate states Ckin GPU High Bandwidth memory (HBM). For this we use a recurrent kernel in certain embodiments. The second level is the parallelization within each chunk, which is achieved in certain embodiments by tiling the intra chunk attention matrix along the chunk dimension. This second level of parallelism enables large chunk sizes and hence reduces the memory consumption for the intermediate states as we have to store and load Nc= ⌈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, certain embodiments of TFLA also parallelize over the embedding dimension. This enables arbitrary large head dimensions and results in a massive parallelization over five dimensions, which is crucial for achieving high performance on modern GPUs.

[0144] Forward Pass.

[0145] 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 gate computations and normalization, as 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).Kernel LhqLkvdqkdhvThread Block Grid

[0146] Ipl dhvLhq\

[0147] H(k)P L

[0148] y &dhv “Lhq ]

[0149] ( dqk ^hq \

[0150] <5Q(fe)P L.... I

[0151] nBd,qk, ’ ° nLTh,q ’ f l

[0152] 8K^k)L P p I | <^qk ^kv |

[0153] y ^dqk °Lkv ]

[0154] L P [_ p | ^-“kv |

[0155] \ Bdhv ^Lkv I

[0156]

[0157] 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 dimensions of the 3D kernel launch grid. The last dimension of all kernels is Nchunk· Nhead· Nbatch.

[0158] In simplified form, the intra-chunk parallel forward pass of the mLSTM (and other linear RNNs) for a chunk k can be written as three matrix multiplications, which we fuse into a single kernel:H(k)=( Q< M ) y(fc)+Q (k}Ck_! (1 2)

[0159] (. LhqXdfrv') (Lhq*dqk')(dqkXLkv) T / cvxd / !v)

[0160] H(fc)HW

[0161]

[0162] nintraninter

[0163] 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 equation (12), we introduce the block

[0164] sizes BLhq, BLkv, Bdqkand Bdflvfor the attention matrix, query, key, value and hidden state dimensions Lhq, Lkv, dqkand dhv, along which we either parallelize or accumulate by using a loop inside the kernel.

[0165] FIG. 4 shows one embodiment of our TFLA tiling strategy for the forward pass H(k)kernel. We parallelize across the outer sequence dimension Lhqwith NLhq= Lhq / BLhqprograms, and across the outer embedding dimension dhvwith Ndhv= dhv / Bdhvprograms. We loop over the inner dimensions Lkvand dqk, which are tiled by the block sizes BLkvand Bdqkrespectively.

[0166] Tiled Computation.

[0167] For the mLSTM we cannot simply accumulate the results of the matrix

[0168] multiplications H(k)intraalong the Lkvdimension and

[0169]

[0170] due to the stabilization of the exponential input gate with the max state mt. The max state tracks the maximum of the forget and input gates over time and is used to stabilize the exponential input gate similar to the safe softmax computation. Since we compute the hidden state output H(k)in blocks along the chunk size (i.e. time) dimension Lkv, we need to rescale during accumulation of the block results

[0171]

[0172] for and the overall results into H(k)in the same way as FlashAttention. 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.Backward Pass.

[0173] 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. In certain embodiments, the individual gradients are 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 1 summarizes the work partitioning of our TFLA kernels.

[0174] Faster mLSTM with Sigmoid Input Gate

[0175] The mLSTM with exponential gating (i.e. exponential input gate) requires to compute and keep track of two additional states, the normalizer state ntand max state mt. Both will increase kernel runtime: The normalizer must be computed through summations, and tracking the max state throughout the tiled computation in TFLA prevents efficient fusing of loops within the kernel. Additionally, our analysis suggests to initialize the input gate biases at 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 0 during training. 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 to use the sigmoid function σ(x) = 1 / (1+exp(−x)) =

[0176]

[0177] exp(x)+i’ converges to exp(x) for x → −∞ and 1 for x → ∞.

[0178]

[0179] mLSTM with Sigmoid Input Gate

[0180] The sigmoid function can be computed in two ways as given above. Depending on the sign of x it can be ensured that the argument of exp is always smaller than 0 to avoid numerical overflow. Therefore, we do not need to control the magnitude of x externally with a max state and as a consequence also drop the normalizer state. This yields the mLSTM with sigmoid input gate (henceforth referred to as mLSTMsig) in its recurrent formulation as

[0181] Ct= σ(f̃t)Ct-1+ σ(ĩt)ktvtT(13)

[0182] h̃t= CtT(qt / √dqk) (14)

[0183]

[0184] = σ(õt) ⊙ NORM(h̃t) (15)

[0185] where the query, key, and value vectors qt,kt,vt, and the gate preactivations

[0186]

[0187] remain the same as for the mLSTM with exponential input gate (from now on referred to as mLSTMexp). We confirm that our TFLA mLSTMsig forward kernel is over 30% faster than the mLSTMexp forward, and show that mLSTMsig performs equally well compared to mLSTMexp in our language modeling experiments up to 1.4B parameters.Normalization of mLSTM and Linear RNNs

[0188] Motivated by the performance of mLSTMsig, we seek to understand the differences between mLSTMsig 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. In our case, we analyze the transfer behavior of mLSTMsig and mLSTMexp for random

[0189] inputs qt,ktand vtand different input gate and forget gate preactivations itand ft.

[0190] 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 with input vector input vector x e Wdand output vector y e Wdis defined as y = x / RMS(x) ⊙ γ, where RMS(x) = √(1 / d Σxi2+ ε) with

[0191]

[0192] with γ ∈ ℝdbeing a learnable scale parameter. The epsilon parameter ε ∈ ℝ is a small constant typically set to 1e-6 to avoid division by zero.

[0193] Transfer Behavior of the mLSTM.

[0194] We analyze the transfer behavior by computing the gain of the mLSTM cells from random inputs sampled from N (0,1) to hidden states before and after the norm layer for varying input and forget gate values. More specifically, we compute the gains Gbeforeand Gafteras

[0195] Gbefore= ‖h̃t‖max / ‖vt‖maxand Gafter= ‖NORM(h̃t)‖max / ‖vt‖max, (10)

[0196]

[0197] where ‖x‖max:= max(|x1|, ..., |xd|) 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.

[0198] Normalization Layers in other Gated Linear RNNs.

[0199] Interestingly, almost all other gated linear RNN variants also place a normalization layer after the RNN cell. Often this is justified with improved training stability, but a more thorough discussion is missing. Some research analyzes the effect of the norm layer after a non-gated, kernel-based linear attention layer and shows that this effectively prevents unbounded gradients. We also confirm that the norm layer has a significant impact on training stability and the gradient norm during training. Initializing the input gate bias at larger negative values, as suggested by our transfer behavior analysis in FIGs. 5A-5B, 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.Effect of Normalization on Gating in Linear RNNs.

[0200] 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 in the information routing or gating mechanism of the linear RNN. For example, if the squared norm of CtTqt, which is controlled by input and forget gates through CtTqt, is smaller than the epsilon, the denominator in the NORM(x) layer is dominated by e and the output moves towards zero (indicated by the lower areas in FIGs. 5A-5B). Hence, by moving through the x-y plane in FIGs. 5A-5B, the gates could learn to suppress or amplify any input in the sequence.

[0201] Experiments

[0202] 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_chunk: A kernel that is limited in chunk size L. (2) xl_chunk: Our Tiled Flash Linear Attention (TFLA) kernels with unlimited chunk size. We assess the performance of mLSTMsig compared to mLSTMexp and benchmark the runtime of our kernels against other baselines.

[0203] Language Modeling with mLSTM

[0204] We train three different model sizes (160M, 400M, 1.4B parameters) with context lengths 4096 and 8192 on the DCLM dataset. We include Llama2 style Transformer models as reference in our comparison.

[0205] Software and Hardware Setup.

[0206] We run our language modeling experiments in JAX 0.4.34 and use FLAX 0.9.0 to implement our models. We implement our kernels in Triton 3.1.0 and use JAX-Triton 0.2.0 to integrate the kernels into JAX. Our kernel benchmark experiments are run in PyTorch 2.5.1, because most kernel baselines are available in PyTorch. All experiments are run on NVIDIA H100 80GB GPUs.

[0207] Model Architecture.

[0208] The model architecture for mLSTMexp and mLSTMsig follows the design of most dense Transformer decoder only large language models.

[0209] An embedding layer, is followed by a stack of blocks and a language model head that produces the output logits (i.e. the values before softmax), which typically consists of a normalization layer and a linear (unembedding) layer. We apply logit soft-capping, such that the value of the logits stay between —c and c for a specific cap value c. We choose c = 30. The logits are capped with the following function:

[0210] softcap(x) = c ■ tanh(x / c)We use the GPT-NeoX tokenizer with vocabulary size 50257 and do not tie the weights for the embedding layers and the last (unembedding) layer.

[0211] Each block consists of two layers, where each layer has skip a connection and a normalization layer before the layer input (i.e. we use the pre-norm block architecture). As normalization layer we use the RMS-norm with epsilon e = 1e-6.

[0212] 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. 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 is a feedforward linear layer that mixes the tokens per timestep channelwise. We use the SwiGLU feedforward linear layers.

[0213] For the mLSTM we set the head dimension for the queries and keys to be half of the values, i.e. dqk= 0.5dflv. We use Layernorm as NORM(x) operation with epsilon e = 1e-6 in our experiments. We apply soft-capping from the above equation on the input and forget gate preactivations, as we found that this improves training stability. For the gate preactivations we set c = 15.

[0214] Training Recipe.

[0215] We train our models with the AdamW optimizer with

[0216]

[0217] = 0.9, β2= 0.95 and e = 1e-8. We apply a weight decay of 0.1 to 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 a 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. 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 a large negative values at the beginning of each new document.

[0218] Performance in Language Modeling.

[0219] We compare mLSTMsig and mLSTMexp models on next-token prediction with different number of heads or head dimensions. Table 2 shows the results for context length 4096.EXP EXP SIG

[0220] Size Tokens Heads Llama

[0221] LIMIT XL XL

[0222] 160M 19B

[0223]

[0224] 6 21.03 21.18 21.03

[0225] 12 20.89 21.03 21.06 21.05

[0226] 400M 24B 4 16.66 16.66 16.67

[0227] 8 16.55 16.80 16.67

[0228] 16 16.85 16.60 16.61 16.61

[0229]

[0230] 1.4B 33B 4 | 13.31 13.35 13.34

[0231] 8 13.20 13.22 13.21

[0232] 16 13.64 13.20 13.87* 13.22

[0233]

[0234] Table 2: Validation Perplexity at context length 4096. EXP and SIG denote mLSTMexp and mLSTMsig. LIMIT and XL correspond to limit_chunk and xl_chunk kernels. We find that our limit_chunk and xl_chunk kernels yield the same loss (up to 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 a main result we find that mLSTMsig performs equally well compared to mLSTMexp.

[0235] Effect of Input Gate Bias Initialization.

[0236] We analyze the effect of the input gate bias initialization on training stability and performance of our mLSTM models. We observe 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, but also improves training stability, if initialized correctly.

[0237] Effect of Norm Layer Epsilon.

[0238] We investigate the effect of the norm layer epsilon on language modeling performance for mLSTMexp. Our transfer behavior analysis in FIGs. 5A-5B suggests, that there exists an interplay between norm layer epsilon and input gate bias initialization. We confirm this in a grid search and find that the best performing configuration is the default epsilon e = 1e-6 with input gate biases initialized to -10.

[0239] Kernel Benchmark

[0240] 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, Mamba and GLA. We run our benchmarks on NVIDIA H100 GPUs.Runtime Benchmark.

[0241] 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. We keep the number of tokens constant at 65,536 and vary sequence length and batch size accordingly. FIG. 6 shows the runtime benchmark results for inference, i.e. forward pass only, (left) and for training, i.e. forward-backward pass (right). Our mLSTMexp TFLA xl_chunk kernels with two level sequence parallelism is about 25% faster than our limit_chunk kernels. Through targeted modifications of the input gate of the mLSTM we save computation and enable more efficient kernel implementations for the forward pass of mLSTMsig. This yields another speedup of over 30% for the forward pass of the mLSTMsig TFLA kernel over the mLSTMexp TFLA kernel.

[0242] In training, our TFLA kernels are faster than FlashAttention 3 for longer sequences and more than 2x faster than Mamba 2 kernels for all sequence lengths.

[0243] Runtime vs. Memory Trade-off.

[0244] The chunk size parameter L balances the computation between the two levels of sequence parallelism. Smaller chunk sizes increase memory consumption, because 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 consumption, but increase quadratic compute FLOPs. In FIG. 7, we measure this trade-off for our mLSTMsig TFLA xl_chunk kernels.

[0245] Theoretical Runtime and Arithmetic Intensity

[0246] In FIG. 7, we empirically observe that there exists an optimal chunk size (between 128 and 256) at which the runtime is minimized. We find that the theoretical runtime optimum exceeds the empirically observed one, suggesting that our current kernel implementations may not yet fully exploit the available performance potential. We analyze the effect of the chunk size parameter L on runtime, FLOPs, and arithmetic intensity and summarize our findings as follows: (1) The chunk size L mediates a trade-off between runtime and GPU memory usage [FIG. 7], (2) L determines the total compute in FLOPs: L = 1 matches the recurrent formulation, while L = T matches the parallel one. (3) There exists an optimal chunk size L ∈ [1, T] that minimized the total FLOP count. (4) Increasing L raises the arithmetic intensity of TFLA kernels. (5) The chunk size determines whether the kernel is memory-bound or compute-bound on a given hardware. (6) FLOPs / s alone can be misleading; the optimal chunk size should be chosen based on total runtime. (7) The runtime-optimal chunk size scales proportionally with the square root of the head dimension and the accelerator’s computational intensity. (8) Newer hardware generations require larger chunk sizes to approach peak performance.

[0247] Certain embodiments have been described to provide a hierarchical partitioning (chunks > tiles) combined with the selective materialization of states (only at chunk boundaries). This approachresolves the physical contradiction between processing speed (which requires parallelism and thus small, independent units) and memory bandwidth (which is bottlenecked by frequent writes of large states). By nesting a fine-grained tiling structure (for speed) within a coarse-grained chunking structure (for state management), the method minimizes the volume of data written to slow memory while maximizing the utilization of fast compute units. This effectively decouples the degree of parallelization from the memory I / O cost, a non-obvious combination of the principles of “Segmentation” (dividing the sequence) and " Local Quality" (optimizing behavior differently for on-chip vs. off-chip memory).

[0248] Implementation in Hardware and Software

[0249] Embodiments of the disclosed technology can be implemented in hardware, software, or both. FIG. 8 illustrates a schematic block diagram of computer hardware on which embodiments of the present disclosure can be implemented. As can be seen, a data processing apparatus 802 comprises one or more processors 804 and one or more memories 806 or storage devices. In embodiments, the data processing apparatus 802 may be a local computer device (e.g., personal computer, laptop, tablet computer or mobile phone), a remote computer device (e.g., remote server or server cluster), or a distributed computer system (e.g., a cloud computing system with one or more processors 804 and one or more memories 806 distributed at various locations, for example, at one or more local clients and / or one or more remote servers). The data processing apparatus 802 may comprise any circuit or combination of circuits.

[0250] In embodiments, the one or more processors 804 may comprise any type of computational circuit, such as but not limited to a microprocessor, a microcontroller, a complex instruction set computing (CISC) microprocessor, a reduced instruction set computing (RISC) microprocessor, a very long instruction word (VLIW) microprocessor, a graphics processor, a digital signal processor (DSP), multiple core processor, a field programmable gate array (FPGA), or any other type of processor or processing circuit. Other types of circuits that may be included in the computer system may be a custom circuit, an application-specific integrated circuit (ASIC), or the like, such as, for example, one or more circuits (such as a communication circuit) for use in wireless devices like mobile telephones, tablet computers, laptop computers, two-way radios, and similar electronic systems.

[0251] In embodiments, the one or more memories 806 or storage devices may include one or more memory elements suitable to the particular application, such as a main memory in the form of random access memory (RAM), one or more hard drives, and / or one or more drives that handle removable media such as compact disks (CD), flash memory cards, digital video disk (DVD), and the like.

[0252] In embodiments, the data processing apparatus 802 may also include a display device 812, one or more speakers, and one or more input devices such as a keyboard 814, a computer mouse 816, trackball, touch screen, voice-recognition device, or any other device that permits a user to input information into and receive information from the data processing apparatus 802.In embodiments, the memory 806 stores a computer program 808. The computer program 808 may also be referred to as a program, software, a software application, an app, a module, a software module, a script, or code. The computer program 808 may be written in a programming language, including compiled or interpreted languages. The computer program 808 may be deployed in any form, including as a stand-alone product or as a module, component, subroutine, or other unit suitable for use in a computing environment, such as the data processing apparatus 802. The computer program 808 may implement some or all aspects of the methods and functionalities of the disclosed technology. Accordingly, embodiments of the disclosed technology can be implemented as a computer program (product) with program code, the program code being operative for performing one of the methods when the computer program runs on a computer. The program code may be stored on a machine-readable carrier. In embodiments, a data stream or a signal sequence may be provided, representing the computer program 808 for performing one of the methods and functionalities described herein. The data stream or the signal sequence may, for example, be configured to be transferred via a data communication connection, for example, via the internet.

[0253] In embodiments, a data processing system, device or apparatus may be provided, e.g., in the form of a file server, configured to transfer the computer program 808 for performing one of the methods and functionalities described herein to a receiver. The receiver may itself comprise any type of data processing system, device or apparatus.

[0254] In certain embodiments, the off-chip memory is a High Bandwidth Memory (HBM) or a GPU global memory, and the on-chip memory is a Static Random Access Memory (SRAM) or a shared memory. FIG. 8 illustrates a data processing apparatus 802 comprising such a memory hierarchy. The apparatus includes DRAM (Off-chip memory), which may be integrated with memory 806, storing inputs (QKV), states (Ck), and outputs (H), and a graphics processing unit 810 (GPU) with streaming multiprocessors (SM) containing L1 / Shared Memory (On-chip memory).

[0255] Optionally, the off-chip memory may refer to GDDR6, HBM2e, HBM3, or system RAM in unified memory architectures. The on-chip memory may refer to the L1 cache, the programmable shared memory (as in NVIDIA CUDA architectures), or a scratchpad memory. In some implementations, the movement of data between these memories is explicitly managed by the software kernel (e.g., using asynchronous copy instructions). Alternatively, the distinction may apply to different levels of cache in a TPU or other AI accelerator.

[0256] Implementing embodiments of the method 200 using these memory types grounds the method 200 in the physical reality of modern high-performance computing. The "partitioning" steps are not merely abstract mathematical operations but are physical data management strategies designed to exploit the specific bandwidth and latency characteristics of HBM versus SRAM. This resolves the contradiction between the need for large model context (requiring large HBM) and low-latency computation (requiring fast SRAM).A GPU (Graphics Processing Unit) is a 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 is important to understand the software hierarchy of the GPU, which closely follows its physical hardware hierarchy.

[0257] 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 are grouped into Thread blocks which together execute a 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, a 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-bandwidth memory (HBM), which has both the largest latency and least bandwidth of all GPU memories. For efficient kernels it is important to minimize memory read and writes from and to HBM.

[0258] 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 tensor core utilization.

[0259] Triton Language: Triton is a GPU kernel programming language with an associated compiler, that provides a Python-based environment for GPU programming. The user can load data from HBM via a tl.load instruction and store data to HBM via tl.store, tl.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 a 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 on the Thread block level and hide register and thread management from the user.

[0260] Further Examples

[0261] Example 1. A system comprising a machine-learning model implemented on a data processing apparatus, comprising a linear recurrent neural network (LRNN) as the machinelearning model, wherein the LRNN is configured for Tiled Flash Linear Attention (TFLA).

[0262] Example 2. The system of example 1, wherein the LRNN comprises a recurrent kernel and a parallel kernel configured to process an input sequence in chunks.

[0263] Example 3. The system of example 2, wherein the recurrent kernel is configured to materialize a memory state for each chunk.Example 4. The system of example 2 or 3, wherein the parallel kernel is configured to compute an output state for all chunks.

[0264] Example 5. The system of any one of examples 2-4, wherein the parallel kernel is configured to use tiling for matrix multiplications, preferably within each chunk.

[0265] Example 6. The system of any one of the preceding examples, wherein the LRNN is an xLSTM with matrix memory (mLSTM).

[0266] Example 7. The system of example 6, wherein the mLSTM comprises an input gate activation function which is similar to the exponential function in the negative range but bounded in the positive range.

[0267] Example 8. The system of example 6 or 7, wherein the mLSTM is an mLSTM with sigmoid input gate (mLSTMsig).

[0268] Example 9. The system of any one of examples 2-8, wherein the recurrent kernel and the parallel kernel are graphics processing unit (GPU) compute kernels.

[0269] Example 10. A method of operating the system of any one of examples 1 -9.

[0270] Example 11. A data processing apparatus comprising means for carrying out the method of example 10.

[0271] Example 12. A computer program comprising instructions which, when the program is executed by a computer, cause the computer to carry out the method of example 10.

[0272] Example 13. A computer-readable medium having stored thereon the computer program of example 12.

[0273] While various aspects, embodiments and examples have been illustrated and described in the foregoing description and the drawings, such illustration and description is illustrative or exemplary and not restrictive. Variations to the disclosed aspects, embodiments and examples can be understood and effected by those skilled in the art in practicing the disclosed technology. Insofar as aspects, embodiments and examples are described, explicitly or implicitly, in the context of a product, apparatus, device or system, such description should be understood as also describing a corresponding process, method or use, and vice versa. Accordingly, aspects, embodiments and examples described in the context of a block, component or feature of a product, apparatus, device or system also represent a description of a corresponding method step or feature of a method or method step. Analogously, aspects, embodiments and examples described in the context of a method or method step also represent a description of a corresponding block, component or feature of a product, apparatus, device or system.

[0274] Insofar as a method is described, explicitly or implicitly, as comprising actions of a particular entity or component, such description should be understood as also describing acomplementary method of a counterpart entity or component. For example, a method described in terms of operations carried out by a client device also discloses a complementary method carried out by a server device that interacts with the client device, and vice versa. In this manner, the disclosed technology encompasses all perspectives of an interaction, regardless of whether a particular embodiment is drafted from one selected perspective.

[0275] The order of execution of the operations in the described embodiments is not essential, unless otherwise specified. That is, the operations may be performed in any order, unless otherwise specified, and embodiments may include additional or fewer operations than those mentioned. In the claims, the word “comprising” does not exclude other elements or steps, and the indefinite article “a” or “an” does not exclude a plurality. A single entity may fulfill the functions of several entities recited in the claims. The mere fact that certain measures are recited in mutually different dependent claims does not indicate that a combination of these measures cannot be used to advantage.

[0276] Particular and preferred aspects of the disclosed technology are set out in the accompanying claims. Features from the dependent claims may be combined with features of the independent claims and with features of other dependent claims as technically appropriate and not merely as explicitly set out in the claims.

Claims

CLAIMS1. A method for processing an input sequence using a linear recurrent neural network, LRNN, implemented on a data processing apparatus having an off-chip memory and an on-chip memory, the method comprising:partitioning the input sequence into a plurality of chunks, each chunk having a chunk size;performing a recurrent processing step across the plurality of chunks to materialize a recurrent memory state for each chunk in the off-chip memory;partitioning each of the plurality of chunks into a plurality of tiles, each tile having a tile size smaller than the chunk size, such that a tile fits within the on-chip memory; and performing a parallel processing step for the plurality of tiles, wherein the parallel processing step comprises loading the materialized recurrent memory state of the corresponding chunk from the off-chip memory into the on-chip memory and computing output states for the tiles within said chunk in parallel using the on-chip memory.

2. The method of claim 1, wherein the off-chip memory is a High Bandwidth Memory, HBM, or a GPU global memory, and wherein the on-chip memory is a Static Random Access Memory, SRAM, or a shared memory.

3. The method of claim 1 or 2, wherein the recurrent processing step is executed by a recurrent compute kernel configured to process the chunks sequentially.

4. The method of any one of the preceding claims, wherein the parallel processing step is executed by a parallel compute kernel configured to process the tiles in parallel to compute the output states for all chunks.

5. The method of any one of the preceding claims, wherein the parallel processing step further comprises parallelizing the computing of the output states over an embedding dimension of the linear recurrent neural network, such that computations for different embedding blocks are distributed across different compute units of the data processing apparatus.

6. The method of any one of the preceding claims, further comprising a backward pass for training the linear recurrent neural network, wherein the backward pass utilizes the partitioning into chunks and tiles, and wherein a parallelization dimension and a looping dimension of the backward pass are swapped relative to the parallel processing step of the forward pass.

7. The method of claim 6, wherein the backward pass further comprises recomputing intermediate hidden states within the on-chip memory from the chunks, rather than loading stored intermediate hidden states from the off-chip memory.

8. The method of any one of the preceding claims, wherein the linear recurrent neural network is an xLSTM with matrix memory, mLSTM, and wherein the recurrent memory state comprises a matrix memory state.

9. The method of claim 8, wherein the mLSTM utilizes an input gate activation function configured to behave similarly to an exponential function in a negative value range and to be bounded in a positive value range.

10. The method of claim 8 or 9, wherein the linear recurrent neural network is an mLSTM with sigmoid input gate, mLSTMsig, wherein one or more of an input gate and a forget gate of the network are determined using a sigmoid activation function.

11. The method of claim 10, wherein the recurrent processing step and the parallel processing step are performed without maintaining a normalizer state or a stabilizer state.

12. The method of claim 11, wherein the parallel processing step comprises executing a fused loop for matrix multiplications within the on-chip memory.

13. The method of any one of the preceding claims, wherein computing the output states for the tiles comprises loading data from the on-chip memory into a register file of a processor, and performing matrix multiplication operations using said register file.

14. The method of any one of claims 10 to 13, further comprising initializing a bias of the input gate to a negative value, preferably a value less than or equal to -5.

15. The method of any one of the preceding claims, wherein the chunk size is selected based on a computational intensity of the data processing apparatus, preferably such that the chunk size scales proportionally to the square root of said computational intensity.

16. A data processing apparatus comprising an off-chip memory, an on-chip memory, and one or more processors configured to carry out the method of any one of claims 1 to 15.

17. A computer program comprising instructions which, when executed by the data processing apparatus of claim 16, cause the data processing apparatus to carry out the method of any one of claims 1 to 15.

18. A computer-readable medium storing the computer program of claim 17.

Images