本文档是用于加快 Web AI 速度的 WebAssembly 和 WebGPU 增强功能（第 1 部分）的延续。我们建议您先阅读这篇博文或先观看 IO 24 上的演讲，然后再继续。
WebGPU
WebGPU 使 Web 应用能够访问客户端的 GPU 硬件，以执行高效且高度并行的计算。自在 Chrome 中推出 WebGPU 以来，我们见证了人工智能 (AI) 和机器学习 (ML) 在网络领域的精彩演示。
例如，Web Stable Diffusion 证明了能够直接在浏览器中使用 AI 技术根据文本生成图片。今年早些时候，Google 自己的 Mediapipe 团队发布了针对大语言模型推断的实验性支持。
下面的动画显示了 Google 的开源大语言模型 (LLM) Gemma，它完全在设备上在 Chrome 中实时运行。
以下 Hugging Face 演示 Meta 的 Segment Anything Model 可完全在客户端生成高质量的对象蒙版。
以上只是几个优秀的项目，它们展示了 WebGPU 在 AI 和机器学习中的强大功能。WebGPU 允许这些模型和其他模型在 CPU 上运行速度明显快于其在 CPU 上运行速度。
与同一模型的 CPU 实现相比，Huging Face 的文本嵌入的 WebGPU 基准表现出了显著的速度。在 Apple M1 Max 笔记本电脑上，WebGPU 速度提高了 30 倍以上。其他人报告称，WebGPU 将基准速度提高超过 120 倍。
针对 AI 和机器学习改进 WebGPU 功能
WebGPU 非常适合 AI 和机器学习模型，得益于对计算着色器的支持，它们可以拥有数十亿个参数。计算着色器在 GPU 上运行，有助于对大量数据执行并行数组操作。
过去一年，我们对 WebGPU 进行了众多改进，并不断添加更多功能，以提高 Web 上的机器学习和 AI 性能。最近，我们推出了两项新功能：16 位浮点和打包的整数点产品。
16 位浮点
请记住，机器学习工作负载不需要精度。shader-f16 功能支持在 WebGPU 着色语言中使用 f16 类型。此浮点类型占用 16 位，而非通常的 32 位。f16 的范围较小且不太精确，但对许多机器学习模型而言，这已经足够。
此功能可从以下几个方面提高效率：
内存使用量减少：具有 f16 元素的张量占据一半空间，内存使用量减半。GPU 计算通常在内存带宽上存在瓶颈，因此一半的内存通常意味着着色器的运行速度快一倍。从技术上讲，您不需要使用 f16 即可节省内存带宽。可以采用低精度格式存储数据，然后在着色器中将其展开至全 f32 进行计算。但是，GPU 会使用额外的计算能力来打包和解压缩数据。
减少数据转换：f16 通过最大限度地减少数据转换，从而减少使用的计算量。可以存储低精度数据，然后直接使用，无需转换。
提高并行性：现代 GPU 能够在 GPU 的执行单元中同时容纳更多值，因而能够执行更多并行计算。例如，支持每秒最多 5 万亿 f32 浮点操作的 GPU 可能支持每秒 10 万亿次 f16 浮点操作。
shader-f16 时，Hugging Face 的文本嵌入 WebGPU 基准基准在 Apple M1 Max 笔记本电脑上的运行速度比 f32 快 3 倍。
WebLLM 是一个可以运行多个大型语言模型的项目。它使用开源机器学习编译器框架 Apache TVM。
我让 WebLLM 使用 Llama 3 80 亿参数模型来规划一次巴黎之旅。结果表明，在模型的预填充阶段，f16 的速度是 f32 的 2.1 倍。在解码阶段，速度提升 1.3 倍以上。
应用必须先确认 GPU 适配器是否支持 f16，并在请求 GPU 设备时明确启用 f16（如果可用）。如果 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 的一半）的情况下仍能正常工作。该方法在 LLM 和用于分割和对象识别的图像模型中很受欢迎。尽管如此，模型的输出质量会降低，精度也会降低，因此 8 位量化并不适合所有应用。
本身支持 8 位值的 GPU 数量相对较少。这就是打包整数点积的用武之地。我们在 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 倍。
请使用 wgslLanguageFeatures 属性检查浏览器是否支持相应浏览器。如果 GPU 本身不支持打包点积，则浏览器会对其自己的实现执行 polyfill 操作。
// main.js
if (navigator.gpu.wgslLanguageFeatures.has('packed_4x8_integer_dot_product')) {
// Use dot4U8Packed, dot4I8Packed builtin
// functions in the shaders.
}
以下代码段“diff”（差异）突出显示了在 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 着色器，编写为使用打包的整数点积。主要区别在于，此着色器会加载单个 32 位整数，而不是从矢量和矩阵加载 4 个浮点值。该 32 位整数可存储 4 个 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);
}
}
16 位浮点产品和打包整数点产品都是 Chrome 中随附的可加快 AI 和机器学习速度的功能。16 位浮点数在硬件支持的情况下可用，并且 Chrome 会在所有设备上实现打包的整数点产品。
您现在可以在 Chrome 稳定版中使用这些功能，以获得更好的性能。
推荐的功能
展望未来，我们正在研究另外两个特征：子组和合作矩阵乘法。
子组功能可让 SIMD 级别的并行处理进行通信或执行集数学运算，例如对超过 16 个数字求和。这样可以实现高效的跨线程数据共享。现代 GPU API 支持子群组，它们的名称各不相同，形式也略有不同。
我们已将通用集提炼成一个提案，并提交给 WebGPU 标准化组。此外，我们已根据实验性标志在 Chrome 中对子群组进行了原型设计，并将初步结果纳入讨论中。主要问题在于如何确保可移植的行为。
协同矩阵乘法是 GPU 的最新功能。一个较大的矩阵乘法可以分解为多个较小的矩阵乘法。合作矩阵乘法可在一个逻辑步中对这些较小的固定大小块执行乘法运算。在该步骤中，一组线程高效协作以计算结果。
我们调查了对底层 GPU API 的支持情况，并计划向 WebGPU 标准化组提出提案。与子小组一样，我们希望大部分的讨论都围绕可移植性展开。
为了评估子群组操作的性能，在真实应用中，我们将对子群组的实验性支持集成到 MediaPipe 中，并使用 Chrome 的子群组操作原型对其进行了测试。
我们使用了大语言模型的预填充阶段 GPU 内核中的子群组，因此我仅报告预填充阶段的加速情况。在 Intel GPU 上，我们看到子群组的执行速度是基准的两倍半。但是，这些改进对于不同的 GPU 不一致。
下一张图表显示了应用子群组以跨多个消费类 GPU 优化矩阵乘法微基准的结果。矩阵乘法是大型语言模型中较繁重的运算之一。数据显示，在许多 GPU 上，子群组的速度是基准的 2 倍、5 倍甚至 13 倍。但请注意，在第一个 GPU 上，子群组的效果并没有太大改善。
GPU 优化很难
归根结底，优化 GPU 的最佳方法取决于客户端提供的 GPU。使用出色的新 GPU 功能并不总能带来预期回报，因为其中可能会涉及到许多复杂的因素。一个 GPU 上的最佳优化策略可能不是另一个 GPU 上的最佳策略。
您希望最大限度地降低内存带宽，同时充分利用 GPU 的计算线程。
内存访问模式也非常重要。当计算线程以对硬件最有利的模式访问内存时，GPU 往往性能更好。 重要提示：不同的 GPU 硬件应该具有不同的性能特征。您可能需要根据 GPU 运行不同的优化。
在下图中，我们采用了相同的矩阵乘法算法，但添加了另一个维度以进一步展示各种优化策略的影响，以及不同 GPU 之间的复杂程度和方差。我们在这里引入了一项新技术，称为“Swizzle”。Swizzle 优化内存访问模式，使其更适合硬件。
您可以看到，内存调配具有重大影响：有时甚至比子群组更具影响力。在 GPU 6 上，swizzle 提供 12 倍的速度，而子群组则可提供 13 倍的速度。加起来，它们的运行速度达到惊人的 26 倍。对于其他 GPU，有时调配和子组的组合比单独使用其中之一时效果更好。在其他 GPU 上，仅使用调配的效果最佳。
调整和优化 GPU 算法以使其在每个硬件上都能正常运行，可能需要大量的专业知识。不过值得庆幸的是，在更高级别的库框架（例如 Mediapipe、Transformers.js、Apache TVM、ONNX Runtime Web 等）中投入了大量才华横溢的工作。
库和框架可以很好地处理以下复杂问题：管理各种 GPU 架构，以及生成可在客户端上良好运行的特定于平台的代码。
要点总结
Chrome 团队将继续帮助改进 WebAssembly 和 WebGPU 标准，以改进面向机器学习工作负载的 Web 平台。我们正在投资开发更快的计算基元、更好地跨 Web 标准实现更好的互操作性，并确保各种大小的模型都能跨设备高效运行。
我们的目标是在保留最佳网络功能（覆盖面、易用性和可移植性）的同时，最大限度地发挥平台的功能。而且，我们并非孤军奋战。我们正在与 W3C 的其他浏览器供应商以及许多开发合作伙伴展开合作。
我们希望您在使用 WebAssembly 和 WebGPU 时牢记以下几点：
- AI 推断现在可在网络上跨设备使用。这带来了在客户端设备上运行的优势，例如降低服务器费用、缩短延迟时间以及加强隐私保护。
- 尽管所讨论的许多功能主要与框架作者相关，但您的应用也可以受益于没有太多开销。
- Web 标准非常灵活且在不断发展完善，我们随时欢迎用户提供反馈。分享您的 WebAssembly 和 WebGPU 模型。
致谢
我们要感谢 Intel 网络图形团队，他们在推动 WebGPU f16 和整合整数点产品功能方面发挥了重要作用。我们要感谢 W3C WebAssembly 和 WebGPU 工作组的其他成员，包括其他浏览器供应商。
感谢 Google 和开源社区的 AI 和机器学习团队成为出色的合作伙伴。当然，还有我们的团队成员让这一切成为可能。