Оглавление

TorchAO внедряет высокопроизводительные низкобитные линейные операторы и операторы эмбеддинга для Arm CPU. В этом обновлении представлены три ключевых улучшения: динамический выбор ядра, интеграция с библиотекой KleidiAI от Arm и поддержка квантованных связанных эмбеддингов — всё для ускорения низкобитного вывода в PyTorch и ExecuTorch, решении PyTorch для эффективного исполнения на устройствах.

С ядрами KleidiAI мы наблюдаем более чем двукратное ускорение предварительного заполнения для 4-битной квантованной Llama1B на M1 Mac (373 токена/сек)!

Динамический выбор ядра

Низкобитные операторы TorchAO теперь автоматически выбирают оптимальное ядро на основе:

  • Формата упакованных весов
  • Возможностей CPU (например, has_arm_neon_dot или has_arm_i8mm)
  • Формы тензора активаций

Как это работает?

Квантованные веса упаковываются в формат, оптимизированный под конкретное ядро. При вызове оператора система анализирует формат весов и возможности CPU, определяя совместимые ядра и кэшируя их в таблице регистрации.

Например, веса в format1 могут использовать GEMV или GEMM ядра, а в format2 — только GEMV. Выбор финального ядра происходит по форме активаций.

Screenshot of kernel registration table

Интеграция с KleidiAI

KleidiAI — библиотека с оптимизированными микроядрами для Arm CPU. Её ядра теперь интегрированы в систему динамического выбора. При отсутствии поддержки (например, для несвязанных эмбеддингов) автоматически используются собственные GEMV-ядра TorchAO.

Гибридный подход даёт:

  • Пиковую производительность от KleidiAI
  • Широкую поддержку операторов через ядра TorchAO

На ExecuTorch это даёт 373 токена/сек на M1 Mac!

Performance comparison chart

Квантованные связанные эмбеддинги

Связанные эмбеддинги (веса embedding и lm_head совмещены) используются в компактных LLM вроде LLaMA 3.2. Проблема: на мобильных устройствах веса часто дублируются из-за разных форматов упаковки.

Diagram of tied embeddings architecture

Решение — SharedEmbeddingQuantizer с унифицированным форматом весов для:

  • 8-битных динамических активаций
  • X-битных весов (X=1-8)

Новшества доступны через API квантизации TorchAO и интегрированы в ExecuTorch для мобильных устройств. Присоединяйтесь к сообществу в Discord и GitHub.

Эти улучшения — не просто апдейт, а стратегический ход в гонке за эффективность LLM на edge-устройствах. Динамическая диспетчеризация ядер решает боль фрагментации Arm-архитектур, а KleidiAI даёт тот самый 2x прирост, который делает локальный запуск Llama на MacBook оправданным. Особенно ценна оптимизация связанных эмбеддингов: в компактных моделях они съедают до 40% памяти. Но главное — теперь это работает «из коробки» в ExecuTorch, снижая порог входа для квантификации. Остаётся вопрос: когда аналогичные оптимизации появятся для Android-чипов среднего класса? Пока это прерогатива флагманов.