WASM + WebGPU：浏览器端大模型推理的 Rust 加速方案-编程阁

WASM + WebGPU：浏览器端大模型推理的 Rust 加速方案

一、浏览器端 AI 推理的瓶颈：CPU 太慢，GPU 难用

在浏览器中运行 AI 推理，最大的瓶颈是计算性能。一个 7B 参数的大模型，单次推理需要数十亿次浮点运算。浏览器的 JavaScript 引擎（V8）在纯 CPU 模式下，推理速度约为每秒 1-2 个 token——用户等一个回答要 30 秒以上，体验不可接受。

WebGPU 是浏览器端访问 GPU 的标准 API，提供了计算着色器（Compute Shader）能力，可以在 GPU 上并行执行大规模矩阵运算。但 WebGPU 的 API 是 JavaScript 接口，直接用 JS 编写计算着色器和管理 GPU 缓冲区的代码复杂且易错。

Rust + WASM + WebGPU 的组合方案，用 Rust 编写推理逻辑和 GPU 管理代码，编译为 WASM 在浏览器中运行，通过 WebGPU API 调用 GPU 加速。Rust 的类型系统保证内存安全，WASM 提供接近原生的执行速度，WebGPU 提供 GPU 并行计算能力——三者结合，让浏览器端大模型推理从"概念验证"走向"可用体验"。

二、WASM + WebGPU 推理的底层机制

2.1 整体架构

flowchart TD A[浏览器页面] --> B[WASM 模块] B --> C[Rust 推理引擎] C --> D[WebGPU 计算管线] D --> E[GPU Compute Shader] E --> F[矩阵乘法 / 注意力计算] F --> G[GPU 缓冲区] G --> C C --> H[Token 解码] H --> B B --> A subgraph 数据流 I[模型权重: ArrayBuffer] --> B J[输入 Token: Uint32Array] --> B B --> K[输出 Token: Uint32Array] end subgraph GPU 执行 L[权重上传到 GPU Buffer] M[Dispatch Compute Pipeline] N[Readback 结果到 CPU] end

2.2 WebGPU 计算管线

WebGPU 的计算管线由三个核心组件构成：

Shader Module：WGSL（WebGPU Shading Language）编写的计算着色器，定义 GPU 上的并行计算逻辑。
Bind Group：将 GPU 缓冲区绑定到着色器的资源槽位，类似于函数参数传递。
Compute Pipeline：将 Shader Module 和 Bind Group Layout 组合为可执行的管线。

矩阵乘法的计算着色器示例（WGSL）：

// 矩阵乘法 C = A × B // A: M×K, B: K×N, C: M×N @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> c: array<f32>; @group(0) @binding(3) var<uniform> dims: vec3<u32>; // M, K, N @compute @workgroup_size(16, 16) fn main(@builtin(global_invocation_id) id: vec3<u32>) { let m = id.x; let n = id.y; if (m >= dims.x || n >= dims.y) { return; } var sum: f32 = 0.0; for (var k: u32 = 0u; k < dims.z; k = k + 1u) { sum += a[m * dims.z + k] * b[k * dims.y + n]; } c[m * dims.y + n] = sum; }

2.3 WASM 与 WebGPU 的桥接

WASM 本身无法直接调用 WebGPU API——WebGPU 是浏览器的 JavaScript API，WASM 需要通过wasm-bindgen桥接到 JS 层调用。具体流程：Rust 代码调用web_syscrate（wasm-bindgen的 Web API 绑定），web_sys在编译时生成 JS 胶水代码，运行时 WASM 通过 JS 胶水代码调用浏览器的 WebGPU API。

这个桥接层有一定的性能开销：每次 GPU 调用都需要从 WASM 切换到 JS 再到浏览器引擎。但对于大模型推理，GPU 计算时间远大于桥接开销（毫秒级 vs 微秒级），影响可以忽略。

三、Rust 生产级代码实现

3.1 GPU 缓冲区管理

use wasm_bindgen::prelude::*; use web_sys::{ GpuDevice, GpuBuffer, GpuBufferDescriptor, GpuBufferUsage, }; /// GPU 缓冲区封装 pub struct GpuTensor { buffer: GpuBuffer, size: usize, shape: Vec<usize>, } impl GpuTensor { /// 创建 GPU 缓冲区并上传数据 pub fn from_data( device: &GpuDevice, data: &[f32], shape: Vec<usize>, usage: u32, ) -> Result<Self, JsValue> { let size = (data.len() * std::mem::size_of::<f32>()) as u64; let descriptor = GpuBufferDescriptor::new(); descriptor.set_size(size); descriptor.set_usage( usage | GpuBufferUsage::CopyDst as u32, ); let buffer = device.create_buffer(&descriptor); // 上传数据到 GPU let js_data = unsafe { js_sys::Float32Array::view(data) }; device.queue().write_buffer_with_f32_array_and_offset( &buffer, 0, &js_data, )?; Ok(Self { buffer, size: data.len(), shape, }) } /// 创建空的 GPU 缓冲区（用于输出） pub fn zeros( device: &GpuDevice, size: usize, shape: Vec<usize>, ) -> Result<Self, JsValue> { let byte_size = (size * std::mem::size_of::<f32>()) as u64; let descriptor = GpuBufferDescriptor::new(); descriptor.set_size(byte_size); descriptor.set_usage( (GpuBufferUsage::Storage as u32) | (GpuBufferUsage::CopySrc as u32) | (GpuBufferUsage::CopyDst as u32), ); let buffer = device.create_buffer(&descriptor); Ok(Self { buffer, size, shape, }) } pub fn buffer(&self) -> &GpuBuffer { &self.buffer } pub fn shape(&self) -> &[usize] { &self.shape } }

3.2 计算管线封装

use web_sys::{ GpuDevice, GpuComputePipeline, GpuPipelineLayout, GpuShaderModuleDescriptor, GpuBindGroupLayout, GpuBindGroupDescriptor, GpuBindGroup, }; /// 矩阵乘法计算管线 pub struct MatmulPipeline { pipeline: GpuComputePipeline, device: GpuDevice, } impl MatmulPipeline { /// 创建矩阵乘法管线 pub fn new(device: &GpuDevice) -> Result<Self, JsValue> { let shader_source = include_str!("shaders/matmul.wgsl"); let shader_descriptor = GpuShaderModuleDescriptor::new(); shader_descriptor.set_code(shader_source); let shader_module = device.create_shader_module(&shader_descriptor); let pipeline = device.create_compute_pipeline(&js_sys::Object::new()); // 简化：实际需要设置 pipeline layout 和 bind group layout // 完整实现需要配置 bind group layout 描述符 Ok(Self { pipeline, device: device.clone(), }) } /// 执行矩阵乘法: C = A × B pub async fn execute( &self, a: &GpuTensor, b: &GpuTensor, m: u32, k: u32, n: u32, ) -> Result<GpuTensor, JsValue> { // 创建输出缓冲区 let c = GpuTensor::zeros( &self.device, (m * n) as usize, vec![m as usize, n as usize], )?; // 创建 uniform 缓冲区（矩阵维度） let dims_data: [f32; 3] = [m as f32, k as f32, n as f32]; let dims_buffer = GpuTensor::from_data( &self.device, &dims_data, vec![3], GpuBufferUsage::Uniform as u32, )?; // 创建 bind group（绑定输入输出缓冲区） // 实际实现需要构造 GpuBindGroupDescriptor // 创建 command encoder 并 dispatch let encoder = self.device.create_command_encoder(); let compute_pass = encoder.begin_compute_pass(); // 设置管线和 bind group // compute_pass.set_pipeline(&self.pipeline); // compute_pass.set_bind_group(0, &bind_group, &[]); // compute_pass.dispatch_workgroups( // (m + 15) / 16, // workgroup 数量 // (n + 15) / 16, // 1, // ); // 结束 compute pass 并提交 // compute_pass.end(); // self.device.queue().submit(&[encoder.finish()]); // 等待 GPU 执行完成 // 实际需要通过 buffer map 或 readback 获取结果 Ok(c) } }

3.3 简易推理引擎

/// 浏览器端简易推理引擎 pub struct WasmLlmEngine { device: GpuDevice, matmul: MatmulPipeline, // 模型权重（GPU 缓冲区） q_weight: Option<GpuTensor>, k_weight: Option<GpuTensor>, v_weight: Option<GpuTensor>, o_weight: Option<GpuTensor>, hidden_dim: usize, num_heads: usize, head_dim: usize, } impl WasmLlmEngine { /// 初始化引擎：请求 GPU 设备 pub async fn new() -> Result<Self, JsValue> { let window = web_sys::window().unwrap(); let navigator = window.navigator(); let gpu = navigator.gpu(); let request_options = web_sys::GpuRequestAdapterOptions::new(); request_options.set_power_preference( web_sys::GpuPowerPreference::HighPerformance, ); let adapter = gpu.request_adapter(&request_options).await?; let device_descriptor = web_sys::GpuDeviceDescriptor::new(); let device = adapter.request_device(&device_descriptor).await?; let matmul = MatmulPipeline::new(&device)?; Ok(Self { device, matmul, q_weight: None, k_weight: None, v_weight: None, o_weight: None, hidden_dim: 0, num_heads: 0, head_dim: 0, }) } /// 加载模型权重 pub async fn load_weights( &mut self, weights_data: &[u8], hidden_dim: usize, num_heads: usize, ) -> Result<(), JsValue> { self.hidden_dim = hidden_dim; self.num_heads = num_heads; self.head_dim = hidden_dim / num_heads; // 解析权重数据并上传到 GPU // 简化：假设权重是连续的 f32 数组 let float_view = unsafe { std::slice::from_raw_parts( weights_data.as_ptr() as *const f32, weights_data.len() / std::mem::size_of::<f32>(), ) }; let weight_size = hidden_dim * hidden_dim; let storage_usage = GpuBufferUsage::Storage as u32; self.q_weight = Some(GpuTensor::from_data( &self.device, &float_view[0..weight_size], vec![hidden_dim, hidden_dim], storage_usage, )?); Ok(()) } /// 执行单步推理 pub async fn forward( &self, input_ids: &[u32], ) -> Result<Vec<u32>, JsValue> { // 简化实现：实际需要完整的 Transformer forward pass // 包括 embedding → self-attention → FFN → logits // 1. Embedding lookup // 2. Q/K/V 投影（矩阵乘法） // 3. 注意力计算 // 4. 输出投影 // 5. FFN // 6. Logits 采样 // 此处仅展示矩阵乘法调用 let _ = input_ids; Ok(vec![]) } }

四、Trade-offs：WASM + WebGPU 方案的局限

4.1 浏览器兼容性

WebGPU 截至 2025 年在 Chrome 113+ 和 Edge 113+ 中可用，Firefox 和 Safari 的支持仍在开发中。这意味着使用 WebGPU 的 WASM 应用无法在所有浏览器中运行。降级方案是使用 WebGL Compute（通过wgpucrate 的 WebGL 后端），但性能会大幅下降。

4.2 GPU 内存限制

浏览器的 WebGPU 实现对 GPU 内存有严格限制——通常不允许单个缓冲区超过 256MB，总 GPU 内存使用量也有限制。对于 7B 参数的模型（约 14GB FP16 权重），无法完整加载到浏览器 GPU 内存中。解决方案是模型量化（INT4/INT8）和分层加载（按层加载权重，计算完一层后释放）。