Hugging Face ra mắt thư viện WebGPU kernels: Đưa AI chạy mượt mà ngay trên trình duyệt

Một trong những mục tiêu lớn nhất của nhóm WebAI tại Hugging Face là làm cho quá trình suy luận (inference) trên trình duyệt trở nên nhanh chóng và thân thiện với người dùng nhất có thể. Để đạt được điều đó đòi hỏi một nỗ lực đa tầng: các mô hình cần những biểu diễn phù hợp với trình duyệt, các runtime cần xây dựng kế hoạch thực thi hiệu quả, và các thao tác GPU riêng lẻ ở tầng dưới cùng cần tận dụng tối đa nhiều thiết bị và cách triển khai trình duyệt khác nhau.

Hôm nay, chúng tôi phát hành tầng đầu tiên của nỗ lực đó: @huggingface/kernels, một thư viện tối giản để tải và chạy các WebGPU kernel đã được tối ưu hóa từ Hugging Face Hub, cùng với bộ sưu tập ban đầu gồm 207 kernel tại huggingface.co/webgpu-kernels.

Bộ sưu tập này bao gồm các thao tác được sử dụng trong nhiều kiến trúc và khối lượng công việc học máy (machine learning) khác nhau. Quan trọng hơn, mỗi kernel được xuất bản dưới dạng một gói hoàn chỉnh, có đánh số phiên bản: giao diện, các mẫu shader, các trường hợp kiểm thử tính đúng đắn, các trường hợp benchmark và hướng dẫn sử dụng đều nằm cùng nhau trên Hub.

Chúng tôi cũng ra mắt Fleet, một bộ công cụ benchmark và kiểm thử GPU trên trình duyệt giúp chạy và chấm điểm các kernel trên phần cứng của bạn. Ngoài kết quả cho máy tính cá nhân, Fleet còn cung cấp cho cộng đồng một cách để đóng góp bằng chứng về hiệu năng và tính đúng đắn từ các thiết bị mà chúng tôi không thể bao phủ trong một phòng thí nghiệm kiểm thử thông thường. Với sự đồng ý của bạn, mỗi lần chạy sẽ bổ sung các bằng chứng riêng tư giúp chúng tôi tìm ra lỗi (kết quả không chính xác, các trường hợp chậm bất thường, v.v.), cải thiện các biến thể kernel và đưa ra các quyết định tối ưu hóa tốt hơn trên phần cứng thực tế.

Tóm tắt (TL;DR)

Tại sao lại bắt đầu với các kernel?

Một mô hình chạy trên trình duyệt cuối cùng sẽ trở thành một chuỗi các thao tác GPU: nhân ma trận, chuẩn hóa, tích chập, các nguyên hàm attention, thao tác lượng tử hóa, biến đổi bố cục dữ liệu và nhiều thao tác khác. WebGPU cung cấp các thao tác này trên các trình duyệt hiện đại thông qua một API di động, trong khi WGSL cung cấp ngôn ngữ chung cho các shader thực thi chúng.

Tuy nhiên, tính di động không tự động đồng nghĩa với hiệu năng. Hai shader có thể thực hiện cùng một thao tác và tạo ra cùng một đầu ra nhưng lại hoạt động hoàn toàn khác nhau trên các bộ tăng tốc khác nhau. Kích thước nhóm công việc (workgroup sizes), mô hình truy cập bộ nhớ, vector hóa, kiểu dữ liệu và chiến lược hợp nhất đều có thể ảnh hưởng đến hiệu năng. Lựa chọn tốt nhất cũng có thể thay đổi tùy theo hình dạng đầu vào (input shape), thiết bị, trình duyệt và các tính năng WebGPU khả dụng.

Đây là lý do tại sao các kernel tạo thành tầng nền tảng cho việc suy luận nhanh trên trình duyệt. Các runtime cấp cao hơn chỉ có thể hiệu quả bằng các thao tác mà chúng điều phối. Bằng cách làm cho các thao tác đó có thể khám phá, kiểm thử, benchmark và đánh số phiên bản một cách riêng biệt, chúng tôi có thể cải thiện nền tảng một cách độc lập trong khi vẫn duy trì một hợp đồng ổn định cho các tầng bên trên.

Một kho lưu trữ kernel, không chỉ là một shader

Mỗi kernel trong bộ sưu tập đều có kho lưu trữ và thẻ kernel (kernel card) riêng. Thẻ này ghi lại ngữ nghĩa của thao tác, đầu vào, đầu ra, thuộc tính, các kiểu dữ liệu được hỗ trợ, tệp nguồn và một ví dụ @huggingface/kernels sẵn sàng để chạy.

Ví dụ, ai.onnx.Add thực hiện phép cộng từng phần tử với cơ chế phát sóng (broadcasting) đa hướng. Đây là một trong những thao tác đơn giản nhất trong mạng thần kinh, được sử dụng ở khắp mọi nơi từ các kết nối dư (residual connections) đến việc cộng bias. Thẻ của nó ghi lại hai đầu vào, hình dạng đầu ra sau khi phát sóng, các kiểu dữ liệu được hỗ trợ và các biến thể khả dụng cho các hình dạng và thiết bị khác nhau.

Phía sau thẻ, kho lưu trữ chứa các thành phần cần thiết để hiểu và đánh giá quá trình triển khai:

Cấu trúc này biến một shader thành một thành phần phần mềm có thể tái sử dụng. Giao diện có thể kiểm tra được mà không cần đọc WGSL, các trường hợp kiểm thử tính đúng đắn và hiệu năng đi kèm với quá trình triển khai, và các phiên bản đã xuất bản có thể được tải một cách rõ ràng thay vì phụ thuộc vào một URL tệp không có phiên bản. Các kernel của chúng tôi cũng có thể đóng vai trò là các bản triển khai tham chiếu cho các nhà phát triển đang xây dựng các WebGPU kernel tùy chỉnh hoặc tích hợp các thao tác này vào runtime của riêng họ.

Tải một kernel từ Hub

Cài đặt gói từ npm:

Việc chạy các kernel này yêu cầu trình duyệt có hỗ trợ WebGPU. Khả năng khả dụng của WebGPU phụ thuộc vào trình duyệt, hệ điều hành, GPU và trình điều khiển (driver). Bạn có thể kiểm tra trong JavaScript với "gpu" in navigator.

@huggingface/kernels cung cấp cầu nối giữa kho lưu trữ kernel và ứng dụng của bạn. Gọi getKernel với ID kho lưu trữ trên Hub và phiên bản hợp đồng, sau đó gọi hàm trả về với dữ liệu đầu vào đã định kiểu và hình dạng tensor. Dưới đây là một ví dụ nhỏ về bias-add:

Đầu vào thứ hai được phát sóng trên chiều đầu tiên, tạo ra đầu ra có hình dạng [2, 3]. Trình tải suy ra hình dạng đầu ra và kiểu dữ liệu logic đó từ hợp đồng manifest và các đầu vào, sau đó tự động cấp phát c.

Phép cộng trên sáu số thực (float) là bản demo nhỏ nhất có thể. Ở kích thước này, chi phí truyền tải GPU lớn hơn nhiều so với phép tính. Điểm mấu chốt là mô hình gọi hàm: nó vẫn giữ nguyên chính xác đối với các thao tác nặng, nơi các kernel được tối ưu hóa thực sự phát huy tác dụng, chẳng hạn như nhân ma trận (ai.onnx.MatMul). Chỉ có ID kho lưu trữ và các đầu vào là thay đổi.

Ngay cả thao tác cơ bản này cũng minh họa lý do tại sao các kernel cần các biến thể. Phép cộng có hình dạng bằng nhau có thể sử dụng đường dẫn vector hóa trực tiếp, trong khi các đầu vào được phát sóng cần logic lập chỉ mục khác. Kernel Add đã xuất bản bao gồm các biến thể cho hình dạng bằng nhau, phát sóng vector hóa, xử lý vô hướng và phát sóng tổng quát. Runtime có thể chọn một bản triển khai phù hợp với lệnh gọi hiện tại và thiết bị mà không cần thay đổi API hướng ứng dụng.

Tùy chọn version: 1 chọn phiên bản 1 của hợp đồng kernel đã xuất bản. Nó tách biệt với ONNX opset, since_version của toán tử hoặc bản sửa đổi mô hình. Việc giữ các khái niệm đó tách biệt cho phép các ứng dụng phụ thuộc vào một hợp đồng hướng JavaScript ổn định trong khi các bản triển khai kernel phát triển phía sau nó.

Các kernel nhanh đến mức nào?

Vậy, các kernel được tối ưu hóa thực sự tạo ra sự khác biệt bao nhiêu? Chúng tôi đã đặt bộ sưu tập của mình đối đầu với ORT WebGPU trên GPU Apple M4, sử dụng ONNX Runtime Web 1.30.0-dev.20260826-b1f76d586a. Chúng tôi bắt đầu với 1.756 trường hợp kiểm thử trên tất cả 207 thao tác và giữ lại 809 trường hợp mà cả hai bên đều tạo ra kết quả khớp nhau và thời gian đo đạc đáng tin cậy.

Qua các so sánh đó, các kernel của chúng tôi nhanh hơn 2,57 lần theo trung bình nhân và 1,90 lần ở giá trị trung vị, với 629 trận thắng, 176 trận thua và 4 trận hòa. Dưới đây là cái nhìn cận cảnh về bốn thao tác quen thuộc:

Một số chiến thắng cá biệt lớn hơn nhiều. Một trường hợp Einsum song tuyến tính đặc biệt khó (i,ij,j với kích thước 4096) chạy trong 0,136 ms với kernel của chúng tôi so với 1.396 ms với ORT WebGPU: nhanh hơn hơn 10.000 lần. Một phép CumSum theo hàng trên [256, 4096] nhanh hơn 301 lần, ở mức 0,016 ms so với 4,784 ms. Đây là những trường hợp bất thường thay vì là mức tăng tốc bạn nên mong đợi ở mọi nơi, nhưng chúng cho thấy một kernel chuyên biệt có thể giúp ích như thế nào khi một bản triển khai tổng quát gặp phải đường dẫn chậm.

Chúng tôi đã đo thời gian công việc thực hiện trên chính GPU, loại bỏ các thiết lập như tải kernel, tạo phiên, tải lên đầu vào, biên dịch shader và đọc đầu ra. Các khối lượng công việc rất ngắn tự nhiên khó đo lường hơn và các trường hợp nhỏ có thể hưởng lợi từ bộ nhớ đệm GPU, vì vậy những con số này tốt nhất nên được đọc như một sự so sánh hữu ích thay vì một lời hứa cho mọi ứng dụng.

Chúng cũng là kết quả cho các thao tác riêng lẻ, không phải các mô hình hoàn chỉnh. Hiệu năng chính xác sẽ thay đổi trên các GPU và trình duyệt khác nhau, đó là lý do tại sao Fleet rất quan trọng để xây dựng một bức tranh rộng lớn hơn.

Chúng tôi cũng đang làm việc với nhóm ONNX Runtime để đưa các cải tiến này lên upstream để chúng có thể mang lại lợi ích cho hệ sinh thái ONNX Runtime Web rộng lớn hơn.

Từ một thiết bị đến một đội ngũ (fleet)

Hiệu năng WebGPU thay đổi tùy theo GPU, trình duyệt và trình điều khiển, vì vậy kết quả từ một máy chỉ kể một phần của câu chuyện. Fleet cho phép bất kỳ ai chạy các kiểm tra tính đúng đắn và hiệu năng trong trình duyệt và xem các kernel hoạt động như thế nào trên phần cứng của họ.

Với sự đồng ý, mỗi lần chạy sẽ đóng góp riêng tư các bằng chứng giúp chúng tôi phát hiện các lỗi cụ thể trên thiết bị, so sánh các biến thể và cải thiện các quy tắc lựa chọn. Mục tiêu rất đơn giản: sử dụng phạm vi bao phủ rộng rãi, thực tế để làm cho các kernel nhanh hơn và đáng tin cậy hơn cho mọi người.

Xây dựng nền tảng chung cho WebAI

207 kernel ban đầu là một điểm khởi đầu, không phải trạng thái cuối cùng. Việc xuất bản các kernel một cách độc lập trên Hub mang lại cho chúng tôi một nơi chung để kiểm tra các hợp đồng, so sánh các bản triển khai, tái tạo các kiểm tra tính đúng đắn và cải thiện hiệu năng mà không cần nhúng trực tiếp mọi shader vào mọi runtime.

Bộ sưu tập này cũng là một phần của hệ sinh thái kernel rộng lớn hơn của Hub: trên trang Kernels, các WebGPU kernel nằm cùng với các kernel cho CUDA, ROCm, Metal và các nền tảng khác, đồng thời có thể được lọc, sắp xếp và khám phá như bất kỳ thành phần nào khác trên Hub.

Các mảnh ghép củng cố lẫn nhau:

Đây là nền tảng cấp thấp cho các bước tiếp theo trong ngăn xếp suy luận trên trình duyệt của chúng tôi. Chúng tôi rất hào hứng khi kết nối các kernel này với các công cụ mô hình cấp cao hơn, tiếp tục mở rộng phạm vi bao phủ thao tác và làm cho việc suy luận cục bộ nhanh chóng trở nên dễ sử dụng hơn trên toàn bộ hệ sinh thái WebAI.