OpenOneRec/KSA-V2-4B-base

KSA V2 4B

Summary tokens as long-range memory and sparse indexes

English | 中文

📖 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.

From Hybrid KSA to KSA V2: shared Summaries, layer-local KV compression and sparse original-token access.

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 released.
• 2026-05-08 — KSA-4B-base weights released.
• 2026-04-28 — KSA technical report 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.

C-KSA hierarchical compression, illustrated with chunk size 4 and two Summaries per Meta entry.

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.

S-KSA block routing, illustrated with chunk size 4 and block size 2. Released model configurations are listed below.

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

1.9B results: long-context and general capabilities (left), theoretical per-token inference FLOPs and KV cache size (right).

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:

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:

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

License and citation

Weights and our runtime: Apache-2.0. Qwen-derived files retain their copyright notices. Bundled FlashAttention CuTe retains its BSD-3-Clause license and AUTHORS; its initializer exposes CuTe without the separate FlashAttention 2 extension.

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