Skip to main content
This guide will help you run your first Linear RNN model with lrnnx.

​
Install lrnnx

First, install the library: 
pip install lrnnx
We recommend installing PyTorch first to match your CUDA version, then install lrnnx with --no-build-isolation.

​
Basic usage

lrnnx provides two main categories of models:
  • LTI (Linear Time-Invariant): Models with fixed dynamics (S4, S4D, S5, LRU, Centaurus)
  • LTV (Linear Time-Varying): Models with input-dependent dynamics (Mamba, S6, S7, RG-LRU)

​
LTI model example

from lrnnx.models.lti import LRU
import torch

# Create model
model = LRU(d_model=64, d_state=64).cuda()

# Create input
batch_size, seq_len, d_model = 2, 128, 64
x = torch.randn(batch_size, seq_len, d_model, dtype=torch.float32, device="cuda")

# Forward pass
output = model(x)
print(output.shape)  # torch.Size([2, 128, 64])

​
LTV model example

from lrnnx.models.ltv import Mamba
import torch

# Create model
model = Mamba(d_model=64, d_state=16).cuda()

# Create input
batch_size, seq_len, d_model = 2, 128, 64
x = torch.randn(batch_size, seq_len, d_model, dtype=torch.float32, device="cuda")

# Forward pass
output = model(x)
print(output.shape)  # torch.Size([2, 128, 64])

​
Available models

​
LTI models

  • S4 - Structured State Space model with DPLR parameterization
  • S4D - Diagonal variant of S4 (simpler, faster)
  • S5 - Simplified State Space model
  • LRU - Linear Recurrent Unit (minimal diagonal SSM)
  • Centaurus - Multi-mode architecture with intra-state mixing

​
LTV models

  • Mamba - Selective State Space model with input-dependent dynamics
  • S6 - Time-varying variant of S4 (included in Mamba)
  • S7 - All matrices input-dependent (A, B, C, D)
  • RG-LRU - Gated LRU from Griffin architecture
  • STREAM - Event-based variant (use with integration_timesteps parameter)
Import models from lrnnx.models.lti for LTI models or lrnnx.models.ltv for LTV models.

​
Next steps

Training guide

Learn how to train models for your tasks

Inference guide

Optimize inference with CUDA graphs for 10x speedup

Model overview

Explore all available models in detail

Architectures

Use pre-built architectures for language modeling, U-Net, and classification

Build docs developers (and LLMs) love

Get started for freeTalk to us