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_sys crate（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（通过 wgpu crate 的 WebGL 后端），但性能会大幅下降。
4.2 GPU 内存限制
浏览器的 WebGPU 实现对 GPU 内存有严格限制——通常不允许单个缓冲区超过 256MB，总 GPU 内存使用量也有限制。对于 7B 参数的模型（约 14GB FP16 权重），无法完整加载到浏览器 GPU 内存中。解决方案是模型量化（INT4/INT8）和分层加载（按层加载权重，计算完一层后释放）。
4.3 适用边界
WASM + WebGPU 推理适用于以下场景：小模型（<1B 参数）、对隐私要求高（数据不出浏览器）、需要离线推理能力。不适用于：大模型（>7B 参数，GPU 内存不足）、对推理速度要求极高（原生 GPU 推理快 5-10 倍）、需要跨浏览器兼容。
五、总结
WASM + WebGPU 让浏览器端 AI 推理从"概念验证"走向"可用体验"，但离"生产级"还有距离。核心落地步骤如下：
浏览器端推理的价值不在于替代服务端推理，而在于提供一种"数据不出浏览器"的隐私保护方案。对于小模型和特定场景，这个方案已经可用。
