Представляем @huggingface/kernels: более 200 ядер WebGPU для локального ИИ

Обновить на GitHub

Одна из главных целей команды WebAI в Hugging Face — сделать выполнение моделей в браузере максимально быстрым и удобным для пользователей. Для этого требуется работа на нескольких уровнях: моделям нужны представления, подходящие для браузера, средам выполнения необходимо строить эффективные планы выполнения, а отдельные GPU-операции на нижнем уровне стека должны максимально использовать возможности множества различных устройств и реализаций браузеров.

Сегодня мы выпускаем первый уровень этой работы: @huggingface/kernels — минималистичную библиотеку для загрузки и запуска оптимизированных ядер WebGPU из Hugging Face Hub, а также первоначальную коллекцию из 207 ядер по адресу huggingface.co/webgpu-kernels.

Коллекция охватывает операции, используемые в самых разных архитектурах и рабочих нагрузках машинного обучения. Что ещё важнее, каждое ядро публикуется как полный пакет с версией: его интерфейс, шаблоны шейдеров, тесты корректности, тестовые случаи для бенчмарков и инструкции по использованию находятся вместе на Hub.

Мы также запускаем Fleet — браузерный набор инструментов для тестирования и бенчмаркинга GPU, который запускает ядра на вашем оборудовании и оценивает их. Помимо результатов для вашего устройства, Fleet даёт сообществу возможность предоставлять данные о производительности и корректности для устройств, которые мы никогда не смогли бы охватить в обычной тестовой лаборатории. С вашего согласия каждый запуск добавляет конфиденциальные данные, которые помогают находить сбои (некорректные результаты, патологически медленные случаи и т. д.), улучшать варианты ядер и принимать более эффективные решения по оптимизации на основе реального оборудования.

Кратко

• 207 ядер WebGPU, опубликованных в виде отдельных репозиториев в организации webgpu-kernels. Лицензия Apache-2.0.
• Загрузчик JavaScript @huggingface/kernels, который скачивает, подготавливает и запускает ядра непосредственно из Hub.
• Явные контракты и воспроизводимые данные для каждого ядра, включая манифесты, тесты корректности, тестовые случаи для бенчмарков и шаблоны шейдеров WGSL.
• Fleet — браузерный инструмент бенчмаркинга, собирающий данные о корректности и производительности на реальных GPU, чтобы помогать нам улучшать ядра и их варианты.

Почему мы начинаем с ядер?

Модель, работающая в браузере, в конечном итоге превращается в последовательность операций GPU: матричные умножения, нормализации, свёртки, примитивы внимания, операции квантования, преобразования расположения данных и многое другое. WebGPU делает эти операции доступными в современных браузерах через переносимый API, а WGSL предоставляет общий язык для шейдеров, которые их выполняют.

Однако переносимость не означает автоматического обеспечения производительности. Два шейдера могут реализовывать одну и ту же операцию и выдавать одинаковый результат, но вести себя совершенно по-разному на разных ускорителях. На производительность могут влиять размеры рабочих групп, схемы доступа к памяти, векторизация, типы данных и стратегии объединения операций. Оптимальный выбор также может меняться в зависимости от формы входных данных, устройства, браузера и доступных возможностей WebGPU.

Именно поэтому ядра образуют фундаментальный слой быстрого выполнения моделей в браузере. Среды выполнения более высокого уровня могут быть эффективны лишь настолько, насколько эффективны операции, которые они вызывают. Сделав эти операции по отдельности доступными для обнаружения, тестирования, бенчмаркинга и версионирования, мы можем независимо улучшать основу, сохраняя стабильный контракт для вышестоящих уровней.

Репозиторий ядра, а не просто шейдер

У каждого ядра коллекции есть собственный репозиторий и карточка ядра. В карточке описаны семантика операции, входные и выходные данные, атрибуты, поддерживаемые типы данных, исходные файлы и готовый к запуску пример с @huggingface/kernels.

Например, ai.onnx.Add реализует поэлементное сложение с разнонаправленным broadcasting. Это одна из простейших операций в нейронной сети, используемая повсюду — от остаточных соединений до сложения смещения. В карточке описаны два входа, форма результата после broadcasting, поддерживаемые типы данных и варианты для разных форм и устройств.

За карточкой находится репозиторий с артефактами, необходимыми для понимания и оценки реализации:

• manifest.json является источником истины для контракта операции. Он определяет входы, выходы, атрибуты, ограничения типов и правила вывода форм.
• metadata.json содержит идентификатор ядра, дайджесты и сведения о происхождении.
• test.json содержит тестовые случаи корректности, позволяющие проверить реализацию на соответствие ожидаемому поведению.
• bench.json содержит тестовые случаи для бенчмаркинга и настройки, отражающие рабочие нагрузки, используемые для оценки ядра.
• *.wgsl.jinja содержит параметризованные реализации WGSL, используемые для создания шейдеров под конкретный запрос и устройство.

Такая структура превращает шейдер в повторно используемый программный артефакт. Интерфейс можно изучить без чтения WGSL, тестовые случаи корректности и производительности хранятся вместе с реализацией, а опубликованные версии можно загружать явно, не полагаясь на URL файла без указания версии. Наши ядра также могут служить эталонными реализациями для разработчиков, создающих собственные ядра WebGPU или интегрирующих эти операции в свои среды выполнения.

Загрузка ядра из Hub

Установите пакет из npm:

npm install @huggingface/kernels@preview

Для запуска этих ядер требуется браузер с поддержкой WebGPU. Доступность WebGPU зависит от браузера, операционной системы, GPU и драйвера. Проверить её в JavaScript можно с помощью "gpu" in navigator.

@huggingface/kernels обеспечивает связь между репозиторием ядра и вашим приложением. Вызовите getKernel, указав идентификатор репозитория Hub и версию контракта, а затем вызовите возвращённую функцию с типизированными входными данными и формами тензоров. Вот небольшой пример сложения со смещением:

import { getKernel } from "@huggingface/kernels";

const add = await getKernel("webgpu-kernels/ai.onnx.Add", { version: 1 });

const { c } = await add({
  a: {
    data: new Float32Array([1, 2, 3, 4, 5, 6]),
    shape: [2, 3],
  },
  b: {
    data: new Float32Array([10, 20, 30]),
    shape: [3],
  },
});

Второй вход распространяется по первому измерению, формируя результат с формой [2, 3]. Загрузчик выводит форму результата и логический тип данных из контракта манифеста и входных данных, а затем автоматически выделяет память под c.

Сложение шести чисел с плавающей точкой намеренно является самым маленьким демонстрационным примером. При таком размере затраты на обмен с GPU значительно превышают стоимость вычислений. Суть заключается в схеме вызова: она остаётся абсолютно такой же для тяжёлых операций, в которых оптимизированные ядра действительно дают преимущество, например для матричного умножения (ai.onnx.MatMul). Меняются только идентификатор репозитория и входные данные.

Даже эта элементарная операция показывает, зачем ядрам нужны варианты. При сложении тензоров одинаковой формы можно использовать прямой векторизованный путь, тогда как для broadcasting входных данных требуется другая логика индексации. Опубликованное ядро Add включает варианты для одинаковых форм, векторизованного broadcasting, обработки скаляров и общего broadcasting. Среда выполнения может выбрать реализацию, подходящую для текущего вызова и устройства, не изменяя API, доступный приложению.

Параметр version: 1 выбирает версию 1 опубликованного контракта ядра. Он отделён от opset ONNX, значения since_version оператора или ревизии модели. Разделение этих понятий позволяет приложениям зависеть от стабильного контракта, ориентированного на JavaScript, пока реализации ядер развиваются за его пределами.

Насколько быстры эти ядра?

Итак, насколько оптимизированные ядра действительно повышают производительность? Мы сравнили нашу коллекцию с ORT WebGPU на GPU Apple M4, используя ONNX Runtime Web 1.30.0-dev.20260826-b1f76d586a. Мы начали с 1 756 тестовых случаев для всех 207 операций и оставили 809 случаев, в которых обе стороны выдавали совпадающие результаты и надёжные замеры времени.

В этих сравнениях наши ядра были в 2,57 раза быстрее по геометрическому среднему и в 1,90 раза быстрее по медиане: 629 побед, 176 поражений и 4 ничьи. Вот более подробный обзор четырёх знакомых операций:

Некоторые отдельные результаты были значительно впечатляющими. Особенно сложный случай билинейного Einsum (i,ij,j с размером 4096) выполнялся за 0.136 мс с нашим ядром против 1 396 мс с ORT WebGPU — более чем в 10 000 раз быстрее. Поэлементный CumSum по [256, 4096] был в 301 раз быстрее: 0.016 мс против 4.784 мс. Это необычные случаи, а не ускорение, которого следует ожидать повсеместно, но они показывают, насколько специализированное ядро может помочь, когда универсальная реализация попадает на медленный путь.

Мы измеряли работу, выполняемую непосредственно на GPU, исключив подготовительные операции: загрузку ядер, создание сессий, отправку входных данных, компиляцию шейдеров и считывание результатов. Очень короткие рабочие нагрузки естественным образом сложнее измерять, а небольшие случаи могут получать преимущество от кэша GPU, поэтому эти числа следует воспринимать как полезное сравнение, а не как обещание для каждого приложения.

Кроме того, это результаты для отдельных операций, а не полноценных моделей. Точная производительность будет меняться в зависимости от GPU и браузера, поэтому Fleet так важен для формирования более полной картины.

Мы также работаем с командой ONNX Runtime над добавлением этих улучшений в основной проект, чтобы они принесли пользу всей экосистеме ONNX Runtime Web.

От одного устройства к целому флоту

Производительность WebGPU различается в зависимости от GPU, браузеров и драйверов, поэтому результаты одной машины показывают лишь часть картины. Fleet позволяет любому пользователю запускать проверки корректности и производительности в браузере и видеть, как ядра ведут себя на его оборудовании.

С согласия пользователя каждый запуск конфиденциально добавляет данные, которые помогают выявлять сбои, зависящие от устройства, сравнивать варианты и улучшать правила выбора. Цель проста: использовать широкий охват реальных устройств, чтобы сделать ядра более быстрыми и надёжными для всех.

Создание общей основы для WebAI

Первые 207 ядер — это начало, а не конечное состояние. Независимая публикация ядер на Hub даёт нам общее место для изучения контрактов, сравнения реализаций, воспроизведения проверок корректности и повышения производительности без встраивания каждого шейдера непосредственно в каждую среду выполнения.

Коллекция также является частью более широкой экосистемы ядер Hub: на странице ядер ядра WebGPU находятся рядом с ядрами для CUDA, ROCm, Metal и других платформ; их можно фильтровать, сортировать и изучать так же, как любой другой артефакт на Hub.

Эти компоненты усиливают друг друга:

1. Репозитории ядер определяют прозрачные контракты операций с версиями.
2. @huggingface/kernels упрощает загрузку и запуск этих операций из JavaScript.
3. Fleet собирает данные из реального мира на гораздо более широком спектре устройств, чем может охватить обычная лаборатория для бенчмаркинга.
4. Каждый предоставленный запуск может выявить сбои, помочь настроить параметры, улучшить выбор вариантов и подтвердить будущие версии ядер.

Это низкоуровневая основа для следующих шагов в нашем стеке выполнения моделей в браузере. Мы рады связать эти ядра с инструментами работы с моделями более высокого уровня, продолжать расширять охват операций и сделать быстрое локальное выполнение моделей более удобным во всей экосистеме WebAI.

Изучите коллекцию ядер WebGPU, попробуйте @huggingface/kernels и присоединяйтесь к Fleet, чтобы предоставить данные со своего устройства и помочь нам сделать ядра лучше для всех.

Другие статьи в нашем блоге