KronLSTM-SWA-1.5B-100B

A 1.5B parameter base language model with 20 KronLSTM layers and 4 sliding-window softmax attention layers (window 2048, no positional encoding), trained on 100B tokens at 8K context. It is one of the four models of the paper Kronecker Delta Rule for Long Context Linear Attention (arXiv:2610.26834); the code is at NX-AI/KronLSTM.

The same architecture at 32K context: KronLSTM-SWA-1.5B-100B-32k.

KronLSTM is a linear-attention recurrence, the repository README has the equations.

Usage

The model runs on the KronLSTM repository, a snapshot of FLA with the KronLSTM kernels, layer and model; import fla registers the model with transformers' Auto classes. The kernels are written in Triton and need an NVIDIA GPU (developed and tested on H100/A100).

git clone https://github.com/NX-AI/KronLSTM.git
cd KronLSTM
pip install -e ".[cuda]"
pip install flash-attn --no-build-isolation  # the attention layers run on FlashAttention

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

import fla  # noqa: F401

tokenizer = AutoTokenizer.from_pretrained('NX-AI/KronLSTM-SWA-1.5B-100B')
model = AutoModelForCausalLM.from_pretrained('NX-AI/KronLSTM-SWA-1.5B-100B', dtype=torch.bfloat16).cuda().eval()

This is a base model: it is not instruction-tuned or aligned, it is published for better reproducibility and is not intended to be used as a production language model!

Architecture


Parameters	1,395,020,608 (embeddings tied with the output layer)
Layers	24: KronLSTM, with softmax attention in place of KronLSTM at layers 5, 11, 17 and 23 (0-based)
Hidden size	2048
KronLSTM heads	16, key and value head dimension 128
Attention	16 heads of dimension 128, sliding window over the token and the 2048 before it, no positional encoding (`use_rope=False`, `window_size=2049`)
MLP	SwiGLU, intermediate size 5632
Normalisation	pre-norm RMSNorm, eps 1e-5
Vocabulary	32,000 (Llama-2 tokenizer)
Training context	8,192 tokens

Training

Trained with lm-engine and converted to FLA's format.


Data	FineWeb-Edu `sample-100BT`
Tokens	100.0B: 95,368 steps of 1,048,576 tokens (128 sequences of 8,192)
Optimizer	AdamW, betas (0.9, 0.95), eps 1e-10, weight decay 0.1, gradient clipping 1.0
Learning rate	3e-4, 9,537 (10B token) warmup steps, cosine decay to 3e-5
Precision	bf16 mixed precision
Hardware	64 NVIDIA GH200

Evaluation

With lm-evaluation-harness in bf16. HellaSwag 10-shot, PIQA 0-shot, ARC-Easy 0-shot, ARC-Challenge 25-shot, Winogrande 5-shot, LAMBADA (OpenAI) 0-shot; accuracy normalised by length for HellaSwag, PIQA and ARC-Challenge. RULER on 500 samples/length for 4k, 8k and 16k. Recall was evaluated on the Based/JRT real-world recall suite (SWDE, FDA, SQuAD, TriviaQA, NQ, DROP). The setup was zero-shot with greedy decoding in bf16, generating up to 48 tokens. Each prompt was left-truncated to 2,000 tokens so that prompt plus generation fits in 2,048 tokens. Scoring used the contains metric, which checks whether the reference answer appears in the generated text.

Citation

@article{hartl2026kronlstm,
  title   = {Kronecker Delta Rule for Long Context Linear Attention},
  author  = {Hartl, Anamaria-Roberta and Z{\'o}lyomi, Levente and List, Michael and Klambauer, G{\"u}nter and B{\"o}ck, Sebastian and Hochreiter, Sepp},
  journal = {arXiv preprint arXiv:2610.26834},
  year    = {2026}
}