WebAssembly 和 WebGPU 強化項目，加快 Web AI 速度，第 2 部分

本文將延續 WebAssembly 和 WebGPU 強化功能，提昇 Web AI 速度 (第 1 部分) 的內容。建議您詳閱這篇文章，或先觀看 IO 24 的演講內容再繼續操作。

Austin Eng
Austin Eng
GitHub
Deepti Gandluri
Deepti Gandluri
GitHub
François Beaufort
François Beaufort
GitHub

WebGPU

WebGPU 可讓網頁應用程式存取用戶端的 GPU 硬體，執行高效率且高平行的運算作業。自在 Chrome 中推出 WebGPU 以來，已經在網路上看過令人驚豔的人工智慧 (AI) 和機器學習 (ML) 示範。

舉例來說，Web Stable Diffusion 可以顯示，直接在瀏覽器中使用 AI 從文字生成圖片。今年稍早，Google 自家的 Mediapipe 團隊發布了實驗支援大型語言模型推論。

以下動畫展示了 Gemma，這是 Google 推出的開放原始碼大型語言模型 (LLM)，該模型完全在 Chrome 裝置上即時執行。

下方是 Meta's Segment Anything Model 的 Hugging Face 示範，該模型完全在用戶端產生了高品質的物件遮罩。

這些只是一些出色的專案，展示了 WebGPU 在 AI 和機器學習方面的強大威力。WebGPU 讓這些模型和其他模型的執行速度遠比在 CPU 上更快。

Hugging Face 針對文字嵌入的 WebGPU 基準，展示了相較於同一模型的 CPU 實作項目，速度可大幅提升。在 Apple M1 Max 筆記型電腦上，WebGPU 的執行速度是以前的 30 倍以上。也有部分已回報 WebGPU 將基準測試的速度加快超過 120 倍。

改善 AI 和機器學習的 WebGPU 功能

WebGPU 支援運算著色器，因此非常適合應用數十億個參數的 AI 和機器學習模型。運算著色器會在 GPU 上執行，有助於對大量資料執行平行陣列作業。

過去一年來，WebGPU 已有許多改進項目，而我們持續新增更多功能，提升網路機器學習和 AI 技術在網路上的效能。我們最近推出了兩項新功能：16 位元浮點產品，以及封裝的整數點產品。

16 位元浮點

提醒您，機器學習工作負載不需要精確度。shader-f16 是一項可在 WebGPU 陰影語言中使用 f16 類型的功能。這個浮點類型會耗用 16 位元，而非一般的 32 位元。f16 的範圍較小且較不精確，但對於許多機器學習模型來說，這已足夠。

這項功能可透過以下幾種方式提高效率：

  • 減少記憶體：含有 f16 元素的 Tensor 會佔用一半的空間，因此記憶體用量減半。GPU 運算作業往往會在記憶體頻寬上出現瓶頸，因此，一半記憶體通常代表著色器執行速度是快上兩倍。嚴格說來，您不需要 f16 即可節省記憶體頻寬。此外，您也可採用低精確度格式儲存資料，然後在著色器中將其展開為 Full f32 以進行運算。但 GPU 會花費額外的運算能力來封裝及拆分資料。

  • 減少資料轉換：f16 可將資料轉換降到最低，減少使用運算成本。Google 可儲存低精確度的資料，然後在沒有轉換的情況下直接使用。

  • 提高平行處理能力：新型 GPU 可在 GPU 的執行單位中同時容納更多值，以便執行更多平行運算。舉例來說，GPU 每秒支援 5 兆次浮點運算，每秒可能會支援 10 兆 f16 浮點作業。

，瞭解如何調查及移除這項存取權。
文字嵌入的 WebGPU 基準螢幕截圖
有了 shader-f16，在 Apple M1 Max 筆電上，Hugging Face 的WebGPU 文字嵌入基準基準執行基準測試，速度比 f32 筆電快 3 倍。

WebLLM 是可執行多個大型語言模型的專案。並使用 Apache TVM 這個開放原始碼機器學習編譯器架構。

我要求 WebLLM 使用 Llama 300 億個參數模型規劃前往巴黎的行程。結果顯示，在模型的預填階段，f16 的速度比 f32 快 2.1 倍。在解碼階段，處理速度的處理速度加快超過 1.3 倍。

應用程式必須先確認 GPU 轉接器支援 f16，如果可用，請在要求 GPU 裝置時明確啟用。如果 f16 不受支援，就無法在 requiredFeatures 陣列中要求。

// main.js

const adapter = await navigator.gpu.requestAdapter();
const supportsF16 = adapter.features.has('shader-f16');
if (supportsF16) {
  // Use f16.
  const device = await adapter.requestDevice({
    requiredFeatures: ['shader-f16'],
  });
  initApp(device);
}

然後在 WebGPU 著色器中明確啟用 f16。之後，您就可以像任何其他浮點資料類型一樣自由使用著色器。

// my-shader.wgsl

enable f16;

struct Data {
  values : array<vec4<f16>>
}
@group(0) @binding(0) var<storage, read> data : Data;
@compute @workgroup_size(64) fn main(@builtin(global_invocation_id) gid : vec3u) {
  let value : vec4<f16> = data.values[gid.x];
  ...
}

包裝整數點產品

許多模型的精確度只有 8 位元 (f16 的一半)，分割及辨識物件時，大型語言模型和圖片模型相當普遍。儘管如此，模型的輸出品質也會降低精確度降低程度，因此不適用於所有應用程式。

相對少數的 GPU 原生支援 8 位元值。這也是裝滿整數內點產品的地方。我們在 Chrome 123 版中推出 DP4a。

現代 GPU 有特殊指示可取用兩個 32 位元整數，並將其解讀為 4 個連續封裝的 8 位元整數，並計算其元件之間的內積。

這對 AI 和機器學習特別有用，因為矩陣乘積核心是由許多內積類產品所組成。

舉例來說，我們可以將 4 x 8 的矩陣乘以 8 x 1 向量。計算過程包括花 4 點乘積來計算輸出向量中的每個值；A、B、C 和 D。

矩陣與向量乘法的圖表

計算這些輸出內容的程序皆相同；我們將說明計算其中一個步驟進行任何計算之前，我們必須先將 8 位元整數資料轉換為可執行算術的型別，例如 f16。接著，我們會執行元素乘法，最後將所有產品加在一起。總而言之，針對整個矩陣向量乘法，我們會執行 40 個整數來為浮動轉換來進行拆分資料、32 浮點乘數以及 28 個浮動計算。

對於執行較多作業的大型矩陣，封裝的整數點產品有助於減少工作量。

針對結果向量中的每個輸出內容，我們使用 WebGPU 著色語言內建的 dot4U8Packed 執行兩個打包的內積運算，然後把結果相加。就整個矩陣向量乘法的總而言，我們不會執行任何資料轉換。我們執行 8 個裝箱產品和 4 個整數加值。

封裝整數矩陣與向量乘法的圖表

我們在各種消費者 GPU 上，測試含有 8 位元資料的整組整數內插式產品。與 16 位元浮點相比，可看到 8 位元快 1.6 到 2.8 倍。如果額外使用封裝的整數點積，則效能會更加出色。速度是 1.7 到 2.9 倍。

矩陣向量乘以速度：f16 與 u8 的螢幕截圖
圖表 1：矩陣向量加速，比較 f16 與 U8 及 U8 及點 4U8Packed。

請使用 wgslLanguageFeatures 屬性確認瀏覽器支援。如果 GPU 本身不支援封裝的內點產品，則瀏覽器會自行執行實作。

// main.js

if (navigator.gpu.wgslLanguageFeatures.has('packed_4x8_integer_dot_product')) {
  // Use dot4U8Packed, dot4I8Packed builtin
  // functions in the shaders.
}

以下程式碼片段差異 (差異) 特別指出在 WebGPU 著色器中使用封裝整數產品所需要的變更。

之前 - 將部分內積積積累計到「sum」變數的 WebGPU 著色器。迴圈結尾時，`sum` 保存了向量和輸入矩陣的一列之間的全點積。

// my-dot-product.wgsl

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3u) {
  var sum : f16;
  let start = gid.x * uniforms.dim;
  for (var i = 0u; i < uniforms.dim; i++) {
    let v1 : vec4<f16> = vector.values[i];
    let v2 : vec4<f16> = matrix.values[start + i];
    sum += dot(v1, v2);
  }
}

更新後 - 為使用封裝的整數點產品撰寫的 WebGPU 著色器。主要差異在於，這個著色器不會從向量和矩陣中載入 4 個浮點值，而是載入單一 32 位元整數。這個 32 位元整數保存了四個 8 位元整數值的資料。接著，我們會呼叫 dot4U8Packed 來計算兩個值的內積。

// my-dot-product.wgsl

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3u) {
  var sum : f32;
  let start = gid.x * uniforms.dim;
  for (var i = 0u; i < uniforms.dim; i++) {
    let v1 : u32 = vector.values[i];
    let v2 : u32 = matrix.values[start + i];
    sum += dot4U8Packed(v1, v2);
  }
}

Chrome 提供 16 位元浮點產品和封裝整數點產品，能加快 AI 和機器學習速度。只要硬體支援 16 位元浮點產品，即可使用，且 Chrome 在所有裝置上實作已封裝的整數點產品。

您現在可以在 Chrome 穩定版中使用這些功能來獲得更好的效能。

提議的功能

日後，我們還在調查第二組功能：子群組及合作矩陣相乘。

子群組功能可讓 SIMD 層級的平行處理進行通訊或執行集體數學運算，例如超過 16 個數字的總和。這可讓您有效率地跨執行緒共用資料。新型 GPU API 支援子群組，名稱可能不同，形式也略有不同。

我們整理了一份提案，並將一般大眾重點整理成 WebGPU 標準化群組。我們在 Chrome 中設計了子群組原型，以實驗性旗幟，並將初始結果帶入討論中。主要問題在於如何確保可攜式行為。

協作矩陣乘法是 GPU 的近況。一個大型矩陣乘法可分為多個較小的矩陣乘積。協作矩陣乘法會在單一邏輯步驟中，對這些較小尺寸的區塊執行乘法。在該步驟中，一組執行緒能夠有效合作計算結果。

我們針對基礎 GPU API 的問卷調查支援，預計向 WebGPU 標準化群組發表提案。和子小組一樣，我們預期大部分的討論都會以可攜性為主軸。

為了評估子群組作業的成效，我們在實際應用程式中將子群組的實驗性支援整合至 MediaPipe，並使用 Chrome 的原型進行子群組操作。

我們在大型語言模型預填階段的 GPU 核心中使用子群組，因此我只回報預填階段的速度速度。在 Intel GPU 上，我們發現子群組的執行速度比基準測試快 2.5 倍。不過，不同 GPU 間的這些改善事項並不一致。

MediaPipe LLM 推論中的子群組速度提升的螢幕截圖
圖表 2.子群組讓 Intel Tiger Lake GT2 GPU 的預填作業速度加快 2.5 倍，並在 Chrome 和 Mediapipe 提供實驗性支援。

下一張圖表顯示，在將矩陣最佳化矩陣後，在多個消費者 GPU 上將矩陣相乘後，應套用子群組的結果。矩陣乘法是大型語言模型中需要處理的其中一項作業。資料顯示，在許多 GPU 中，子群組可使速度加快二、5，甚至十三倍。但請注意，在第一個 GPU 上，子群組基本上的成效不盡理想。

矩陣乘上子群組加速的螢幕截圖
圖表 3.套用子群組處理矩陣乘法可能進一步提高效能。

GPU 最佳化難以進行

最後，將 GPU 最佳化的最佳做法取決於用戶端提供的 GPU。使用高效率的全新 GPU 功能不一定能達到預期的效果，因為可能會有許多複雜的因素。對不同 GPU 而言，最佳最佳化策略對其他 GPU 來說可能不是最佳策略。

您希望盡可能降低記憶體頻寬，同時充分利用 GPU 的運算執行緒。

記憶體存取模式也很重要。當運算執行緒以最適合硬體的模式存取記憶體時，GPU 的效能會大幅提高。 重要事項：不同 GPU 硬體的效能特性會有所不同。您可能需要根據 GPU 執行不同的最佳化作業。

在下圖中，我們採用了相同的矩陣乘法，再加上另一個維度，進一步示範各種最佳化策略的影響，以及不同 GPU 的複雜度與差異。我們在這裡引進了新的技術，也就是「Swizzle」。Swizzle 會將記憶體存取模式最佳化，提升硬體的最佳效能。

可以看到記憶體配置會產生重大影響效果有時甚至比子群組更出色在 GPU 6 上，swizzle 的速度提升 12 倍，子群組則加快 13 倍。並帶來 26 倍的驚人速度。對其他 GPU 而言，有時快速靈活的架構和子群組組合的組合成效會比單獨使用其中一體更好。在其他 GPU 上，只使用 Swizzle 可以達到最佳效果。

矩陣乘上策略速度的螢幕截圖
圖表 4.

如要調整及最佳化 GPU 演算法，讓每部硬體都能順利運作，需要具備大量專業知識。不過值得慶幸的是，在 Mediapipe、Transformers.js、Apache TVM、ONNX Runtime Web 等高階程式庫架構而言，仍有許多才華。

程式庫和架構非常適合用來處理複雜的管理不同 GPU 架構，以及產生能在用戶端順暢運作的平台專屬程式碼。

重點整理

Chrome 團隊持續協助改善 WebAssembly 和 WebGPU 標準，改善用於機器學習工作負載的網路平台。我們持續投入資源開發速度更快的運算基元，提升各種網路標準之間的互通性，並確保各種大小的模型能在各種裝置上有效率地執行。

我們的目標是盡可能發揮平台的最大效益，同時保留最出色的網路體驗，包括便利性、可用性和可攜性。而且，我們並非單憑一個己力，我們與 W3C 的其他瀏覽器廠商及許多開發合作夥伴攜手合作，

使用 WebAssembly 和 WebGPU 時，我們希望您記住下列項目：

  • 你現在可以在各種裝置上透過網頁使用 AI 推論功能。這樣就能在用戶端裝置上執行，例如降低伺服器成本、低延遲及加強隱私保護。
  • 討論的許多功能大多與架構作者相關，但您的應用程式在不耗費太多的負擔的情況下就能派上用場。
  • 網路標準瞬息萬變且不斷演進，我們也一直在尋求意見回饋。提供 WebAssembly 和 WebGPU 的共用權限。

特別銘謝

我們由衷感謝 Intel 網頁繪圖團隊，他們是推動 WebGPU f16 及大量整數點產品功能的重要推手。我們由衷感謝 W3C 的 WebAssembly 和 WebGPU 工作團隊的其他成員，包括其他瀏覽器廠商。

感謝 Google 和開放原始碼社群的 AI 和機器學習團隊感謝他們的支持。當然，我們的團隊成員，往往能實踐這一切。