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 推理从"概念验证"走向"可用体验"，但离"生产级"还有距离。核心落地步骤如下：

1. 初始化 WebGPU 设备：通过 navigator.gpu 请求适配器和设备，优先选择高性能 GPU。
2. 上传模型权重：将量化后的权重数据从 JS ArrayBuffer 上传到 GPU 缓冲区。
3. 构建计算管线：用 WGSL 编写矩阵乘法和注意力计算着色器，创建 Compute Pipeline。
4. 执行推理循环：Embedding → Attention → FFN → Logits，每步通过 GPU Compute Shader 加速。
5. 结果回读：通过 buffer map 或 readback 将 GPU 计算结果拷贝回 CPU/WASM。

浏览器端推理的价值不在于替代服务端推理，而在于提供一种"数据不出浏览器"的隐私保护方案。对于小模型和特定场景，这个方案已经可用。