| --- |
| license: apache-2.0 |
| language: |
| - en |
| - zh |
| library_name: transformers |
| pipeline_tag: text-generation |
| tags: |
| - ksa |
| - ksa-v1 |
| - custom_code |
| - long-context |
| --- |
| |
|
|
| <div align="center"> |
| <h1>KSA V1 1.9B</h1> |
| <p><strong>Kwai Summary Attention for long-context modeling</strong></p> |
| <p> |
| <a href="https://arxiv.org/abs/2604.24432">Technical report</a> · |
| <a href="https://github.com/Kuaishou-OneRec/KSA">Code</a> · |
| <a href="LICENSE">Apache-2.0</a> |
| </p> |
| <p><a href="README.md">English</a> | <a href="README_zh.md">ä¸æ–‡</a></p> |
| </div> |
| |
| ## 📖 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. |
|
|
| <p align="center"><img src="assets/figures/mainmodel.png" width="80%" alt="Hybrid KSA: Summary tokens interleaved with text, and KSA / Full Attention layers in a 3:1 layout." /></p> |
|
|
| ## 🔥 News |
|
|
| - **2026-10-13** — [KSA V2 technical report](https://arxiv.org/abs/2610.18654) released. |
| - **2026-05-08** — [KSA-4B-base](https://huggingface.co/OpenOneRec/KSA-4B-base), our Qwen3-4B CPT variant, released. |
| - **2026-04-28** — [KSA technical report](https://arxiv.org/abs/2604.24432) and training code released. |
|
|
| ## 🤖 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: |
|
|
| ```bash |
| 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: |
|
|
| ```bash |
| export PYTHONPATH="$PWD/kernel:$PWD/kernel/deps${PYTHONPATH:+:$PYTHONPATH}" |
| ``` |
|
|
| ```python |
| 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](LICENSE). Qwen-derived files retain their |
| copyright notices. Bundled FlashAttention CuTe retains its |
| [BSD-3-Clause license](kernel/deps/flash_attn/LICENSE) and |
| [AUTHORS](kernel/deps/flash_attn/AUTHORS). |
|
|
| ```bibtex |
| @techreport{kwai2026ksa, |
| title = {Kwai Summary Attention Technical Report}, |
| author = {OneRec Team}, |
| year = {2026}, |
| institution = {Kuaishou Technology}, |
| url = {https://arxiv.org/abs/2604.24432} |
| } |
| ``` |
|
|
