浏览器内的推理引擎：WASM 端侧 AI 推理的架构与实现

浏览器内的推理引擎：WASM 端侧 AI 推理的架构与实现

一、云端推理的延迟困境：为什么 AI 需要跑到用户设备上

大模型推理依赖云端 GPU 集群，每次请求需要经过网络传输、排队等待和 GPU 计算三个阶段。在弱网环境下，仅网络往返延迟就可能超过 500ms，对于实时交互场景（语音助手、AR 滤镜、智能输入法）而言，这个延迟不可接受。此外，将用户数据上传到云端推理存在隐私风险——医疗影像、金融文档等敏感数据不应离开用户设备。

端侧推理是解决延迟和隐私问题的直接方案，但面临两个工程挑战：模型体积和计算性能。一个 7B 参数的量化模型仍需 4GB 存储，浏览器环境无法承受；而 JavaScript 的动态类型和 JIT 编译无法提供足够的计算吞吐。WebAssembly 提供了第三条路径：接近原生的执行性能、沙箱化的安全隔离、跨浏览器的标准化运行时。通过 WASM + WebGPU 的组合，可以在浏览器中运行量化后的轻量模型，实现端侧 AI 推理。

二、WASM 线性内存与计算沙箱：端侧推理的运行时基础

2.1 WASM 执行模型

WASM 模块运行在虚拟机提供的线性内存（Linear Memory）中，这是一块可增长的字节数组，所有数据通过偏移量访问。WASM 代码无法直接访问宿主环境（浏览器）的内存，所有交互通过导入/导出函数完成。这种沙箱模型保证了 WASM 代码的安全性——即使推理引擎存在漏洞，也无法越权访问用户数据。

graph TB subgraph 浏览器宿主环境 A[JavaScript 运行时] -->|调用导出函数| B[WASM 模块] B -->|写入线性内存| C[Linear Memory] A -->|读取线性内存| C B -->|调用导入函数| D[WebGL/WebGPU API] D -->|GPU 计算| E[GPU Shader] end subgraph WASM 沙箱内部 B --> F[模型权重<br/>f16 量化] B --> G[推理计算图] B --> H[中间激活 Buffer] F --> C G --> C H --> C end subgraph 数据流 I[用户输入] -->|JS 传入| C C -->|WASM 计算| G G -->|结果写回| C C -->|JS 读出| J[推理结果] end

2.2 Rust 到 WASM 的编译链

Rust 通过wasm32-unknown-unknown目标直接编译为 WASM 字节码，无需 C 中间层。wasm-bindgen工具生成 JS 绑定代码，处理类型映射（Rust 的Vec<T>对应 JS 的Float32Array）、内存管理和函数导出。wasm-pack将编译产物打包为 npm 包，支持直接在 Web 项目中引用。

2.3 内存管理的关键约束

WASM 线性内存的增长通过memory.grow指令实现，每次增长一个 Page（64KB）。频繁的内存增长会导致性能抖动，因此推理引擎应在初始化时预分配足够的内存。Rust 的默认分配器（dlmalloc）在 WASM 中可用但性能一般，推荐使用wee_alloc减小代码体积，或使用lol_alloc获得更快的分配速度。

三、端侧推理引擎的 WASM 实现

3.1 模型加载与量化推理

use wasm_bindgen::prelude::*; /// WASM 端侧推理引擎 /// 编译目标：wasm32-unknown-unknown #[wasm_bindgen] pub struct WasmInferenceEngine { // 模型权重，使用 f16 量化减少内存占用 weights: Vec<u16>, // 推理中间激活 Buffer，预分配避免运行时增长 activations: Vec<f32>, // 模型配置 hidden_size: usize, num_layers: usize, } #[wasm_bindgen] impl WasmInferenceEngine { /// 从二进制权重数据创建推理引擎 /// 权重在 JS 端通过 fetch 加载后传入 #[wasm_bindgen(constructor)] pub fn new( weights_data: &[u8], hidden_size: usize, num_layers: usize, ) -> Result<WasmInferenceEngine, JsValue> { // 预分配激活 Buffer，避免推理过程中的内存增长 let activation_size = hidden_size * 4; // 4 倍隐藏层大小用于中间结果 let activations = vec![0.0f32; activation_size * num_layers]; // 将字节流解码为 f16 权重 if weights_data.len() % 2 != 0 { return Err(JsValue::from_str("权重数据长度必须是偶数字节")); } let weights: Vec<u16> = weights_data .chunks_exact(2) .map(|chunk| u16::from_le_bytes([chunk[0], chunk[1]])) .collect(); Ok(Self { weights, activations, hidden_size, num_layers, }) } /// 执行单步推理 /// input_ids: 输入 token ID 序列 /// 返回下一个 token 的概率分布 pub fn forward(&mut self, input_ids: &[u32]) -> Result<Vec<f32>, JsValue> { if input_ids.is_empty() { return Err(JsValue::from_str("输入序列不能为空")); } let seq_len = input_ids.len(); let vocab_size = self.weights.len() / (self.hidden_size * self.num_layers); // Token Embedding：从权重中查找对应向量 // 使用 f16 到 f32 的转换进行计算 let mut hidden_state = vec![0.0f32; self.hidden_size]; for &token_id in input_ids { let embed_offset = token_id as usize * self.hidden_size; if embed_offset + self.hidden_size > self.weights.len() { return Err(JsValue::from_str("Token ID 超出词表范围")); } for i in 0..self.hidden_size { hidden_state[i] += f16_to_f32(self.weights[embed_offset + i]); } } // 逐层 Transformer 计算 for layer in 0..self.num_layers { let layer_offset = layer * self.hidden_size * 4; // QKV+O 四组权重 self.attention_forward( &mut hidden_state, layer_offset, )?; } // 输出投影：隐藏状态 -> 词表概率分布 let mut logits = vec![0.0f32; vocab_size]; // 矩阵乘法：hidden_state @ output_weight for v in 0..vocab_size.min(32000) { // 限制词表大小防止内存溢出 let weight_offset = v * self.hidden_size; let mut sum = 0.0f32; for h in 0..self.hidden_size { let w = f16_to_f32(self.weights[weight_offset + h]); sum += hidden_state[h] * w; } logits[v] = sum; } // Softmax 归一化 let max_logit = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max); let exp_sum: f32 = logits.iter().map(|&x| (x - max_logit).exp()).sum(); for logit in logits.iter_mut() { *logit = (*logit - max_logit).exp() / exp_sum; } Ok(logits) } /// 获取引擎的内存占用（字节） pub fn memory_usage(&self) -> usize { self.weights.len() * 2 + self.activations.len() * 4 } } impl WasmInferenceEngine { /// 单头注意力计算 fn attention_forward( &self, hidden: &mut [f32], weight_offset: usize, ) -> Result<(), JsValue> { let hs = self.hidden_size; if weight_offset + hs * 4 > self.weights.len() { return Err(JsValue::from_str("权重偏移量超出范围")); } // Q = hidden @ W_q let mut q = vec![0.0f32; hs]; for i in 0..hs { q[i] = hidden[i] * f16_to_f32(self.weights[weight_offset + i]); } // K = hidden @ W_k let mut k = vec![0.0f32; hs]; for i in 0..hs { k[i] = hidden[i] * f16_to_f32(self.weights[weight_offset + hs + i]); } // V = hidden @ W_v let mut v = vec![0.0f32; hs]; for i in 0..hs { v[i] = hidden[i] * f16_to_f32(self.weights[weight_offset + hs * 2 + i]); } // Attention: softmax(Q * K^T / sqrt(d)) * V let scale = (hs as f32).sqrt().recip(); let mut attn_weights = vec![0.0f32; hs]; let mut max_val = f32::NEG_INFINITY; for i in 0..hs { attn_weights[i] = q[i] * k[i] * scale; max_val = max_val.max(attn_weights[i]); } let exp_sum: f32 = attn_weights.iter().map(|&x| (x - max_val).exp()).sum(); for w in attn_weights.iter_mut() { *w = (*w - max_val).exp() / exp_sum; } // 输出 = attn_weights * V for i in 0..hs { hidden[i] = attn_weights[i] * v[i]; } Ok(()) } } /// IEEE 754 半精度浮点数转单精度 /// 在 WASM 中使用软件模拟，因为 WASM 规范不原生支持 f16 fn f16_to_f16(half: u16) -> f32 { let sign = (half >> 15) & 1; let exponent = (half >> 10) & 0x1F; let mantissa = half & 0x3FF; match exponent { 0 => { if mantissa == 0 { // 零 if sign == 1 { -0.0 } else { 0.0 } } else { // 非正规数 let f = (mantissa as f32) * 2.0f32.powi(-24); if sign == 1 { -f } else { f } } } 31 => { if mantissa == 0 { // 无穷 f32::INFINITY.copysign(if sign == 1 { -1.0 } else { 1.0 }) } else { // NaN f32::NAN } } _ => { // 正规数 let f = 2.0f32.powi(exponent as i32 - 15) * (1.0 + mantissa as f32 / 1024.0); if sign == 1 { -f } else { f } } } } // 修正函数名 fn f16_to_f32(half: u16) -> f32 { f16_to_f16(half) }

3.2 JS 端集成与内存交互

// JavaScript 端加载与调用 import init, { WasmInferenceEngine } from './pkg/inference_engine.js'; async function runInference() { // 初始化 WASM 模块 await init(); // 从服务器加载量化模型权重 const response = await fetch('/models/tinyllama-q4.wasm'); const weightsBuffer = await response.arrayBuffer(); const weightsData = new Uint8Array(weightsBuffer); // 创建推理引擎 const engine = new WasmInferenceEngine( weightsData, 512, // hidden_size 4, // num_layers ); // 执行推理 const inputIds = new Uint32Array([1, 1543, 566, 1024]); const logits = engine.forward(inputIds); // 解码输出：取概率最高的 token let maxIdx = 0; let maxVal = -Infinity; for (let i = 0; i < logits.length; i++) { if (logits[i] > maxVal) { maxVal = logits[i]; maxIdx = i; } } console.log(`预测下一个 token: ${maxIdx}, 概率: ${maxVal}`); console.log(`引擎内存占用: ${engine.memory_usage()} 字节`); }

四、WASM 推理的性能边界与架构妥协

WASM 端侧推理存在几个硬性限制，需要在架构设计时充分考虑。

计算性能天花板。WASM 的计算性能约为原生代码的 70%-90%（纯 CPU 计算密集型场景），但缺乏 SIMD 的完整支持（WASM SIMD 128 目前已广泛支持，但 256/512 位宽尚在提案阶段）。对于 Transformer 模型中的矩阵乘法，SIMD 向量化是性能的关键，128 位 SIMD 意味着每次处理 4 个 f32，而 AVX-512 可以处理 16 个。这导致 WASM 推理的吞吐量约为原生代码的 1/3 到 1/2。

内存限制。浏览器对 WASM 线性内存有上限（Chrome 默认 4GB，Firefox 约 2GB）。一个 1.5B 参数的 Q4 量化模型需要约 1GB 权重内存，加上推理中间激活，总内存接近 2GB。更大的模型在浏览器中无法运行。解决方案是模型切片（将权重分片按需加载）或更激进的量化（Q2、1.5-bit），但都会影响推理质量。

多线程限制。WASM 的SharedArrayBuffer要求页面配置 COOP/COEP 安全头，许多网站无法满足。没有共享内存，WASM 只能使用单线程计算，无法利用多核 CPU。WebGPU 可以部分弥补这个缺陷，将矩阵运算卸载到 GPU，但 WebGPU 的计算 Shader 性能也低于原生 CUDA。

适用边界。WASM 端侧推理最适合以下场景：模型参数量 < 1.5B 的轻量级任务（文本分类、NER、小规模生成）、对延迟极度敏感且无法容忍网络往返的交互场景、隐私敏感数据不能离开用户设备的合规要求。不适合的场景包括：大参数量模型的生成任务、需要 GPU 级吞吐量的批量推理、对计算精度有严格要求的科学计算。

五、总结

WASM 端侧推理通过浏览器沙箱和线性内存模型，在保护用户隐私的前提下实现了接近原生的推理性能。本文从 WASM 执行模型出发，实现了基于 Rust 编译的端侧推理引擎，包含 f16 量化权重管理、注意力计算和 JS 端集成。落地路线建议：第一步，使用wasm-pack build --target web编译推理引擎为 npm 包，在本地浏览器中验证基础推理功能；第二步，对计算热点（矩阵乘法、Softmax）使用 WASM SIMD intrinsics 优化，通过wasm-simd特性门启用；第三步，引入 WebGPU 后端将矩阵运算卸载到 GPU，使用wgpucrate 的 WASM 后端；第四步，对模型进行 Q4 量化压缩，使用safetensors格式存储权重，通过fetch+ReadableStream实现流式加载。