ParaRNN: Parallel Training for Non-Linear RNNs (sLSTM, RWKV-7, CfC, Titans) via Triton Newton Scans [P]

I've built and open-sourced ParaRNN, a PyTorch/Triton library that parallelizes the training of non-linear recurrent architectures across time using Newton-Raphson iterations mapped to associative prefix scans. While non-linear RNNs provide O(1) inference without KV-cache explosion, training them has historically been bottlenecked by slow sequential unrolls. ParaRNN solves this with custom fused Triton solvers: ~200x wall-clock speedup: Runs orders of magnitude faster than sequential BPTT on long sequences (e.g., CfC at T=2048 takes 2.79ms vs 643ms on an RTX 3060). Flat convergence up to 131k context: Empirically verified that Newton iteration budgets stay flat (K) le 3) matching sequential unroll (au le 10{-4}) up to 131,072 tokens. Supported cells: Fused kernels for sLSTM (4x4 block Jacobians & Beck-style head-mixing), RWKV-7 (Goose matrix state monoid), CfC (Liquid continuous-time), Titans (associative surprise GD), Modern Hopfield, and M²RNN. Production hooks: torch.compile(fullgraph=True) with 0 graph breaks, deterministic closed-form VJPs (no atomics), and an out-of-tree vLLM plugin. GitHub: https://github.com/bugkira/pararnn-torch Quick Install: pip install pararnn-torch I’d love to get your feedback, bug reports on custom shapes, or benchmark numbers on different hardware! submitted by /u/bugkira [link] [comments]

Related

• Generating Bad Apple autonomously from a single initial state using a tiny recurrent dynamical system (417k params) [P]
• I released a softmax-free attention model at GPT-2 Medium scale (~354M params, 11.5B tokens): structural sparsity + tile-skipping kernels for long-context VRAM savings. Open weights + custom Triton kernels [R]

Source: r/MachineLearning | 2026-09-07