WebGPU 实战：在浏览器中运行 AI 模型推理

Author Avatar
via
发表：2026-06-26 09:01:53
修改：2026-06-26 09:01:52

引言

WebGPU 作为下一代 Web 图形 API，不仅替代了 WebGL 用于渲染，更重要的是它打开了在浏览器中直接进行高性能 GPU 计算的大门。2026 年的今天，Chrome、Edge 和 Safari 都已全面支持 WebGPU，Firefox 也已默认启用。这意味着前端开发者终于可以利用 GPU 的并行计算能力，在浏览器中直接运行 AI 模型推理，而无需依赖服务器端。

本文将从 WebGPU 的基础概念出发，逐步带你理解如何在浏览器中加载并运行一个轻量级 AI 模型，涵盖从计算着色器编写到推理流水线构建的完整流程。

一、WebGPU 与 WebGL 的核心区别

很多人第一反应是：WebGL 也能做计算（通过 WebGL2 的 Transform Feedback），为什么要搞 WebGPU？区别在于设计理念。

WebGL 的设计初衷是图形渲染，计算能力是"附赠"的。它的状态机模型（绑定纹理、设置混合模式、配置 VAO…）对于纯计算任务来说过于冗余。而 WebGPU 从一开始就将计算作为一等公民：

  • 显式的命令编码：WebGPU 引入了 Command Encoder 模式，你先录制一系列命令，再一次性提交给 GPU 队列执行。这与现代图形 API（Vulkan、Metal、D3D12）一致，减少了 CPU-GPU 之间的同步开销。
  • 计算着色器（Compute Shader）原生支持：不需要用顶点着色器"欺骗"GPU 来做通用计算，直接写 compute shader 即可。
  • Storage Buffer：WebGPU 的 Storage Buffer 支持读写，不像 WebGL 的 Uniform Buffer 只能读。这对 AI 推理中的权重加载和中间激活值的读写至关重要。
  • Bind Group 模型：资源绑定更加灵活，可以在管线创建时预定义布局，运行时只需切换 Bind Group 即可改变输入输出，减少了状态切换开销。

二、在浏览器中运行 AI 推理的整体架构

一个典型的 WebGPU AI 推理流水线包含以下几个层次：

  1. 模型层：将训练好的模型（如 ONNX 格式）解析为 WebGPU 可用的张量结构和算子图。
  2. 算子层：每个算子（Conv2D、MatMul、Softmax 等）对应一个或多个计算着色器。
  3. 调度层：按照拓扑顺序执行算子图，管理中间张量的内存分配。
  4. 缓冲层：管理 GPU Buffer 的创建、复用和垃圾回收。

目前已有成熟的开源库帮你做了大部分工作：transformers.js（Hugging Face 出品）已经集成了 WebGPU 后端，ONNX Runtime Web 也支持 WebGPU EP（Execution Provider）。但理解底层原理，才能在性能瓶颈出现时知道去哪里优化。

三、实战：用计算着色器实现矩阵乘法

矩阵乘法是 AI 推理中最核心的算子。全连接层、注意力机制中的 QKV 计算，本质上都是 MatMul。我们先从最基础的 WebGPU 计算着色器开始，实现一个 4×4 矩阵乘法。

3.1 创建 Pipeline

首先需要创建一个 Compute Pipeline。关键代码如下：

const shaderModule = device.createShaderModule({
  code: `
    @group(0) @binding(0) var<storage, read> a: array<f32>;
    @group(0) @binding(1) var<storage, read> b: array<f32>;
    @group(0) @binding(2) var<storage, read_write> result: array<f32>;

    @compute @workgroup_size(4, 4)
    fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
      let row = gid.x;
      let col = gid.y;
      let N: u32 = 4u;
      var sum: f32 = 0.0;
      for (var k: u32 = 0u; k < N; k = k + 1u) {
        sum = sum + a[row * N + k] * b[k * N + col];
      }
      result[row * N + col] = sum;
    }
  `
});

const pipeline = device.createComputePipeline({
  layout: 'auto',
  compute: { module: shaderModule, entryPoint: 'main' }
});

3.2 创建 Buffer 和 Bind Group

着色器写好后，需要创建 GPU Buffer 来存储数据，并将它们绑定到 Bind Group：

// 创建 Storage Buffer
const bufferA = device.createBuffer({
  size: 4 * 4 * 4, // 4x4 矩阵，每个元素 4 字节 (f32)
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
});
const bufferB = device.createBuffer({
  size: 4 * 4 * 4,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
});
const bufferResult = device.createBuffer({
  size: 4 * 4 * 4,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
});

// 写入数据
device.queue.writeBuffer(bufferA, 0, new Float32Array([...]).buffer);
device.queue.writeBuffer(bufferB, 0, new Float32Array([...]).buffer);

// 创建 Bind Group
const bindGroup = device.createBindGroup({
  layout: pipeline.getBindGroupLayout(0),
  entries: [
    { binding: 0, resource: { buffer: bufferA } },
    { binding: 1, resource: { buffer: bufferB } },
    { binding: 2, resource: { buffer: bufferResult } },
  ],
});

3.3 提交计算命令

const encoder = device.createCommandEncoder();
const pass = encoder.beginComputePass();
pass.setPipeline(pipeline);
pass.setBindGroup(0, bindGroup);
pass.dispatchWorkgroups(1, 1); // 4x4 矩阵只需 1 个 workgroup
pass.end();

// 将结果拷回可读 buffer
const readBuffer = device.createBuffer({
  size: 4 * 4 * 4,
  usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
});
encoder.copyBufferToBuffer(bufferResult, 0, readBuffer, 0, 4 * 4 * 4);

device.queue.submit([encoder.finish()]);

// 读取结果
await readBuffer.mapAsync(GPUMapMode.READ);
const result = new Float32Array(readBuffer.getMappedState());
console.log('MatMul 结果:', result);
readBuffer.unmap();

四、性能优化关键技巧

上面的例子是最简单的实现。在实际 AI 模型推理中，矩阵尺寸往往很大（如 4096×4096），需要考虑以下优化：

4.1 Workgroup Size 选择

Workgroup 的大小直接影响 GPU 占用率和共享内存利用率。对于矩阵乘法，通常选择 8×8 或 16×16 的 workgroup size，配合共享内存（Shared Memory）做分块计算（Tiling），可以大幅减少全局内存访问。

4.2 内存对齐与合并访问

GPU 内存访问的效率很大程度上取决于访问模式。如果一个 workgroup 内的所有线程访问连续的内存地址（合并访问），GPU 可以一次性读取整块数据；如果是随机访问，则会产生大量内存事务，性能急剧下降。

在 WGSL 中，使用 vec4<f32> 而不是 4 个独立的 f32 可以帮助编译器生成更好的内存访问模式。

4.3 异步 Pipeline 与双缓冲

对于多算子的模型推理，可以使用双缓冲（Double Buffering）技术：当前算子正在 GPU 执行时，CPU 已经在编码下一个算子的命令。WebGPU 的 Command Encoder 天然支持这种模式——你可以录制多帧命令，然后交替提交。

4.4 量化推理

FP32 在 GPU 上并不总是最快的。许多现代 GPU 对 FP16 甚至 INT8 有专门的硬件支持。在 WebGPU 中，你可以使用 f16 类型（需要启用 chrome://flags/#enable-unsafe-webgpu 或使用 WGSL 的 enable f16; 扩展），将推理速度提升 2-4 倍，同时减半内存占用。

五、生态现状与实战工具链

2026 年的 WebGPU AI 生态已经相当成熟：

  • transformers.js v3：Hugging Face 的 JS 库，支持 WebGPU 后端，可以直接在浏览器中运行 BERT、Whisper、LLaMA 等模型的量化版本。配合 IndexedDB 缓存模型权重，首次加载后后续推理完全离线。
  • ONNX Runtime Web：微软的 ONNX Runtime 有了 WebGPU EP，支持直接在浏览器中运行 ONNX 模型。适合已有 ONNX 模型Pipeline 的团队迁移。
  • MediaPipe Web：Google 的 MediaPipe 也提供了 WebGPU 后端，适合实时视觉任务（人脸检测、手势识别、姿态估计）。
  • TensorFlow.js WebGPU Backend：虽然 TF.js 式微，但它的 WebGPU backend 仍然是一个不错的参考实现。

六、实用建议：何时该用 WebGPU 推理

不是所有场景都适合在浏览器中跑 AI 推理。以下是一些决策参考：

适合 WebGPU 推理的场景：

  • 隐私敏感场景：用户数据不出本地，如面部特征提取、语音转文字。
  • 低延迟交互：实时风格迁移、实时滤镜、手势控制等对延迟敏感的场景。
  • 离线场景：PWA 应用或网络条件差的环境。
  • 成本控制：客户端推理不消耗服务器 GPU 资源。

不适合的场景：

  • 大模型推理（70B 参数级别）：浏览器内存和 GPU 显存有限，目前移动端设备能跑的最大模型大约在 1-3B 参数级别。
  • 需要高精度的场景：量化推理会损失精度。
  • 老旧设备兼容性：WebGPU 需要较新的浏览器和 GPU 驱动。

七、展望：从 WebGPU 到 WebNN

值得一提的是，W3C 正在推进 Web Neural Network API（WebNN） 标准，它位于 WebGPU 之上，提供更高层的 AI 推理接口。WebNN 的目标是让前端开发者不需要理解着色器和内存管理，直接用几行 JavaScript 就能调用硬件加速的 AI 推理。

目前 WebNN 在 Chrome 中已经可以通过 Origin Trial 试用。未来 WebGPU 和 WebNN 的关系类似于 Canvas 2D 和 WebGL 的关系——一个易用，一个灵活，各司其职。

总结

WebGPU 让浏览器成为了一等公民的 AI 推理平台。从计算着色器实现矩阵乘法，到利用 transformers.js 直接在浏览器中运行 Hugging Face 模型，前端开发者第一次拥有了在用户设备上直接进行 AI 计算的能力。

核心要点回顾：

  • WebGPU 的计算着色器是浏览器 AI 推理的基础
  • Storage Buffer 和 Bind Group 模型是理解 WebGPU 资源管理的关键
  • 性能优化的核心在于 workgroup size、内存合并访问和量化
  • 实战中优先考虑 transformers.js 或 ONNX Runtime Web 等成熟库
  • WebNN API 是未来趋势，值得持续关注

2026 年是 WebGPU AI 推理从实验走向生产的一年。无论你是前端工程师还是 AI 工程师，掌握这套技术栈都将为你的工具箱增添一把利器。

评论