Представляємо @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 реалізує поелементне додавання з багатоспрямованим транслюванням. Це одна з найпростіших операцій у нейронній мережі, яка використовується всюди — від залишкових з’єднань до додавання зміщення. У картці описано два входи, форму виходу після транслювання, підтримувані типи даних і доступні варіанти для різних форм і пристроїв.

За карткою репозиторій містить артефакти, необхідні для розуміння та оцінювання реалізації:

• 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). Змінюються лише ідентифікатор репозиторію та входи.

Навіть ця елементарна операція демонструє, навіщо ядрам потрібні варіанти. Додавання тензорів однакової форми може використовувати прямий векторизований шлях, тоді як транслювання входів потребує іншої логіки індексації. Опубліковане ядро Add містить варіанти для однакових форм, векторизованого транслювання, скалярної обробки та загального транслювання. Середовище виконання може вибрати реалізацію, що відповідає поточному виклику та пристрою, не змінюючи API, видимий застосунку.

Параметр version: 1 вибирає версію 1 опублікованого контракту ядра. Він не залежить від набору операторів 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: на сторінці Kernels ядра WebGPU розміщені поруч із ядрами для CUDA, ROCm, Metal та інших платформ, і їх можна фільтрувати, сортувати та переглядати так само, як будь-який інший артефакт на Hub.

Ці складові підсилюють одна одну:

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

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

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

Більше статей із нашого блогу