Worker 里的推理引擎：消息怎么来，Token 怎么回

浏览器跑大模型（四）：终于有人把 Worker 流式推理的消息设计讲清楚了

本文为原创。基于真实项目 webgpu-deepseek（WebGPU + Transformers.js 在浏览器本地跑 DeepSeek-R1-Distill-Qwen-1.5B）讲解。
系列第 4 篇 / 共 5 篇 · 上篇：《[navigator.gpu 存在就够了？WebGPU 检测的两层陷阱](#navigator.gpu 存在就够了？WebGPU 检测的两层陷阱 "#")》 · 下篇预告：《封装 useLLM Hook：把 Worker 推理收进一个 React Hook》

先看一段对话。你在浏览器的聊天框里发了一条消息，几秒钟后，AI 一个字一个字地"吐"出回答------和 ChatGPT 网页版一模一样。

但这段对话的背后发生了什么？消息是怎么从 React 组件传到 Web Worker 的？Token 又是怎么从模型里流回来的？中断按钮点了之后，模型为什么真的停了？

这篇文章把整个消息驱动的推理引擎从头拆开给你看。读完你就能自己写一个浏览器端的大模型推理消息系统。

sequenceDiagram participant R as React 主线程 participant W as Web Worker participant T as Tokenizer participant M as ONNX Model (WebGPU) R->>W: postMessage({ type: "generate", data: messages }) W->>W: stopping_criteria.reset() W->>T: apply_chat_template(messages) T-->>W: input_ids + attention_mask W->>M: model.generate({...inputs, streamer, stopping_criteria}) loop 每个 token M-->>W: token_callback (计算 TPS) W->>W: callback_function W->>R: postMessage({ status: "update", output, tps, state }) end M-->>W: { past_key_values, sequences } W->>W: 缓存 past_key_values W->>R: postMessage({ status: "complete", output: decoded })

上面这张图就是全文的地图。下面我们逐一拆解图中的每一个环节。

第一关：消息怎么从 React 传到 Worker

Worker 实例化------整个推理引擎的底座

所有推理都跑在 Web Worker 里。为什么？因为模型推理是 CPU/GPU 密集型操作------如果跑在主线程，任何一次 model.generate() 都会让整个页面冻住，滚动不了、按钮点不动，浏览器甚至会弹出"页面无响应"。

typescript 复制代码

// App.tsx --- Worker 只创建一次，用 useRef 管住引用
const worker = useRef(null);

useEffect(() => {
  if (!worker.current) {
    // 🔑 关键：Vite 下用 new URL() 导入 Worker，支持 ESM 模块
    worker.current = new Worker(
      new URL("./worker.js", import.meta.url),
      { type: "module" }  // ⚠️ 不加这个，Worker 里 import 第三方库会报错
    );
    worker.current.postMessage({ type: "check" }); // 第一步：检测 WebGPU 是否可用
  }
  // ...监听消息
}, []);

为什么 useRef 而不是 useState？

Worker 实例是一个活着的线程引用 ，不是 UI 状态。每次渲染都重新 new Worker() 会导致内存泄漏（旧 Worker 不会被 GC），而且之前下载的模型进度全丢了。用 useRef 保证整个组件生命周期内只有一个 Worker 实例------这和第一篇讲的单例模式是同一个思路，只是管的是 Worker 而不是 Pipeline。

消息协议设计------6 种消息类型

Worker 和主线程之间通过 postMessage 通信。消息协议是整个推理引擎的神经系统，消息类型设计好了，后面的一切才能井井有条：

javascript 复制代码

// worker.js --- 消息路由（整个推理引擎的入口）
self.addEventListener("message", async (e) => {
  const { type, data } = e.data;

  switch (type) {
    case "check":   check();    break;   // WebGPU 能力检测
    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;       // 🔑 清空 KV 缓存，开启全新对话
      stopping_criteria.reset();
      break;
  }
});

消息方向	type	含义	携带数据
React → Worker	`check`	检测 WebGPU 是否可用	无
React → Worker	`load`	开始下载模型	无
React → Worker	`generate`	发起一次推理	`data`: messages 数组
React → Worker	`interrupt`	中断当前生成	无
React → Worker	`reset`	重置对话上下文	无
Worker → React	`loading`	模型下载进度	`data`: 进度文本
Worker → React	`initiate`	单个文件开始下载	`file`, `progress`, `total`
Worker → React	`progress`	单个文件下载进度	`file`, `progress`, `total`
Worker → React	`done`	单个文件下载完成	`file`
Worker → React	`ready`	模型就绪	无
Worker → React	`update`	流式生成中，新 token 到达	`output`, `tps`, `numTokens`, `state`
Worker → React	`complete`	生成完毕	`output`: 完整解码文本
Worker → React	`error`	Worker 内部报错	`data`: 错误信息

消息类型从 5 种演进到 13 种，不是因为"设计先行"，而是因为"问题驱动"------每当你发现"这个状态主线程不知道"，就加一种消息。设计消息协议的最好方式不是预先画 UML，而是让代码跑起来，缺什么补什么。

第二关：generate()------推理流水线的核心

消息到了 Worker，路由到 generate(data)。这个函数是整个推理引擎的心脏------它把 messages 数组变成 token，再让模型生成，最后通过回调把结果逐字推回主线程。

javascript 复制代码

// worker.js --- 推理流水线（核心）
async function generate(messages) {
  // 第一步：拿到单例的 tokenizer 和 model
  const [tokenizer, model] = await TextGenerationPipeline.getInstance();

  // 第二步：messages 数组 → 模型能吃的 token ids
  const inputs = tokenizer.apply_chat_template(messages, {
    add_generation_prompt: true,  // 🔑 自动追加 <|im_start|>assistant\n，让模型续写回答
    return_dict: true,            // 返回 {input_ids, attention_mask} 而不是纯数组
  });

  // 第三步：解析 DeepSeek 的 <think> 标签，区分"思考"和"回答"
  const [START_THINKING_TOKEN_ID, END_THINKING_TOKEN_ID] = tokenizer.encode(
    "<think></think>",
    { add_special_tokens: false }
  );

  let state = "thinking";  // 初始状态：思考中
  let startTime;
  let numTokens = 0;
  let tps;  // tokens per second
  // ...（回调函数见下一节）
}

为什么用 `apply_chat_template` 而不是手动拼接字符串？

javascript 复制代码

// ❌ 新手常写的错误做法：
const prompt = `User: ${messages[0].content}\nAssistant:`;
const inputs = tokenizer(prompt);

// ✅ 正确做法：
const inputs = tokenizer.apply_chat_template(messages, {
  add_generation_prompt: true,
  return_dict: true,
});

DeepSeek-R1 训练时用的是特定的 chat template 格式（<|im_start|>user\n...<|im_end|>\n<|im_start|>assistant\n）。如果你手动拼字符串，格式和训练时对不上，模型输出质量会严重下降------它收到的"信号"和训练时完全不一样。apply_chat_template 保证格式和训练时严格一致。永远不要手动拼接 LLM 的 prompt 字符串，让 tokenizer 处理。

第三关：TextStreamer------Token 怎么"流"回来

这是整个推理引擎最精妙的部分。model.generate() 不是一次性返回全部结果的------那样用户要等很久。而是每生成一个 token，就触发一次回调。

javascript 复制代码

// worker.js --- 流式回调的设置
const token_callback_function = (tokens) => {
  // 🔑 用 performance.now() 实现高性能计时（比 Date.now() 精度高 100 倍）
  startTime ??= performance.now();

  if (numTokens++ > 0) {
    // 🔑 TPS = token 数 / 耗时（秒）× 1000
    tps = (numTokens / (performance.now() - startTime)) * 1000;
  }

  // ⚠️ 检测 DeepSeek 的思考结束标记
  // tokens[0] 是模型刚生成的那个 token
  if (tokens[0] == END_THINKING_TOKEN_ID) {
    state = "answering";  // 从"思考模式"切换到"回答模式"
  }
};

const callback_function = (output) => {
  // 🔑 每个 token 生成后，立刻通过 postMessage 推回主线程
  self.postMessage({
    status: "update",
    output,        // 当前累积的文本（每次都在增长）
    tps,           // 实时生成速度
    numTokens,     // 已生成 token 总数
    state,         // "thinking" | "answering"
  });
};

const streamer = new TextStreamer(tokenizer, {
  skip_prompt: true,           // 不在输出里重复 prompt
  skip_special_tokens: true,   // 过滤掉 <|im_end|> 等特殊 token
  callback_function,           // 🔑 每个 token 触发一次
  token_callback_function,     // 🔑 比 callback_function 更底层：拿到的是原始 token id
});

两个回调的区别很重要：

回调	触发时机	拿到的数据	用途
`token_callback_function`	模型预测完一个 token（解码之前）	原始 token id（整数）	性能统计（TPS）、状态切换
`callback_function`	token 解码为文本之后	累积的文本字符串	推送给 UI 显示

先 token ID 后文本 这个顺序很关键。如果是性能监控需求（比如 TPS、检测特殊 token），在 token_callback 里做------它拿 token ID 是零成本的；如果等解码成文本再去解析，白白多了一次字符串操作。

流式回调是怎么被调用的？

整个流程如下：

scss 复制代码

model.generate() 内部循环：
  for (let i = 0; i < max_new_tokens; i++) {
    1. 模型预测下一个 token → 得到 token_id
    2. streamer.token_callback([token_id])  ← 算 TPS、检测 </think>
    3. tokenizer.decode(token_id) → 文本片段
    4. streamer.callback(累积文本)          ← postMessage 推回主线程
    5. 检查 stopping_criteria → 是否被中断？
  }

TextStreamer 是 Transformers.js 内置的流式输出工具。它不需要你手动写循环------你只需要把回调函数传进去，其他的（解码、拼接、跳过特殊 token）它都帮你处理好了。

第四关：KV-Cache------为什么第二次对话比第一次快

你可能注意到：同一轮对话里，第二轮回答比第一轮快。这不是错觉，是 past_key_values_cache 在起作用。

javascript 复制代码

// worker.js --- KV-Cache 缓存
let past_key_values_cache = null;  // 🔑 模块级变量，跨 generate() 调用复用

async function generate(messages) {
  // ...
  const { past_key_values, sequences } = await model.generate({
    ...inputs,
    // past_key_values: past_key_values_cache,  // 当前版本暂时注释掉
    max_new_tokens: 2048,
    streamer,
    stopping_criteria,
    return_dict_in_generate: true,  // 🔑 必须设为 true，才能拿到 past_key_values
  });

  past_key_values_cache = past_key_values;  // 🔑 缓存本次的 KV 计算结果，下次复用
  // ...
}

KV-Cache 是什么？

Transformer 模型每生成一个新 token，都要对所有历史 token 做 Attention 计算（每个 token 的 Key 和 Value 矩阵相乘）。不用 KV-Cache 的话：

erlang 复制代码

生成第 1 个 token：计算 1 次 Attention
生成第 2 个 token：重新计算前 2 个 token 的 Attention
生成第 3 个 token：重新计算前 3 个 token 的 Attention
...
生成第 N 个 token：重新计算前 N 个 token 的 Attention

总计算量是 O(N²) 。用 KV-Cache 之后：每轮只计算新 token 的 Attention，历史 token 的 Key/Value 直接从缓存读。总计算量降到 O(N)。

复制代码

❌ 不用 KV-Cache：
  第1轮对话：计算 1+2+3+...+2048 次 Attention ≈ 200 万次
  第2轮对话：再从头计算一遍 → 又 200 万次

✅ 用 KV-Cache：
  第1轮对话：同上
  第2轮对话：只计算新 token 的 Attention → 几千次

这就是为什么"继续对话"比"新开对话"快。不是模型学聪明了，是之前的计算结果没丢。

reset 消息为什么要清空 KV-Cache？

javascript 复制代码

case "reset":
  past_key_values_cache = null;  // 🔑 开启全新对话，之前的上下文全部丢弃
  stopping_criteria.reset();
  break;

如果不清空，模型会把上一轮对话的所有历史 token 的 Key/Value 都带上------结果就是：

新对话的回答里掺杂了上一轮的上下文
缓存越来越大，推理越来越慢
显存/内存占用无限增长

第五关：中断------点了停止按钮，模型为什么真的停了？

生成是一个同步的 for 循环（虽然 WebGPU 计算是异步的），你不能直接"杀"掉它。但可以在循环的每一步检查一个标记：

javascript 复制代码

// worker.js
const stopping_criteria = new InterruptableStoppingCriteria();

// 用户点了停止按钮：
case "interrupt":
  stopping_criteria.interrupt();  // 把内部 _interrupted 设为 true
  break;

// model.generate() 内部：
// 每生成一个 token 后，检查 stopping_criteria
// 如果 _interrupted === true → 立即停止，不再生成下一个 token

stateDiagram-v2 [*] --> 空闲: reset 空闲 --> 生成中: generate 生成中 --> 空闲: complete 生成中 --> 被中断: interrupt 被中断 --> 空闲: reset note right of 被中断: 保留已生成的 token

InterruptableStoppingCriteria 本质是一个协作式中断------它不强制杀线程（Worker 里也杀不了），而是在每次生成下一个 token 之前检查。这就像你在跑马拉松，每个水站都有人举牌问"要停吗？"------你看到牌子就可以停，但不需要在每个脚步之间都检查。

中断后 key 的设计 ：调用 interrupt() 之前已经生成的 token 不会丢失 ------它们已经通过 callback_function 推回主线程显示了。中断只是不再生成新的 token。

第六关：Think/Answer 状态机------DeepSeek 的"内心独白"

DeepSeek-R1 是一种 reasoning model （推理模型）。它在给出最终回答之前，会先在 <think> 标签里进行"内心推理"：

ini 复制代码

<think>
方程 x² - 3x + 2 = 0，我需要用求根公式...
判别式 = 9 - 8 = 1 > 0，所以有两个实根
x = (3 ± 1) / 2 = 2 或 1
</think>

方程 x² - 3x + 2 = 0 的解是 **x = 1** 和 **x = 2**。

为了让 UI 能区分"思考过程"和"最终回答"（比如思考过程显示在可折叠区域里），worker 里实现了一个简单的状态机：

javascript 复制代码

// worker.js --- Think/Answer 状态机
const [START_THINKING_TOKEN_ID, END_THINKING_TOKEN_ID] = tokenizer.encode(
  "<think></think>",
  { add_special_tokens: false }
);

let state = "thinking";  // 两种状态：'thinking' | 'answering'

const token_callback_function = (tokens) => {
  if (tokens[0] == END_THINKING_TOKEN_ID) {
    state = "answering";  // 🔑 检测到 </think> token，切换状态
  }
};

const callback_function = (output) => {
  self.postMessage({
    status: "update",
    output,
    tps,
    numTokens,
    state,  // 🔑 UI 根据这个字段判断当前文本是"思考"还是"回答"
  });
};

这个状态机只有 2 个状态、1 个转换条件，但非常可靠。因为 token ID 的匹配是精确的 ------不存在正则表达式可能匹配错误的问题。模型生成的 token 流里，</think> token 一定是那个特定的 ID。

UI 端可以根据 state 字段做不同的渲染（比如 thinking 时文字用灰色斜体，answering 时用正常样式），但这部分留到下一篇（React Hook 封装）再展开。

为什么把推理放在 Web Worker？不只是"不卡 UI"

放在主线程	放在 Web Worker
`model.generate()` 阻塞 DOM 渲染	主线程始终空闲，用户可滚动、点击
浏览器可能弹出"页面无响应"	完全无感
无法利用多核 CPU	Worker 跑在独立线程
组件卸载时无法优雅清理	`worker.terminate()` 干净利落
大文件下载阻塞用户交互	Worker 里下载不影响 UI

有人问"React 的 useTransition 能不能替代 Worker？"答案是不能。useTransition 只是把状态更新标记为低优先级，但 JS 计算本身还是在主线程------遇到 model.generate() 这种同步阻塞操作，该卡还是卡。Worker 是真正的多线程，计算和 UI 渲染互不干扰。

总结：推理引擎的 5 个设计决策

回头看这张全景图：

每个设计决策对应一个问题：

#	决策	解决的问题
1	推理放 Worker 线程	主线程不阻塞，页面不卡死
2	`apply_chat_template`	prompt 格式和训练时一致，模型输出质量不下降
3	TextStreamer 双回调	每个 token 实时推送给 UI，用户看到"打字机效果"
4	past_key_values_cache	多轮对话不重复计算 Attention，第二轮起速度翻倍
5	InterruptableStoppingCriteria	用户可随时中断生成，不浪费算力

金句：一个推理引擎的消息设计，不是在画 UML 图------而是在每个"主线程不知道"的瞬间，加一条 postMessage。

本系列其他文章

第一篇：浏览器跑大模型，为什么你一刷新页面就卡死？ --- 单例模式管住 Pipeline 和 Worker
[第二篇：点了 Load 之后，那个 800MB 的下载进度条去哪了？](#第二篇：点了 Load 之后，那个 800MB 的下载进度条去哪了？ "#") --- 模型下载进度状态机
[第三篇：navigator.gpu 存在就够了？WebGPU 检测的两层陷阱](#第三篇：navigator.gpu 存在就够了？WebGPU 检测的两层陷阱 "#") --- 兼容性检测
← 本文：Worker 里的推理引擎：消息怎么来，Token 怎么回 --- 你在这里
[下一篇：封装 useLLM Hook](#下一篇：封装 useLLM Hook "#") --- 把整个 Worker 推理收进一个 React Hook

下次写 React + Web Worker + LLM 时，记住这张消息路由表。它比任何架构图都值钱。

Worker 里的推理引擎：消息怎么来，Token 怎么回

浏览器跑大模型（四）：终于有人把 Worker 流式推理的消息设计讲清楚了

第一关：消息怎么从 React 传到 Worker

Worker 实例化------整个推理引擎的底座

消息协议设计------6 种消息类型

第二关：generate()------推理流水线的核心

为什么用 apply_chat_template 而不是手动拼接字符串？

第三关：TextStreamer------Token 怎么"流"回来

流式回调是怎么被调用的？

第四关：KV-Cache------为什么第二次对话比第一次快

KV-Cache 是什么？

reset 消息为什么要清空 KV-Cache？

第五关：中断------点了停止按钮，模型为什么真的停了？

第六关：Think/Answer 状态机------DeepSeek 的"内心独白"

为什么把推理放在 Web Worker？不只是"不卡 UI"

总结：推理引擎的 5 个设计决策

本系列其他文章

为什么用 `apply_chat_template` 而不是手动拼接字符串？