llm-jp-4.1-8b-thinking-lsa

Экспериментальная версия llm-jp/llm-jp-4.1-8b-thinking на 8,59 млрд параметров: во всех 32 слоях обычное attention заменено на Latent Sparse Attention (LSA). Веса базовой модели не менялись — добавлены только послойные индексаторы W_K_idx и W_Q_idx размером 4096×64. Языки: японский и английский.

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

• ✅ LSA выбирает, какие key читать: 4 sink-key, окно 512 key и до 2048 длинных key через Ward-кластеры.
• 💥 На 64k контексте одна query теоретически читает около 3,9% key, но KV cache не уменьшается — экономия только на вычислениях.
• 🦾 В комплекте три индексора: PCA (по умолчанию), recovered v2 (лучший long-range) и distilled v1 (справочный).
• ✨ PPL у PCA ниже dense на 4k/16k/32k: 13,95/14,51/12,93 против 14,17/15,41/14,04.
• ⬇ Короткие последовательности <512 token уходят в dense attention.

💥 Пример

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    'ChiTako/llm-jp-4.1-8b-thinking-lsa',
    trust_remote_code=True, dtype='bfloat16', device_map='auto'
)
tokenizer = AutoTokenizer.from_pretrained('ChiTako/llm-jp-4.1-8b-thinking-lsa')

Чтобы включить recovered-индексор, замените config.lsa['maps_file'] на 'lsa_maps_recovered.pt' до загрузки модели.

Подробнее

Бенчмарки: PPL у PCA ниже dense на 4k/16k/32k — 13,95/14,51/12,93 против 14,17/15,41/14,04; needle 16k — 5/9 у dense и recovered v2, 3/9 у PCA.
На H200 64k prefill peak: 80,5 GiB у dense → 49,2 GiB у LSA, но KV cache не сокращается.
Скорость: torch-референс, prefill медленнее dense (64k — 8,2×), decode ~0,103 с против ~0,055 с. Ускорение из статьи требует SGLang fused top-k и CUDA graph.
Ограничения: batch 1 без padding, контекст до 65 536 token. Лицензия базы — Apache-2.0.