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.

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.

4B: Scaling to a GQA backbone

C-KSA uses group32 and S-KSA uses block4, preserving the backbone's GQA structure.

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