FlashRNN

FlashRNN provides different backends optimized for different hardware and constraints:

• cuda_fused: This is the high-performance option. It fuses matrix multiplication with pointwise non-linearity into a single CUDA kernel using wmma instructions and custom caching on SRAM/registers. This provides a $2\times$ to $5\times$ speedup over the standard cuda backend by reducing I/O overhead for the recurrent matrix $\mathbf{R}$ and biases $\mathbf{b}$.
• cuda: A basic CUDA implementation that alternates between matrix multiplication and pointwise non-linearity. It has no restrictions on head dimension $d_{head}$.
• triton: Uses Triton kernels. It is highly efficient but imposes a constraint on the head dimension: $d_{head} \leq 64$.
• vanilla: A standard, non-optimized implementation.

To ensure correct compilation and library linking, you may need to set the following environment variables:

1. Set Torch CUDA Architecture List: To ensure torch uses compatible compilation flags, set TORCH_CUDA_ARCH_LIST to your specific hardware architectures (e.g., 8.0;8.6;9.0).
2. Inject CUDA Include Paths: If you have a custom CUDA setup, use FLASHRNN_EXTRA_INCLUDE_PATHS to provide the correct include directories. This can be set in the shell or within Python using os.environ.

# Set architecture list
export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"

# Inject include paths
export FLASHRNN_EXTRA_INCLUDE_PATHS='/usr/local/include/cuda/:/usr/include/cuda/'

FlashRNN can be installed via pip.

Hardware Requirements:

• CUDA Compute Capability 8.0 or later.
• An up-to-date g++ compiler.

Recommended Setup: It is recommended to use conda with the provided environment file to ensure compatible versions of PyTorch and CUDA:

conda env create -n flashrnn -f environment_pt240cu124.yaml

pip install flashrnn

When working with FlashRNN, ensure your input tensors follow these expected shapes:

• Wx (Input weights): [B, T, NG, NH, D]
• R (Recurrent weights): [NG, NH, D, D]
• b (Bias): [NG, NH, D]
• states_initial (Initial states): [NS, B, NH, D]

Where:

• B: Batch size
• T: Sequence length
• NG: Number of gates
• NH: Number of heads
• D: Embedding/Hidden dimension
• NS: Number of states (e.g., 2 for LSTM: c and h)

When initializing the states_initial tensor for sLSTM, it is critical to ensure that all states except the hidden state (h) are set to zero. Failure to do so will result in numerical mismatches between the optimized kernels and standard PyTorch implementations.

The states_initial tensor has the shape [NS, B, NH, D], where NS is the number of states (typically 4 for sLSTM, representing c, h, and others).

# NS = 4 (e.g., c, h, and others)
# B = batch size, NH = number of heads, D = dimension
states_initial = torch.zeros([NS, B, NH, D], device=device, dtype=dtype)

# IMPORTANT: Only the h state (index 1) should be non-zero if you have specific initial values.
# All other states (c, etc.) must be zero for numerical consistency.
states_initial[1, :, :, :] = 0.0 

To ensure your configuration of flashrnn (specifically the triton_fused backend) is numerically accurate compared to a baseline, use the model_test utility. This performs a comparison between the optimized backend and a reference backend (like vanilla_fwbw) including backward pass gradients.

from flashrnn.tests.utils import model_test

model_test(
    batch_size=B,
    sequence_size=T,
    num_heads=NH,
    head_dim=D,
    backend="triton_fused",
    backend_cmp="vanilla_fwbw",
    function="slstm",
    dtype=dtype,
    include_backward=True,
    tensor_compare_kwargs={"atol": 0.5, "rtol": 1.0},
)

To ensure a custom implementation (like Triton) matches PyTorch's expected behavior, you can compare the outputs and gradients against a forward_sequence baseline using lstm_pointwise_fw.

1. Compute the forward pass with the baseline.
2. Compute the forward pass with the target implementation (e.g., lstm_tr_fwbw).
3. Compare the maximum absolute difference between outputs: (h_target - h_baseline).abs().max().
4. Perform a backward pass on both and compare gradients: (grad_target - grad_baseline).abs().max().

The flashrnn function supports several backends for executing RNN operations. Choose the backend based on your performance and compatibility requirements:

• vanilla_fwbw: Uses standard PyTorch forward-backward implementations.
• triton_fused: Uses optimized Triton kernels for fused operations.
• cuda_fused: Uses optimized CUDA kernels.

You can also specify the dtype (e.g., "float32", "bfloat16") to control precision.

The flashrnn function accepts two primary string arguments to control its behavior:

function (RNN Architecture)

Specifies the type of recurrent network to implement:

• lstm
• gru
• elman
• slstm

backend (Execution Engine)

Specifies the computational backend used for the operation:

• cuda_fused: Uses a fused CUDA kernel with wmma instructions and custom SRAM/register caching. This is the fastest version but requires specific hardware.
• cuda: A standard CUDA backend that alternates between recurrent matrix multiplication and non-linear pointwise function application. It is not limited by head dimension.
• triton: A Triton-based backend. For this backend, heads are limited to small head dimensions ($d_{head} \leq 64$).
• vanilla: Standard implementation.