KSA V1 1.9B

Kwai Summary Attention for long-context modeling

Technical report · Code · Apache-2.0

English | 中文

📖 Introduction

KSA V1 1.9B is our Hybrid KSA base model trained from scratch, with 128K context. Learnable Summary tokens compress distant text into persistent long-range memory. Three KSA layers alternate with one Full Attention layer; each Summary represents eight original text tokens.

This is the Hybrid KSA 1.9B backbone used to initialize KSA V2. It has no C-KSA compression or S-KSA Indexer and is not instruction-tuned.

Hybrid KSA: Summary tokens interleaved with text, and KSA / Full Attention layers in a 3:1 layout.

🔥 News

🤖 Model configuration

This directory contains KSA V1 1.9B.

Configuration KSA V1 1.9B KSA V1 4B
Training From scratch CPT from Qwen3-4B
Number of layers 24 36
Hidden size 2048 2560
Intermediate size 6144 9728
Attention heads (Q / KV) 16 / 16 (MHA) 32 / 8 (GQA)
Head dimension 128 128
Hybrid ratio (KSA : Full) 3 : 1 3 : 1
Maximum context (text tokens) 131072 131072
Summary chunk size 8 8
KSA sliding chunk number 128 128
Summary Q/K/V projections Shared with text Shared with text
RoPE base 1000000 1000000
Tied embeddings False True
Weight / inference precision BF16 / BF16 BF16 / BF16

📈 Long-context retrieval

RULER results from the KSA V1 report:

1.9B: From scratch

Model 4K 8K 16K 32K 64K 128K Mean
Full 76.08 72.85 73.24 69.06 65.32 48.75 67.55
Hybrid SWA 74.54 71.69 69.54 67.86 63.03 56.64 67.22
Hybrid SCA 77.72 75.22 72.55 67.74 63.54 58.01 69.13
Hybrid GDN 79.83 76.01 74.04 70.41 69.39 59.87 71.59
Pure KSA 70.44 65.91 66.74 62.54 57.13 39.29 60.34
Hybrid KSA 80.65 73.35 74.07 72.30 69.95 65.35 72.61

4B: Continual pretraining

Model 4K 8K 16K 32K 64K 128K Mean
Full 92.88 91.38 89.12 84.74 78.16 65.86 83.69
Hybrid SWA 91.30 88.03 82.87 78.94 73.88 66.27 80.22
Hybrid SCA 86.02 84.28 80.67 76.89 68.88 60.94 76.28
Hybrid Linear 86.39 83.86 78.06 76.48 73.50 67.98 77.71
Pure KSA 91.55 86.78 84.78 80.30 76.09 66.81 81.05
Hybrid KSA 92.97 90.53 88.86 86.65 76.04 71.67 84.45

Quick start

Linux, Python 3.12 and H100/H800. Download the entire model directory, including kernel/, and run from it:

python -m pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu126
python -m pip install -r requirements.txt

CUDA_VISIBLE_DEVICES=0 python generate.py \
    --prompt "The main advantage of long-context language models is" \
    --max-new-tokens 128

Omit --prompt for interactive completion. The optional --apply-chat-template only applies the tokenizer template; it does not add instruction tuning.

Transformers 5.0.0 and the bundled ksa_attn / FlashAttention CuTe sources are used. No training checkout, kernel build or FlashAttention 2 CUDA extension is required.

Transformers usage

Set the bundled kernel paths before importing Transformers:

export PYTHONPATH="$PWD/kernel:$PWD/kernel/deps${PYTHONPATH:+:$PYTHONPATH}"
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(".", local_files_only=True)
model = AutoModelForCausalLM.from_pretrained(
    ".",
    trust_remote_code=True,
    local_files_only=True,
    dtype=torch.bfloat16,
    device_map={"": "cuda:0"},
    attn_implementation="sdpa",
).eval()

inputs = tokenizer("The main advantage of long-context models is", return_tensors="pt")
inputs = inputs.to(model.device)
with torch.inference_mode():
    output = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
    )
print(tokenizer.decode(output[0, inputs.input_ids.shape[1]:], skip_special_tokens=True))

License and citation

Weights and our runtime: Apache-2.0. Qwen-derived files retain their copyright notices. Bundled FlashAttention CuTe retains its BSD-3-Clause license and AUTHORS.

@techreport{kwai2026ksa,
  title       = {Kwai Summary Attention Technical Report},
  author      = {OneRec Team},
  year        = {2026},
  institution = {Kuaishou Technology},
  url         = {https://arxiv.org/abs/2604.24432}
}
Downloads last month
28
Safetensors
Model size
2B params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for OpenOneRec/KSA-1.9B-base

Quantizations
1 model

Collection including OpenOneRec/KSA-1.9B-base

Papers for OpenOneRec/KSA-1.9B-base