Представяме @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 opset, since_version на оператор или ревизия на модел. Разделянето на тези понятия позволява на приложенията да разчитат на стабилен договор, видим за JavaScript, докато реализациите на ядрата се развиват зад него.

Колко бързи са ядрата?

И така, доколко оптимизираните ядра действително променят нещата? Сравнихме колекцията си директно с ORT WebGPU на GPU Apple M4, използвайки ONNX Runtime Web 1.30.0-dev.20260826-b1f76d586a. Започнахме с 1756 тестови случая за всичките 207 операции и оставихме 809 случая, при които и двете страни произведоха съвпадащи резултати и надеждни измервания.

В тези сравнения нашите ядра бяха 2,57 пъти по-бързи по геометрична средна и 1,90 пъти по-бързи по медиана, с 629 победи, 176 загуби и 4 равенства. Ето по-подробен поглед към четири познати операции:

Някои отделни победи бяха много по-големи. Особено труден двулинеен случай на Einsum (i,ij,j с размер 4096) се изпълни за 0,136 ms с нашето ядро, спрямо 1396 ms с ORT WebGPU: над 10 000 пъти по-бързо. CumSum по редове върху [256, 4096] беше 301 пъти по-бърз — 0,016 ms спрямо 4,784 ms. Това са необичайни случаи, а не ускорения, които трябва да очаквате навсякъде, но показват колко много може да помогне специализираното ядро, когато общата реализация попадне на бавен път.

Измервахме работата, извършена непосредствено на 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, за да допринесете с данни от вашето устройство и да ни помогнете да направим ядрата по-добри за всички.

Още статии от нашия блог