Триты, таблицы сумм и почти десятикратное ускорение

Мы ускорили инференс модели на упакованных троичных весах: в 8,13 раза для одиночного запроса и в 9,86 раза для пачки из 32 запросов. Всё это — на одной RTX 5070 Ti.

При этом веса модели с 1,47 млрд параметров занимают около 795 МБ вместо 5,88 ГБ в FP32.

Но самое интересное — как именно удалось получить этот результат.

Триты сами по себе не гарантируют скорость

Трит принимает три значения: −1, 0, +1. Вклад такого веса можно получить, прибавив активацию, вычтя её или пропустив.

Звучит как очевидное ускорение: убрали умножения — стало быстрее. На практике GPU устроен сложнее.

Первый прямой троичный reader экономил память, но работал медленно. Более того, контрольный тест на одинаковых данных показал: простая замена умножения на выбор +x / −x / 0 сама по себе преимущества не дала.

Поэтому мы пошли дальше.

Не складывать заново, а переиспользовать готовые суммы

Пять тритов помещаются в один байт. Для пяти активаций можно заранее вычислить суммы, соответствующие различным сочетаниям знаков.

Тогда при обработке весов вместо повторного вычисления пяти вкладов ядро получает готовую сумму по коду из упакованного байта.

Именно здесь сложение становится полезным для ускорения: сумма вычисляется один раз, а используется многократно.

Дополнительно мы:

• перестроили хранение весов и таблиц под совместную обработку пачки;
• использовали симметрию положительных и отрицательных сочетаний, сократив таблицы;
• перенесли выбор следующего токена на GPU;
• применили CUDA-графы, сократив накладные расходы запуска операций;
• объединили часть промежуточных суммирований внутри GPU-ядер.

Что получилось по скорости

Сравнение с нашей первой реализацией прямого инференса на упакованных тритах:

Это замеры после прогрева, включающие обработку входного запроса и генерацию. Загрузка модели, упаковка весов и подготовка CUDA-графа в эти цифры не входят.

Для пачки указана общая производительность всех 32 запросов.

Что получилось по памяти

Здесь важно разделять веса модели и всю память выполнения.

Весовой пакет, включая масштабы и оставшиеся FP32-параметры, сократился:

5,88 ГБ → 0,795 ГБ — экономия 86,48%.

Мы используем несколько троичных плоскостей на вес: преимущественно две или три. Поэтому это не утверждение, что вся модель хранится ровно в 1,58 бита на параметр.

Отдельная оптимизация промежуточных сумм уменьшила резерв памяти при обработке пачки:

3654 → 3404 МиБ — ещё 250 МиБ экономии без заметного ухудшения скорости.

Эта оптимизация включена в постоянный механизм.

Не каждая идея оказалась полезной

В старшей троичной плоскости выходного слоя оказалось почти 77% нулей. Мы проверили, можно ли выиграть за счёт их пропуска и более компактного хранения.

Результат оказался неоднозначным:

• списки ненулевых тритов увеличили память и замедлили ядро;
• пропуск нулевых обращений к таблице тоже не помог;
• bitmap-хранение сэкономило 4,66 МБ на одной плоскости, но ухудшило скорость.

Эти варианты мы не включили в основной путь. Экономия байтов полезна только тогда, когда её не съедает стоимость доступа и декодирования.

Что с точностью

Все оптимизированные версии сохранили 600 из 600 полных последовательностей ответов исходного троичного reader на нашем проверочном наборе.

Это сохранение поведения уже квантованной модели. Само квантование ранее снизило арифметический результат с 507/600 у FP32 до 345/600 у троичной версии. Ускорение эту потерю не устраняет, но и не добавляет новых расхождений на проверенном наборе.

Главный вывод

Триты дают компактное представление весов. Предвычисленные суммы позволяют переиспользовать вычисления. Правильная раскладка памяти и организация GPU-работы превращают это в реальное ускорение.

Полученный результат — около 140 токенов/с одиночно, 584 токена/с пачкой и менее 800 МБ на веса модели.

Не потому, что «сложение всегда быстрее умножения», а потому, что мы научились меньше вычислять заново и эффективнее использовать уже вычисленное.