Hardware-optimized recurrent neural network system

Abstract

A system includes a machine-learning model implemented on a data processing apparatus, which features a parallel processor with a memory hierarchy. The machine-learning model is a recurrent neural network (RNN) with a multi-head architecture, comprising multiple sub- vectors that process parallel data streams. The RNN's weight matrix is structured as a block-diagonal matrix, allowing for parallel processing of the sub-vectors. A fused computational kernel executes an entire time-series processing loop for the multi-head RNN, maintaining the weight matrix blocks in on-chip memory and performing matrix multiplications and element-wise operations for each sub-vector in a single kernel execution.

Description

TECHNICAL FIELDThe 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.BACKGROUNDSequence modeling is a fundamental task in machine learning with applications ranging from natural language processing and speech recognition to time-series forecasting and video analysis. As the complexity of sequence-based tasks has grown, so too has the diversity of neural network architectures designed to process sequential data.Recurrent Neural Networks (RNNs) have historically played a central role in sequence modeling due to their ability to maintain internal state information across sequential inputs. Traditional RNN variants, such as Long Short-Term Memory (LSTM) networks and Gated Recurrent Units (GRUs), were designed to address challenges in capturing long-range dependencies within sequential data. These architectures maintain explicit state representations that can be updated at each time step, enabling them to track information across potentially long sequences.In recent years, attention-based models, such as the Transformer, and state space models have gained prominence in many sequence modeling tasks. While Transformers and other sequence-parallelizable neural network architectures seem like the current state of the art in sequence modeling, they specifically lack state-tracking capabilities. Thus, there remain numerous applications where the state-tracking capabilities inherent in traditional RNNs may be advantageous, such as in time-series tasks requiring explicit memory of previous states and logical reasoning problems where information must be carried forward in a structured manner. Traditional RNNs like LSTMs and GRUs, as well as modern variants like sLSTM do have these capabilities at the cost of strictly sequential processing. This is often seen as a strong limitation.The computational characteristics of RNNs also present distinct challenges on modern hardware. Unlike some alternative architectures, the sequential nature of RNN computation-where each time step depends on the previous one-limits opportunities for parallelization across the time dimension. This sequential dependency can lead to processing inefficiencies on parallel computing hardware such as Graphics Processing Units (GPUs).Modern GPUs offer complex memory hierarchies, including high-bandwidth memory (HBM), on-chip static random-access memory (SRAM), and registers, each with different capacity and access speed characteristics. Utilizing these memory resources effectively while managing data movement between them remains challenging for sequence processing tasks.The implementation of RNNs on modern hardware typically involves complex trade-offs between computational efficiency, memory utilization, and numerical precision. These trade-offs become increasingly important as model sizes grow and as applications demand faster processing of longer sequences.Challenges remain in developing efficient implementations of RNN architectures that can leverage the full capabilities of modern parallel processors while maintaining the flexibility needed for research and practical applications.BRIEF DESCRIPTION OF THE DRAWINGSThe disclosed technology may be better understood by reference to the following drawings:FIG. 1 illustrates a block diagram of a machine-learning model system with a recurrent neural network and fused computational kernel according to embodiments.FIG. 2 illustrates a schematic diagram of a fused kernel leveraging caching options for speed optimization in a recurrent neural network according to embodiments.FIG. 3 illustrates a schematic diagram of alternating kernels for maximum hidden sizes according to embodiments.FIG. 4 illustrates a method 400 for processing a sequence of input data using a recurrent neural network according to embodiments.FIG. 5 illustrates a performance comparison graph of LSTM runtime performance with varying head dimensions and number of heads according to embodiments.FIG. 6 illustrates a comparative bar graph showing LSTM runtime performance for different batch sizes across multiple implementation methods according to embodiments.FIG. 7 illustrates a graph of LSTM runtime performance for different batch sizes with a single head according to embodiments.FIG. 8 illustrates a schematic block diagram of computer hardware for implementing the disclosed technology according to embodiments.DETAILED DESCRIPTIONIn 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.Certain embodiments of the present disclosure demonstrate how fast recurrent neural networks can get through hardware-optimization. In one specific implementation referred to as "FlashRNN" in Triton and CUDA, kernels are optimized to the register level on modern GPUs. Certain embodiments extend traditional RNNs with a parallelization variant that processes multiple RNNs of smaller hidden state in parallel, similar to the head-wise processing in Transformers. To enable flexibility on different GPU variants, certain embodiments introduce a new optimization framework for hardware-internal cache sizes, memory and compute handling. In certain embodiments, it models the hardware in a setting using polyhedral-like constraints, including the notion of divisibility. This speeds up the solution process in a library for general integer constraint satisfaction problems (integer CSPs) referred to as " ConstrINT". We show that, in certain embodiments, our kernels can achieve 50x speedups over a vanilla PyTorch implementation and allow 32x larger hidden sizes compared to our Triton implementation.TermsThe 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:As used herein, and unless otherwise specified, the term "machine-learning model" refers to any computational model constructed or trained to identify patterns, make decisions, or generate predictions from data, typically involving learned parameters that are adjusted during a training process. Examples of specific embodiments include neural networks (such as feedforward networks, convolutional neural networks, and recurrent neural networks), support vector machines, decision trees, random forests, and ensembles of the foregoing.As used herein, and unless otherwise specified, the term "parallel processor" refers to a processor or processing unit capable of executing multiple operations concurrently, rather than strictly sequentially. Examples of specific embodiments include multi-core central processing units (CPUs), graphics processing units (GPUs), tensor processing units (TPUs), and other specialized hardware architectures optimized for parallel execution.As used herein, and unless otherwise specified, the term "memory hierarchy" refers to an arrangement of memory or data storage elements within a computing system that are organized according to access speed, capacity, or proximity to processing units. Examples of specific embodiments include caches (L1, L2, and L3), on-chip memory, off-chip dynamic random-access memory (DRAM), and external storage devices such as solid-state drives or hard disk drives. As an example, a memory hierarchy in a GPU may comprise main CPU memory, GPU HMB, GPU SRAM and GPU registers.As used herein, and unless otherwise specified, the term "on-chip memory" refers to a memory component physically integrated into the same chip or module as the processor or computational elements, allowing faster access compared to off-chip memory. Examples of specific embodiments include Static Random-Access Memory (SRAM) regions, processor register files, embedded DRAM, or specialized cache blocks included in system-on-chip designs.As used herein, and unless otherwise specified, the term "Static Random-Access Memory (SRAM)" refers to a type of semiconductor memory that uses bistable latching circuitry to store each bit, typically providing faster access times and lower latency than dynamic RAM. Examples of specific embodiments include on-chip cache lines in CPU or GPU architectures, register files in digital signal processors, and high-speed embedded memory blocks in field-programmable gate arrays.As used herein, and unless otherwise specified, the term "processor registers" refers to small, fast-access storage locations within a processor, used for temporary storage of instructions, addresses, and data during active computation. Examples of specific embodiments include general-purpose registers for integer or floating-point arithmetic, dedicated registers for vector or tensor operations, and special-purpose registers for storing program counters or status flags.As used herein, and unless otherwise specified, the term “recurrent neural network," sometimes referred to as "RNN," refers to a class of artificial neural networks where connections between nodes form directed cycles, enabling the network to maintain internal states that capture temporal or sequential dependencies. Examples of specific embodiments include vanilla RNNs that apply a simple recurrence formula, Long Short-Term Memory (LSTM) networks, Gated Recurrent Unit (GRU) networks, and other architectures that preserve or propagate hidden states across time steps.As used herein, and unless otherwise specified, the term "parallelization mechanism" refers to any structural or algorithmic feature enabling simultaneous or concurrent execution of computations within a model or processing pipeline. Examples of specific embodiments include splitting input data across multiple processing units, dividing computations for different heads within a multi-head architecture, or leveraging batch-level parallelism in training or inference.As used herein, and unless otherwise specified, the term "multi-head architecture" refers to a design in which a model's internal representation is divided into multiple portions or "heads," each of which processes a subset of the input or hidden state in parallel, often followed by a mechanism to combine the outputs from each head. Examples of specific embodiments include multi-head self-attention modules in transformer-based models, multi-branch convolutional layers in convolutional neural networks, or multi-head recurrent layers that operate on separate sub-vectors of a shared hidden state.As used herein, and unless otherwise specified, the term "hidden state vector" refers to the internal state representation maintained by a recurrent neural network or a similar sequential model at any given point in time. Examples of specific embodiments include the hidden state output of LSTM cells (often combined with a cell state), the state vector used in GRU cells, or intermediate hidden representations in a custom recurrent layer design.As used herein, and unless otherwise specified, the term "recurrent weight matrix" refers to a parameter matrix within a recurrent neural network that is multiplied by the hidden state vector from a previous time step to update the hidden state at a current time step. Examples of specific embodiments include the weight matrix in a basic RNN cell, the recurrent gate matrices within an LSTM, or the update and reset gate matrices in a GRU cell.As used herein, and unless otherwise specified, the term "block-diagonal matrix" refers to a matrix that is partitioned into separate square blocks along its diagonal, with zeros elsewhere in the off-diagonal regions, enabling independent operations on each block. Examples of specific embodiments include matrices composed of several square sub-matrices placed along the diagonal for multi-head operations, or matrices split by feature subsets for specialized concurrency.As used herein, and unless otherwise specified, the term "fused computational kernel" refers to a software or hardware routine that combines multiple computational steps into a single execution pass, aiming to reduce overhead or improve efficiency by reusing intermediate data kept in fast-access memory. Examples of specific embodiments include a GPU kernel that performs both matrix multiplication and activation functions, a CPU routine that fuses layer normalization with subsequent transformation steps, or a specialized accelerator kernel that merges convolution and pooling operations.As used herein, and unless otherwise specified, the term "time-series processing loop" refers to the iterative procedure through which sequential data points are processed over consecutive time steps, typically in a recurrent neural network. Examples of specific embodiments include the iterative unrolling of an RNN for natural language processing tasks, the stepwise processing of sensor data in a control system, or the sequential handling of financial time-series data in forecasting.As used herein, and unless otherwise specified, the term "matrix multiplication" refers to the standard linear algebra operation of multiplying two matrices, where the entry in the ith row and jth column of the product is computed as the sum of element-wise products between the ith row of one matrix and the jth column of the other. Examples of specific embodiments include general-purpose dense matrix multiplication (GEMM), batched matrix multiplication for multiple input sets, and specialized high-throughput operations using hardware-accelerated units.As used herein, and unless otherwise specified, the term "element-wise operations" refers to operations that are applied independently to each element of a vector, matrix, or tensor without inter-element interaction in that operation. Examples of specific embodiments include activation functions like sigmoid, tanh, ReLU, or leaky ReLU, as well as arithmetic operations (e.g., addition, subtraction, scaling) that act on each element individually.As used herein, and unless otherwise specified, the term "Long Short-Term Memory (LSTM) network" refers to a type of recurrent neural network designed to capture long-term dependencies through a gated architecture, typically including input, output, and forget gates that control the flow of information. Examples of specific embodiments include standard LSTM cells used in language modeling, stacked LSTM layers in speech recognition systems, and bidirectional LSTMs for sequence-to-sequence tasks.As used herein, and unless otherwise specified, the term "Gated Recurrent Unit (GRU) network" refers to a type of recurrent neural network that utilizes fewer gates than an LSTM but maintains a gating mechanism to regulate how much of the past information to keep. Examples of specific embodiments include standard GRU layers for machine translation, stacked GRU layers in time-series forecasting models, and bidirectional GRUs for classification tasks.As used herein, and unless otherwise specified, the term "scalar Long Short-Term Memory (sLSTM) network" refers to a variant of an LSTM network where certain operations or state representations may be maintained at a scalar granularity, typically optimizing memory or computational usage. Examples of specific embodiments include sLSTMs that reduce the dimensionality of gating operations, or specialized hardware implementations designed for minimal resource consumption. Technological background on sLSTM and other LTSM variants (collectively referred to as xLSTM) is disclosed in WO 2025 / 168512 titled "EXTENDED LONG SHORT-TERM MEMORY NEURAL NETWORKS", the content of which is incorporated herein by reference in its entirety.As used herein, and unless otherwise specified, the term "bias vectors" refers to sets of parameters in neural network layers or cells that are added to weighted inputs before or after non-linear activation functions in order to shift the input domain. Examples of specific embodiments include bias terms in fully connected layers, bias offsets in convolution kernels, and gate-specific biases in LSTM and GRU cells.As used herein, and unless otherwise specified, the term "bfloat16 or float16 precision" refers to numerical formats that represent floating-point values using 16 bits, typically reducing memory bandwidth requirements and computational overhead while retaining a dynamic range sufficient for many neural-network computations. Examples of specific embodiments include the brain floating-point format (bfloat16) commonly used in tensor processing units, the IEEE 754 half-precision floating-point format (float16) implemented in some GPUs, and other hardware-supported half-precision representations.As used herein, and unless otherwise specified, the term "back-propagation through time" refers to the training procedure for recurrent neural networks in which gradients are calculated by unrolling the network over the full sequence of time steps and then applying back-propagation, accumulating errors to update weights associated with each time step. Examples of specific embodiments include truncated back-propagation through time for very long sequences, full-sequence training for shorter inputs, and real-time recurrent learning variants.As used herein, and unless otherwise specified, the term "gradient" refers to a partial derivative of a loss or cost function with respect to model parameters, representing the direction and magnitude of updates needed to minimize or optimize the loss. Examples of specific embodiments include weight gradients computed during back-propagation for neural networks, parameter updates in stochastic gradient descent, and gradient vectors maintained for each trainable parameter in momentum-based optimizers or adaptive optimizers.As used herein, and unless otherwise specified, the term "clip gradient values" refers to a technique used to limit the magnitude of gradients to prevent excessively large parameter updates, typically by applying a threshold beyond which gradient values are scaled down. Examples of specific embodiments include norm-based clipping (where gradients are scaled if they exceed a certain L2 norm) and value-based clipping (where each gradient component is individually limited).As used herein, and unless otherwise specified, the term "optimization module" refers to software, hardware, or a combination thereof that automates or simplifies the process of configuring parameters or settings for a computational routine, aiming to improve resource usage or runtime performance. Examples of specific embodiments include compiler backends specialized in neural network operations, automated optimization frameworks that select kernel launch configurations, or heuristic-based modules that tune block sizes and thread counts for efficient parallel execution.As used herein, and unless otherwise specified, the term "execution parameters" refers to configurable settings or variables that determine how a computational routine is performed on a particular hardware platform. Examples of specific embodiments include block and grid dimensions in a GPU kernel launch, memory layout choices for matrix multiplication, thread scheduling policies in multi-core CPUs, or tiling parameters for breaking down large tensors into smaller chunks.As used herein, and unless otherwise specified, the term "memory tiling sizes" refers to dimensions or shapes used to partition data into smaller blocks that fit within certain memory constraints, such as cache lines or register files, thereby optimizing data locality. Examples of specific embodiments include dividing large matrices into tile-shaped sub-matrices for GPU-based multiplication, splitting tensors into blocks that fit into on-chip SRAM, or subdividing convolutional feature maps in a tiled fashion for specialized hardware.As used herein, and unless otherwise specified, the term "integer constraint satisfaction problem" refers to any formulation of a problem in which the goal is to find integer solutions to a set of numerical constraints, commonly involving equality or inequality constraints. Examples of specific embodiments include model configuration problems that assign integer values to block sizes and tile shapes, job scheduling problems that allocate integer resources among tasks, or hardware resource allocation problems that must respect capacity limits.As used herein, and unless otherwise specified, the term "physical constraints" refers to real-world or hardware-imposed limits on computational resources, such as memory capacity, bandwidth, or the number of parallel processing units available. Examples of specific embodiments include maximum allowable memory usage in registers or SRAM on a processor, limitations on how many threads can run concurrently on a GPU, and the total number of multiply-accumulate units available for matrix operations.As used herein, and unless otherwise specified, the term "divisibility constraints" refers to constraints requiring that certain parameters or dimensions be integer multiples of given values, often to align with hardware block sizes or warp sizes. Examples of specific embodiments include requiring matrix dimensions to be multiples of a processor's vector register size, enforcing particular block dimensions for a GPU's thread group, or aligning data structure sizes with cache line boundaries.As used herein, and unless otherwise specified, the term "hardware-accelerated matrix multiplication unit" refers to a specialized computational block within processors dedicated to performing matrix multiplication operations with high throughput and efficiency. Examples of specific embodiments include tensor cores in GPUs, systolic array units in machine-learning accelerators, or matrix-multiply-and-accumulate pipelines in digital signal processors.As used herein, and unless otherwise specified, the term "computational super-unit" refers to a group or subset of a processor's functional blocks or computational cores that collectively act as a higher-level block for processing. Examples of specific embodiments include streaming multiprocessors in a GPU, compute clusters in a system-on-chip, or multi-core complexes in a CPU that share a common cache subsystem.As used herein, and unless otherwise specified, the term “threads per computational super-unit" refers to the number of individual threads of execution that can be concurrently scheduled or multiplexed on one such group of computational resources. Examples of specific embodiments include the maximum number of threads per warp or wavefront on a GPU, the concurrency limit per CPU core in simultaneous multithreading, or the total hardware threads assigned to a heterogeneous processing cluster.As used herein, and unless otherwise specified, the term "equality constraints" refers to constraints specifying that certain parameters must be exactly equal to one another or to a specified value in the context of a resource allocation or optimization problem. Examples of specific embodiments include requiring that two tiling dimensions are the same, matching the number of heads to a sub-matrix block count, or setting an exact limit on the amount of on-chip memory allocated per tile.As used herein, and unless otherwise specified, the term "inequality constraints" refers to constraints specifying that certain parameters must remain greater than or equal to, or less than or equal to a given value, in the context of a resource allocation or optimization problem. Examples of specific embodiments include memory usage being less than or equal to the total SRAM size, the number of threads not exceeding hardware limits, or tile dimensions not exceeding a maximum allowable block size.As used herein, and unless otherwise specified, the term "heuristics" refers to rules of thumb or strategies that guide an algorithm toward a good, though not necessarily optimal, solution, often used to reduce computational complexity in large search or optimization problems. Examples of specific embodiments include greedy selection of the largest tile size that fits in available memory, iterative deepening strategies, or machine-learning-based approaches that predict optimal configurations based on past runs.As used herein, and unless otherwise specified, the term "sequence of input data" refers to an ordered collection of data elements that are processed in succession, particularly relevant when maintaining dependencies between data points. Examples of specific embodiments include time-sequenced sensor readings, words in a sentence, or frames in a video stream.As used herein, and unless otherwise specified, the term "time steps" refers to discrete points in an ordered sequence at which a system updates its state or processes data. Examples of specific embodiments include clock ticks in time-series processing, character-by-character processing in natural language tasks, and sample-by-sample processing in audio signal analysis.As used herein, and unless otherwise specified, the term "point-wise operations" refers to operations that independently transform or combine corresponding elements in arrays, vectors, or matrices without mixing data across indices. Examples of specific embodiments include adding a bias term to every output neuron, multiplying each element of a vector by a scalar, or applying an activation function element by element.As used herein, and unless otherwise specified, the term "backward pass" refers to the computational phase in which gradients of a model's loss function with respect to its parameters are computed, typically proceeding in reverse order from the final output back through intermediate layers or states. Examples of specific embodiments include back-propagation through time for RNN sequences, reverse-mode differentiation in feedforward networks, or distributed gradient computation across multiple devices.As used herein, and unless otherwise specified, the term "gradient clipping" refers to the procedure of enforcing an upper limit on the magnitude of gradients during back-propagation or parameter updates to stabilize training. Examples of specific embodiments include rescaling gradients when their total norm exceeds a threshold, bounding individual gradient components, or applying more sophisticated strategies that adapt clipping thresholds during training.System OverviewFIG. 1 illustrates a block diagram of a machine-learning system 100 implemented on a data processing apparatus 102. The data processing apparatus 102 comprises a parallel processor 104 with a memory hierarchy that includes an on-chip memory 106. The system 100 comprises a recurrent neural network 108 as the machine-learning model, a fused computational kernel 110, and an optimization module 112.The recurrent neural network 108 comprises a parallelization mechanism, in that the recurrent neural network 108 comprises a multi-head architecture where a hidden state vector is partitioned into a plurality of head sub-vectors, and a recurrent weight matrix of the recurrent neural network 108 is structured as a block-diagonal matrix, each block corresponding to one of the plurality of head sub-vectors, enabling parallel processing of the heads.The fused computational kernel 110 is configured to execute an entire time-series processing loop for the multi-head recurrent neural network 108. This kernel is configured to maintain the blocks of the block-diagonal recurrent weight matrix in the on-chip memory 106 for the duration of the time-series loop and, for each timestep within the loop, perform both a matrix multiplication for each head using its corresponding block of the recurrent weight matrix and subsequent element-wise operations, within the same kernel execution.This fusion of operations within a single kernel invocation reduces the overhead associated with launching multiple kernels and allows for more efficient use of the on-chip memory 106 and computational resources. The system 100 enables high-throughput, low-latency execution of multi-head recurrent neural networks 108 on parallel processors 104 with hierarchical memory architectures. By structuring the recurrent weight matrix as a block-diagonal matrix and fusing the time-series processing loop into a single kernel, the system 100 achieves efficient parallelization and minimizes memory bottlenecks, making it suitable for demanding machine-learning applications involving sequential data.In embodiments, the processor 104 may be a graphics processing unit (GPU), a tensor processing unit (TPU), or another form of parallel computing hardware. This provides massive parallelism, high memory bandwidth, and optimized support for matrix and tensor operations, thereby significantly accelerating the execution of recurrent neural network workloads.In embodiments, the on-chip memory 106 may comprise Static Random-Access Memory (SRAM). The use of SRAM as the on-chip memory 106 can provide advantages in terms of access speed and energy efficiency, as SRAM typically offers lower latency and higher bandwidth compared to other types of memory such as Dynamic Random-Access Memory (DRAM). The parallel processor 104 may optionally incorporate SRAM to store the blocks of the block-diagonal recurrent weight matrix during execution of the time-series processing loop. It is also possible for the on-chip memory 106 to include a combination of SRAM and other memory types, depending on the architecture of the data processing apparatus 102. For example, certain architectures might employ SRAM for frequently accessed data, such as the recurrent weight matrix blocks, while utilizing other memory technologies for less time-critical data.In embodiments, the on-chip memory may comprise processor registers of the parallel processor 104, allowing even faster memory access and boosting computational efficiency. Processor registers can serve as a form of on-chip memory 106 that is directly accessible by the computational units within the parallel processor 104. Utilizing processor registers for storing data, such as the blocks of the block-diagonal recurrent weight matrix, can enable extremely low-latency access during execution of the time-series processing loop. This arrangement may be particularly advantageous in scenarios where the recurrent weight matrix blocks are sufficiently small to fit within the available register space, thereby maximizing computational efficiency and minimizing memory access delays. The data processing apparatus 102 can be configured so that processor registers are used in conjunction with other types of on-chip memory, such as shared memory or cache, to optimize data locality and throughput. In some implementations, critical data required for each timestep-such as the current head sub-vectors and their associated weight blocks-can be loaded into processor registers at the start of the kernel execution, allowing all subsequent operations within the time-series loop to proceed without further off-chip or even shared memory accesses.In embodiments, the recurrent neural network 108 implemented in the system 100 can be a Long Short-Term Memory (LSTM) network. An LSTM network is a type of recurrent neural network that is capable of learning long-term dependencies by incorporating memory cells and gating mechanisms, such as input, output, and forget gates. This architecture may be advantageous in scenarios where the sequential data exhibits complex temporal relationships or where it is desirable to mitigate the vanishing gradient problem commonly encountered in standard recurrent neural networks.Alternatively, the recurrent neural network 108 may be realized as a Gated Recurrent Unit (GRU) network. A GRU network is another variant of recurrent neural network that utilizes gating mechanisms to control the flow of information, but with a simplified structure compared to LSTM. GRU networks can offer comparable performance to LSTM networks while requiring fewer parameters and computational resources, making them suitable for applications where efficiency is a priority.In other possible configurations, the recurrent neural network 108 may be implemented as a scalar LSTM (sLSTM). A sLSTM network is a variant of the standard LSTM architecture in which certain operations or parameters are represented as scalars rather than vectors or matrices. This approach can reduce the computational complexity and memory footprint of the network, which may be beneficial in resource-constrained environments or when deploying the system on hardware with limited on-chip memory.Fused Computational KernelFIG. 2 illustrates a schematic diagram of an example implementation of the fused computational kernel 110 leveraging caching options for speed optimization in the recurrent neural network 108 according to embodiments. In the diagram, a circular arrow on top of the main processing block indicates the sequence-loop nature of the RNN processing. The dashed arrows represent read actions from HBM, and the solid arrows represent write actions to HBM.In embodiments, the fused computational kernel 110 may be configured to maintain bias vectors associated with the recurrent neural network 108 in the on-chip memory 106 for the duration of the time-series loop. The bias vectors can correspond to the various gates or units within the recurrent neural network 108, such as input, output, forget, or update gates, depending on the specific architecture employed. By retaining these bias vectors in the on-chip memory 106 throughout the processing of the time-series data, it is possible to reduce memory access latency and improve overall computational efficiency, as the kernel can access the required bias values directly without repeatedly fetching them from off-chip or slower memory tiers.In embodiments, the bias vectors are loaded into the on-chip memory 106 at the start of the time-series loop and remain resident for the entire duration of the loop execution. This approach can be particularly advantageous in parallel processing environments, where minimizing memory transfers and maximizing data locality are important for achieving high throughput. In some implementations, the bias vectors may be stored alongside the blocks of the block-diagonal recurrent weight matrix, allowing both sets of parameters to be accessed efficiently within the same kernel execution.Alternatively, the system 100 can be arranged so that only a subset of the bias vectors, such as those most frequently accessed or those associated with specific heads in the multi-head architecture, are maintained in the on-chip memory 106, while others are fetched as needed. The decision regarding which bias vectors to retain in on-chip memory 106 may depend on the available memory resources, the size of the bias vectors, and the requirements of the particular recurrent neural network configuration.In embodiments, the fused computational kernel 110 may keep the recurrent weight blocks and bias vectors in processor registers for the entire sequence loop. For example, the kernel may allocate registers such that the block-diagonal weight tiles of one head, along with their biases (5), remain register-resident. Intermediate results may be written to shared memory, and only final hidden states may be committed to off-chip memory. This register-resident caching approach may eliminate nearly all high-bandwidth-memory traffic associated with recurrent parameters, reducing latency and overall energy consumption.In embodiments, a generic multi-state, multi-gate cell template may be adopted, wherein the numbers of states (N₂) and gates (Ng) are supplied as run-time parameters, rather than hard-coded. The fused computational kernel 110 may be generated with loop counts defined by N, and Ng, allowing modifications to gate activation functions or state-update formulas through templated functors without altering the overall kernel structure. This approach may let a single compiled kernel handle different architectures-such as an Elman RNN (N=1, N=1), a conventional LSTM (N=2, Ng=4) (4), or an sLSTM (N=3, Ng=4) (4)-using the same executable but different parameter sets. As a result, maintenance is simplified, and rapid experimentation is promoted.In embodiments, the system 100 may be configured to pre-compute an input-to-hidden matrix product outside the recurrent loop. Specifically, if the complete input sequence is known in advance, a batched GEMM may be executed before any recurrent computation, so that the fused computational kernel 110 only processes recurrent terms and biases for each timestep. By caching the results of W-X for all T timesteps, the fused computational kernel 110 may omit loading W altogether during the loop, thereby removing an entire matrix multiplication from every timestep and improving overall throughput.In certain embodiments, the fused computational kernel 110 may be configured to operate on data represented with bfloat16 or float16 precision. The use of reduced-precision data formats, such as bfloat16 (brain floating point 16) or float16 (half-precision floating point), can offer significant advantages in terms of computational efficiency and memory utilization. By employing bfloat16 or float16 representations, it is possible to decrease the memory footprint of the recurrent neural network 108 parameters and intermediate activations, thereby enabling the processing of larger models or longer sequences within the constraints of the available on-chip memory 106.Alternating Computational KernelIn embodiments, the data processing apparatus 102 may be configured for adaptive kernel mode switching between the persistent fused fused computational kernel 110 and an alternating computational kernel. For instance, the system 100 may monitor a workload attribute such as hidden-state width, batch size, or available on-chip memory and select the persistent fused mode if the attribute value stays below a certain threshold, while switching to the alternating mode if that threshold is exceeded. By way of example, for an LSTM with 12 heads, the controller may select the fused mode for batch sizes up to 32 and switch to the alternating mode for batch sizes above 128, thus preserving high throughput on small workloads without sacrificing the ability to run very large models.FIG. 3 illustrates a schematic diagram of an alternating computational kernel for maximum hidden sizes according to embodiments.Second Fused Computational KernelIn some embodiments, the system 100 may comprise a second fused computational kernel that is configured to execute a backward pass for back-propagation through time. This second kernel can be arranged to perform the gradient computations required for training the recurrent neural network 108, such as those arising during the optimization of model parameters via stochastic gradient descent or related algorithms.Optionally, the second fused computational kernel may be designed to maintain gradients in the on-chip memory 106 across multiple time steps of the backward pass. By retaining the gradients in the on-chip memory 106 throughout the backward traversal of the time-series data, it is possible to reduce memory access latency and improve the efficiency of the training process.In certain configurations, the gradients corresponding to the recurrent weight matrix blocks, bias vectors, and other trainable parameters may be stored in the on-chip memory 106 for the duration of the backward pass. This approach may allow the second kernel to access and update the gradients efficiently at each time step, without incurring the overhead of repeated memory fetches from slower memory tiers. Alternatively, the system could be arranged so that only a subset of the gradients-such as those associated with the most frequently updated parameters are maintained in the on-chip memory 106, while others are transferred as needed based on memory availability and computational requirements.In embodiments, the second fused computational kernel may also be implemented to process the backward pass in a manner analogous to the forward pass, leveraging parallelism across the multiple heads of the recurrent neural network 108. In this way, the kernel can compute gradients for each head in parallel, further accelerating the back-propagation through time procedure. It is also possible for the system 100 to support flexible allocation of on-chip memory resources between the forward and backward computational kernels, depending on the phase of operation (inference or training), the size of the model, and the available hardware resources. This flexibility can help to optimize the utilization of the on-chip memory and maximize overall system performance during both inference and training of the recurrent neural network.In a further possible embodiment, the second fused computational kernel may be configured to perform gradient clipping after each time step within the time-series loop of the backward pass. Gradient clipping can be implemented as an optional technique to address the issue of exploding gradients, which may arise during the training of the recurrent neural network 108, particularly when processing long sequences. In such an arrangement, after the computation of gradients at each time step during the backward traversal, the kernel can optionally evaluate the magnitude of the gradient values and, if necessary, limit them to a predefined threshold.Optimization ModuleIn embodiments, the optimization module 112 of the system 100 serves for configuring the fused computational kernel 110 and / or other kernels. The optimization module 112 can be arranged to automatically determine execution parameters for the kernel, such as memory tiling sizes, by solving an integer constraint satisfaction problem that models the physical constraints of the parallel processor 104. This approach allows the system 100 to adapt the computational kernel to the specific hardware characteristics of the data processing apparatus 102, thereby improving computational efficiency and resource utilization.The optimization module 112 may, for example, analyze the available on-chip memory 106, the number and arrangement of processing elements, and the architectural features of the parallel processor 104 in order to select suitable tiling sizes and other execution parameters. By formulating the selection of these parameters as an integer constraint satisfaction problem, the optimization module 112 can systematically explore the space of possible configurations while ensuring that all relevant hardware limitations are respected.Among the constraints considered by the optimization module 112, divisibility constraints may be included to ensure compatibility with the required input dimensions of hardware-accelerated matrix multiplication units present in the parallel processor 104. For instance, certain matrix multiplication engines or tensor cores may require that input matrices have dimensions that are integer multiples of a specific value, such as 8, 16, or 32. The optimization module 112 can take such requirements into account when determining the memory tiling sizes and other kernel parameters, thereby enabling efficient use of the hardware acceleration features.It is also possible for the optimization module 112 to incorporate additional constraints, such as limits on the total memory footprint, alignment requirements, or restrictions imposed by the scheduling and synchronisation mechanisms of the processor 104.In embodiments, the optimization module 112 may employ various algorithms for solving the integer constraint satisfaction problem, including but not limited to branch-and-bound, integer linear programming, or heuristic search methods. Optionally, the optimization module 112 can operate at runtime, dynamically adapting the kernel configuration in response to changes in workload, input data size, or available hardware resources. Alternatively, the optimization module 112 may be performed offline during deployment or compilation, with the resulting execution parameters being stored for use during subsequent kernel invocations.By providing an optimization module 112 capable of automatically determining execution parameters for the computational kernel through the resolution of an integer constraint satisfaction problem, the system 100 can achieve improved performance, portability, and scalability across a range of parallel processor architectures. This capability may be particularly beneficial in environments where hardware resources are heterogeneous or subject to change, or where optimal utilization of specialized matrix multiplication units is critical for achieving high throughput in machine-learning workloads.In embodiments, the physical constraints of the parallel processor 104 that are modeled by the integer constraint satisfaction problem may include, for example, a total Static Random-Access Memory (SRAM) size per computational super-unit, a total register size per computational super-unit, and / or a number of threads per computational super-unit. The optimization module 112 can optionally take into account the available SRAM within each computational super-unit, which may represent a cluster of processing elements sharing a common on-chip memory resource. By considering the total SRAM size, the optimization module 112 may ensure that the selected memory tiling sizes and other execution parameters do not exceed the memory capacity available to each super-unit, thereby preventing memory overflows and maximizing data locality.In embodiments, the integer constraint satisfaction problem used by the optimization module 112 may comprise equality constraints and / or inequality constraints. These types of constraints can be incorporated as part of the mathematical formulation that determines suitable execution parameters for the computational kernel. For example, equality constraints may be employed to ensure that certain variables, such as tiling sizes or memory allocations, exactly match required values dictated by hardware specifications or algorithmic requirements. In other cases, inequality constraints might be introduced to specify that particular parameters must remain within defined upper or lower bounds, such as limiting memory usage to not exceed the available on-chip resources or ensuring that thread counts do not surpass hardware-supported limits.In embodiments, the optimization module 112 may be configured to utilize heuristics in order to prioritize solutions that maximize utilization of the hardware-accelerated matrix multiplication units. The use of heuristics in this context can enable the optimization module to efficiently explore the space of possible execution parameters and to favor those configurations that are likely to achieve higher throughput or computational efficiency on the available hardware.FIG. 4 illustrates a method 400 for processing a sequence of input data using the recurrent neural network 108 according to embodiments.In block 402, method 400 executes the fused computational kernel 110 for a plurality of time steps of the recurrent neural network 108, which may comprise several sub-steps: In block 404, method 400 loads recurrent weight parameters of the recurrent neural network 108 into the on-chip memory 106. In block 406, method 400 performs, for each of the plurality of time steps, a recurrent matrix multiplication using a hidden state from a previous time step and the recurrent weight parameters stored in the on-chip memory 106. In block 408, method 400 performs one or more point-wise operations using a result of the matrix multiplication to generate a new hidden state for a current time step. In block 410 of method 400, the recurrent weight parameters remain in the on-chip memory 106 for the duration of the processing of the plurality of time steps within the fused computational kernel 110.The method 400 is designed to efficiently process sequential input data by leveraging the parallel processing capabilities and memory hierarchy of the data processing apparatus 102 of FIG. 1. By maintaining the recurrent weight parameters in the on-chip memory 106 for the entire duration of the time-series processing loop, the method 400 avoids repeated transfers of these parameters from slower, off-chip memory. This approach reduces memory access overhead, improves computational efficiency, and enables the recurrent neural network 108 to process long sequences or large batches of data with high throughput. B y fusing the time-series loop into a single kernel execution and optimally utilizing the on-chip memory 106, the method 400 achieves efficient and scalable processing of sequential data, making it suitable for demanding machine-learning applications such as speech recognition, natural language processing, and time-series forecasting.FlashRNN ImplementationAn example implementation of embodiments of the disclosed technology is referred to as "FlashRNN", which will be described next.Conventional recurrent neural networks are typically limited in speed due to a sequence of alternating calls of matrix multiplication and point-wise kernels, as well as their limitation to higher (but slower) precision. FlashRNN overcomes the limitation by fusing the recurrent matrix multiplication with the pointwise operations into a single persistent kernel with custom caching of the recurrent weights in registers. FlashRNN also supports the bfloat16 dtype and block-diagonal recurrent matrices.Generic Recurrent Neural Network Architecture with Memory MixingA generic RNN architecture that we aim to optimize has Ns states s(i) ∈ Rd, and Ng gates (or pre-activations) g(1) ∈ Rd, with d being the embedding dimension or hidden size of the RNN. For example, the LSTM has Ns = 2 states and Ng = 4 gates.Each gate receives an input x(i) ∈ Rd. As learnable parameters, the gates have a recurrent matrix RJ) ∈ Rdxd that models the dependency on the previous hidden state s), st and a bias b ∈ Rd. The state sequence of the RNN is then defined as:g) = x(0) + RO)s(0) + b) (1)St-1s(i) = p(i) ((sin)in (g);) (2)Stwith a point-wise / element-wise function P(i) that does not mix different cells along the vector dimension (unlike the recurrent weight).Usually for these networks, the input is modified with another weight matrix W. We omit this here as it can be moved outside of the basic kernels. In the common training setting, where the whole sequence is given as input, the weight matrix W can be applied in parallel to all timesteps before processing a sequence in the RNN. Our runtime experiments show that this operation has only marginal impact on the overall runtime.Generic Gradient for Back-Propagation Through TimeIn back-propagation through time, the backward pass of this RNN architecture can be recursively defined as well. The backward pass reads:δε(i) = ap(1)((s)k(g)j) ds(1) (3)St-1 t (1) 89t= ap(1) ((s1kg);) 8s (1) + (RJ)T8g), ifi = 0) (4)(i) as1 9t-1The structure of the gradient shows that, also for the backward pass, we have an alternation of point-wise operations (left) and matrix multiplication (right).The input gradient is equal to the gate gradients, the bias gradient is the sum of the input gradients and the recurrent weight matrix gradient is the time-wise sum of the outer product of gate gradients with the state values:8x) = 8g) (5)8b(i) = Στ δg) (6)SR) = Στ 8g)s(0)T (7)StVanishing and Exploding Gradients and Gradient ModificationsFor a neural network to be stably trainable, there must not be exploding gradients, also vanishing gradients should be prohibited for long context sequence modeling. Still, for the generic structure of Equation 3, there can be exploding components: Firstly, one or more eigenvalues of the point-wise function Jacobian can be greater than one in magnitude. This can be mitigated by a proper choice of the point-wise function. Secondly, the combination of recurrent matrix and gate gradients with the gradient RJ)T could have singular(0) agvalues of magnitude > 1. This case cannot be excluded directly, as the recurrent matrix consists of trainable weights with usually unconstrained magnitude. However, for practical training this is rarely a limitation.In our library, we implement a simple approach for mitigating this at the cost of additional gradient noise, clipping the gradient values on a scalar level after each time step. Specifically, we clip the term containing the recurrent matrix to within a pre-defined magnitude. The gradients can even be cut to zero, leading to typically worse convergence at the benefit of faster training, as the recurrent matrix part in Equation 3 is cut to zero for the backward pass.Head-Wise ParallelizationWhen increasing the size of a neural network, typically the width, i.e. the embedding dimension or hidden size is increased. Research has found that for the attention operation it is beneficial to linearly project the input embedding vectors into multiple smaller input vectors, the so called heads, and then perform attention on each of these small vectors in parallel. This parallelization primitive enables also efficient implementations on GPUs, since each head can be computed in different thread blocks of the GPU in parallel, as will be explained further below.The embedding or hidden vector of dimension d can be split into Nheads heads of smaller dimension dhead = d / Nhead, each of which is processed independently inside the sequential part. In FlashRNN, we apply this primitive to traditional RNNs by dividing the recurrent matrix R into multiple blocks or heads Rhead ∈ Rdhead dhead rendering the recurrent matrix R as a block-diagonal matrix.Hardware-Efficient ImplementationGPU-Acclerated ComputingModern compute hardware in the form of GPUs enables massive parallelization and accelerated matrix multiplication. This means that both point-wise (scalar) operations can be parallelized and matrix multiplications have good support via BLAS-like libraries, as used for RNN training workloads as defined above.Execution Model: Specifically, a modern GPU consists of larger computational super-units (i.e. streaming multiprocessors (SMs)) that have some faster memory attached to them. There are typically three levels of memory, the large HBM which allows global random access from all computational units at the cost of low speed (still fast compared to CPU RAM access), the SRAM which is attached to one computational super-unit and the registers which are tied to a smallest computational unit (i.e. thread). One super-unit usually supports up to 1024 threads in parallel (with varying register sizes) which are typically referred to as a block or thread block. Multiple blocks executed in parallel on multiple super-units are called the grid. An NVIDIA H100, for example, consists of 132 streaming multiprocessors, with 256 KB SRAM per SM and a SRAM bandwidth of around 33 TB / s, compared to the up to 3 TB / s for access to the 80 GB of HBM. Starting from the NVIDIA Ampere Architecture and newer, there is hardware acceleration for asynchronous loading and SRAM interconnection, which we did not utilize in this work. Beyond the memory levels, a computational super-unit allows for hardware-accelerated matrix multiplication (e.g. via TensorCores, "wmma" operation). Typically, it is divided into sub-units (warps) of a certain number of threads (NVIDIA: 32) that act as one for a matrix multiplication. There are certain size limitations for this acceleration, which have to be considered in the kernel optimization process. For a NVIDIA H100, this means that only minimal matrices of sizes 32x16x8, 16x16x16 or 8x16x32 can be multiplied for the low-precision bfloat16 or float16 dtypes, larger matrix multiplications have to be composed of those, by parallelization along the outer dimensions and summation along the accumulating dimension.Performance measures: The specific limitation of a computational load falls into two regimes: Being compute-bound or being memory-bound. In the former case, the arithmetic intensity is high, there are many compute operations per loaded byte and therefore, the main limitation is the computational part. In the latter case, arithmetic intensity is low and the bottleneck is the memory access to load inputs and store outputs. Small operations, like applying an activation function in parallel are typically memory bound and should be grouped together into a fused kernel.Fused Kernels: Το minimize HBM memory accesses, one combines multiple arithmetic operations in one GPU kernel. A kernel is a set of instructions on the GPU which is executed in parallel on its parts. Only within the execution of one kernel SRAM and registers are kept and can serve as a cache. Therefore, for memory-bound operations it is helpful to fuse multiple arithmetic operations into one kernel to leverage these lower cache levels. While compilers can fuse point-wise operations, an alternation of both point-wise computations and matrix multiplication is non-trivial.FlashRNN KernelsAs the RNN operations of Equations 1 and 3 are a sequential alternation between matrix multiplication and pointwise non-linearities, there is a simple speed up variant that optimizes these two primitives separately. Our library implements this variant, in the alternating backend. This enables arbitrarily large head dimensions (to the limits of HBM GPU memory). Also, a vanilla PyTorch implementation relying on auto-grads will work in this primitive, but for every time step a separate state is saved for the backward pass, leading to inefficiencies beyond memory accesses. We show that moving the time-loop into CUDA can already give large speedups over the vanilla PyTorch implementation.The downside of the alternating implementation is that there are no I / O optimizations beyond a single time step. For every time step, the current input and last state, as well as the recurrent matrix and the biases have to be re-loaded. However, both the recurrent matrix R and the biases remain the same for the whole time loop and the previous states can stay in memory as they were computed in the previous time step. Since the structure of the computation remains the same over the time steps, one can even store most of these values in registers. Registers have the highest memory bandwidth and, while they can only be used within the lowest computation unit (threads), their total size on a GPU is comparable to the SRAM (both 256 KB per SM on H100).To reach the maximum speed, we implement FlashRNN fused kernels that store the recurrent matrix R and the biases b in registers (and SRAM if register memory is exceeded), as was already described in connection with the fused computational kernel 110 of FIG. 1. The matrix multiplication results are stored and accumulated in shared memory (or HBM if SRAM sizes are exceeded). In the forward pass, the computations are mainly tiled along the gate dimension (or the dimension of the new hidden states). This way, we use the maximum amount of memory along the previous state dimension. This dimension is the accumulating dimension of the recurrent matrix multiplication. For the backward pass, the computations are typically tiled along the previous state gradient dimension, such that the gate dimension, which is accumulated over, is minimally tiled. The following algorithm shows a simplified representation of the forward pass in pseudo-code:Algorithm 1: FlashRNN-fused forward passAll states are tiled along threads (single ALU) in Warps (for e.g. Matrix Multiplication) in a block (SRAM level, streaming multiprocessor) and blocks in the grid (multiple streaming multiprocessors) - additionally there can be looping levels where the parallelization is resolved to a simple loop. Dimensions are: b: batch, t: time, g: gates, s / s': previous / new state.Recurrent matrix Rgs, inputs Xtbg, biases bgInitial states SobsLoad Rgs,bg to registers and SRAMfor lo in LB doLoad Sobs to registersforte 0..T-1 dofor Matrix Tiles in Registers doCalculate and Accumulate Matrix product ytbg = Rgs Stbsalong send forfor Matrix Tiles in SRAM doLoad Matrix Tile of RgsCalculate and Accumulate Matrix product ytbg = Rgs gs Stbs(0)along send forAccumulate MatMul results Ytbg along s in shared memory (Write, Load and Sum)if state dimension too big for SRAM thenAccumulate MatMul results ytbg along s in HBM (Write, Grid Sync, Load, Sum)end ifSum Gate inputs Xtbg with Ytbg and biases by to gates gtbgCompute Point-wise Function St+1bs' = P(Stbs', gtbg) with aligned states s' and gates gWrite out gates for backward pass and new states to HBMGrid-Level Sync (for new states to be available across the whole grid)end forend forTriton ImplementationWith FlashRNN we also implement a Triton variant of the fused FlashRNN kernel. Triton is a domain specific language and compiler for parallel programming that provides a Python-based environment for writing custom GPU kernels.For the Triton kernel we parallelize the computation over two dimensions the batch dimension and the head dimension. As described further above, we partition the embedding dimension into multiple heads and compute each head in parallel in different programs (or thread blocks) with no synchronization in between these programs. In Triton each program (which corresponds to a thread block in CUDA) will hold its recurrent weightmatrix Rhead and bias bhead in SRAM. In contrast to CUDA, Triton gives no access to registers on the GPU. Therefore, we cannot apply the custom caching strategy of the fused CUDA kernels and instead rely on Triton for managing the shared memory and register cache. Additionally, there is no (grid) synchronization between programs in Triton, which makes it impossible to communicate values between different programs over HBM. This poses a limitation on the maximum head dimension of 128 for the forward pass and 64 for the backward pass on a NVIDIA H100 GPU.The recurrent matrix multiply in Equation 1 and 3 is implemented with Triton's matrix multiply operation tl.dot which gives an interface to the Tensor Core units on GPUs. In Triton minimum block size of these matrix multiplies is 16x16, which gives a limit on the minimum batch size. In practice, we enable smaller batch sizes by padding zeros at the cost of efficiency.Automatic Tuning of Tiling and Looping DimensionsWhile Algorithm 1 above describes the algorithmic behavior, the tile, block and grid sizes and loop iterations depend on the specific hardware architecture, i.e. the number of computational super-units (streaming multiprocessors), the SRAM per super-unit, the sizes of matrix-multiplication units, threads (warps and threads) per super-unit and the number of registers per thread. On NVIDIA H100s (and most other NVIDIA GPUs), there is a varying amount of registers per thread, depending on the block size used. The total number of registers on chip per streaming multiprocessor is physically fixed.These physical constraints can now be reformulated as equalities, inequalities and divisibility constraints inside an integer constraint satisfaction problem (integer CSP). Typically this optimization is done via polyhedral constraint optimization in compilers (Baghdadi et al., 2018). For solving these constraints in FlashRNN, we implement an efficient solver ConstrINT in Python for general integer CSPs going over large number ranges and including the notion of divisibility constraints, which are needed to model the minimal matrix sizes.ExperimentsIn the following, we benchmark the runtime of our FlashRNN kernels and compare against the LSTM and Attention implementations provided in PyTorch. We measure training time with FlashRNN kernels on language modeling, and we confirm that traditional RNNs like LSTM and more recent variants like sLSTM implemented in FlashRNN can solve state tracking problems.Runtime BenchmarkWe evaluate the runtime of all backends of our FlashRNN library that implement the LSTM operation:CUDA fused: CUDA implementation that fuses matrix multiplication and pointwise operations of the LSTM in a single kernel that is persistent over all time iterations.CUDA alternating: CUDA implementation that implements the time loop in C++ and alternates between a matrix multiply kernel and a LSTM pointwise kernel.Triton fused: Triton implementation that fuses matrix multiplication and pointwise operations similar to CUDA fused.Vanilla PyTorch: PyTorch implementation of the LSTM operation with our custom backward pass implementation, which is faster than the PyTorch autograd backward pass. We do not use torch.compile due to very long compile times.We compare our backends to two references from PyTorch and the haste library:FlashAttention2: PyTorch Attention with FlashAttention2 backend. Note that FlashAttention2 is not a recurrent operation and can be parallelized across batch, head, and sequence dimension on the GPU. FlashAttention2 does not fall into the category of RNNs, which FlashRNN aims to speed up, and is not able to solve state tracking tasks. Therefore, in our benchmarks it should be seen as a widely adopted reference to better interpret the runtimes instead of a direct baseline that we aim to outperform.nn.LSTM: PyTorch LSTM with NVIDIA cuDNN as backend. In contrast to our FlashRNN LSTM, nn.LSTM also integrates the gate pre-activation computation into the function call (not kernel call), which we do not. Moreover, nn.LSTM does not support multiple heads on the embedding dimension. nn.LSTM always uses a single head.haste: The haste library is an implementation of LSTM and GRU and variations in CUDA, using alternating kernels between pointwise and matrix multiplication operations. Its last release was in 2020, with no compilation support for Ampere or later architectures in the standard setting. It solely supports float32 and float64 precision and does not have a multi-head option.Setup: We assess the impact of the input dimensions batch size (B), sequence length (T) and head dimension (DH) and number of heads (NH). The number of heads together with the head dimension give the embedding dimension d = NH × DH. Except for PyTorch nn.LSTM we run all runtime experiments with bfloat16 precision. For nn.LSTM we use float16 precision, since this precision yielded the fastest runtimes. For every runtime measurement we do 25 warmup iterations and then report the average across 1000 iterations on NVIDIA H100 GPUs. We use PyTorch 2.4 and with CUDA version 12.4 for our experiments.Head dimension: We measure the runtime of all of our FlashRNN kernels and our two references FlashAttention2 and PyTorch nn.LSTM for different head dimensions. We fix the embedding dimension d = NH × DH to 768 and vary the head dimension from 16 to 768. We use batch size 16 and sequence length 1024. In Figure 2 we report the runtime of each the forward pass only on the left and the forward combined with the backward pass. FlashAttention2 does not allow for head dimension larger than 256, due shared memory limitation. The PyTorch nn.LSTM does not support multiple heads or blockdiagonal recurrent matrices. Therefore, we only report the runtime for a single head of dimension 768, including the gate pre-activation computation. At this dimension, nn.LSTM is about 3 times faster than CUDA fused. The Triton kernels are limited to head dimension 128 and 64, but are about two times faster than CUDA fused for small head dimensions 16 and 32. The fused CUDA kernels support all head dimensions up to 768 (actually more) and are about two to three times faster than the alternating kernels.FIG. 5 shows LSTM runtime performance for different head dimensions (DH) and number of heads (NH) on a NVIDIA H100. Overall embedding dimension is fixed at 768. We use batch size 16 and sequence length 1024. Left: Forward pass. Right: Forward + backward pass.Batch size: We measure the runtime of all LSTM kernels while varying the batch size (B) from 2 to 256 at sequence length 1024. FIG. 6 shows the results for NH=12 heads with head dimension DH=64. The CUDA fused backend is optimized for smaller batch sizes and shows a 2x speed up over the alternating backend for batch sizes up to 32. For larger batch sizes than 128 CUDA alternating is faster. FIG. 7 shows the results for a single head with head dimension DH=768. At this head dimension CUDA fused is still faster than CUDA alternating up to batch size 32. For larger batch sizes, CUDA alternating is more than two times faster. Comparing to the PyTorch nn.LSTM, we find for medium batch sizes from 8 to 64 it is about 2-3 times faster than and CUDA fused and for larger batch sizes about about 30% faster than CUDA alternating.Language ModelingEven though we do no expect traditional RNNs to outperform Transformers, the language modeling setting serves as an important benchmark for speed on larger scales. Here, we train models at the 165M parameter scale for a Llama-style Transformer without weight tying, i.e. 12 Transformer blocks with Pre-LayerNorm and a Swish-Gated MLP after the attention layer. We replace attention with FlashRNN LSTM and sLSTM layers for a speed comparison. The results show a slowdown of roughly 25% over attention for equal head dimensions or 140% for one RNN head, see Table 1 below. In our experiments, we also compare to the cuDNN implementation of LSTM integrated in PyTorch (torch.nn.LSTM). While it's integration into PyTorch is considerably faster, there are numerical differences to the FlashRNN implementation. With same initialization, FlashRNN LSTMs converge faster in our language experiments (both bfloat16 and float32), even though the differences in a single kernel call are at the expected levels of numerical precision. This deviation should be investigated further and suggests the use of FlashRNN even for the established LSTM architecture. For larger models, we expect local batch sizes to be smaller and the effective speed difference for fused kernels to be higher compared to the alternating version.(Table 1: 165M Model training on 15B tokens of SlimPajama on 8xH100s with two gradient accumulation steps.)State Tracking TaskTo show state tracking capabilities of traditional RNNs in contrast to Transformers and State Space Models experimentally, we train our implementation on the Parity task and evaluate on longer sequences to measure extrapolation capabilities. This serves as a litmus test for state tracking capabilities:(Table 2: Parity Task in Sequence Extrapolation: Transformers, State Space Models and mLSTM fails at this task (close to random chance at 0.5), while traditional recurrent models can learn to extrapolate. Extrapolation accuracies are averaged over three seeds for the best respective learning rate.)Implementation in Hardware and SoftwareEmbodiments 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) or may be 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 a local client and / or one or more remote server farms and / or data centers). The data processing apparatus 802 may comprise any circuit or combination of circuits.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.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.In embodiments, the data processing apparatus 802 may also include a display device 810, one or more speakers, and one or more input devices such as a keyboard 812, a computer mouse 814, 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 present disclosure can be implemented as a computer program (product) with a program code, the program code being operative for performing one of the methods when the computer program product runs on a computer. The program code may, for example, be stored on a machine-readable carrier.In embodiments, a data stream or a sequence of signals may be provided, representing the computer program 808 for performing one of the methods described herein. The data stream or the sequence of signals may, for example, be configured to be transferred via a data communication connection, for example, via the internet.In embodiments, an apparatus or a system may be provided, configured to transfer (e.g., electronically or optically) a computer program 808 for performing one of the methods described herein to a receiver. The receiver may, for example, comprise a computer, a mobile device, a memory device or and other type of data processing apparatus. The apparatus or system may, for example, comprise a file server for transferring the computer program to the receiver.ConclusionThe FlashRNN library serves as a fast and extendable implementation of traditional RNNs with a recurrent connection or memory mixing. It extends RNNs with the multi-head paradigm for sLSTM. FlashRNN provides a speed-up of up to 50x over vanilla PyTorch implementations of RNNs and may serve as a backbone for future RNN architectures that have a recurrent connection.FlashRNN implements two variants, an alternating version switching between point-wise and matrix-multiplication kernels and a fused implementation - optimizing memory transfers, while using hardware-optimized matrix-multiplication. The second leads to a further 3-4x speed-up over the alternating option for small batch sizes. The implementation auto-optimizes its internal sizes for different cache levels via the ConstrINT library a custom library solving general integer constraint satisfaction problems with equality, inequality and divisibility constraints. This library may be re-used for other optimization problems regarding cache sizes on hardware platforms and beyond.We show that with FlashRNN, traditional RNNs are not too far in speed from Transformers in practice, even though they are not parallelizable along the sequence dimension. In the future, it may be optimized to leverage asynchronous memory operations and inter-SRAM connections - recent hardware features that promise further speed ups not realized in this work.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 also represents a description of the 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 a 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.Insofar as a disclosed method is described, explicitly or implicitly, as comprising actions performed by a particular entity or component, such description is likewise to be understood as describing a corresponding complementary method performed by a counterpart entity or component. For example, a method described in terms of operations carried out by a client device or system may alternatively be construed as a method carried out by a server device or system that interacts with the client device or system, and vice versa. In this manner, the disclosed technology encompasses both perspectives of an interaction, regardless of whether a particular embodiment is drafted from the standpoint of one entity or its counterpart.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. Any reference signs in the claims should not be construed as limiting the scope.Particular and preferred aspects of the disclosed technology are set out in the accompanying independent and dependent claims. Features from the dependent claims may be combined with features of the independent claims and with features of other dependent claims as appropriate and not merely as explicitly set out in the claims.

Claims

1. A system comprising a machine-learning model implemented on a data processing apparatus, the data processing apparatus comprising a parallel processor with a memory hierarchy that includes an on-chip memory, the system comprising:a recurrent neural network as the machine-learning model, wherein the recurrent neural network comprises a parallelization mechanism, wherein the recurrent neural network comprises a multi-head architecture wherein a hidden state vector is partitioned into a plurality of head sub-vectors, and a recurrent weight matrix of the recurrent neural network is structured as a block-diagonal matrix where each block corresponds to one of the plurality of head sub-vectors, enabling parallel processing of the heads; anda fused computational kernel configured to execute an entire time-series processing loop for the multi-head recurrent neural network, wherein the kernel is configured to:maintain the blocks of the block-diagonal recurrent weight matrix in the on-chip memory for the duration of the time-series loop; andfor each timestep within the loop, perform both a matrix multiplication for each head using its corresponding block of the recurrent weight matrix and subsequent element-wise operations, within the same kernel execution.

2. The system of claim 1, wherein the on-chip memory comprises at least one of Static Random-Access Memory (SRAM) and processor registers of the parallel processor.

3. The system of claim 1 or 2, wherein the processor is a graphics processing unit (GPU) or a tensor processing unit (TPU).

4. The system of any one of claims 1 to 3, wherein the fused computational kernel is further configured to maintain bias vectors associated with the recurrent neural network in the on-chip memory for the duration of the time-series loop.

5. The system of any one of claims 1 to 4, wherein the fused computational kernel is configured to operate on data represented with bfloat16 or float16 precision.

6. The system of any one of claims 1 to 5, wherein the data processing apparatus is configured for adaptive kernel mode switching between the persistent fused computational kernel and an alternating computational kernel.

7. The system of any one of claims 1 to 6, further comprising a second fused computational kernel configured to execute a backward pass for back-propagation through time, wherein the second kernel is configured to maintain gradients in the on-chip memory across multiple time steps of the backward pass.

8. The system of claim 7, wherein the second fused computational kernel is configured to clip gradient values after each time step within the time-series loop of the backward pass.

9. The system of any one of claims 1 to 8, further comprising an optimization module for configuring the fused computational kernel, the optimization module configured to automatically determine execution parameters for the fused computational kernel, including memory tiling sizes, by solving an integer constraint satisfaction problem that models physical constraints of the parallel processor, wherein the constraints include divisibility constraints corresponding to required input dimensions of hardware-accelerated matrix multiplication units of the parallel processor.

10. The system of claim 9, wherein the physical constraints of the parallel processor modeled by the integer constraint satisfaction problem comprise at least one of: a total SRAM size per computational super-unit, a total register size per computational super-unit, and a number of threads per computational super-unit.

11. The system of any one of claims 1 to 10, wherein the integer constraint satisfaction problem further includes equality constraints and inequality constraints.

12. The system of any one of claims 1 to 11, wherein the optimization module is configured to utilize heuristics to prioritize solutions that maximize utilization of the hardware-accelerated matrix multiplication units.

13. The system of any one of claims 1 to 12, wherein the recurrent neural network is selected from the group consisting of a Long Short-Term Memory (LSTM) network, a Gated Recurrent Unit (GRU) network, and a scalar Long Short-Term Memory (sLSTM) network.

14. A method for processing a sequence of input data using a recurrent neural network on a data processing apparatus, the data processing apparatus comprising a parallel processor with a memory hierarchy including an on-chip memory, the method comprising:executing a fused computational kernel for a plurality of time steps of the recurrent neural network, wherein executing the fused kernel comprises:loading recurrent weight parameters of the recurrent neural network into the on-chip memory;for each of the plurality of time steps:performing a recurrent matrix multiplication using a hidden state from a previous time step and the recurrent weight parameters stored in the on-chip memory; andperforming one or more point-wise operations using a result of the matrix multiplication to generate a new hidden state for a current time step;wherein the recurrent weight parameters remain in the on-chip memory for the duration of the processing of the plurality of time steps within the fused computational kernel.

15. The method of claim 14, wherein loading the recurrent weight parameters into the on-chip memory comprises loading the parameters into Static Random-Access Memory (SRAM).

16. The method of claim 14, wherein loading the recurrent weight parameters into the on-chip memory comprises loading the parameters into processor registers.

17. The method of claim 14, wherein the recurrent neural network comprises a multi-head architecture, and wherein the method further comprises:partitioning a hidden state vector into a plurality of head sub-vectors;structuring the recurrent weight parameters as a block-diagonal matrix where each block corresponds to one of the plurality of head sub-vectors; andprocessing each of the plurality of heads in parallel within the fused computational kernel.

18. The method of claim 14, further comprising, prior to executing the fused computational kernel, automatically determining execution parameters for the kernel by solving an integer constraint satisfaction problem that models physical constraints of the parallel processor.

19. The method of claim 18, wherein the physical constraints include divisibility constraints corresponding to required input dimensions of hardware-accelerated matrix multiplication units of the parallel processor.

20. The method of claim 14, further comprising performing a backward pass using a second fused computational kernel for back-propagation through time, wherein the second fused computational kernel maintains gradients in the on-chip memory across multiple time steps of the backward pass.

21. The method of claim 20, further comprising clipping gradient values after each time step within the execution of the second fused computational kernel.

22. The method of claim 14, wherein the processing is performed using data with bfloat16 or float16 precision.

23. A computer program or a computer-readable medium storing a computer program, the computer program comprising instructions which, when executed by at least one processor, cause the at least one processor to carry out the method of any one of claims 14 to 22.

Images