Instructions to use OpenOneRec/KSA-V2-4B-base with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use OpenOneRec/KSA-V2-4B-base with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="OpenOneRec/KSA-V2-4B-base", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("OpenOneRec/KSA-V2-4B-base", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("OpenOneRec/KSA-V2-4B-base", trust_remote_code=True, device_map="auto")
messages = [
    {"role": "user", "content": "Who are you?"},
]
inputs = tokenizer.apply_chat_template(
	messages,
	add_generation_prompt=True,
	tokenize=True,
	return_dict=True,
	return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=256)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use OpenOneRec/KSA-V2-4B-base with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "OpenOneRec/KSA-V2-4B-base"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "OpenOneRec/KSA-V2-4B-base",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/OpenOneRec/KSA-V2-4B-base

SGLang

How to use OpenOneRec/KSA-V2-4B-base with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "OpenOneRec/KSA-V2-4B-base" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "OpenOneRec/KSA-V2-4B-base",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "OpenOneRec/KSA-V2-4B-base" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "OpenOneRec/KSA-V2-4B-base",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use OpenOneRec/KSA-V2-4B-base with Docker Model Runner:
```
docker model run hf.co/OpenOneRec/KSA-V2-4B-base
```

KSA-V2-4B-base / README_zh.md

OpenOneRec

Release KSA-V2-4B-base weights and inference runtime

a1111f3 verified 2 days ago

preview code

Raw History Blame Contribute Delete

9.62 kB

---
license: apache-2.0
language:
- en
- zh
library_name: transformers
pipeline_tag: text-generation
tags:
- ksa
- ksa-v2
- custom_code
- long-context
---


<div align="center">
  <h1>KSA V2 4B</h1>
  <p><strong>用 Summary 同时构建长程记忆与稀疏索引</strong></p>
  <p>
    <a href="https://arxiv.org/abs/2610.18654"><img alt="KSA V2 Paper" src="https://img.shields.io/badge/KSA_V2-arXiv%3A2610.18654-b31b1b?logo=arxiv" /></a>
    <a href="https://arxiv.org/abs/2604.24432"><img alt="KSA Paper" src="https://img.shields.io/badge/KSA-arXiv%3A2604.24432-b31b1b?logo=arxiv" /></a>
    <a href="https://github.com/Kuaishou-OneRec/KSA"><img alt="GitHub" src="https://img.shields.io/badge/GitHub-Kuaishou--OneRec-black?logo=github" /></a>
    <a href="LICENSE"><img alt="License" src="https://img.shields.io/badge/License-Apache%202.0-green" /></a>
  </p>
  <p><a href="README.md">English</a> | <a href="README_zh.md">中文</a></p>
</div>
## 📖 简介

**Kwai Summary Attention（KSA）** 是我们面向长上下文建模的高效注意力机制。
KSA 在固定 chunk 边界插入可学习的 Summary Token，以压缩后的状态保留对远端历史的访问。

**KSA V2** 进一步让 Summary 承担两个角色：

- **Compressed KSA（C-KSA）**：将多份 Summary KV 聚合为 Meta KV，进一步降低长程记忆的存储开销。
- **Sparse KSA（S-KSA）**：以 Summary 聚合原文块的选路表示，用轻量 Indexer 选择相关块，直接读取这些块的原始 KV。

我们保留 Hybrid KSA 的 3:1 层间布局，将 KSA 层替换为 C-KSA、Full 层替换为 S-KSA，
分别优化长程记忆和原文读取。

<p align="center"><img src="assets/figures/ksa_v2_overview.png" width="100%" alt="从 Hybrid KSA 到 KSA V2：保持 Summary 布局，引入 C-KSA 压缩与 S-KSA 原文块选路。" /></p>
<p align="center"><em>从 Hybrid KSA 到 KSA V2：复用 Summary，在注意力层内压缩 KV，并稀疏读取原文。</em></p>

本模型基于 **Qwen3-4B 继续预训练得到的 4B Hybrid KSA 主干**进行适配，是未经过指令微调的基础模型。

## 🔥 新闻

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

## ✨ 方法

**C-KSA：Summary → Meta KV。** 在注意力层内进行可学习的逐维聚合。
每组 Summary 生成一份更紧凑的 Meta KV；Summary Token 本身仍按原 chunk 粒度传递，
不必为增加压缩率重新训练更大 chunk 的主干。查询同时访问远程 Meta、边界 Summary
和局部原文 KV；Meta 仅在当前层内使用，不作为新 token 传递。

<p align="center"><img src="assets/figures/ksa_v2_compression.png" width="80%" alt="C-KSA：将完整的远程 Summary 分组压缩为 Meta KV，保留边界 Summary 与局部原文 KV。" /></p>
<p align="center"><em>C-KSA 两级压缩示意，图中 chunk 大小为 4，每 2 个 Summary 聚合为一个 Meta。</em></p>

**S-KSA：Summary → 块级索引。** 以 Summary 为 query 聚合原文块，再映射到低维索引空间。
主注意力读取选中块的原始 KV，并保留 sink 与局部窗口。Indexer 用按块取 max 的
dense attention teacher 监督训练。

<p align="center"><img src="assets/figures/ksa_v2_sparse.png" width="80%" alt="S-KSA：Summary 引导块内聚合，Indexer 选择远程原文块，主注意力读取其原始 KV 以及 sink 和局部窗口。" /></p>
<p align="center"><em>S-KSA 块级选路示意，图中 chunk 大小为 4、block 大小为 2；发布模型的实际配置见下表。</em></p>

**三个适配阶段。** 压缩模块 warmup → C-KSA 全参数 CPT → S-KSA 聚合模块与 Indexer warmup。
最后一阶段冻结继承参数，完成后直接导出用于稀疏推理。

## 🤖 模型配置

当前目录为 **KSA-v2-4B**。

| 配置项 | KSA-v2-1.9B | KSA-v2-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 比例（C-KSA : S-KSA） | 3 : 1 | 3 : 1 |
| 最大上下文（原文 token） | 131072 | 131072 |
| Summary chunk size | 8 | 8 |
| C-KSA sliding chunk number | 128 | 128 |
| C-KSA group size | 128 | 32 |
| S-KSA block size | 4 | 4 |
| S-KSA 远程 token 预算 | 4096 | 4096 |
| Sink tokens | 8 | 8 |
| S-KSA 局部窗口（原文 token） | 1024 | 1024 |
| Indexer heads × head dimension | 8 × 128 | 8 × 128 |
| Indexer query rank | 512 | 512 |
| Indexer RoPE dimension | 64 | 64 |
| Indexer Q/K RMSNorm | False | True |
| Teacher 聚合 | max | max |
| Tied embeddings | False | True |

远程预算包含 8 个 sink token，局部窗口另计。

权重以 FP32 保存，推理时以 BF16 加载。

## 📈 实验结果

<p align="center"><img src="assets/figures/ksa_v2_results.png" width="85%" alt="技术报告第一页的 1.9B 主结果图：长上下文检索与通用能力，以及单 token 推理 FLOPs 和 KV cache 大小的理论值。" /></p>
<p align="center"><em>1.9B 主结果：长上下文与通用能力（左），单 token 推理 FLOPs 与 KV cache 大小的理论值（右）。</em></p>

我们在两种 Hybrid KSA 主干上验证 KSA V2：从零训练的 1.9B，以及基于 Qwen3-4B
继续预训练的 4B。两者均经过三阶段适配，原始 baseline 结果沿用 KSA V1。

### 1.9B：长上下文检索

C-KSA 使用 group128，S-KSA 使用 block4；Hybrid KSA 是适配前的父模型。

| 模型 | 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 |
| C-KSA + Full | 81.52 | 72.24 | 74.44 | 71.93 | 69.61 | 66.31 | 72.68 |
| **C-KSA + S-KSA** | 81.48 | 72.42 | 75.67 | 70.78 | 68.99 | 65.01 | 72.39 |

### 4B：扩展到 GQA 主干

C-KSA 使用 group32，S-KSA 使用 block4，保留主干的 GQA 结构。

| 模型 | 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 |
| Hybrid KSA + CPT | 92.22 | 89.85 | 87.20 | 83.09 | 74.74 | 72.03 | 83.19 |
| C-KSA + Full | 92.77 | 88.12 | 86.27 | 82.68 | 78.58 | 73.92 | 83.72 |
| **C-KSA + S-KSA** | 93.38 | 87.78 | 85.45 | 83.60 | 80.28 | 74.41 | 84.15 |

Hybrid KSA + CPT 是不改模型结构的续训对照。

更多消融实验和推理效率分析见技术报告。

## 快速开始

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)；入口仅加载 CuTe，不依赖单独的 FlashAttention 2 扩展。

```bibtex
@techreport{kwai2026ksav2,
  title       = {Kwai Summary Attention V2 Technical Report},
  author      = {OneRec Team},
  year        = {2026},
  institution = {Kuaishou Technology},
  url         = {https://arxiv.org/abs/2610.18654}
}
```

	---
	license: apache-2.0
	language:
	- en
	- zh
	library_name: transformers
	pipeline_tag: text-generation
	tags:
	- ksa
	- ksa-v2
	- custom_code
	- long-context
	---



	<div align="center">
	<h1>KSA V2 4B</h1>
	<p><strong>用 Summary 同时构建长程记忆与稀疏索引</strong></p>
	<p>
	<a href="https://arxiv.org/abs/2610.18654"><img alt="KSA V2 Paper" src="https://img.shields.io/badge/KSA_V2-arXiv%3A2610.18654-b31b1b?logo=arxiv" /></a>
	<a href="https://arxiv.org/abs/2604.24432"><img alt="KSA Paper" src="https://img.shields.io/badge/KSA-arXiv%3A2604.24432-b31b1b?logo=arxiv" /></a>
	<a href="https://github.com/Kuaishou-OneRec/KSA"><img alt="GitHub" src="https://img.shields.io/badge/GitHub-Kuaishou--OneRec-black?logo=github" /></a>
	<a href="LICENSE"><img alt="License" src="https://img.shields.io/badge/License-Apache%202.0-green" /></a>
	</p>
	<p><a href="README.md">English</a> \| <a href="README_zh.md">中文</a></p>
	</div>

	## 📖 简介

	Kwai Summary Attention（KSA）是我们面向长上下文建模的高效注意力机制。
	KSA 在固定 chunk 边界插入可学习的 Summary Token，以压缩后的状态保留对远端历史的访问。

	KSA V2 进一步让 Summary 承担两个角色：

	- Compressed KSA（C-KSA）：将多份 Summary KV 聚合为 Meta KV，进一步降低长程记忆的存储开销。
	- Sparse KSA（S-KSA）：以 Summary 聚合原文块的选路表示，用轻量 Indexer 选择相关块，直接读取这些块的原始 KV。

	我们保留 Hybrid KSA 的 3:1 层间布局，将 KSA 层替换为 C-KSA、Full 层替换为 S-KSA，
	分别优化长程记忆和原文读取。

	<p align="center"><img src="assets/figures/ksa_v2_overview.png" width="100%" alt="从 Hybrid KSA 到 KSA V2：保持 Summary 布局，引入 C-KSA 压缩与 S-KSA 原文块选路。" /></p>
	<p align="center"><em>从 Hybrid KSA 到 KSA V2：复用 Summary，在注意力层内压缩 KV，并稀疏读取原文。</em></p>

	本模型基于 Qwen3-4B 继续预训练得到的 4B Hybrid KSA 主干进行适配，是未经过指令微调的基础模型。

	## 🔥 新闻

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

	## ✨ 方法

	C-KSA：Summary → Meta KV。在注意力层内进行可学习的逐维聚合。
	每组 Summary 生成一份更紧凑的 Meta KV；Summary Token 本身仍按原 chunk 粒度传递，
	不必为增加压缩率重新训练更大 chunk 的主干。查询同时访问远程 Meta、边界 Summary
	和局部原文 KV；Meta 仅在当前层内使用，不作为新 token 传递。

	<p align="center"><img src="assets/figures/ksa_v2_compression.png" width="80%" alt="C-KSA：将完整的远程 Summary 分组压缩为 Meta KV，保留边界 Summary 与局部原文 KV。" /></p>
	<p align="center"><em>C-KSA 两级压缩示意，图中 chunk 大小为 4，每 2 个 Summary 聚合为一个 Meta。</em></p>

	S-KSA：Summary → 块级索引。以 Summary 为 query 聚合原文块，再映射到低维索引空间。
	主注意力读取选中块的原始 KV，并保留 sink 与局部窗口。Indexer 用按块取 max 的
	dense attention teacher 监督训练。

	<p align="center"><img src="assets/figures/ksa_v2_sparse.png" width="80%" alt="S-KSA：Summary 引导块内聚合，Indexer 选择远程原文块，主注意力读取其原始 KV 以及 sink 和局部窗口。" /></p>
	<p align="center"><em>S-KSA 块级选路示意，图中 chunk 大小为 4、block 大小为 2；发布模型的实际配置见下表。</em></p>

	三个适配阶段。压缩模块 warmup → C-KSA 全参数 CPT → S-KSA 聚合模块与 Indexer warmup。
	最后一阶段冻结继承参数，完成后直接导出用于稀疏推理。

	## 🤖 模型配置

	当前目录为 KSA-v2-4B。

	\| 配置项 \| KSA-v2-1.9B \| KSA-v2-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 比例（C-KSA : S-KSA） \| 3 : 1 \| 3 : 1 \|
	\| 最大上下文（原文 token） \| 131072 \| 131072 \|
	\| Summary chunk size \| 8 \| 8 \|
	\| C-KSA sliding chunk number \| 128 \| 128 \|
	\| C-KSA group size \| 128 \| 32 \|
	\| S-KSA block size \| 4 \| 4 \|
	\| S-KSA 远程 token 预算 \| 4096 \| 4096 \|
	\| Sink tokens \| 8 \| 8 \|
	\| S-KSA 局部窗口（原文 token） \| 1024 \| 1024 \|
	\| Indexer heads × head dimension \| 8 × 128 \| 8 × 128 \|
	\| Indexer query rank \| 512 \| 512 \|
	\| Indexer RoPE dimension \| 64 \| 64 \|
	\| Indexer Q/K RMSNorm \| False \| True \|
	\| Teacher 聚合 \| max \| max \|
	\| Tied embeddings \| False \| True \|

	远程预算包含 8 个 sink token，局部窗口另计。

	权重以 FP32 保存，推理时以 BF16 加载。

	## 📈 实验结果

	<p align="center"><img src="assets/figures/ksa_v2_results.png" width="85%" alt="技术报告第一页的 1.9B 主结果图：长上下文检索与通用能力，以及单 token 推理 FLOPs 和 KV cache 大小的理论值。" /></p>
	<p align="center"><em>1.9B 主结果：长上下文与通用能力（左），单 token 推理 FLOPs 与 KV cache 大小的理论值（右）。</em></p>

	我们在两种 Hybrid KSA 主干上验证 KSA V2：从零训练的 1.9B，以及基于 Qwen3-4B
	继续预训练的 4B。两者均经过三阶段适配，原始 baseline 结果沿用 KSA V1。

	### 1.9B：长上下文检索

	C-KSA 使用 group128，S-KSA 使用 block4；Hybrid KSA 是适配前的父模型。

	\| 模型 \| 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 \|
	\| C-KSA + Full \| 81.52 \| 72.24 \| 74.44 \| 71.93 \| 69.61 \| 66.31 \| 72.68 \|
	\| C-KSA + S-KSA \| 81.48 \| 72.42 \| 75.67 \| 70.78 \| 68.99 \| 65.01 \| 72.39 \|

	### 4B：扩展到 GQA 主干

	C-KSA 使用 group32，S-KSA 使用 block4，保留主干的 GQA 结构。

	\| 模型 \| 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 \|
	\| Hybrid KSA + CPT \| 92.22 \| 89.85 \| 87.20 \| 83.09 \| 74.74 \| 72.03 \| 83.19 \|
	\| C-KSA + Full \| 92.77 \| 88.12 \| 86.27 \| 82.68 \| 78.58 \| 73.92 \| 83.72 \|
	\| C-KSA + S-KSA \| 93.38 \| 87.78 \| 85.45 \| 83.60 \| 80.28 \| 74.41 \| 84.15 \|

	Hybrid KSA + CPT 是不改模型结构的续训对照。

	更多消融实验和推理效率分析见技术报告。

	## 快速开始

	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)；入口仅加载 CuTe，不依赖单独的 FlashAttention 2 扩展。

	```bibtex
	@techreport{kwai2026ksav2,
	title = {Kwai Summary Attention V2 Technical Report},
	author = {OneRec Team},
	year = {2026},
	institution = {Kuaishou Technology},
	url = {https://arxiv.org/abs/2610.18654}
	}
	```