lrnnx 1.0.0

A library for Linear RNNs

lrnnx: A library for Linear RNNs

A unified PyTorch library providing easy access to state-of-the-art Linear RNN architectures for sequence modeling. The technical report of this system was accepted to EACL Student Research Workshop 2026. We recommend reading the report before using / contributing to the library.

Installation

From PyPI

Since this library compiles custom CUDA kernels upon installation, we recommend using --no-build-isolation to avoid downloading a duplicate version of PyTorch during the build process.

# standard installation
pip install lrnnx --no-build-isolation
# with optional causal-conv1d
pip install "lrnnx[conv1d]" --no-build-isolation
# for development
pip install "lrnnx[dev]" --no-build-isolation

From Source

We recommend installation with uv for fast, reliable dependency management, though standard pip is fully supported.

Using uv

git clone https://github.com/SforAiDl/lrnnx.git
cd lrnnx
# standard installation
uv sync
# with optional causal-conv1d
uv sync --extra conv1d
# for development
uv sync --extra dev

Using pip

git clone https://github.com/SforAiDl/lrnnx.git
cd lrnnx
# standard installation
pip install -e . --no-build-isolation
# with optional causal-conv1d
pip install -e ".[conv1d]" --no-build-isolation
# for development
pip install -e ".[dev]" --no-build-isolation

Note that since our library builds several custom CUDA kernels, it can take time for this installation to finish. Along with causal-conv1d the full installation can take about 30 minutes, depending on the number of CPUs available.

Model Zoo

Our library provides implementations of the following Linear RNN architectures:

• S4
• S4D
• S5
• Event-SSM (inside S5, use by passing integration_timesteps)
• LRU
• S6 (we implemented other discretizations)
• STREAM (inside S6, use by passing integration_timesteps)
• RG-LRU
• S7
• aTENNuate

We expose several levels of API for each model, including a scan, a recurrent step, and a full layer API matching the paper. For S5 we implement both a convolution based approach and a parallel scan approach. The latter is more stable and faster for most use cases, but the convolution based approach can be faster for very long sequences.

Usage

Training

It is easy to instantiate a model from our library

from lrnnx.models.lti import LRU
from lrnnx.models.ltv import Mamba

model_lti = LRU(d_model, d_state).cuda()
x = torch.randn(
	batch_size, seq_len, d_model, dtype=torch.float32, device="cuda"
)
output = model_lti(x)

model_ltv = Mamba(d_model, d_state).cuda()
x = torch.randn(
	batch_size, seq_len, d_model, dtype=torch.float32, device="cuda"
)
output = model_ltv(x)

Inference

Linear RNNs in torch require special handling during inference, following mamba, we also implement CUDA graphs based inference which reduces CPU overheads, this leads to > 10x speedup compared to using a simple for loop over the sequence length. The main file is generation.py which provides a simple API for autoregressive generation with any of the models in our library. You can see a simple way to use it in our benchmarking script.

Reproducing the Benchmarks from the paper

This script will run both training and inference benchmarks.

python -m benchmarks.run_all

Architectures

We also implement some common architectures based on the models in our library, such as a U-Net (inspired from aTENNuate ) and a hierarchical classifier (inspired from Event-SSM). Additionally, there is a Language Model architecture inspired from Mamba and RG-LRU which can be used for language modeling tasks, with replaceable LRNN and attention layers. This can be used as

from lrnnx.models.language_model import LRNNLMHeadModel

model = LRNNLMHeadModel(
	d_model, d_state, num_layers, vocab_size, mixer_types=["s5", "s6", "attn"]
)
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
logits = model(input_ids)

Tutorial Overview

Based on the architectures, there are tutorials on how to use them for 2 very popular use cases:

1. U-Net Seq2Seq for audio denoising Tutorial
2. Hierarchical Classification Tutorial

Contributing

Please check out our Contributing Guide for details on how to contribute to this project.

Citation

If you use lrnnx in your research, please cite:

@misc{bania2026textttlrnnxlibrarylinearrnns,
	title={$\texttt{lrnnx}$: A library for Linear RNNs}, 
	author={Karan Bania and Soham Kalburgi and Manit Tanwar and Dhruthi and Aditya Nagarsekar and Harshvardhan Mestha and Naman Chibber and Raj Deshmukh and Anish Sathyanarayanan and Aarush Rathore and Pratham Chheda},
	year={2026},
	eprint={2602.08810},
	archivePrefix={arXiv},
	primaryClass={cs.LG},
	url={https://arxiv.org/abs/2602.08810}, 
}

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

This library builds upon the excellent work of researchers who developed the individual LRNN models. Please see individual model documentation for proper citations of the original papers.