TileMix: ускорение LLM без потери качества через смешанную точность

Проблема длинного контекста и памяти

Префиллинг (prefill) в больших языковых моделях (LLM) с длинным контекстом требует огромных вычислительных ресурсов. Стандартное плотное самовнимание (dense self-attention) вычисляет квадратичное количество очков взаимодействия запрос-ключ (query-key scores), что создает узкое место по пропускной способности памяти. Существующие методы либо используют единый путь низкой точности (что снижает качество), либо выбирают отдельные взаимодействия токенов, нарушая плотность связей.

Решение: Tile-Centric Routing

Новый метод TileMix предлагает пространственную маршрутизацию точности на уровне аппаратно-выровненных «плиток» (tiles) матрицы внимания. Вместо глобального снижения точности, система:

• Разбивает матрицу внимания на плитки, выровненные под архитектуру GPU.
• Упаковывает решения о маршрутизации в компактные битовые маски.
• Направляет каждую группу плиток либо через вычисление в FP16, либо через INT8.
• Оба пути обновляют общее состояние онлайн-softmax, сохраняя математическую корректность.

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

TileMix не требует дообучения моделей (training-free) и сохраняет плотную связность токенов. Метод поддерживает:

• Grouped-Query Attention (GQA).
• Пакеты переменной длины (variable-length batches).
• Кэширование ключей/значений в INT8.

Результаты бенчмарков

Тестирование проводилось на моделях LLaMA, Qwen и Vicuna с использованием GPU NVIDIA A100. Метод оценивался на наборах данных LongEval и LV-Eval. TileMix восстанавливает качество длинного контекста, которое обычно теряется при использовании унифицированного INT8, и одновременно увеличивает пропускную способность префилла по сравнению с чистым FP16.

Реализация метода доступна в открытом доступе, что позволяет разработчикам внедрить более эффективный префиллинг в свои LLM-инференс решения без компромиссов в точности.

Источник: arXiv cs.AI ↗