Hugging Face выложила 207 ядер WebGPU, чтобы нейросети в браузере работали быстрее

Hugging Face выпустила 207 ядер WebGPU и JavaScript-библиотеку @huggingface/kernels. Так называются мелкие программы для видеокарты, из которых складывается работа нейросети в браузере. Библиотека скачивает ядро с Hub, готовит его и запускает на видеокарте пользователя. Все ядра открыты под лицензией Apache-2.0. Команда WebAI объясняет выбор отправной точки просто: любая модель в браузере в итоге сводится к цепочке операций на видеокарте. Среди них умножение матриц, нормализация, свёртки. Среда выполнения не может работать быстрее операций, которые она вызывает.

У каждого ядра свой репозиторий. В нём лежат описание входов и выходов, шаблоны шейдеров на языке WGSL, тесты на правильность и замеры скорости. Версии ядер фиксируются, поэтому приложение подключает конкретную версию, а не безымянный файл по ссылке. Под одну операцию часто пишут несколько вариантов. Даже у простого сложения есть отдельные пути для массивов одинаковой формы и для случаев, когда меньший массив нужно растянуть до формы большего. Библиотека сама выбирает подходящий вариант под данные и устройство.

Hugging Face сравнила свои ядра с ONNX Runtime Web на видеокарте Apple M4. Из 1 756 тестовых случаев в сравнение попали 809: в них обе реализации выдали одинаковый ответ и дали надёжные замеры. В среднем геометрическом ядра Hugging Face оказались быстрее в 2,57 раза, по медиане — в 1,9 раза. Счёт побед — 629 против 176, ещё 4 случая закончились вничью. Встречались и крайние случаи: одна операция Einsum выполнилась более чем в 10 000 раз быстрее. Авторы сами называют такие результаты нетипичными. Замерялась только работа видеокарты по отдельным операциям, а не целые модели. Компания вместе с командой ONNX Runtime переносит эти улучшения в сам движок.

Вместе с ядрами вышел Fleet — сервис, который прогоняет тесты и замеры прямо в браузере на железе пользователя. Если пользователь согласен, результаты анонимно уходят разработчикам. Так они надеются найти ошибки и медленные случаи на видеокартах, браузерах и драйверах, которых нет в обычной тестовой лаборатории.

Читатели в комментариях обратили внимание на 947 случаев, которые в сравнение не вошли. Снаружи не понятно, в чём там дело: ошибка в ONNX Runtime, в новом ядре или в самом тестовом стенде. Пока нет такой разбивки, цифру 2,57 стоит читать с этой оговоркой.