Проблема «разбавления» в длинных контекстах
Инференс больших языковых моделей (LLM) с длинным контекстом (long-context) упирается в этап префилла (prefill), где плотное самовнимание (dense self-attention) обрабатывает весь промпт перед генерацией. Традиционные методы разреженного выбора блоков часто страдают от эффекта mean dilution (разбавления средним): если в блоке есть один критически важный токен, но много шумовых, усредненная оценка релевантности может скрыть этот токен, что приведет к потере информации.
Решение: двухветвевой выбор RBS-Attention
Авторы (Chuxu Song, Jiuqi Wei, Zhencan Peng) предлагают метод RBS-Attention (Radius-Bounded Sparse Prefill), который не требует дообучения модели. Алгоритм использует две комплементарные ветви для выбора блоков:
- Centroid base branch: захватывает среднюю релевантность (стандартный подход).
- Rescue branch: использует максимальный радиус ключ-блока и распределение, зависящее от промпта, слоя и головы, чтобы выявить блоки, которые рискуют быть пропущенными из-за «разбавления».
Маски от обеих ветвей порогово фильтруются независимо, а затем комбинируются. Это позволяет сохранить регулярное выполнение FlashAttention для основных блоков, добавляя «спасательные» блоки только там, где это необходимо.
Результаты на Qwen3-30B и Qwen3-32B
Эксперименты проводились на GPU NVIDIA H100. Метод демонстрирует колоссальный прирост скорости на модели Qwen3-30B-A3B-Instruct-2507-FP8 с контекстом 128K токенов. При этом качество генерации на модели Qwen3-32B остается практически на уровне плотного внимания.
| Метрика | Результат RBS-Attention | Примечание |
|---|---|---|
| Ускорение standalone prefill-attention | 20.65× | На H100, контекст 128K |
| Ускорение prefill-attention в vLLM | 11.92× | Интеграция с популярным движком |
| Ускорение Time-to-First-Token (TTFT) | 5.97× | Ключевой метрик UX для чат-ботов |
| Точность RULER (Qwen3-32B) | 88.65 | Сравнение с плотным вниманием: 89.52 |
Почему это важно
Снижение задержки первого токена (TTFT) в 6 раз делает LLM с длинным контекстом значительно более отзывчивыми в реальных приложениях. При этом падение точности на бенчмарке RULER составляет менее 1% (88.65 против 89.52), что подтверждает эффективность подхода. Дополнительные оценки на LongBench-v2, InfiniteBench и Video-MME также показывают стабильные результаты, что делает RBS-Attention перспективным решением для развертывания тяжелых моделей в продакшене без переплаты за GPU-мощности.
Источник: arXiv cs.AI ↗