1. 项目概述浏览器里的“本地大模型”最近在折腾一个前端项目想给用户加点智能交互比如让页面能理解用户输入的自然语言然后给出点个性化的反馈。一提到这个大家第一反应可能就是去调 OpenAI 的 API或者用国内的一些大模型服务。但转念一想这条路有几个绕不开的坎一是 API Key 得管着不能在前端代码里裸奔二是按 token 计费用户量一大成本就有点吓人三是网络请求有延迟体验上总感觉差那么点意思。就在琢磨有没有更“轻巧”的办法时我看到了一个挺有意思的思路直接在用户的浏览器里跑一个轻量级的大语言模型LLM。这听起来有点天方夜谭毕竟我们印象中的大模型动辄几十上百亿参数需要强大的 GPU 集群。但别忘了现在浏览器有了WebGPU这个新武器它能让网页应用直接调用设备的 GPU 进行计算。再加上模型量化技术的成熟一些经过高度压缩和优化的轻量级模型比如几亿参数的小模型已经具备了在浏览器环境中实时推理的潜力。这个项目的核心目标就是用 React 框架结合 WebGPU在大概 50 行核心代码的体量内实现一个在用户本地浏览器中运行、无需网络请求、没有 API 密钥和 token 费用的 LLM 交互 demo。它解决的不仅仅是成本和隐私问题更是一种全新的交互范式——将 AI 能力真正“下发”到终端实现零延迟、高隐私的智能体验。无论是做个人知识库助手、文档摘要工具还是嵌入到在线教育、创意写作等场景中都能开辟出新的可能性。2. 核心思路与技术选型解析2.1 为什么是 WebGPU 而不是 WebGL要实现浏览器内的 GPU 计算过去我们可能首先想到 WebGL。但 WebGL 本质上是为图形渲染设计的用它来做通用计算GPGPU就像用螺丝刀当锤子虽然也能凑合但效率低下且写法别扭。WebGPU 则是专门为现代 GPU 的通用计算和图形渲染设计的下一代 Web API。选择 WebGPU 的核心原因有三点计算着色器原生支持WebGPU 提供了更直接、高效的计算管线Compute Pipeline和计算着色器Compute Shader其编程模型如工作组、共享内存更贴近 CUDA 或 Metal 的计算范式能极大发挥 GPU 的并行计算能力这对于矩阵运算密集的 LLM 推理至关重要。更好的性能与能效WebGPU 的底层抽象更接近现代 GPU如 Vulkan, Metal, DirectX 12减少了驱动层的开销能更高效地利用硬件资源。在模型推理这种需要大量重复计算的任务上性能提升是数量级的。未来的标准WebGPU 是 W3C 正在推进的标准得到了主流浏览器Chrome, Edge, Firefox, Safari的积极支持。虽然目前可能还在实验阶段或需要手动开启但它代表了未来方向提前布局是值得的。注意截至当前WebGPU 在部分浏览器如 Chrome/Edge 113中已默认启用但在 Safari 或某些版本中可能需要用户在chrome://flags或about:config中手动开启#enable-unsafe-webgpu等标志。在生产环境中使用需要做好特性检测和降级方案。2.2 模型选择如何在浏览器中“装下”一个 LLM让一个完整的 GPT-3 或 LLaMA 在浏览器里跑是不现实的。我们的目标是选择一个足够小、但能力又不错的模型。这里的关键技术是模型量化Quantization。量化是指将模型参数从高精度如 32 位浮点数FP32转换为低精度如 8 位整数INT8 甚至 4 位表示。一个 FP32 模型量化成 INT8理论上模型大小能减少至 1/4内存占用和计算量也大幅下降而精度损失在可控范围内。对于这个项目一个理想的选择是类似Google 的 Gemma 2B或Meta 的 Llama 3.1 8B的量化版本如 GGUF 格式。但即使是 2B 参数的 INT4 量化模型文件也可能有几百 MB。因此我们还需要借助模型分片Model Sharding技术。我们可以将一个大模型文件按层或按权重矩阵切分成多个小文件在浏览器中运行时按需加载当前推理所需的片段而不是一次性加载整个模型。这需要模型格式和推理引擎的支持。在本项目中为了极致简化50行代码的约束我们实际上不会去加载一个完整的、多层的 Transformer 模型。我们会选择一个极简的、演示性质的“模型”它可能只是一个简单的词袋模型、一个微型的神经网络如单层 LSTM 或 TinyBERT 的变体或者甚至是一个预先计算好的、固化权重的计算图。我们的核心目的是验证 WebGPU 进行模型推理的技术链路而非复现一个完整的 ChatGPT。因此我们会准备一个非常小的、权重已内置在代码或一个微小二进制文件中的计算图。2.3 整体架构设计整个应用的架构可以清晰地分为三层React 前端层提供用户界面输入框、按钮、输出区域处理用户交互事件。WebGPU 计算层这是核心。负责初始化 WebGPU 设备、创建计算管线、将模型权重和数据用户输入编码后的向量送入 GPU执行计算着色器进行推理并读回结果。模型与数据处理层包含一个极简的“模型”权重和计算图定义以及一个简单的分词器Tokenizer将文本转换为模型能理解的数字 IDToken IDs。由于模型极小分词器可能只是一个简单的字典映射。数据流如下用户输入文本 - 前端调用分词器得到 Token IDs - 将 Token IDs 和模型权重数据通过 WebGPU 提交到 GPU - GPU 执行计算着色器完成前向传播 - 取回计算结果下一个 token 的概率分布- 前端根据概率采样得到下一个 token并循环此过程生成完整回复。3. 环境准备与核心依赖3.1 创建 React 项目我们使用 Vite 来快速搭建一个 React TypeScript 项目它更轻量配置更简单。npm create vitelatest browser-llm-demo -- --template react-ts cd browser-llm-demo npm install3.2 关键依赖分析在这个 demo 中我们刻意避免使用庞大的机器学习框架如 TensorFlow.js 或 ONNX Runtime Web以保持代码的纯粹性和精简性。我们将直接使用 WebGPU 的原生 API。webgpu/types(可选但推荐)提供 TypeScript 类型定义让编码时有更好的提示。npm install --save-dev webgpu/types无其他主要依赖是的我们不需要tensorflow/tfjs或onnxruntime-web。所有计算逻辑都将通过我们手写的 WebGPU 代码完成。3.3 启用 WebGPU 支持确保你的浏览器支持 WebGPU。最新版的 Chrome/Edge 通常已支持。在地址栏输入chrome://flags搜索 “WebGPU”确保其状态为 “Enabled”。4. 核心实现50行代码拆解接下来我们进入最核心的部分。我们将创建一个WebGPULLM组件它包含了从初始化到推理的完整逻辑。我会将代码分成几个逻辑块并详细解释。4.1 初始化 WebGPU 设备与计算管线这是与 GPU 通信的起点。我们需要获取适配器、设备并创建计算管线。// WebGPULLM.tsx 核心部分 import React, { useState, useRef } from react; const WebGPULLM: React.FC () { const [input, setInput] useState(); const [output, setOutput] useState(); const [isLoading, setIsLoading] useState(false); const computePipelineRef useRefGPUComputePipeline | null(null); // 初始化 WebGPU const initWebGPU async (): PromiseGPUDevice { if (!navigator.gpu) { throw new Error(WebGPU not supported on this browser.); } const adapter await navigator.gpu.requestAdapter(); if (!adapter) { throw new Error(No appropriate GPUAdapter found.); } const device await adapter.requestDevice(); return device; }; // 创建计算管线这里定义了GPU上运行的计算内核Shader const createComputePipeline async (device: GPUDevice) { // 这是一个极度简化的“模型”计算着色器。 // 实际模型会是复杂的矩阵乘法和激活函数。 // 这里模拟一个操作 output input * weight bias const shaderModule device.createShaderModule({ code: group(0) binding(0) varstorage, read input: arrayf32; group(0) binding(1) varstorage, read weight: arrayf32; group(0) binding(2) varstorage, read bias: arrayf32; group(0) binding(3) varstorage, read_write output: arrayf32; compute workgroup_size(64) fn main(builtin(global_invocation_id) global_id: vec3u32) { let index global_id.x; if (index arrayLength(output)) { // 简化计算实际是复杂的矩阵运算 output[index] input[index] * weight[index] bias[index]; // 这里可以加入激活函数例如output[index] max(0.0, output[index]); } } , }); const pipeline device.createComputePipeline({ layout: auto, compute: { module: shaderModule, entryPoint: main, }, }); computePipelineRef.current pipeline; return pipeline; };代码解读initWebGPU: 标准流程获取 GPU 适配器和设备。requestDevice()是我们与 GPU 对话的主要接口。createComputePipeline: 这里定义了计算着色器WGSL 语言。我们声明了四个存储缓冲区Storage Buffer分别绑定到输入、权重、偏置和输出。workgroup_size(64)定义了一个工作组有 64 个线程并行执行。main函数是每个线程的入口global_id.x是当前线程的全局索引。这个简化内核只是做了逐元素的乘加运算。在一个真实的 LLM 推理中这里会是复杂的矩阵乘法MatMul、LayerNorm、Softmax 等操作的组合。4.2 准备模型数据与执行推理假设我们有一个“模型”它只有一层权重和偏置都是已知的小数组。我们将它们、用户输入编码后的数据一起送到 GPU 进行计算。// 执行推理 const runInference async (device: GPUDevice, pipeline: GPUComputePipeline, userInput: string) { setIsLoading(true); try { // 1. 模拟分词和嵌入将输入文本转换为数字向量。这里极度简化。 // 实际会使用分词器得到 token ids再通过嵌入层查找表得到向量。 const inputText userInput.toLowerCase(); const vocab: Recordstring, number { hello: 0.1, world: 0.2, ai: 0.3 }; // 模拟嵌入 const inputArray new Float32Array(64).fill(0); // 假设我们的“模型”接受长度为64的向量 inputText.split( ).forEach(word { if (vocab[word]) inputArray[0] vocab[word]; // 简单累加毫无道理仅为演示 }); inputArray[0] Math.tanh(inputArray[0]); // 模拟一个激活 // 2. 定义“模型权重”和偏置这里全是伪造的随机数 const weightArray new Float32Array(64); const biasArray new Float32Array(64); for (let i 0; i 64; i) { weightArray[i] Math.random() * 0.1; biasArray[i] Math.random() * 0.05; } const outputArray new Float32Array(64); // 3. 创建GPU缓冲区并写入数据 const inputBuffer device.createBuffer({ size: inputArray.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, }); device.queue.writeBuffer(inputBuffer, 0, inputArray); const weightBuffer device.createBuffer({ size: weightArray.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, }); device.queue.writeBuffer(weightBuffer, 0, weightArray); const biasBuffer device.createBuffer({ size: biasArray.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, }); device.queue.writeBuffer(biasBuffer, 0, biasArray); const outputBuffer device.createBuffer({ size: outputArray.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, }); // 4. 创建绑定组Bind Group将缓冲区与着色器中的 binding 点位关联 const bindGroup device.createBindGroup({ layout: pipeline.getBindGroupLayout(0), entries: [ { binding: 0, resource: { buffer: inputBuffer } }, { binding: 1, resource: { buffer: weightBuffer } }, { binding: 2, resource: { buffer: biasBuffer } }, { binding: 3, resource: { buffer: outputBuffer } }, ], }); // 5. 录制并提交命令到命令队列 const commandEncoder device.createCommandEncoder(); const passEncoder commandEncoder.beginComputePass(); passEncoder.setPipeline(pipeline); passEncoder.setBindGroup(0, bindGroup); passEncoder.dispatchWorkgroups(Math.ceil(64 / 64)); // 计算需要多少个工作组 (数据大小 / 工作组大小) passEncoder.end(); // 6. 从输出缓冲区读取数据回CPU const readbackBuffer device.createBuffer({ size: outputArray.byteLength, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ, }); commandEncoder.copyBufferToBuffer(outputBuffer, 0, readbackBuffer, 0, outputArray.byteLength); device.queue.submit([commandEncoder.finish()]); await readbackBuffer.mapAsync(GPUMapMode.READ); const resultData new Float32Array(readbackBuffer.getMappedRange().slice(0)); readbackBuffer.unmap(); // 7. 后处理将结果向量“解码”成文本这里极度简化直接取第一个值并映射 const simulatedResponse resultData[0] 0.5 ? I sense positive vibes about ${userInput}. : Hmm, ${userInput} is interesting.; setOutput(simulatedResponse); } catch (error) { console.error(Inference error:, error); setOutput(Error: ${error instanceof Error ? error.message : Unknown error}); } finally { setIsLoading(false); } };代码解读数据准备我们模拟了 NLP 流程的冰山一角。vocab模拟了一个嵌入表将单词映射为数字。inputArray模拟了经过简单处理的输入向量。weightArray和biasArray是我们的“模型参数”。GPU 缓冲区createBuffer创建了在 GPU 上存储数据的区域。usage标志非常重要它决定了缓冲区可以用于什么操作存储、复制目的地、复制源等。绑定组这是 WebGPU 中连接 CPU 端数据缓冲区和 GPU 端着色器通过 binding 点位的桥梁。getBindGroupLayout(0)从管线获取布局信息。调度计算dispatchWorkgroups(1)启动计算。参数1表示我们启动 1 个工作组。因为我们数据大小是 64工作组大小也是 64所以刚好一个工作组处理所有数据。在真实场景中数据量巨大这里会调度成百上千个工作组。读回结果这是一个典型的三步操作创建一个可映射读取的缓冲区readbackBuffer- 使用命令编码器将outputBuffer的数据复制到readbackBuffer- 提交命令并异步映射读取数据。这是 WebGPU 中 CPU 与 GPU 同步的关键点。4.3 整合与UI交互最后我们将初始化、推理和 React 的 UI 事件绑定起来。const handleSubmit async (e: React.FormEvent) { e.preventDefault(); if (!input.trim() || isLoading) return; try { const device await initWebGPU(); let pipeline computePipelineRef.current; if (!pipeline) { pipeline await createComputePipeline(device); } await runInference(device, pipeline, input); } catch (error) { console.error(Failed to run inference:, error); setOutput(WebGPU initialization or inference failed. Check console and browser support.); } }; return ( div style{{ padding: 20px, fontFamily: sans-serif }} h1Browser LLM via WebGPU (Demo)/h1 pstrongNote:/strong This is a simulation. The “model” is a single, tiny tensor operation./p form onSubmit{handleSubmit} input typetext value{input} onChange{(e) setInput(e.target.value)} placeholderType something (e.g., hello ai) disabled{isLoading} style{{ padding: 10px, width: 300px, marginRight: 10px }} / button typesubmit disabled{isLoading} {isLoading ? Thinking... : Ask} /button /form {output ( div style{{ marginTop: 20px, padding: 15px, border: 1px solid #ccc, borderRadius: 5px }} strongResponse:/strong {output} /div )} /div ); }; export default WebGPULLM;在App.tsx中引入并使用这个组件即可。至此一个完整的、在浏览器中利用 WebGPU 进行“模型推理”的 React 应用就完成了。核心逻辑初始化、管线创建、数据准备、调度、读取确实被浓缩在了一个主要函数中算上必要的状态和 UI整体思路在 50 行左右的概念框架内是清晰的。5. 从演示到真实应用差距与进阶路径上面的代码是一个“概念验证”它跑通了 WebGPU 计算管线。但要运行一个真正的、有用的轻量级 LLM如 Phi-2, TinyLlama 的量化版我们还需要跨越以下几个鸿沟5.1 引入专业的推理引擎手动编写整个 Transformer 的 WGSL 着色器是不现实的。我们需要借助成熟的推理引擎它们已经实现了高效的 GPU 内核。ONNX Runtime Web with WebGPU backend: ONNX Runtime 是一个高性能推理引擎其 Web 版本正在积极集成 WebGPU 后端。我们可以将模型导出为 ONNX 格式然后使用onnxruntime-web库在浏览器中加载并运行它底层会调用 WebGPU API。Transformers.js: 这是一个直接在浏览器中运行 Hugging Face 模型的 JavaScript 库。它使用 WebGL 或 WebGPU实验性作为后端。随着 WebGPU 的成熟它的支持会越来越好。专门为 Web 优化的运行时例如web-llm等项目它提供了完整的端到端方案包括模型加载、分词、推理基于 WebGPU支持 Llama 等系列模型的量化版本。进阶步骤将项目依赖改为onnxruntime-web。准备一个 ONNX 格式的量化小模型如phi-2-int4.onnx。使用 ONNX Runtime 的 API 加载模型并创建 WebGPU 会话。将用户输入通过分词器处理后喂给模型进行推理。5.2 处理真实的模型文件与分词真实模型可能很大即使量化后也有几十到几百 MB。我们需要模型分片与懒加载将模型文件切成多个小块在推理过程中动态加载所需的分片。这需要模型格式如 GGUF和推理引擎的支持。集成分词器需要一个 JavaScript 实现的分词器如bert-tokenizer或gpt-tokenizer将文本转换为 token IDs。许多推理引擎会内置或提供配套的分词器。5.3 性能优化与内存管理流水线优化将数据上传、计算、结果读取等操作重叠进行减少 GPU 空闲时间。缓冲区复用避免频繁创建和销毁 GPU 缓冲区创建缓冲区池进行复用。精度与速度权衡根据模型和硬件能力选择 FP16 甚至 INT8 的推理以进一步提升速度。6. 常见问题与调试技巧“navigator.gpu is undefined” 错误原因浏览器不支持或未启用 WebGPU。解决检查浏览器版本Chrome 113并前往chrome://flags确保#enable-unsafe-webgpu已启用。在代码中做好特性检测提供友好的降级提示。“Failed to create buffer” 或内存不足错误原因尝试分配超过设备限制的缓冲区大小。不同 GPU 的设备限制不同。解决查询设备限制device.limits根据maxStorageBufferBindingSize等限制来规划模型大小和数据分块。对于大模型分片加载是必须的。着色器编译错误原因WGSL 代码语法错误或者使用了目标设备不支持的特性。解决仔细检查 WGSL 代码。使用device.createShaderModule()后可以调用shaderModule.compilationInfo()异步获取详细的编译信息和错误消息这是调试着色器的关键工具。推理结果全是 0 或 NaN原因最常见的原因是数据没有正确地从 CPU 传到 GPU或者着色器中的计算逻辑有误如访问越界。调试使用 RenderDoc 或 WebGPU 开发者工具这些工具可以捕获一帧的 WebGPU 调用检查缓冲区的实际内容是终极调试手段。简化与验证首先用极小的、已知的数据在着色器中做一个简单的加法验证整个数据通路是否正确。再逐步替换为真实的模型计算。检查绑定组确保绑定组Bind Group的binding索引与着色器中binding(X)的声明一一对应。性能不及预期原因工作组大小设置不合理内存访问模式低效如未利用好局部内存或者 CPU-GPU 同步开销太大。优化工作组大小尝试不同的workgroup_size如 64, 128, 256这是一个经验值需要针对特定硬件和算法进行微调。减少读回尽量避免每一层推理后都从 GPU 读回数据。尽可能在 GPU 上完成所有计算只在最终需要结果时才读回。使用计算通道Compute Pass高效录制命令将多个计算任务录制到同一个计算通道中一次性提交减少提交开销。这个项目就像打开了一扇新世界的大门。虽然我们起步于一个仅 50 行核心代码的“玩具”但它清晰地勾勒出了未来 Web AI 应用的蓝图去中心化、高隐私、低延迟。真正的挑战和乐趣在于如何将那个简陋的output[index] input[index] * weight[index] bias[index]替换成真正的 Transformer 层并看着它在你自己的浏览器里无需连接任何远程服务器流畅地生成文字。这不仅仅是技术上的尝试更是一种对 AI 应用形态的重新想象。
基于WebGPU的浏览器端轻量级大语言模型推理实践
1. 项目概述浏览器里的“本地大模型”最近在折腾一个前端项目想给用户加点智能交互比如让页面能理解用户输入的自然语言然后给出点个性化的反馈。一提到这个大家第一反应可能就是去调 OpenAI 的 API或者用国内的一些大模型服务。但转念一想这条路有几个绕不开的坎一是 API Key 得管着不能在前端代码里裸奔二是按 token 计费用户量一大成本就有点吓人三是网络请求有延迟体验上总感觉差那么点意思。就在琢磨有没有更“轻巧”的办法时我看到了一个挺有意思的思路直接在用户的浏览器里跑一个轻量级的大语言模型LLM。这听起来有点天方夜谭毕竟我们印象中的大模型动辄几十上百亿参数需要强大的 GPU 集群。但别忘了现在浏览器有了WebGPU这个新武器它能让网页应用直接调用设备的 GPU 进行计算。再加上模型量化技术的成熟一些经过高度压缩和优化的轻量级模型比如几亿参数的小模型已经具备了在浏览器环境中实时推理的潜力。这个项目的核心目标就是用 React 框架结合 WebGPU在大概 50 行核心代码的体量内实现一个在用户本地浏览器中运行、无需网络请求、没有 API 密钥和 token 费用的 LLM 交互 demo。它解决的不仅仅是成本和隐私问题更是一种全新的交互范式——将 AI 能力真正“下发”到终端实现零延迟、高隐私的智能体验。无论是做个人知识库助手、文档摘要工具还是嵌入到在线教育、创意写作等场景中都能开辟出新的可能性。2. 核心思路与技术选型解析2.1 为什么是 WebGPU 而不是 WebGL要实现浏览器内的 GPU 计算过去我们可能首先想到 WebGL。但 WebGL 本质上是为图形渲染设计的用它来做通用计算GPGPU就像用螺丝刀当锤子虽然也能凑合但效率低下且写法别扭。WebGPU 则是专门为现代 GPU 的通用计算和图形渲染设计的下一代 Web API。选择 WebGPU 的核心原因有三点计算着色器原生支持WebGPU 提供了更直接、高效的计算管线Compute Pipeline和计算着色器Compute Shader其编程模型如工作组、共享内存更贴近 CUDA 或 Metal 的计算范式能极大发挥 GPU 的并行计算能力这对于矩阵运算密集的 LLM 推理至关重要。更好的性能与能效WebGPU 的底层抽象更接近现代 GPU如 Vulkan, Metal, DirectX 12减少了驱动层的开销能更高效地利用硬件资源。在模型推理这种需要大量重复计算的任务上性能提升是数量级的。未来的标准WebGPU 是 W3C 正在推进的标准得到了主流浏览器Chrome, Edge, Firefox, Safari的积极支持。虽然目前可能还在实验阶段或需要手动开启但它代表了未来方向提前布局是值得的。注意截至当前WebGPU 在部分浏览器如 Chrome/Edge 113中已默认启用但在 Safari 或某些版本中可能需要用户在chrome://flags或about:config中手动开启#enable-unsafe-webgpu等标志。在生产环境中使用需要做好特性检测和降级方案。2.2 模型选择如何在浏览器中“装下”一个 LLM让一个完整的 GPT-3 或 LLaMA 在浏览器里跑是不现实的。我们的目标是选择一个足够小、但能力又不错的模型。这里的关键技术是模型量化Quantization。量化是指将模型参数从高精度如 32 位浮点数FP32转换为低精度如 8 位整数INT8 甚至 4 位表示。一个 FP32 模型量化成 INT8理论上模型大小能减少至 1/4内存占用和计算量也大幅下降而精度损失在可控范围内。对于这个项目一个理想的选择是类似Google 的 Gemma 2B或Meta 的 Llama 3.1 8B的量化版本如 GGUF 格式。但即使是 2B 参数的 INT4 量化模型文件也可能有几百 MB。因此我们还需要借助模型分片Model Sharding技术。我们可以将一个大模型文件按层或按权重矩阵切分成多个小文件在浏览器中运行时按需加载当前推理所需的片段而不是一次性加载整个模型。这需要模型格式和推理引擎的支持。在本项目中为了极致简化50行代码的约束我们实际上不会去加载一个完整的、多层的 Transformer 模型。我们会选择一个极简的、演示性质的“模型”它可能只是一个简单的词袋模型、一个微型的神经网络如单层 LSTM 或 TinyBERT 的变体或者甚至是一个预先计算好的、固化权重的计算图。我们的核心目的是验证 WebGPU 进行模型推理的技术链路而非复现一个完整的 ChatGPT。因此我们会准备一个非常小的、权重已内置在代码或一个微小二进制文件中的计算图。2.3 整体架构设计整个应用的架构可以清晰地分为三层React 前端层提供用户界面输入框、按钮、输出区域处理用户交互事件。WebGPU 计算层这是核心。负责初始化 WebGPU 设备、创建计算管线、将模型权重和数据用户输入编码后的向量送入 GPU执行计算着色器进行推理并读回结果。模型与数据处理层包含一个极简的“模型”权重和计算图定义以及一个简单的分词器Tokenizer将文本转换为模型能理解的数字 IDToken IDs。由于模型极小分词器可能只是一个简单的字典映射。数据流如下用户输入文本 - 前端调用分词器得到 Token IDs - 将 Token IDs 和模型权重数据通过 WebGPU 提交到 GPU - GPU 执行计算着色器完成前向传播 - 取回计算结果下一个 token 的概率分布- 前端根据概率采样得到下一个 token并循环此过程生成完整回复。3. 环境准备与核心依赖3.1 创建 React 项目我们使用 Vite 来快速搭建一个 React TypeScript 项目它更轻量配置更简单。npm create vitelatest browser-llm-demo -- --template react-ts cd browser-llm-demo npm install3.2 关键依赖分析在这个 demo 中我们刻意避免使用庞大的机器学习框架如 TensorFlow.js 或 ONNX Runtime Web以保持代码的纯粹性和精简性。我们将直接使用 WebGPU 的原生 API。webgpu/types(可选但推荐)提供 TypeScript 类型定义让编码时有更好的提示。npm install --save-dev webgpu/types无其他主要依赖是的我们不需要tensorflow/tfjs或onnxruntime-web。所有计算逻辑都将通过我们手写的 WebGPU 代码完成。3.3 启用 WebGPU 支持确保你的浏览器支持 WebGPU。最新版的 Chrome/Edge 通常已支持。在地址栏输入chrome://flags搜索 “WebGPU”确保其状态为 “Enabled”。4. 核心实现50行代码拆解接下来我们进入最核心的部分。我们将创建一个WebGPULLM组件它包含了从初始化到推理的完整逻辑。我会将代码分成几个逻辑块并详细解释。4.1 初始化 WebGPU 设备与计算管线这是与 GPU 通信的起点。我们需要获取适配器、设备并创建计算管线。// WebGPULLM.tsx 核心部分 import React, { useState, useRef } from react; const WebGPULLM: React.FC () { const [input, setInput] useState(); const [output, setOutput] useState(); const [isLoading, setIsLoading] useState(false); const computePipelineRef useRefGPUComputePipeline | null(null); // 初始化 WebGPU const initWebGPU async (): PromiseGPUDevice { if (!navigator.gpu) { throw new Error(WebGPU not supported on this browser.); } const adapter await navigator.gpu.requestAdapter(); if (!adapter) { throw new Error(No appropriate GPUAdapter found.); } const device await adapter.requestDevice(); return device; }; // 创建计算管线这里定义了GPU上运行的计算内核Shader const createComputePipeline async (device: GPUDevice) { // 这是一个极度简化的“模型”计算着色器。 // 实际模型会是复杂的矩阵乘法和激活函数。 // 这里模拟一个操作 output input * weight bias const shaderModule device.createShaderModule({ code: group(0) binding(0) varstorage, read input: arrayf32; group(0) binding(1) varstorage, read weight: arrayf32; group(0) binding(2) varstorage, read bias: arrayf32; group(0) binding(3) varstorage, read_write output: arrayf32; compute workgroup_size(64) fn main(builtin(global_invocation_id) global_id: vec3u32) { let index global_id.x; if (index arrayLength(output)) { // 简化计算实际是复杂的矩阵运算 output[index] input[index] * weight[index] bias[index]; // 这里可以加入激活函数例如output[index] max(0.0, output[index]); } } , }); const pipeline device.createComputePipeline({ layout: auto, compute: { module: shaderModule, entryPoint: main, }, }); computePipelineRef.current pipeline; return pipeline; };代码解读initWebGPU: 标准流程获取 GPU 适配器和设备。requestDevice()是我们与 GPU 对话的主要接口。createComputePipeline: 这里定义了计算着色器WGSL 语言。我们声明了四个存储缓冲区Storage Buffer分别绑定到输入、权重、偏置和输出。workgroup_size(64)定义了一个工作组有 64 个线程并行执行。main函数是每个线程的入口global_id.x是当前线程的全局索引。这个简化内核只是做了逐元素的乘加运算。在一个真实的 LLM 推理中这里会是复杂的矩阵乘法MatMul、LayerNorm、Softmax 等操作的组合。4.2 准备模型数据与执行推理假设我们有一个“模型”它只有一层权重和偏置都是已知的小数组。我们将它们、用户输入编码后的数据一起送到 GPU 进行计算。// 执行推理 const runInference async (device: GPUDevice, pipeline: GPUComputePipeline, userInput: string) { setIsLoading(true); try { // 1. 模拟分词和嵌入将输入文本转换为数字向量。这里极度简化。 // 实际会使用分词器得到 token ids再通过嵌入层查找表得到向量。 const inputText userInput.toLowerCase(); const vocab: Recordstring, number { hello: 0.1, world: 0.2, ai: 0.3 }; // 模拟嵌入 const inputArray new Float32Array(64).fill(0); // 假设我们的“模型”接受长度为64的向量 inputText.split( ).forEach(word { if (vocab[word]) inputArray[0] vocab[word]; // 简单累加毫无道理仅为演示 }); inputArray[0] Math.tanh(inputArray[0]); // 模拟一个激活 // 2. 定义“模型权重”和偏置这里全是伪造的随机数 const weightArray new Float32Array(64); const biasArray new Float32Array(64); for (let i 0; i 64; i) { weightArray[i] Math.random() * 0.1; biasArray[i] Math.random() * 0.05; } const outputArray new Float32Array(64); // 3. 创建GPU缓冲区并写入数据 const inputBuffer device.createBuffer({ size: inputArray.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, }); device.queue.writeBuffer(inputBuffer, 0, inputArray); const weightBuffer device.createBuffer({ size: weightArray.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, }); device.queue.writeBuffer(weightBuffer, 0, weightArray); const biasBuffer device.createBuffer({ size: biasArray.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, }); device.queue.writeBuffer(biasBuffer, 0, biasArray); const outputBuffer device.createBuffer({ size: outputArray.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC, }); // 4. 创建绑定组Bind Group将缓冲区与着色器中的 binding 点位关联 const bindGroup device.createBindGroup({ layout: pipeline.getBindGroupLayout(0), entries: [ { binding: 0, resource: { buffer: inputBuffer } }, { binding: 1, resource: { buffer: weightBuffer } }, { binding: 2, resource: { buffer: biasBuffer } }, { binding: 3, resource: { buffer: outputBuffer } }, ], }); // 5. 录制并提交命令到命令队列 const commandEncoder device.createCommandEncoder(); const passEncoder commandEncoder.beginComputePass(); passEncoder.setPipeline(pipeline); passEncoder.setBindGroup(0, bindGroup); passEncoder.dispatchWorkgroups(Math.ceil(64 / 64)); // 计算需要多少个工作组 (数据大小 / 工作组大小) passEncoder.end(); // 6. 从输出缓冲区读取数据回CPU const readbackBuffer device.createBuffer({ size: outputArray.byteLength, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ, }); commandEncoder.copyBufferToBuffer(outputBuffer, 0, readbackBuffer, 0, outputArray.byteLength); device.queue.submit([commandEncoder.finish()]); await readbackBuffer.mapAsync(GPUMapMode.READ); const resultData new Float32Array(readbackBuffer.getMappedRange().slice(0)); readbackBuffer.unmap(); // 7. 后处理将结果向量“解码”成文本这里极度简化直接取第一个值并映射 const simulatedResponse resultData[0] 0.5 ? I sense positive vibes about ${userInput}. : Hmm, ${userInput} is interesting.; setOutput(simulatedResponse); } catch (error) { console.error(Inference error:, error); setOutput(Error: ${error instanceof Error ? error.message : Unknown error}); } finally { setIsLoading(false); } };代码解读数据准备我们模拟了 NLP 流程的冰山一角。vocab模拟了一个嵌入表将单词映射为数字。inputArray模拟了经过简单处理的输入向量。weightArray和biasArray是我们的“模型参数”。GPU 缓冲区createBuffer创建了在 GPU 上存储数据的区域。usage标志非常重要它决定了缓冲区可以用于什么操作存储、复制目的地、复制源等。绑定组这是 WebGPU 中连接 CPU 端数据缓冲区和 GPU 端着色器通过 binding 点位的桥梁。getBindGroupLayout(0)从管线获取布局信息。调度计算dispatchWorkgroups(1)启动计算。参数1表示我们启动 1 个工作组。因为我们数据大小是 64工作组大小也是 64所以刚好一个工作组处理所有数据。在真实场景中数据量巨大这里会调度成百上千个工作组。读回结果这是一个典型的三步操作创建一个可映射读取的缓冲区readbackBuffer- 使用命令编码器将outputBuffer的数据复制到readbackBuffer- 提交命令并异步映射读取数据。这是 WebGPU 中 CPU 与 GPU 同步的关键点。4.3 整合与UI交互最后我们将初始化、推理和 React 的 UI 事件绑定起来。const handleSubmit async (e: React.FormEvent) { e.preventDefault(); if (!input.trim() || isLoading) return; try { const device await initWebGPU(); let pipeline computePipelineRef.current; if (!pipeline) { pipeline await createComputePipeline(device); } await runInference(device, pipeline, input); } catch (error) { console.error(Failed to run inference:, error); setOutput(WebGPU initialization or inference failed. Check console and browser support.); } }; return ( div style{{ padding: 20px, fontFamily: sans-serif }} h1Browser LLM via WebGPU (Demo)/h1 pstrongNote:/strong This is a simulation. The “model” is a single, tiny tensor operation./p form onSubmit{handleSubmit} input typetext value{input} onChange{(e) setInput(e.target.value)} placeholderType something (e.g., hello ai) disabled{isLoading} style{{ padding: 10px, width: 300px, marginRight: 10px }} / button typesubmit disabled{isLoading} {isLoading ? Thinking... : Ask} /button /form {output ( div style{{ marginTop: 20px, padding: 15px, border: 1px solid #ccc, borderRadius: 5px }} strongResponse:/strong {output} /div )} /div ); }; export default WebGPULLM;在App.tsx中引入并使用这个组件即可。至此一个完整的、在浏览器中利用 WebGPU 进行“模型推理”的 React 应用就完成了。核心逻辑初始化、管线创建、数据准备、调度、读取确实被浓缩在了一个主要函数中算上必要的状态和 UI整体思路在 50 行左右的概念框架内是清晰的。5. 从演示到真实应用差距与进阶路径上面的代码是一个“概念验证”它跑通了 WebGPU 计算管线。但要运行一个真正的、有用的轻量级 LLM如 Phi-2, TinyLlama 的量化版我们还需要跨越以下几个鸿沟5.1 引入专业的推理引擎手动编写整个 Transformer 的 WGSL 着色器是不现实的。我们需要借助成熟的推理引擎它们已经实现了高效的 GPU 内核。ONNX Runtime Web with WebGPU backend: ONNX Runtime 是一个高性能推理引擎其 Web 版本正在积极集成 WebGPU 后端。我们可以将模型导出为 ONNX 格式然后使用onnxruntime-web库在浏览器中加载并运行它底层会调用 WebGPU API。Transformers.js: 这是一个直接在浏览器中运行 Hugging Face 模型的 JavaScript 库。它使用 WebGL 或 WebGPU实验性作为后端。随着 WebGPU 的成熟它的支持会越来越好。专门为 Web 优化的运行时例如web-llm等项目它提供了完整的端到端方案包括模型加载、分词、推理基于 WebGPU支持 Llama 等系列模型的量化版本。进阶步骤将项目依赖改为onnxruntime-web。准备一个 ONNX 格式的量化小模型如phi-2-int4.onnx。使用 ONNX Runtime 的 API 加载模型并创建 WebGPU 会话。将用户输入通过分词器处理后喂给模型进行推理。5.2 处理真实的模型文件与分词真实模型可能很大即使量化后也有几十到几百 MB。我们需要模型分片与懒加载将模型文件切成多个小块在推理过程中动态加载所需的分片。这需要模型格式如 GGUF和推理引擎的支持。集成分词器需要一个 JavaScript 实现的分词器如bert-tokenizer或gpt-tokenizer将文本转换为 token IDs。许多推理引擎会内置或提供配套的分词器。5.3 性能优化与内存管理流水线优化将数据上传、计算、结果读取等操作重叠进行减少 GPU 空闲时间。缓冲区复用避免频繁创建和销毁 GPU 缓冲区创建缓冲区池进行复用。精度与速度权衡根据模型和硬件能力选择 FP16 甚至 INT8 的推理以进一步提升速度。6. 常见问题与调试技巧“navigator.gpu is undefined” 错误原因浏览器不支持或未启用 WebGPU。解决检查浏览器版本Chrome 113并前往chrome://flags确保#enable-unsafe-webgpu已启用。在代码中做好特性检测提供友好的降级提示。“Failed to create buffer” 或内存不足错误原因尝试分配超过设备限制的缓冲区大小。不同 GPU 的设备限制不同。解决查询设备限制device.limits根据maxStorageBufferBindingSize等限制来规划模型大小和数据分块。对于大模型分片加载是必须的。着色器编译错误原因WGSL 代码语法错误或者使用了目标设备不支持的特性。解决仔细检查 WGSL 代码。使用device.createShaderModule()后可以调用shaderModule.compilationInfo()异步获取详细的编译信息和错误消息这是调试着色器的关键工具。推理结果全是 0 或 NaN原因最常见的原因是数据没有正确地从 CPU 传到 GPU或者着色器中的计算逻辑有误如访问越界。调试使用 RenderDoc 或 WebGPU 开发者工具这些工具可以捕获一帧的 WebGPU 调用检查缓冲区的实际内容是终极调试手段。简化与验证首先用极小的、已知的数据在着色器中做一个简单的加法验证整个数据通路是否正确。再逐步替换为真实的模型计算。检查绑定组确保绑定组Bind Group的binding索引与着色器中binding(X)的声明一一对应。性能不及预期原因工作组大小设置不合理内存访问模式低效如未利用好局部内存或者 CPU-GPU 同步开销太大。优化工作组大小尝试不同的workgroup_size如 64, 128, 256这是一个经验值需要针对特定硬件和算法进行微调。减少读回尽量避免每一层推理后都从 GPU 读回数据。尽可能在 GPU 上完成所有计算只在最终需要结果时才读回。使用计算通道Compute Pass高效录制命令将多个计算任务录制到同一个计算通道中一次性提交减少提交开销。这个项目就像打开了一扇新世界的大门。虽然我们起步于一个仅 50 行核心代码的“玩具”但它清晰地勾勒出了未来 Web AI 应用的蓝图去中心化、高隐私、低延迟。真正的挑战和乐趣在于如何将那个简陋的output[index] input[index] * weight[index] bias[index]替换成真正的 Transformer 层并看着它在你自己的浏览器里无需连接任何远程服务器流畅地生成文字。这不仅仅是技术上的尝试更是一种对 AI 应用形态的重新想象。
