KSA V2 4B
用 Summary 同时构建长程记忆与稀疏索引
从 Hybrid KSA 到 KSA V2：复用 Summary，在注意力层内压缩 KV，并稀疏读取原文。
本模型基于 **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 传递。
C-KSA 两级压缩示意，图中 chunk 大小为 4，每 2 个 Summary 聚合为一个 Meta。
**S-KSA：Summary → 块级索引。** 以 Summary 为 query 聚合原文块，再映射到低维索引空间。 主注意力读取选中块的原始 KV，并保留 sink 与局部窗口。Indexer 用按块取 max 的 dense attention teacher 监督训练。
S-KSA 块级选路示意，图中 chunk 大小为 4、block 大小为 2；发布模型的实际配置见下表。
**三个适配阶段。** 压缩模块 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 加载。 ## 📈 实验结果
1.9B 主结果：长上下文与通用能力（左），单 token 推理 FLOPs 与 KV cache 大小的理论值（右）。
我们在两种 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} } ```