Hugging Face's logo

OpenOneRec
/
KSA-1.9B-base

Text Generation
Transformers
Safetensors
English
Chinese
qwen3
ksa
ksa-v1
custom_code
long-context
conversational
text-generation-inference
Model card Files
xet
 Community
KSA-1.9B-base
File size: 5,576 Bytes
6c91dd5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
---
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}
}
```