KSA-1.9B-base / README_zh.md
OpenOneRec's picture
Release KSA-1.9B-base weights and inference runtime
6c91dd5 verified
|
Raw History Blame Contribute Delete
5.58 kB
metadata
license: apache-2.0
language:
  - en
  - zh
library_name: transformers
pipeline_tag: text-generation
tags:
  - ksa
  - ksa-v1
  - custom_code
  - long-context

KSA V1 1.9B

Kwai Summary Attention:面向长上下文建模的高效注意力

技术报告 · 代码 · Apache-2.0

English | 中文

📖 简介

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, 未经过指令微调。

Hybrid KSA:Summary 与原文交错排列,KSA 和 Full Attention 层按 3:1 配置。

🔥 新闻

🤖 模型配置

当前目录为 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/),并在该目录下执行:

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 前设置随包内核路径:

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

许可证与引用

模型权重及我们的推理代码采用 Apache-2.0。基于 Qwen 的文件保留原版权声明。 随包的 FlashAttention CuTe 保留其 BSD-3-Clause 许可证 与 AUTHORS。

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