KronLSTM-1.5B-100B-32k

A 1.5B parameter base language model with 24 KronLSTM layers, trained on 100B tokens at 32K 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 8K context: KronLSTM-1.5B-100B.

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]"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

import fla  # noqa: F401

tokenizer = AutoTokenizer.from_pretrained('NX-AI/KronLSTM-1.5B-100B-32k')
model = AutoModelForCausalLM.from_pretrained('NX-AI/KronLSTM-1.5B-100B-32k', 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,414,272,384 (embeddings tied with the output layer)
Layers	24 KronLSTM
Hidden size	2048
KronLSTM heads	16, key and value head dimension 128
MLP	SwiGLU, intermediate size 5632
Normalisation	pre-norm RMSNorm, eps 1e-5
Vocabulary	32,000 (Llama-2 tokenizer)
Training context	32,768 tokens

Training

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

The model continues KronLSTM-1.5B-100B's run: it branches from that run at step 70,000 (73.4B tokens of FineWeb-Edu at 8K context, the recipe on that model's card) and trains the remaining 26.6B tokens at 32K context.

Data (32K stage)	the long-document mixture of ProLong (code 30%, books 30%, textbooks 3%, arXiv 3%), re-tokenized with the Llama-2 tokenizer, and FineWeb-Edu (34%)
Tokens (32K stage)	26.6B: 25,368 steps of 1,048,576 tokens (32 sequences of 32,768)
Optimizer	AdamW, betas (0.9, 0.95), eps 1e-10, weight decay 0.1, gradient clipping 1.0, fresh optimizer state
Learning rate	1,268 (1.3B token) warmup steps to 8.41e-5 (the 8K schedule's value at step 70,000), cosine decay to 0.1x
Precision	bf16 mixed precision
Hardware	32 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, 16k and 32k. 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}
}