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

metadata

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

KSA V2 4B

用 Summary 同时构建长程记忆与稀疏索引

English | 中文

📖 简介

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，分别优化长程记忆和原文读取。

从 Hybrid KSA 到 KSA V2：保持 Summary 布局，引入 C-KSA 压缩与 S-KSA 原文块选路。

从 Hybrid KSA 到 KSA V2：复用 Summary，在注意力层内压缩 KV，并稀疏读取原文。

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

🔥 新闻

2026-10-13 —— KSA V2 技术报告发布。
2026-05-08 —— KSA-4B-base 权重发布。
2026-04-28 —— KSA 技术报告与训练代码发布。

✨ 方法

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

C-KSA：将完整的远程 Summary 分组压缩为 Meta KV，保留边界 Summary 与局部原文 KV。

C-KSA 两级压缩示意，图中 chunk 大小为 4，每 2 个 Summary 聚合为一个 Meta。

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

S-KSA：Summary 引导块内聚合，Indexer 选择远程原文块，主注意力读取其原始 KV 以及 sink 和局部窗口。

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 大小的理论值。

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/），并在该目录下执行：

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

许可证与引用

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