---
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：面向长上下文建模的高效注意力</strong></p>
  <p>
    <a href="https://arxiv.org/abs/2604.24432">技术报告</a> ·
    <a href="https://github.com/Kuaishou-OneRec/KSA">代码</a> ·
    <a href="LICENSE">Apache-2.0</a>
  </p>
  <p><a href="README.md">English</a> | <a href="README_zh.md">中文</a></p>
</div>

## 📖 简介

**KSA V1 1.9B** 是我们**从零训练**的 Hybrid KSA 基础模型，支持 128K 上下文。
可学习的 Summary Token 将远端原文压缩为持久的长程记忆。每 3 层 KSA 配置
1 层 Full Attention，每个 Summary 对应 8 个原文 token。

这也是 KSA V2 使用的 Hybrid KSA 1.9B 初始主干，**不含 C-KSA 二级压缩或 S-KSA Indexer**，
未经过指令微调。

<p align="center"><img src="assets/figures/mainmodel.png" width="80%" alt="Hybrid KSA：Summary 与原文交错排列，KSA 和 Full Attention 层按 3:1 配置。" /></p>

## 🔥 新闻

- **2026-10-13** —— [KSA V2 技术报告](https://arxiv.org/abs/2610.18654)发布。
- **2026-05-08** —— 基于 Qwen3-4B 继续预训练的 [KSA-4B-base](https://huggingface.co/OpenOneRec/KSA-4B-base) 发布。
- **2026-04-28** —— [KSA 技术报告](https://arxiv.org/abs/2604.24432)与训练代码发布。

## 🤖 模型配置

当前目录为 **KSA V1 1.9B**。

| 配置项 | KSA V1 1.9B | KSA V1 4B |
| :--- | :--- | :--- |
| 训练方式 | 从零训练 | 基于 Qwen3-4B 继续预训练 |
| 层数 | 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 比例 (KSA : Full) | 3 : 1 | 3 : 1 |
| 最大上下文（原文 token） | 131072 | 131072 |
| Summary chunk size | 8 | 8 |
| KSA sliding chunk number | 128 | 128 |
| Summary Q/K/V 投影 | 与原文共享 | 与原文共享 |
| RoPE base | 1000000 | 1000000 |
| Tied embeddings | False | True |
| 权重／推理精度 | BF16 / BF16 | BF16 / BF16 |

## 📈 长上下文检索

KSA V1 报告中的 RULER 结果：

### 1.9B：从零训练

| 模型 | 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：继续预训练

| 模型 | 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 |

## 快速开始

Linux、Python 3.12、H100/H800。下载**完整模型目录**（包含 `kernel/`），并在该目录下执行：

```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
```

省略 `--prompt` 可进行交互式文本续写。`--apply-chat-template` 只应用 tokenizer 中的对话模板，不代表模型经过指令微调。

使用 Transformers 5.0.0 和随包提供的 `ksa_attn`／FlashAttention CuTe 源码，无需训练仓库、单独编译内核或安装 FlashAttention 2 CUDA 扩展。

### Transformers 用法

导入 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))
```

## 许可证与引用

模型权重及我们的推理代码采用 [Apache-2.0](LICENSE)。基于 Qwen 的文件保留原版权声明。
随包的 FlashAttention CuTe 保留其 [BSD-3-Clause 许可证](kernel/deps/flash_attn/LICENSE)
与 [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}
}
```