Text Generation
Transformers
Safetensors
English
Chinese
qwen3
ksa
ksa-v2
custom_code
long-context
conversational
text-generation-inference
Instructions to use OpenOneRec/KSA-V2-4B-base with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
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
|
Download README.md from OpenOneRec/KSA-V2-4B-base: direct link, hf CLI and curl.
- Browser
- Download file 9.92 kB
-
https://huggingface.co/OpenOneRec/KSA-V2-4B-base/resolve/main/README.md
- Command line
-
hf download hf://OpenOneRec/KSA-V2-4B-base/README.md
-
curl -L -o README.md https://huggingface.co/OpenOneRec/KSA-V2-4B-base/resolve/main/README.md
9.92 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 tokens as long-range memory and sparse indexes</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> | |
| ## π Introduction | |
| **Kwai Summary Attention (KSA)** is our attention architecture for efficient | |
| long-context modeling. Learnable Summary tokens at fixed chunk boundaries | |
| preserve a compressed representation of distant history. | |
| **KSA V2** gives these Summaries two complementary roles: | |
| - **Compressed KSA (C-KSA)** aggregates Summary KV into compact Meta KV for | |
| more efficient long-range memory. | |
| - **Sparse KSA (S-KSA)** uses Summaries to build block-level routing keys. | |
| A lightweight Indexer selects relevant blocks, whose original KV is read | |
| directly by the main attention. | |
| We retain Hybrid KSA's 3:1 layer layout, replacing KSA layers with C-KSA and | |
| Full layers with S-KSA. | |
| <p align="center"><img src="assets/figures/ksa_v2_overview.png" width="100%" alt="From Hybrid KSA to KSA V2: the same Summary layout, with C-KSA compression and S-KSA block routing." /></p> | |
| <p align="center"><em>From Hybrid KSA to KSA V2: shared Summaries, layer-local KV compression and sparse original-token access.</em></p> | |
| This base model builds on our **4B Hybrid KSA backbone continually pretrained from Qwen3-4B**. It is not an instruction-tuned model. | |
| ## π₯ News | |
| - **2026-10-13** β [KSA V2 technical report](https://arxiv.org/abs/2610.18654) released. | |
| - **2026-05-08** β [KSA-4B-base](https://huggingface.co/OpenOneRec/KSA-4B-base) weights released. | |
| - **2026-04-28** β [KSA technical report](https://arxiv.org/abs/2604.24432) and training code released. | |
| ## β¨ Method | |
| **C-KSA: Summary β Meta KV.** Learnable dimension-wise pooling compresses groups | |
| of Summary KV within each attention layer. Summary tokens continue through | |
| the backbone at their original chunk granularity; stronger KV compression does | |
| not require a new large-chunk backbone. Each query attends to remote Meta, | |
| boundary Summary and local text KV. Meta entries remain local to the layer | |
| that constructs them; they are not propagated as new tokens. | |
| <p align="center"><img src="assets/figures/ksa_v2_compression.png" width="80%" alt="C-KSA compresses complete remote Summary groups into Meta KV while retaining boundary Summaries and local text KV." /></p> | |
| <p align="center"><em>C-KSA hierarchical compression, illustrated with chunk size 4 and two Summaries per Meta entry.</em></p> | |
| **S-KSA: Summary β block index.** A Summary query aggregates each text block | |
| into a representation for low-dimensional routing. Main attention reads the | |
| selected blocks' original KV together with sinks and local context. A | |
| blockwise-max dense-attention teacher supervises the Indexer. | |
| <p align="center"><img src="assets/figures/ksa_v2_sparse.png" width="80%" alt="S-KSA uses Summary-guided block aggregation and an Indexer to select remote text blocks, then attends to their original KV together with sinks and local context." /></p> | |
| <p align="center"><em>S-KSA block routing, illustrated with chunk size 4 and block size 2. Released model configurations are listed below.</em></p> | |
| **Three adaptation stages.** Compression-module warmup β full-parameter C-KSA | |
| CPT β S-KSA aggregation and Indexer warmup. The final stage freezes inherited | |
| parameters; its checkpoint is exported directly for sparse inference. | |
| ## π€ Model configuration | |
| This directory contains **KSA-v2-4B**. | |
| | Configuration | KSA-v2-1.9B | KSA-v2-4B | | |
| | :--- | :--- | :--- | | |
| | Number of layers | 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 ratio (C-KSA : S-KSA) | 3 : 1 | 3 : 1 | | |
| | Maximum context (text tokens) | 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 remote token budget | 4096 | 4096 | | |
| | Sink tokens | 8 | 8 | | |
| | S-KSA local window (text tokens) | 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 pooling | max | max | | |
| | Tied embeddings | False | True | | |
| The remote budget includes eight sink tokens; the local window is additional. | |
| Weights are stored in FP32 and loaded in BF16. | |
| ## π Results | |
| <p align="center"><img src="assets/figures/ksa_v2_results.png" width="85%" alt="The technical report's first-page 1.9B results: long-context retrieval and general capabilities, and theoretical per-token inference FLOPs and KV cache size." /></p> | |
| <p align="center"><em>1.9B results: long-context and general capabilities (left), theoretical per-token inference FLOPs and KV cache size (right).</em></p> | |
| We evaluate KSA V2 on two Hybrid KSA backbones: a 1.9B model trained from | |
| scratch and a 4B model continually pretrained from Qwen3-4B. Both undergo | |
| three-stage adaptation; original baseline results are reused from KSA V1. | |
| ### 1.9B: Long-context retrieval | |
| C-KSA uses group128 and S-KSA uses block4. Hybrid KSA is the parent model before adaptation. | |
| | Model | 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: Scaling to a GQA backbone | |
| C-KSA uses group32 and S-KSA uses block4, preserving the backbone's GQA structure. | |
| | Model | 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 is a continuation control with no architectural changes. | |
| See the technical report for further ablations and inference efficiency. | |
| ## Quick start | |
| Linux, Python 3.12 and H100/H800. Download the **entire model directory**, | |
| including `kernel/`, and run from it: | |
| ```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 | |
| ``` | |
| Omit `--prompt` for interactive completion. The optional `--apply-chat-template` | |
| only applies the tokenizer template; it does not add instruction tuning. | |
| Transformers 5.0.0 and the bundled `ksa_attn` / FlashAttention CuTe sources are | |
| used. No training checkout, kernel build or FlashAttention 2 CUDA extension | |
| is required. | |
| ### Transformers usage | |
| Set the bundled kernel paths before importing 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)) | |
| ``` | |
| ## License and citation | |
| Weights and our runtime: [Apache-2.0](LICENSE). Qwen-derived files retain their | |
| copyright notices. Bundled FlashAttention CuTe retains its | |
| [BSD-3-Clause license](kernel/deps/flash_attn/LICENSE) and | |
| [AUTHORS](kernel/deps/flash_attn/AUTHORS); its initializer exposes CuTe without | |
| the separate FlashAttention 2 extension. | |
| ```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} | |
| } | |
| ``` | |