Llama 1B with Kimi Delta Attention, 6B tokens

A dense ~1B-parameter Llama 3-style decoder where one attention layer in every four is Kimi Delta Attention (KDA), a linear-attention layer from Kimi Linear. It is otherwise identical to the QK-norm baseline and was trained on the same 6B FineWeb tokens.

This is a base model. It is not instruction-tuned or safety-tuned.

Results

Training

Zero-shot benchmarks

Scores from lm-eval on each task's full split. Shared-9 is the unweighted mean of the nine tasks.

The +0.14 Shared-9 gap is within single-seed noise, so treat KDA as matching the baseline, not beating it.

Usage

The architecture class ships with the checkpoint, so load it with trust_remote_code=True. The KDA layers need flash-linear-attention and a CUDA GPU (its kernels are written in Triton). Tested with transformers==5.8.0 and flash-linear-attention==0.5.2.

pip install "transformers==5.8.0" "flash-linear-attention==0.5.2"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "Mercity/pretrain-kda-1b"

tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
    repo,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)

inputs = tokenizer("The capital of France is", return_tensors="pt").to(model.device)
# use_cache=False is required: the KDA layers keep no recurrent state between
# decoding steps, so cached generation would feed them one token at a time.
output = model.generate(**inputs, max_new_tokens=32, do_sample=False, use_cache=False)
print(tokenizer.decode(output[0], skip_special_tokens=True))

Generation without the cache re-reads the full sequence at every step, so it is slow for long outputs. Scoring text in one forward pass has no such cost:

batch = tokenizer("FineWeb is a large web-text dataset.", return_tensors="pt").to(model.device)
with torch.no_grad():
    loss = model(**batch, labels=batch["input_ids"]).loss
print(f"loss={loss.item():.3f}  ppl={loss.exp().item():.1f}")

The first forward pass on a new machine spends about 90 seconds compiling the KDA kernels.

Model details

Training

Related checkpoints

Limitations

Trained on 6B English web tokens only, a small budget for a 1B model. Benchmark scores are single-seed. KDA inference needs a CUDA GPU and does not support cached generation through transformers. The model will repeat or make up facts and has had no alignment training.