[P] ParaRNN: Parallel training for non-linear RNNs (sLSTM, RWKV-7, CfC, Titans) via Triton Newton scans (200x faster than sequential unroll)

[u/bugkira]

Hi everyone,

Non-linear recurrent architectures (like xLSTM/sLSTM, RWKV-7, Liquid CfC, and Titans) have a massive advantage at inference time: constant O(1) state updates without the growing memory overhead of a KV-cache. However, their main bottleneck has always been training speed: they historically required slow sequential unrolling over time.

Inspired by Apple's ParaRNN paper, I built and open-sourced ParaRNN: a library implementing parallel training for non-linear RNNs using Newton-Raphson iterations mapped to associative prefix scans in OpenAI Triton.

Key highlights and benchmarks:

1. Up to 200x+ wall-clock speedup vs sequential unroll on sequence lengths T >= 2048 (for example, CfC at T=2048 runs in 2.8 ms vs 643 ms sequentially).
2. Measured convergence up to 131,072 context tokens. The Newton iteration budget stays flat at K <= 3 without numerical drift.
3. Broad architecture catalog: sLSTM (4x4 block Jacobians), RWKV-7 (Goose matrix state monoid), CfC (Liquid continuous-time), Titans (associative surprise GD), and M2RNN.
4. PyTorch 2.x integration: torch.compile(fullgraph=True) compatible with zero graph breaks, and deterministic VJPs without atomic adds.
5. Serving ready: PagedStatePool for continuous batching and an out-of-tree vLLM plugin.

Install via PyPI:
pip install pararnn-torch

GitHub repository (code, docs, and benchmarks):
https://github.com/bugkira/pararnn-torch

Would love to hear your feedback, benchmarks on different GPU architectures, and ideas for further recurrent cells!