KSA-V2-4B-base

KSA V2 — attention-архитектура Kuaishou/OpenOneRec для эффективного long-context моделирования. Это base-модель на 4,26 млрд параметров, продолжательно предобученная с Qwen3-4B; не instruction-tuned. 36 слоёв, hidden size 2560, GQA 32/8, контекст до 131 072 токенов.

🔍 Ключевые особенности

• 🔍 C-KSA: Summary KV сжимаются в компактные Meta KV внутри слоя для дальней памяти.
• 🔗 S-KSA: Summary строят блочные routing-ключи, Indexer выбирает релевантные блоки, attention читает их исходный KV.
• ⚙ Гибридная схема 3:1 C-KSA : S-KSA, всего 36 слоёв.
• 📌 Контекст 131 072 токена; remote budget 4096 (включая 8 sink), локальное окно 1024.
• 💻 Веса хранятся в FP32 и загружаются в BF16; нужны Linux, Python 3.12, H100/H800 и bundled kernel.
• 🛡 Веса и runtime — Apache-2.0; Qwen-derived файлы сохраняют копирайты.

💻 Пример

Запуск из каталога модели; перед импортом Transformers добавьте kernel/ в PYTHONPATH.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = 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 = tok("The main advantage of long-context models is", return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=128, do_sample=False, pad_token_id=tok.eos_token_id)
print(tok.decode(out[0, inputs.input_ids.shape[1]:], skip_special_tokens=True))

Подробнее

Бенчмарки 4B: среднее по 4K–128K — 84,15; на 128K — 74,41 (C-KSA + S-KSA).
Для сравнения Full: 83,69 и 65,86 соответственно.
Максимальный контекст — 131 072 текстовых токена; удалённый бюджет 4096 токенов включает 8 sink, локальное окно — 1024.
Железо: Linux, Python 3.12, H100/H800; нужен весь каталог модели вместе с kernel/.
Это base-модель без instruction tuning; языки — английский и китайский.