从零在浏览器里跑 DeepSeek-R1：WebGPU + Transformers.js 全链路实战（三）

从零在浏览器里跑 DeepSeek-R1：WebGPU + Transformers.js 全链路实战（三）

本文是系列第三篇，聚焦主线程（React UI）与 Worker 线程（LLM 推理）之间的完整通信链路，从用户按下回车到模型吐出第一个 token，再到流式输出、思考/回答分离、中断重置------把每行代码讲透。

---

一、前情回顾与本文定位

前两篇我们分别讲了：

• （一）环境搭建：WebGPU 浏览器兼容性检测、Transformers.js 的 ONNX Runtime Web 后端、Vite 工程化
• （二）模型加载 ：AutoTokenizer 和 AutoModelForCausalLM 的 from_pretrained 流程、单例模式管理、进度回调

本文（三）聚焦两个文件的核心交互 ------App.tsx（主线程 React UI）和 worker.js（Web Worker 中的 LLM 推理），把整条链路串起来。

---

二、整体架构：两个线程，一个协议

┌──────────────────────────────────────────────────────────────────┐
│                        浏览器标签页                                │
│                                                                   │
│  ┌─────────────────────────┐    postMessage    ┌────────────────┐ │
│  │     主线程 (React)       │ ←──────────────→ │  Worker 线程    │ │
│  │                         │                   │                │ │
│  │  App.tsx                │   {type, data}    │  worker.js     │ │
│  │  - UI 渲染              │                   │  - LLM 推理    │ │
│  │  - 用户交互              │                   │  - GPU 计算    │ │
│  │  - 状态管理              │                   │  - 无 DOM 访问  │ │
│  └─────────────────────────┘                   └────────────────┘ │
└──────────────────────────────────────────────────────────────────┘

为什么要用 Web Worker？因为 LLM 推理是计算密集型任务，如果放在主线程，UI 会完全卡死。Worker 跑在独立线程里，主线程的滚动、点击、动画不受影响。

两端通过 postMessage 通信。消息格式遵循一套状态机协议：

方向	type	含义	携带数据
主→Worker	check	检测 WebGPU 是否可用	---
主→Worker	load	开始下载模型	---
主→Worker	generate	启动 LLM 生成	data: messages[]
主→Worker	interrupt	用户点击停止	---
主→Worker	reset	重置对话	---
Worker→主	status 系列	加载进度/生成更新/完成/错误	见下文

---

三、主线程端：App.tsx 的状态机

3.1 核心状态一览

const [status, setStatus] = useState(null);
// null → "loading" → "ready"  三个状态驱动整个 UI

const [messages, setMessages] = useState([]);
// 对话历史 [{role, content}, ...]  遵循 OpenAI 消息格式

const [input, setInput] = useState("");
// 受控输入框文本

const [isRunning, setIsRunning] = useState(false);
// 是否正在生成中（控制按钮切换和输入禁用）

const worker = useRef(null);
// Worker 实例存在 ref 里，不触发重渲染

status 是全局状态机，驱动三个视图：

status = null      →  欢迎页（Logo + Load model 按钮）
status = "loading" →  进度条页（显示模型文件下载进度）
status = "ready"   →  聊天页（消息列表 + 输入框）

3.2 Worker 的诞生：useRef + useEffect

useEffect(() => {
    if (!worker.current) {
        worker.current = new Worker(
            new URL("./worker.js", import.meta.url),
            { type: "module" }
        );
        worker.current.postMessage({ type: "check" });
    }

    const onMessageReceived = (e) => { /* 处理 Worker 回传消息 */ };

    worker.current.addEventListener("message", onMessageReceived);
    return () => worker.current.removeEventListener("message", onMessageReceived);
}, []);  // ← 空依赖数组，只在组件挂载时执行一次

关键设计点：

• useRef 而非 useState：Worker 实例不需要驱动渲染，放在 ref 里避免不必要的重渲染
• type: "module" ：让 Worker 支持 ES Module（import/export），与 Vite 打包体系兼容
• [] 依赖：整个 effect 只跑一次，Worker 只创建一次，永远不被销毁重建
• 清理函数：组件卸载时移除事件监听，防止内存泄漏

3.3 消息接收：从 Worker 到 UI

const onMessageReceived = (e) => {
    switch (e.data.status) {
        case "loading":   // 模型加载中，更新状态文本
        case "initiate":  // 某个文件开始下载，初始化进度条
        case "progress":  // 某个文件下载中，更新进度百分比
        case "done":      // 某个文件下载完成，移除进度条
        case "ready":     // 全部就绪，进入聊天界面
        case "start":     // 开始生成了（流式第一个信号）
        case "update":    // 流式文本到达，追加到对话
        case "complete":  // 生成完毕
        case "error":     // 出错了
    }
};

loading → initiate → progress → done → ready 是一条链。其中 progress 和 done 为什么要用函数式更新 (prev) =>？

// ❌ 直接引用可能读到旧值
setProgressItems([...progressItems, e.data]);

// ✅ 函数式保证基于最新快照
setProgressItems((prev) => [...prev, e.data]);

因为多个模型文件并发下载 ，进度回调触发极快，React 的批量更新可能导致部分 progress 丢失。函数式写法每次追加都基于最新状态。

---

四、从回车到生成：触发的双保险

4.1 onEnter：用户按下回车

function onEnter(message) {
    setMessages((prev) => [...prev, { role: "user", content: message }]);
    setInput("");          // 清空输入框
    setIsRunning(true);    // 锁定 UI，发送按钮变停止按钮
}

四件事一气呵成：① 立即在聊天列表里追加用户气泡（秒级 UI 反馈），② 清空输入框，③ 锁住按钮防止重复发送，④ 等待 useEffect 触发下一步。

4.2 useEffect 双守卫：精准触发 generate

useEffect(() => {
    // 守卫①: 没有用户消息就不触发
    if (messages.filter((x) => x.role === "user").length === 0) {
        return;
    }
    // 守卫②: 最后一条已经是 AI 说的了就不触发
    if (messages.at(-1).role === "assistant") {
        return;
    }
    // 两个守卫都通过 → 最后一条是用户消息 → 触发生成
    worker.current.postMessage({ type: "generate", data: messages });
}, [messages]);

两个守卫防止了什么？

守卫	拦截场景	典型时刻
filter(user).length === 0	没有任何用户消息	页面刚加载、reset 后
.at(-1).role === "assistant"	最后一条是 AI 消息	流式输出中每一次 setMessages、生成完成后

守卫②是最关键的 。LLM 流式输出时，每收到一个文本片段，主线程都会 setMessages 更新 assistant 消息的内容。messages 是数组引用，每次 setMessages 都会创建新数组 → [messages] 依赖检测到变化 → effect 重新执行。如果没有守卫②，每次流式更新都会再发一次 generate，直接陷入死循环。

---

五、Worker 端：generate() 全链路拆解

5.1 消息路由

self.addEventListener("message", async (e) => {
    const { type, data } = e.data;
    switch (type) {
        case "check":     check();    break;
        case "load":      load();     break;
        case "generate":  stopping_criteria.reset();
                          generate(data);  break;
        case "interrupt": stopping_criteria.interrupt();  break;
        case "reset":     past_key_values_cache = null;
                          stopping_criteria.reset();      break;
    }
});

Worker 就是一个单线程事件循环 。收到主线程消息后根据 type 分发到不同处理函数。注意 generate 前先调一次 reset()------把上一轮可能遗留的中断标记清掉。

5.2 第一步：获取 tokenizer 和 model

const [tokenizer, model] = await TextGenerationPipeline.getInstance();

背后的单例模式：

class TextGenerationPipeline {
    static model_id = "onnx-community/DeepSeek-R1-Distill-Qwen-1.5B-ONNX";

    static async getInstance(progress_callback = null) {
        this.tokenizer ??= AutoTokenizer.from_pretrained(this.model_id, {
            progress_callback,
        });
        this.model ??= AutoModelForCausalLM.from_pretrained(this.model_id, {
            dtype: "q4f16",     // 4-bit 量化，省显存
            device: "webgpu",   // 跑在 GPU 上
            progress_callback,
        });
        return Promise.all([this.tokenizer, this.model]);
    }
}

??= 操作符保证不管调多少次，只下载一次 。第一次调用时 this.tokenizer 和 this.model 都是 undefined（空值），触发下载。后续调用直接返回静态属性上已有的实例。

dtype: "q4f16" 是 ONNX Runtime 的 4-bit 量化格式，把 FP16 的模型权重压缩到 4-bit，显存占用降到原来的 1/4。1.5B 参数模型才能在浏览器 WebGPU 的有限显存里跑起来。

5.3 第二步：Chat Template 转 token IDs

const inputs = tokenizer.apply_chat_template(messages, {
    add_generation_prompt: true,
    return_dict: true,
});

这一步将下面这样的 JSON 消息数组：

[{ "role": "user", "content": "1+1等于几？" }]

转换成模型训练时用的特殊标记格式文本：

<|im_start|>user
1+1等于几？<|im_end|>
<|im_start|>assistant

然后再分词成数字数组：

inputs = {
    input_ids: [151644, 8948, 198, 16, 18, 16, 41118, 17564, 16199, 151645, 198, 151644, 77091, 198],
    attention_mask: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
}

为什么不用 tokenizer("原始文本") 手动拼接？ 因为每个模型的特殊标记不同------Llama 用 <|begin_of_text|>、Qwen 用 <|im_start|>、ChatML 又是另一套。apply_chat_template 从模型的 tokenizer_config.json 读取正确的模板，永远不出错。

三个参数的作用：

| 参数 | 作用 |
|-------------------------------|-------------------------------------------------------|----------|------------------------------|
| add_generation_prompt: true | 末尾追加 `< | im_start | >assistant\n`，告诉模型"该你说话了" |
| return_dict: true | 返回 {input_ids, attention_mask} 对象，直接喂给 generate() |
| 不加 add_generation_prompt | 适用于训练场景，已有完整的 assistant 回复 |

5.4 第三步：提取思考标记的 Token ID

const [START_THINKING_TOKEN_ID, END_THINKING_TOKEN_ID] = tokenizer.encode(
    "<think></think>",
    { add_special_tokens: false },
);

DeepSeek-R1 的推理模式是思考→回答两阶段：

<think>
1+1等于几？这是一个简单的加法问题...       ← 思考过程（模型自言自语）
答案是2。
</think>

1+1等于2。                                 ← 最终回答

为了在流式生成时实时区分 两个阶段，需要知道 <> 和 </> 对应的 token ID：

// 编码结果示意：
// "<", "think", ">", "<", "/", "think", ">"
// [27,  17845,  29,  27, 1526, 17845,  29]
//  ↑                              ↑
//  解构取第一个                    解构取最后一个
//  START = 27                     END = 29

add_special_tokens: false 是关键------不需要 BOS/EOS token，只要纯净的标签 token ID。

5.5 第四步：状态追踪变量 + 双回调

let state = "thinking";    // 'thinking' | 'answering'
let startTime;             // 性能计时起点
let numTokens = 0;         // 已生成的 token 总数
let tps;                   // 每秒 token 数

const token_callback_function = (tokens) => {
    startTime ??= performance.now();              // 首 token 开始计时
    if (numTokens++ > 0) {                        // 第 2 个 token 起算 TPS
        tps = (numTokens / (performance.now() - startTime)) * 1000;
    }
    if (tokens[0] == END_THINKING_TOKEN_ID) {     // 检测到 </> 的最后字符
        state = "answering";
    }
};

const callback_function = (output) => {
    self.postMessage({
        status: "update",
        output,        // 解码后的人类文本片段
        tps,           // 实时速度
        numTokens,     // 计数
        state,         // 当前阶段
    });
};

两个回调的分工：

每生成一个 token ID
  │
  ├─→ token_callback_function
  │     ├─ 记录时间戳（只记一次）
  │     ├─ numTokens++ 并计算 TPS
  │     └─ 检测是否进入 "answering" 阶段
  │
  └─→ TextStreamer 内部累积 → 解码
        └─→ callback_function
              └─ postMessage("update") 推给主线程

token_callback 是 token 粒度 （每个都触发），负责后台统计。callback_function 是文本粒度 （累积一段后触发），负责前台推送。分离的设计让主线程收到的 postMessage 频率可控，不会每个 token 都发一次消息。

5.6 第五步：组装 TextStreamer

const streamer = new TextStreamer(tokenizer, {
    skip_prompt: true,          // 跳过输入部分，只输出新生成的
    skip_special_tokens: true,  // 过滤 <|im_end|> 等控制标记
    callback_function,          // 文本级回调
    token_callback_function,    // Token 级回调
});

TextStreamer 是 Transformers.js 内置的流式解码器，架构上是一个回调驱动的管道：

Token ID → [token_callback 统计] → [内部缓冲] → [tokenizer.decode] → [callback_function 输出]

skip_prompt: true 确保用户不会看到整段对话模板被重新"打"一遍，只看新内容。skip_special_tokens: true 过滤掉 <|im_start|>、<|im_end|> 等特殊控制 token。

5.7 第六步：调用 model.generate()

self.postMessage({ status: "start" });

const { past_key_values, sequences } = await model.generate({
    ...inputs,                    // input_ids + attention_mask
    do_sample: false,             // 贪婪解码（确定性输出）
    max_new_tokens: 2048,         // 最多新生成 2048 个 token
    streamer,                     // 流式回调
    stopping_criteria,            // 可中断停止条件
    return_dict_in_generate: true,// 返回 KV Cache
});

各参数详解：

...inputs ：把 apply_chat_template 返回的 {input_ids, attention_mask} 展开传入。

do_sample: false ：贪婪解码（Greedy Decoding）。每一步直接选概率最高的 token，不做随机采样。对于数学推理任务（DeepSeek-R1 的强项），贪婪解码保证答案的一致性和确定性。如果是创意写作，应该启采样（do_sample: true + temperature + top_p）。

max_new_tokens: 2048 ：只统计新生成的 token，不包含输入的 prompt。达到上限自动停止，方式类似于安全阀。

streamer ：挂载上文创建的 TextStreamer 实例。没有它的话，model.generate() 会全部生成完才一次性返回，用户体验就是"等 10 秒，突然蹦出全文"。

stopping_criteria ：传入 InterruptableStoppingCriteria 实例。model.generate() 内部的生成循环每步都会调用 _call() 方法，检查 this.interrupted 标志位。用户点击停止按钮后，标志位变成 true，循环立即终止。

return_dict_in_generate: true ：除了 sequences（生成的 token 序列），额外返回 past_key_values（KV Cache），存下来供下一轮复用。

5.8 第七步：保存 KV Cache + 最终解码

past_key_values_cache = past_key_values;

const decoded = tokenizer.batch_decode(sequences, {
    skip_special_tokens: true,
});

self.postMessage({
    status: "complete",
    output: decoded,
});

past_key_values 保存到模块级变量，下一轮 generate 时传入，省去重复计算。关于 KV Cache 的原理，下面单独展开讲。

batch_decode 把完整的 token ID 序列一次性转回可读文本，跳过特殊标记。这是最终完整结果，和流式的 update 形成互补：

消息	时机	内容
update	生成中（N 次）	文本片段，一段一段追加
complete	生成结束（1 次）	全部 token 解码后的完整文本

---

六、KV Cache：空间换时间的注意力加速

6.1 背景问题

Transformer 每生成一个新 token，都要让这个 token "关注"前面所有 token。如果不做缓存：

Step 1: 计算 K₁V₁
Step 2: 重算 K₁V₁ + 计算 K₂V₂        ← 重算了！
Step 3: 重算 K₁V₁K₂V₂ + 计算 K₃V₃    ← 又重算了！
...
Step N: 重算前 N-1 对的 KV + 计算 KₙVₙ ← O(n²)！

6.2 KV Cache 的解决方案

前面 token 的 K 和 V 不会变------存起来就行：

Step 1: 算 K₁V₁ → 存入 Cache: [K₁V₁]
Step 2: 只算 K₂V₂ → 追加 Cache: [K₁V₁|K₂V₂]
Step 3: 只算 K₃V₃ → 追加 Cache: [K₁V₁|K₂V₂|K₃V₃]
...
每步只算 1 对 KV，O(n)

6.3 在代码中的体现

// 模块顶层声明
let past_key_values_cache = null;

// generate() 中保存（第 6 步）
const { past_key_values } = await model.generate({ ... });
past_key_values_cache = past_key_values;

// reset 中清空
case "reset":
    past_key_values_cache = null;
    stopping_criteria.reset();
    break;

为什么 reset 时要设为 null？

• 隔离对话：旧 KV Cache 存的是上一段对话的注意力信息，不清会导致新对话"串台"
• 释放显存：对于 WebGPU（浏览器内运行），显存极为宝贵。1.5B 参数模型生成 2048 token 的 KV Cache 可占 350MB+
• 状态复位 ：非空的 past_key_values 会让模型误以为这是"继续生成"而非"从头开始"

---

七、中断机制：InterruptableStoppingCriteria

// 模块顶层创建一次
const stopping_criteria = new InterruptableStoppingCriteria();

// 用户点击停止按钮
case "interrupt":
    stopping_criteria.interrupt();  // this.interrupted = true
    break;

源码（来自 Transformers.js）：

class InterruptableStoppingCriteria extends StoppingCriteria {
    constructor() { super(); this.interrupted = false; }
    interrupt()   { this.interrupted = true; }
    reset()       { this.interrupted = false; }
    _call(input_ids, scores) {
        return new Array(input_ids.length).fill(this.interrupted);
    }
}

运作机制：

model.generate() 内部循环
  │
  ├─ 生成 token
  ├─ _call() → 检查 interrupted 标志
  │     ├─ false → 继续循环
  │     └─ true  → 立即终止    ← interrupt() 触发
  │
  └─ ...重复

为什么放在模块顶层？因为 model.generate() 内部循环和 interrupt 消息处理必须共享同一个对象引用 。放在函数内部的话，每次调用 generate() 都是新实例，interrupt() 改的就不是同一个对象的 interrupted 字段了。

---

八、完整时序：从回车到完成

把整个链路串起来：

用户按 Enter
    │
    ├── onEnter("1+1等于几?")
    │     ├── setMessages(prev => [...prev, {role:"user", content:"1+1等于几?"}])
    │     ├── setInput("")
    │     └── setIsRunning(true)
    │
    ├── React 重渲染
    │     └── useEffect([messages]) 触发
    │           ├── 守卫①: filter(user).length > 0 ✅
    │           ├── 守卫②: at(-1).role !== "assistant" ✅
    │           └── worker.current.postMessage({type:"generate", data:messages})
    │
    ├── ──────────── 跨线程边界 ────────────
    │
    ├── Worker: case "generate"
    │     ├── stopping_criteria.reset()      // 先松刹车
    │     └── generate(messages)
    │           │
    │           ├── ① getInstance() → tokenizer + model
    │           ├── ② apply_chat_template(messages) → input_ids + attention_mask
    │           ├── ③ encode("<think></think>") → START/END ID
    │           ├── ④ 初始化 state/tps/numTokens
    │           ├── ⑤ new TextStreamer(tokenizer, {...})
    │           ├── ⑥ postMessage({status:"start"})
    │           ├── ⑦ await model.generate({...})
    │           │      │
    │           │      ├── [token 循环开始]
    │           │      │    ├── token_callback: 计时/计数/检测state
    │           │      │    └── streamer → callback_function → postMessage("update")
    │           │      │         ↓ 到主线程 → setMessages 更新 UI
    │           │      └── [生成结束]
    │           │
    │           ├── ⑧ past_key_values_cache = past_key_values
    │           └── ⑨ batch_decode → postMessage("complete")
    │
    ├── ──────────── 跨线程边界 ────────────
    │
    └── 主线程处理 complete
          ├── setMessages 追加最终内容
          ├── setIsRunning(false) → 停止按钮变回发送按钮
          └── UI 恢复可用

---

九、总结

本文覆盖了 App.tsx + worker.js 的完整交互链路，核心要点：

1. 双线程架构 ：主线程管 UI，Worker 管推理，postMessage 通信
2. 单例模式 ：??= 操作符保证 tokenizer 和 model 只下载/初始化一次
3. Chat Template ：apply_chat_template 自动处理不同模型的特殊标记格式，add_generation_prompt 是对模型说"该你说话了"
4. 双回调流式输出 ：token_callback 做后台统计（计时/计数/阶段检测），callback_function 做前台推送（文本→主线程→UI）
5. 两个 useEffect 守卫：防止空对话触发生成、防止流式更新时重复触发
6. KV Cache：空间换时间，O(n²)→O(n)，reset 时必须清空防止显存泄漏
7. InterruptableStoppingCriteria ：模块顶层单例，interrupt() 和 model.generate() 共享同一个布尔标志位实现紧急停止