WebGPU DeepSeek项目实战(四):接通聊天生成、流式输出与中断控制
WebGPU DeepSeek项目实战(四):接通聊天生成、流式输出与中断控制
- 前言
- 1. 先对比第三篇留下的接口
- 2. 进入代码前需要理解的生成基础
- 3. App 先把输入变成一条用户消息
- 4. 监听 `messages`,把最新对话发送给 Worker
- 5. Worker 引入流式输出与停止条件
- 6. 用 Chat Template 生成模型输入
- 7. 区分思考阶段和正式回答
- 8. 统计 Token 数量和 TPS
- 9. 用 `TextStreamer` 建立流式输出
- 10. 调用 `model.generate()` 完成正式生成
- 11. App 如何把流式片段合成一条助手消息
- 12. 实现可中断生成
- 13. KV Cache 到底有没有启用
- 14. 本篇功能边界:逻辑闭环与页面展示
- 15. 从打开页面到流式回答的整体流程
- 16. 本篇两份完整注释代码
- 总结
前言
系列第三篇完成了模型下载、WebGPU 预热和聊天页面框架。模型进入 ready 后,输入框可以输入内容,发送图标和停止图标也已经准备好,但当时还留下了几个空入口:onEnter() 没有保存用户消息,onInterrupt() 没有发送中断命令,App 中的 start、update、complete 没有处理 Worker 状态,Worker 中的 generate 也没有真正调用模型。
这一篇继续把这些入口逐个接通,让项目形成一条完整的文本生成链路:用户提交问题后,React 把消息保存为对话历史,再将完整消息数组发送给 Worker;Worker 使用模型配套的 Chat Template 把对话转换成 Token,调用 model.generate() 开始推理;TextStreamer 一边接收新 Token,一边解码文字并回传页面;React 将每个文本片段追加到助手消息中,直到收到 complete。
同时,本篇还会加入 Token 统计、TPS 计算、思考与回答阶段划分、可中断生成、KV Cache 基础。这些代码不是彼此孤立的 API,而是共同服务于同一件事:让一个加载完成的本地模型真正接收多轮对话,并把生成过程持续反馈给前端。
这一篇完成的是聊天生成的数据闭环。消息气泡、思考区域、TPS 展示和自动滚动等前端页面优化,将在系列第五篇继续完成。
1. 先对比第三篇留下的接口
1.1 原来的代码停在哪里
第三篇已经完成输入框和操作图标,但两个交互函数还是空的:
function onEnter(message) {
// 等待接入 generate 命令
}
function onInterrupt() {
// 等待接入 interrupt 命令
}
App 虽然提前写好了 Worker 状态分支,但生成阶段也没有实际行为:
case "start":
break;
case "update":
break;
case "complete":
break;
Worker 中同样只保留命令入口:
case "generate":
break;
case "interrupt":
break;
case "reset":
break;
这些空分支不是没有意义。前三篇先建立了稳定的双向消息协议,本篇只需要按照原来的职责继续填充:
| 方向 | 消息字段 | 本篇用途 |
|---|---|---|
| App → Worker | type: "generate" |
发送完整对话历史并启动生成 |
| App → Worker | type: "interrupt" |
请求停止当前生成 |
| App → Worker | type: "reset" |
清理缓存和中断状态 |
| Worker → App | status: "start" |
通知页面创建助手消息 |
| Worker → App | status: "update" |
返回流式文本和性能数据 |
| Worker → App | status: "complete" |
通知页面本次生成结束 |
前三篇的 check → loading → ready 负责把模型准备好,本篇的 generate → start → update → complete 负责真正使用模型。两条链路通过同一个 Worker 和同一个消息监听器连接起来,不需要另外创建推理线程。
1.2 本篇新增内容怎样连接旧代码
本篇没有创建新的页面组件,而是在原有 App.tsx 和 worker.js 中继续向下实现:
第三篇已有
├── TextGenerationPipeline:提供 tokenizer 和 model
├── status === "ready":允许用户输入
├── messages:预留对话状态
├── onEnter():预留发送入口
├── onInterrupt():预留停止入口
└── Worker switch:预留 generate / interrupt / reset
第四篇接入
├── onEnter():创建 user 消息并进入运行状态
├── messages Effect:发送完整对话历史
├── apply_chat_template():生成模型输入
├── TextStreamer:流式解码输出
├── start / update / complete:维护 assistant 消息
├── InterruptableStoppingCriteria:停止生成
└── past_key_values_cache:保存注意力缓存
这也是项目逐步推进的好处:加载、通信、输入和按钮都已经存在,本篇只关注“数据怎样进入模型,又怎样从模型回到 React”。
2. 进入代码前需要理解的生成基础
2.1 大模型为什么可以一个 Token 一个 Token 地回答
DeepSeek-R1-Distill-Qwen-1.5B 属于因果语言模型。它的核心任务不是一次写出整段答案,而是根据已有 Token 预测下一个 Token:
输入:用户的问题
↓
预测第 1 个新 Token
↓
把新 Token 拼回上下文
↓
预测第 2 个新 Token
↓
不断重复,直到结束标记、长度上限或外部中断
例如模型要生成“你好”,底层过程可能是:
已有输入 → “你” → “好” → 结束 Token
这种生成方式叫 Autoregressive Generation(自回归生成)。它带来两个直接结果:
- 每得到一批新 Token,就可以提前解码并展示,不必等待完整答案。
- 每次预测之间都可以检查停止条件,所以用户能够中断生成。
TextStreamer 和 InterruptableStoppingCriteria 正是建立在这个生成机制之上。
2.2 聊天消息为什么使用 role 和 content
普通字符串只能表示一句文字,却无法告诉模型这句话是谁说的。多轮对话通常使用下面的结构:
[
{
role: "user",
content: "什么是 WebGPU?",
},
{
role: "assistant",
content: "WebGPU 是浏览器中的现代 GPU API……",
},
{
role: "user",
content: "它为什么适合运行模型?",
},
]
其中:
| 字段 | 作用 |
|---|---|
role: "user" |
表示内容来自用户 |
role: "assistant" |
表示内容来自模型 |
content |
保存这一轮消息正文 |
完整数组表示对话历史。每次向 Worker 发送整个 messages,模型才能结合前面的提问和回答理解上下文。
消息数组是应用层的对话数据;模型不能直接理解 JavaScript 对象,还需要 Chat Template 把它转换成训练时熟悉的文本格式和 Token。
2.3 Chat Template 解决了什么问题
聊天模型在训练时不会只看到裸文本,它还需要角色边界和消息边界。例如一段对话在底层可能被组织成:
<|im_start|>user
什么是 WebGPU?<|im_end|>
<|im_start|>assistant
不同模型使用的特殊标记可能不同。如果项目手动拼接,很容易使用错误格式,导致模型不能正确区分用户与助手。
Tokenizer 保存了与模型配套的聊天模板,apply_chat_template() 会完成以下工作:
role/content 消息数组
↓ 套用模型配套 Chat Template
加入角色标记、消息结束标记和生成提示
↓ Tokenizer 编码
input_ids + attention_mask
↓
交给 model.generate()
所以 Chat Template 不是为了美化提示词,而是为了让应用输入符合模型训练时的对话格式。
2.4 Token 回调和文本回调有什么区别
模型生成的是 Token ID,人看到的是文字。TextStreamer 提供两个不同层级的回调:
| 回调 | 接收到的数据 | 本项目中的用途 |
|---|---|---|
token_callback_function |
新生成的原始 Token ID | 计数、计算 TPS、判断思考结束标记 |
callback_function |
Tokenizer 解码后的可显示文本片段 | 通过 update 消息发送给 App |
有些 Token 还不能单独组成完整单词,TextStreamer 会暂时缓存它们。等形成适合输出的文本后,才调用 callback_function。因此,Token 回调的触发粒度和文本回调的触发粒度并不一定完全相同。
3. App 先把输入变成一条用户消息
3.1 新增 TPS 和 Token 数量状态
在原有 input 和 messages 下面加入:
const [input, setInput] = useState("");
const [messages, setMessages] = useState([]);
// 模型平均每秒生成多少 Token。
const [tps, setTps] = useState(null);
// 本轮已经生成的 Token 数量。
const [numTokens, setNumTokens] = useState(null);
tps 和 numTokens 不负责控制模型,它们保存 Worker 回传的性能数据。数据先进入 React State,下一篇再把它们显示到聊天页面中。
这两个 State 与原有 isRunning 共同描述一次生成任务:
| State | 表示什么 | 由谁更新 |
|---|---|---|
isRunning |
模型是否正在生成 | 提交时设为 true,完成时设为 false |
tps |
每秒生成 Token 数 | 每条 update 消息更新 |
numTokens |
已生成 Token 总数 | 每条 update 消息更新 |
3.2 完成 onEnter()
上一节预留的函数现在负责提交前端状态:
function onEnter(message) {
// 在最新消息数组末尾追加一条用户消息。
setMessages((prev) => [
...prev,
{
role: "user",
content: message,
},
]);
// 新一轮生成重新统计 TPS。
setTps(null);
// 立即切换为运行状态,让发送图标变成停止图标。
setIsRunning(true);
// 提交后清空受控输入框。
setInput("");
}
按钮点击和 Enter 键原本都已经调用 onEnter(input),所以函数完成后,两种提交方式会走同一条逻辑。
这里继续使用函数式更新:
setMessages((prev) => [...prev, newMessage]);
原因与进度数组相同:新消息依赖旧消息数组,React 应该在处理更新时把最新数组传给 prev。如果直接使用闭包中的旧 messages,连续提交或其他消息更新可能相互覆盖。
执行后的状态变化如下:
input = "解释一下 WebGPU"
messages = []
isRunning = false
↓ onEnter(input)
input = ""
messages = [{ role: "user", content: "解释一下 WebGPU" }]
isRunning = true
tps = null
onEnter() 只负责更新 React 数据,没有直接调用模型。真正发送 Worker 命令的是下一步的 Effect。
4. 监听 messages,把最新对话发送给 Worker
4.1 为什么不直接在 onEnter() 中读取旧消息
调用 setMessages() 后,当前函数中的 messages 不会立刻原地改变。React 会把更新放入队列,在下一次渲染中提供新的消息数组。
如果在同一个函数中直接这样写:
function onEnter(message) {
setMessages((prev) => [...prev, { role: "user", content: message }]);
// 这里的 messages 仍然是本次渲染的旧快照。
worker.current.postMessage({ type: "generate", data: messages });
}
Worker 可能收不到刚刚提交的那条用户消息。因此,代码增加一个监听 messages 的 Effect:
useEffect(() => {
// 没有任何 user 消息时,不启动生成。
if (messages.filter((x) => x.role === "user").length === 0) {
return;
}
// 最后一条已经是 assistant 时,不再次请求生成。
if (messages.at(-1).role === "assistant") {
return;
}
setTps(null);
// 此时 messages 已经是下一次渲染中的最新对话数组。
worker.current.postMessage({
type: "generate",
data: messages,
});
}, [messages]);
依赖数组 [messages] 表示:只要消息数组引用发生变化,React 就重新执行这个 Effect。这样 Worker 取得的是更新后的完整对话历史。
4.2 两个提前返回分别防止什么
第一个判断:
if (messages.filter((x) => x.role === "user").length === 0) {
return;
}
组件第一次渲染时 messages 是空数组。没有这个判断,页面初始化后就可能向 Worker 发送一次空生成请求。
第二个判断:
if (messages.at(-1).role === "assistant") {
return;
}
Worker 发送 start 后,App 会在数组末尾加入空的 assistant 消息,这同样会改变 messages 并再次触发 Effect。如果不判断最后一条角色,流程会变成:
加入 user
↓ Effect 发送 generate
Worker 发送 start
↓ 加入 assistant
Effect 再次发送 generate
↓ 重复生成
加上角色判断后:
最后一条是 user → 需要模型回答
最后一条是 assistant → 模型回答已开始或已经存在,不再重复请求
messages.at(-1) 表示取得数组最后一个元素。项目的 TypeScript 目标是 ES2023,因此可以直接使用 Array.prototype.at()。
4.3 发送给 Worker 的对象是什么样的
worker.current.postMessage({
type: "generate",
data: messages,
});
假设用户进行到第二轮,Worker 收到的数据大致是:
{
type: "generate",
data: [
{ role: "user", content: "什么是 WebGPU?" },
{ role: "assistant", content: "WebGPU 是浏览器 GPU API……" },
{ role: "user", content: "它和 WebGL 有什么区别?" },
],
}
Worker 继续使用原来的命令解构:
const { type, data } = e.data;
进入 case "generate" 后,data 就是完整消息数组:
case "generate":
stopping_criteria.reset();
generate(data);
break;
生成前调用 reset() 是为了清除上一次中断留下的标记,否则新任务可能一启动就被旧状态停止。
5. Worker 引入流式输出与停止条件
5.1 扩展 Transformers.js 导入
在原有 Tokenizer 和模型类后加入两个工具:
import {
AutoTokenizer,
AutoModelForCausalLM,
// 将生成 Token 持续解码为文本片段。
TextStreamer,
// 提供可以由外部切换的生成停止条件。
InterruptableStoppingCriteria,
} from "@huggingface/transformers";
这次不需要安装新依赖。两个类都来自前面已经安装的 @huggingface/transformers。
TextGenerationPipeline 也不需要重新加载资源。生成函数继续调用:
const [tokenizer, model] =
await TextGenerationPipeline.getInstance();
由于第二篇和第三篇已经用单例思路缓存 Promise,这里取得的是同一份 Tokenizer 和模型,而不是重新下载模型。
5.2 创建全局停止条件
在类外创建实例:
const stopping_criteria =
new InterruptableStoppingCriteria();
它内部维护一个布尔状态,可以简单理解为:
reset() → interrupted = false
interrupt() → interrupted = true
当它被传给 model.generate() 后,模型每进行一轮生成都会询问停止条件。只要 interrupted 变成 true,生成循环就会结束。
实例定义在 generate() 外部,是为了让 Worker 的两个位置访问同一个对象:
case "generate"调用reset()。case "interrupt"调用interrupt()。
如果每次都在函数内部创建新实例,外面的 interrupt 分支拿不到正在被 model.generate() 使用的那个停止条件。
6. 用 Chat Template 生成模型输入
6.1 generate(messages) 接收完整对话
Worker 新增生成函数:
async function generate(messages) {
// 复用已经加载并预热过的资源。
const [tokenizer, model] =
await TextGenerationPipeline.getInstance();
// 后续生成逻辑写在这里。
}
参数名使用复数 messages,因为 App 发送的是消息数组,不是单个字符串。生成函数重新调用 getInstance() 不会重新下载资源,单例缓存会直接返回已有 Promise。
6.2 调用 apply_chat_template()
const inputs = tokenizer.apply_chat_template(messages, {
// 在末尾加入“轮到 assistant 回答”的生成提示。
add_generation_prompt: true,
// 返回包含 input_ids 等具名字段的对象。
return_dict: true,
});
两个参数的作用如下:
| 参数 | 值 | 效果 |
|---|---|---|
messages |
完整对话数组 | 为所有 user/assistant 消息套用模型模板 |
add_generation_prompt |
true |
在末尾提示模型继续生成 assistant 内容 |
return_dict |
true |
返回可展开给 generate() 的输入对象 |
apply_chat_template() 默认会继续执行 Tokenize,所以 inputs 不是普通模板字符串,而是类似下面的模型输入对象:
{
input_ids: /* Token ID Tensor */,
attention_mask: /* 有效 Token 位置 */,
}
后面使用:
model.generate({
...inputs,
});
对象展开会把 input_ids、attention_mask 等字段交给模型。
7. 区分思考阶段和正式回答
7.1 取得 <think> 与 </think> 的 Token ID
DeepSeek 推理模型的输出可以分成思考过程和正式回答。代码先编码两个标记:
const [
START_THINKING_TOKEN_ID,
END_THINKING_TOKEN_ID,
] = tokenizer.encode("<think></think>", {
// 不额外加入 BOS、EOS 等特殊 Token。
add_special_tokens: false,
});
数组解构把编码结果中的前两个 ID 分别保存为开始思考和结束思考标记。
本篇实际使用的是:
END_THINKING_TOKEN_ID
因为生成开始时默认处于思考阶段,只要检测到 </think>,就可以切换到回答阶段。START_THINKING_TOKEN_ID 先被取出,但这一版代码没有继续使用。
7.2 使用 state 保存生成阶段
let state = "thinking";
这里的 state 是 generate() 内部的普通局部变量,不是 React State。一次生成开始时设为 thinking,检测到结束思考 Token 后改为 answering:
if (tokens[0] == END_THINKING_TOKEN_ID) {
state = "answering";
}
代码使用 == 而不是 ===,是为了兼容 Token ID 可能以 BigInt 等数值形式返回的情况。它只比较数值含义,不强制要求两边类型完全相同。
状态变化随后会跟随每条 update 消息发给 App,前端才能记录思考内容与正式回答的分界。
8. 统计 Token 数量和 TPS
8.1 三个统计变量
每次进入 generate() 都重新创建:
let startTime; // 第一个生成 Token 到达的时间
let numTokens = 0; // 本轮已生成 Token 数量
let tps; // 平均每秒生成 Token 数
它们是单次生成任务的局部变量。下一轮用户再次提问时会重新计时,不会把前一次的 Token 数混进来。
8.2 token_callback_function 的执行过程
const token_callback_function = (tokens) => {
// 只在第一个 Token 到达时记录时间。
startTime ??= performance.now();
// 后缀 ++ 先参与判断,再把 numTokens 加一。
if (numTokens++ > 0) {
tps =
(numTokens / (performance.now() - startTime)) *
1000;
}
// 遇到 </think> 后进入正式回答阶段。
if (tokens[0] == END_THINKING_TOKEN_ID) {
state = "answering";
}
};
performance.now() 返回高精度毫秒时间。??= 保证 startTime 只在第一次为空时赋值,后续 Token 不会覆盖起点。
TPS 的公式是:
TPS = 已生成 Token 数 ÷ 已经过的秒数
代码中时间差是毫秒,所以乘以 1000:
tps =
(numTokens / (performance.now() - startTime)) * 1000;
numTokens++ 是后缀自增。第一次执行时,参与 > 0 判断的旧值是 0,条件不成立,但判断结束后 numTokens 已变成 1。从第二个 Token 开始,经过时间更有意义,代码才正式计算 TPS。
这里得到的是从第一个 Token 到当前时刻的平均生成速度,不是某一个 Token 的瞬时速度。
9. 用 TextStreamer 建立流式输出
9.1 文本回调负责发送 update
Token 回调用于底层计数,文本回调则把可以阅读的内容发回 App:
const callback_function = (output) => {
self.postMessage({
status: "update",
output,
tps,
numTokens,
state,
});
};
一条消息大致是:
{
status: "update",
output: "WebGPU ",
tps: 18.7,
numTokens: 42,
state: "answering",
}
output 是本次新解码出来的文本片段,不一定是完整答案。App 要不断把多个 output 追加起来。
9.2 创建 Streamer
const streamer = new TextStreamer(tokenizer, {
// 不把原始提示词再次输出给页面。
skip_prompt: true,
// 解码时跳过模型内部特殊标记。
skip_special_tokens: true,
// 可显示文本形成后执行。
callback_function,
// 新 Token 生成后执行。
token_callback_function,
});
TextStreamer 同时拿到 Tokenizer 和两类回调。模型生成新 Token 后,大致经过:
model.generate() 得到新 Token
↓
token_callback_function(tokens)
├── Token 计数
├── TPS 计算
└── 判断 </think>
↓
TextStreamer 使用 tokenizer 解码
↓
形成适合显示的文本片段
↓
callback_function(output)
↓
self.postMessage({ status: "update", ... })
skip_prompt: true 非常重要。传给模型的输入中包含所有聊天历史,如果不跳过 Prompt,Streamer 可能把原来的用户问题和历史内容也当成输出交给页面。
10. 调用 model.generate() 完成正式生成
10.1 生成前发送 start
self.postMessage({ status: "start" });
Worker 在调用模型前先通知 App。页面收到它后创建一条空助手消息,为后续流式片段预留写入位置。
这个顺序不能反过来。如果第一个 update 已经到达,但数组中还没有助手消息,App 就找不到应该把输出追加到哪里。
10.2 每个生成参数的作用
const { past_key_values, sequences } =
await model.generate({
// Chat Template 返回的 input_ids、attention_mask 等输入。
...inputs,
// 暂未开启跨轮 KV Cache 传入。
// past_key_values: past_key_values_cache,
// 关闭随机采样,使用稳定的确定性解码路线。
do_sample: false,
// 最多生成 2048 个新 Token。
max_new_tokens: 2048,
// 让生成结果持续进入两个回调。
streamer,
// 每一步都检查是否收到外部中断。
stopping_criteria,
// 返回 sequences、past_key_values 等具名结果。
return_dict_in_generate: true,
});
参数汇总如下:
| 参数 | 解决的问题 |
|---|---|
...inputs |
把格式化和分词后的聊天历史传给模型 |
do_sample: false |
关闭随机采样,让解码路线更稳定 |
max_new_tokens: 2048 |
防止生成无限进行 |
streamer |
在生成过程中持续输出文本 |
stopping_criteria |
允许外部中断生成 |
return_dict_in_generate: true |
取得生成序列和注意力缓存等完整结果 |
await 会一直等待生成结束。结束原因可能是:模型生成结束 Token、达到 2048 上限,或者 stopping_criteria 检测到中断。
10.3 最终解码与 complete
生成结束后保存返回的缓存:
past_key_values_cache = past_key_values;
再把完整序列批量解码:
const decoded = tokenizer.batch_decode(sequences, {
skip_special_tokens: true,
});
最后通知 App:
self.postMessage({
status: "complete",
output: decoded,
});
update 和 complete 的职责不同:
| 状态 | 发生时机 | output 的含义 |
|---|---|---|
update |
生成过程中多次触发 | 本次新增的可显示文本片段 |
complete |
整次生成结束后触发一次 | 最终序列的批量解码结果 |
这一版 App 主要依靠 update 拼接助手内容,complete.output 没有再次写入消息,避免把完整结果重复追加。complete 主要负责结束运行状态。
11. App 如何把流式片段合成一条助手消息
11.1 start:先创建空助手消息
case "start": {
setMessages((prev) => [
...prev,
{
role: "assistant",
content: "",
},
]);
}
break;
假设发送前数组是:
[
{
role: "user",
content: "解释一下 WebGPU",
},
]
收到 start 后变成:
[
{
role: "user",
content: "解释一下 WebGPU",
},
{
role: "assistant",
content: "",
},
]
这条空消息就是流式输出的容器。它加入数组后会触发监听 [messages] 的 Effect,但由于最后一条角色是 assistant,Effect 会提前返回,不会再次生成。
11.2 update:复制数组并更新最后一项
case "update": {
const { output, tps, numTokens, state } = e.data;
setTps(tps);
setNumTokens(numTokens);
setMessages((prev) => {
// 创建新数组,避免直接修改旧 State。
const cloned = [...prev];
// start 已经保证最后一条是 assistant。
const last = cloned.at(-1);
// 创建新的 assistant 对象,并追加本次文本片段。
const data = {
...last,
content: last.content + output,
};
// 第一次进入 answering 时记录分界位置。
if (
data.answerIndex === undefined &&
state === "answering"
) {
data.answerIndex = last.content.length;
}
// 用新对象替换数组最后一项。
cloned[cloned.length - 1] = data;
return cloned;
});
}
break;
这里不能每收到一个片段就追加一条新消息,否则模型回答会被拆成大量 assistant 对象。正确思路是:start 只创建一次消息,所有 update 都更新这一个对象。
假设三个流式片段依次是:
"WebGPU "
"是浏览器中的 "
"GPU API。"
最后一条消息会经历:
content = ""
content = "WebGPU "
content = "WebGPU 是浏览器中的 "
content = "WebGPU 是浏览器中的 GPU API。"
11.3 为什么既复制数组又复制对象
下面两步分别保证不同层级不被直接修改:
const cloned = [...prev];
创建新的数组引用。
const data = {
...last,
content: last.content + output,
};
创建新的最后一条消息对象。
如果只复制数组,却直接修改:
last.content += output;
旧 State 快照中的消息对象仍然会被改变。React 状态应该按照不可变数据思路更新:从旧值计算新值,但不修改旧值本身。
11.4 answerIndex 如何记录思考与答案边界
第一次收到 state === "answering" 时:
data.answerIndex = last.content.length;
这里记录的是追加本次 output 之前,已有内容的字符长度。假设:
data = {
role: "assistant",
content: "这是思考过程……这是正式回答……",
answerIndex: 8,
}
下一篇渲染页面时就可以按位置切开:
content.slice(0, answerIndex) → 思考部分
content.slice(answerIndex) → 正式回答
判断中还包含:
data.answerIndex === undefined
它保证分界位置只记录一次。后续所有 answering 片段不会不断覆盖边界。
11.5 complete:统一结束运行状态
case "complete":
setIsRunning(false);
break;
onEnter() 提交时把 isRunning 设为 true,complete 再把它设为 false:
提交问题
↓ isRunning = true
显示 StopIcon
↓ 模型正常结束或被中断
收到 complete
↓ isRunning = false
恢复 ArrowRightIcon
这样正常生成和中断生成都使用同一个收尾入口,页面不会维护两套结束逻辑。
12. 实现可中断生成
12.1 App 发送 interrupt
第三篇的停止图标已经绑定 onInterrupt,现在补上函数:
function onInterrupt() {
// 不立即修改 isRunning,等待 Worker 真正结束后发送 complete。
worker.current.postMessage({
type: "interrupt",
});
}
为什么不在点击时直接写:
setIsRunning(false);
因为 postMessage() 只是发送请求。Worker 收到消息、停止条件被检查、模型退出生成循环,都需要一个过程。如果 App 提前显示已结束,页面状态可能和后台任务不一致。
12.2 Worker 设置中断标记
case "interrupt":
// 与 generate() 使用的是同一个 stopping_criteria 实例。
stopping_criteria.interrupt();
break;
InterruptableStoppingCriteria 不会强制杀死 Worker。它只把内部标志设为 true,让 model.generate() 在下一次生成检查时正常退出。
这种方式比直接终止整个 Worker 更合适:
| 方式 | 结果 |
|---|---|
worker.terminate() |
整个 Worker 被销毁,已加载模型也会随线程失去 |
stopping_criteria.interrupt() |
只结束本次生成,Worker、Tokenizer 和模型继续保留 |
模型退出后仍然执行 batch_decode() 和 complete 消息,App 再恢复可发送状态。
13. KV Cache 到底有没有启用
13.1 KV Cache 是什么
Transformer 在生成每个新 Token 时,需要进行注意力计算。已经处理过的历史 Token 会产生 Key 和 Value。如果每生成一步都从头计算全部历史内容,序列越长,重复计算越多。
KV Cache 会保存已经计算过的 Key 和 Value。下一步预测时,只需要处理新 Token,再复用旧缓存:
没有 KV Cache
第 1 步:计算 Token 1
第 2 步:重新计算 Token 1、2
第 3 步:重新计算 Token 1、2、3
使用 KV Cache
第 1 步:计算 Token 1,保存 KV
第 2 步:复用 Token 1,只计算 Token 2
第 3 步:复用 Token 1、2,只计算 Token 3
这对自回归生成非常重要,因为一次回答可能包含数百甚至上千个 Token。
13.2 当前代码保存了什么
Worker 先准备缓存变量:
let past_key_values_cache = null;
由于生成配置使用:
return_dict_in_generate: true
模型返回值中包含:
const { past_key_values, sequences } =
await model.generate(...);
生成结束后将其保存:
past_key_values_cache = past_key_values;
完成了 KV Cache 的接收与保存。
13.4 reset 清理了什么
Worker 已经预留重置命令:
case "reset":
// 放弃上一轮保存的注意力缓存。
past_key_values_cache = null;
// 同时清除可能残留的中断状态。
stopping_criteria.reset();
break;
这个分支让未来“新建对话”时可以清理推理状态。不过这一版 App 还没有发送 reset 的按钮或函数,页面入口会与下一篇的会话界面一起完善。
14. 本篇功能边界:逻辑闭环与页面展示
生成相关数据已经能够进入 React:
messages → 用户消息与流式助手消息
tps → 平均生成速度
numTokens → 本轮生成 Token 数量
answerIndex → 思考与正式回答的分界
isRunning → 发送/停止按钮状态
但聊天容器依然为空:
{status === "ready" && (
<div
ref={chatContainerRef}
className="overflow-y-auto scrollbar-thin w-full flex flex-col items-center h-full"
>
{/* 下一篇在这里渲染 messages */}
</div>
)}
也就是说,本篇完成的是 数据和生成逻辑:模型输出已经可以被保存到 messages,但页面还没有执行 messages.map() 创建聊天气泡,tps 和 numTokens 也没有显示。
下一篇的前端优化会继续处理:
- 用户消息与助手消息的不同布局。
- 思考过程和正式回答的拆分显示。
- Markdown、代码块与数学公式渲染。
- TPS 和 Token 数量展示。
- 消息增长后的自动滚动。
- 示例问题与会话重置入口。
这不是重新实现生成,而是把本篇已经进入 State 的数据转换成完整页面。
15. 从打开页面到流式回答的整体流程
前三篇和本篇合在一起后,项目整体执行链如下:
1. React 创建 module Worker
↓
2. App 发送 check
↓
3. Worker 通过 navigator.gpu 请求 GPUAdapter
↓
4. 用户点击 Load model
↓
5. 单例 Pipeline 加载 Tokenizer 与 q4f16 ONNX 模型
↓
6. Transformers.js 发送 initiate / progress / done
↓
7. App 使用 progressItems 渲染文件进度条
↓
8. Worker 用最小输入预热 WebGPU 模型
↓
9. Worker 发送 ready,App 解锁输入框
↓
10. 用户提交问题,onEnter 添加 user 消息
↓
11. messages Effect 发送完整对话历史
↓
12. Worker 使用 apply_chat_template() 生成模型输入
↓
13. Worker 创建 TextStreamer 和停止条件
↓
14. Worker 发送 start,App 创建空 assistant 消息
↓
15. model.generate() 在 WebGPU 上自回归生成
↓
16. Token 回调统计数量、TPS 和思考阶段
↓
17. 文本回调持续发送 update
↓
18. App 把 output 追加到最后一条 assistant 消息
↓
19. 正常结束或 interrupt 停止生成
↓
20. Worker 发送 complete,App 恢复发送状态
从职责上再次归纳:
| 模块 | 整体职责 |
|---|---|
App.tsx |
收集输入、保存对话、发送命令、接收状态和维护页面数据 |
worker.js |
加载资源、格式化模型输入、执行推理、流式解码与中断控制 |
TextGenerationPipeline |
单例式管理 Tokenizer 和模型 Promise |
AutoTokenizer |
Chat Template、文本编码和结果解码 |
AutoModelForCausalLM |
在 WebGPU 后端执行自回归文本生成 |
TextStreamer |
把持续生成的 Token 转换成文本片段 |
InterruptableStoppingCriteria |
在不销毁 Worker 和模型的情况下停止本次生成 |
16. 本篇两份完整注释代码
16.1 worker.js 完整版
import {
AutoTokenizer, // 文本编码、解码和 Chat Template
AutoModelForCausalLM, // 因果语言模型
TextStreamer, // 流式解码生成内容
InterruptableStoppingCriteria, // 可由外部触发的停止条件
} from "@huggingface/transformers";
// 检查 Worker 环境是否能取得 WebGPU Adapter。
async function check() {
try {
const adapter = await navigator.gpu.requestAdapter();
if (!adapter) {
throw new Error("WebGPU is not supported (no adapter found)");
}
// 后续还可以检查 shader-f16 等特性。
// fp16_supported = adapter.features.has("shader-f16")
} catch (e) {
self.postMessage({
status: "error",
data: e.toString(),
});
}
}
/**
* 使用单例思路懒加载并复用 Tokenizer 与模型。
*/
class TextGenerationPipeline {
static model_id =
"onnx-community/DeepSeek-R1-Distill-Qwen-1.5B-ONNX";
static async getInstance(progress_callback = null) {
// 第一次调用时加载 Tokenizer,后续复用同一 Promise。
this.tokenizer ??= AutoTokenizer.from_pretrained(
this.model_id,
{
progress_callback,
},
);
console.log(this.tokenizer, "////////");
// 第一次调用时加载 q4f16 模型,并指定 WebGPU 后端。
this.model ??= AutoModelForCausalLM.from_pretrained(
this.model_id,
{
dtype: "q4f16",
device: "webgpu",
progress_callback,
},
);
return Promise.all([this.tokenizer, this.model]);
}
}
// 所有生成任务共用这个实例,外部 interrupt 命令才能停止正在运行的任务。
const stopping_criteria =
new InterruptableStoppingCriteria();
// 保存 generate() 返回的 KV Cache。
// 下一轮传入模型的代码仍被注释,因此跨轮复用尚未开启。
let past_key_values_cache = null;
async function generate(messages) {
// 复用已加载并预热的 Tokenizer 和模型。
const [tokenizer, model] =
await TextGenerationPipeline.getInstance();
// 把 role/content 对话转换成模型配套格式和 Token。
const inputs = tokenizer.apply_chat_template(messages, {
// 在末尾提示模型继续生成 assistant 回答。
add_generation_prompt: true,
// 返回 input_ids、attention_mask 等具名字段。
return_dict: true,
});
// 获取思考开始、结束标记对应的 Token ID。
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;
// 每生成一批新 Token 时执行。
const token_callback_function = (tokens) => {
// 第一个 Token 到达时记录高精度时间。
startTime ??= performance.now();
// 第一个 Token 只计数,从第二个开始计算平均 TPS。
if (numTokens++ > 0) {
tps =
(numTokens / (performance.now() - startTime)) *
1000;
}
// 检测到 </think> 后切换到正式回答阶段。
if (tokens[0] == END_THINKING_TOKEN_ID) {
state = "answering";
}
};
// TextStreamer 解码出可显示文本时执行。
const callback_function = (output) => {
self.postMessage({
status: "update",
output,
tps,
numTokens,
state,
});
};
const streamer = new TextStreamer(tokenizer, {
// 不重复输出完整聊天 Prompt。
skip_prompt: true,
// 不显示模型内部特殊标记。
skip_special_tokens: true,
callback_function,
token_callback_function,
});
// 必须先创建前端 assistant 消息,再发送流式 update。
self.postMessage({ status: "start" });
const { past_key_values, sequences } =
await model.generate({
...inputs,
// 这里只保存了缓存,跨轮复用仍未正式启用。
// past_key_values: past_key_values_cache,
do_sample: false,
max_new_tokens: 2048,
streamer,
stopping_criteria,
return_dict_in_generate: true,
});
// 接收并保存本轮生成返回的 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,
});
}
async function load() {
self.postMessage({
status: "loading",
data: "Loading model...",
});
const [tokenizer, model] =
await TextGenerationPipeline.getInstance((x) => {
console.log(x, "//////////////");
self.postMessage(x);
});
self.postMessage({
status: "loading",
data: "Compiling shaders and warming up model...",
});
// 使用最小输入完成 WebGPU 模型预热。
const inputs = tokenizer("a");
console.log(inputs);
await model.generate({ ...inputs, max_new_tokens: 1 });
self.postMessage({ status: "ready" });
}
// 统一接收 App 发来的命令。
self.addEventListener("message", async (e) => {
const { type, data } = e.data;
switch (type) {
case "check":
check();
break;
case "load":
load();
break;
case "generate":
// 新任务开始前清除旧的 interrupt 标志。
stopping_criteria.reset();
generate(data);
break;
case "interrupt":
// 通知正在运行的 generate() 正常退出。
stopping_criteria.interrupt();
break;
case "reset":
// 新会话可以清空缓存和中断标志。
past_key_values_cache = null;
stopping_criteria.reset();
break;
}
});
16.2 App.tsx 完整版
import { useEffect, useState, useRef } from "react";
import Progress from "./components/Progress.tsx";
import ArrowRightIcon from "./components/icons/ArrowRightIcon.tsx";
import StopIcon from "./components/icons/StopIcon.tsx";
const IS_WEBGPU_AVAILABLE = !!navigator.gpu;
// 下一篇用于聊天区自动滚动判断。
const STICKY_SCROLL_THRESHOLD = 120;
// 下一篇用于欢迎页示例问题。
const EXAMPLES = [
"Solve the equation x^2 - 3x + 2 = 0",
"Lily is three times older than her son. In 15 years, she will be twice as old as him. How old is she now?",
"Write python code to compute the nth fibonacci number.",
];
function App() {
// 保存 Worker 和两个 DOM 节点。
const worker = useRef(null);
const textareaRef = useRef(null);
const chatContainerRef = useRef(null);
// 模型加载、进度与运行状态。
const [status, setStatus] = useState(null);
const [error, setError] = useState(null);
const [loadingMessage, setLoadingMessage] = useState("");
const [progressItems, setProgressItems] = useState([]);
const [isRunning, setIsRunning] = useState(false);
// 对话输入、消息与生成性能数据。
const [input, setInput] = useState("");
const [messages, setMessages] = useState([]);
const [tps, setTps] = useState(null);
const [numTokens, setNumTokens] = useState(null);
function onEnter(message) {
// 使用最新消息数组追加 user 消息。
setMessages((prev) => [
...prev,
{ role: "user", content: message },
]);
// 重置统计、进入运行状态并清空输入框。
setTps(null);
setIsRunning(true);
setInput("");
}
function onInterrupt() {
// 不立即修改 isRunning,等待 Worker 的 complete 消息。
worker.current.postMessage({ type: "interrupt" });
}
useEffect(() => {
if (!worker.current) {
// module Worker 才能在 worker.js 中使用 import。
worker.current = new Worker(
new URL("./worker.js", import.meta.url),
{ type: "module" },
);
worker.current.postMessage({ type: "check" });
}
const onMessageReceived = (e) => {
switch (e.data.status) {
case "loading":
setStatus("loading");
setLoadingMessage(e.data.data);
break;
case "initiate":
// 新文件开始加载:追加到进度数组。
setProgressItems((prev) => [...prev, e.data]);
break;
case "progress":
// 按 file 更新对应文件进度。
setProgressItems((prev) =>
prev.map((item) => {
if (item.file === e.data.file) {
return { ...item, ...e.data };
}
return item;
}),
);
break;
case "done":
// 文件完成后从进度数组删除。
setProgressItems((prev) =>
prev.filter(
(item) => item.file !== e.data.file,
),
);
break;
case "ready":
setStatus("ready");
break;
case "start": {
// 先创建一条空 assistant 消息,承接后续流式内容。
setMessages((prev) => [
...prev,
{ role: "assistant", content: "" },
]);
break;
}
case "update": {
const {
output,
tps,
numTokens,
state,
} = e.data;
setTps(tps);
setNumTokens(numTokens);
setMessages((prev) => {
// 同时创建新数组和新消息对象。
const cloned = [...prev];
const last = cloned.at(-1);
const data = {
...last,
content: last.content + output,
};
// 第一次进入正式回答时记录分界位置。
if (
data.answerIndex === undefined &&
state === "answering"
) {
data.answerIndex = last.content.length;
}
cloned[cloned.length - 1] = data;
return cloned;
});
break;
}
case "complete":
// 正常结束和中断结束都从这里恢复发送状态。
setIsRunning(false);
break;
case "error":
setError(e.data.data);
break;
}
};
const onErrorReceived = (e) => {
console.error("Worker error:", e);
};
worker.current.addEventListener(
"message",
onMessageReceived,
);
worker.current.addEventListener(
"error",
onErrorReceived,
);
// 防止开发环境或组件生命周期造成重复监听。
return () => {
worker.current.removeEventListener(
"message",
onMessageReceived,
);
worker.current.removeEventListener(
"error",
onErrorReceived,
);
};
}, []);
// messages 更新后,把最新完整对话同步给 Worker。
useEffect(() => {
// 初始化阶段没有 user 消息,不生成。
if (
messages.filter((x) => x.role === "user").length === 0
) {
return;
}
// start 加入 assistant 后会再次触发 Effect,必须阻止重复生成。
if (messages.at(-1).role === "assistant") {
return;
}
setTps(null);
worker.current.postMessage({
type: "generate",
data: messages,
});
}, [messages]);
return (
IS_WEBGPU_AVAILABLE ? (
<div className="flex flex-col h-screen mx-auto items justify-end text-gray-800 dark:text-gray-200 bg-white dark:bg-gray-900">
{/* 欢迎页 */}
{status === null && messages.length === 0 && (
<div className="h-full overflow-auto scrollbar-thin flex justify-center items-center flex-col relative">
<div className="flex flex-col items-center mb-1 max-w-[400px] text-center">
<img
src="logo.png"
width="80%"
height="auto"
className="block drop-shadow-lg bg-transparent"
></img>
<h1 className="text-4xl font-bold mb-1">
DeepSeek-R1 WebGPU
</h1>
<h2 className="font-semibold">
A next-generation reasoning model that runs locally in your
browser with WebGPU acceleration.
</h2>
</div>
<div className="flex flex-col items-center px-4">
<p className="max-w-[510px] mb-4">
<br />
You are about to load{" "}
<a
href="https://huggingface.co/onnx-community/DeepSeek-R1-Distill-Qwen-1.5B-ONNX"
target="_blank"
rel="noreferrer"
className="font-medium underline"
>
DeepSeek-R1-Distill-Qwen-1.5B
</a>
, a 1.5B parameter reasoning LLM optimized for in-browser
inference. Everything runs entirely in your browser with{" "}
<a
href="https://huggingface.co/docs/transformers.js"
target="_blank"
rel="noreferrer"
className="underline"
>
🤗 Transformers.js
</a>{" "}
and ONNX Runtime Web, meaning no data is sent to a server. Once
loaded, it can even be used offline. The source code for the demo
is available on{" "}
<a
href="https://github.com/huggingface/transformers.js-examples/tree/main/deepseek-r1-webgpu"
target="_blank"
rel="noreferrer"
className="font-medium underline"
>
GitHub
</a>
.
</p>
{error && (
<div className="text-red-500 text-center mb-2">
<p className="mb-1">
Unable to load model due to the following error:
</p>
<p className="text-sm">{error}</p>
</div>
)}
<button
className="border px-4 py-2 rounded-lg bg-blue-400 text-white hover:bg-blue-500 disabled:bg-blue-100 cursor-pointer disabled:cursor-not-allowed select-none"
onClick={() => {
worker.current.postMessage({ type: "load" });
setStatus("loading");
}}
disabled={status !== null || error !== null}
>
Load model
</button>
</div>
</div>
)}
{/* 模型加载页 */}
{status === "loading" && (
<>
<div className="w-full max-w-[500px] text-left mx-auto p-4 bottom-0 mt-auto">
<p className="text-center mb-1">
{loadingMessage}
</p>
{progressItems.map(
({ file, progress, total }, i) => (
<Progress
key={i}
text={file}
percentage={progress}
total={total}
/>
),
)}
</div>
</>
)}
{/*
生成数据已经进入 messages。
下一篇在此渲染用户消息、思考过程、正式回答和性能指标。
*/}
{status === "ready" && (
<div
ref={chatContainerRef}
className="overflow-y-auto scrollbar-thin w-full flex flex-col items-center h-full"
>
</div>
)}
{/* 输入框 */}
<div className="mt-2 border border-gray-300 dark:bg-gray-700 rounded-lg w-[600px] max-w-[80%] max-h-[200px] mx-auto relative mb-3 flex">
<textarea
ref={textareaRef}
className="scrollbar-thin w-[550px] dark:bg-gray-700 px-3 py-4 rounded-lg bg-transparent border-none outline-hidden text-gray-800 disabled:text-gray-400 dark:text-gray-200 placeholder-gray-500 dark:placeholder-gray-400 disabled:placeholder-gray-200 resize-none disabled:cursor-not-allowed"
placeholder="Type your message..."
rows={1}
value={input}
disabled={status !== "ready"}
title={
status === "ready"
? "Model is ready"
: "Model not loaded yet"
}
onKeyDown={(e) => {
if (
input.length > 0 &&
!isRunning &&
e.key === "Enter" &&
!e.shiftKey
) {
e.preventDefault();
onEnter(input);
}
}}
onInput={(e) =>
setInput(e.currentTarget.value)
}
/>
</div>
{/* 运行时显示停止;有输入时显示可发送;否则显示灰色箭头。 */}
{isRunning ? (
<div
className="cursor-pointer"
onClick={onInterrupt}
>
<StopIcon className="h-8 w-8 p-1 rounded-md text-gray-800 dark:text-gray-100 absolute right-3 bottom-3" />
</div>
) : input.length > 0 ? (
<div
className="cursor-pointer"
onClick={() => onEnter(input)}
>
<ArrowRightIcon className="h-8 w-8 p-1 bg-gray-800 dark:bg-gray-100 text-white dark:text-black rounded-md absolute right-3 bottom-3" />
</div>
) : (
<div>
<ArrowRightIcon className="h-8 w-8 p-1 bg-gray-200 dark:bg-gray-600 text-gray-50 dark:text-gray-800 rounded-md absolute right-3 bottom-3" />
</div>
)}
</div>
) : (
<div className="fixed w-screen h-screen bg-black z-10 bg-opacity-[92%] text-white text-2xl font-semibold flex justify-center items-center text-center">
WebGPU is not supported
<br />
by this browser :(
</div>
)
);
}
export default App;
完整代码中仍然保留 STICKY_SCROLL_THRESHOLD、EXAMPLES、tps 和 numTokens。它们将在下一篇页面实现中被消费。messages Effect 中的 setTps(null) 与 onEnter() 中的重置作用重复,后续整理时可保留一个;它不改变本篇生成链路,但会触发严格 ESLint 对 Effect 内同步 State 更新的提示。
总结
这一篇把第三篇预留的交互框架接成了真实生成链路。App 通过 onEnter() 把输入转换成 role/content 用户消息,监听 messages 的 Effect 再将最新完整对话发送给 Worker。Worker 使用 apply_chat_template() 套用模型配套的聊天格式,通过 model.generate() 在 WebGPU 上执行自回归生成,并使用 TextStreamer 分别处理 Token 回调和文本回调。
Token 回调负责统计 numTokens、计算平均 TPS,并检测 </think> 将状态从 thinking 切换为 answering;文本回调则发送 update。App 收到 start 后先建立空助手消息,后续通过不可变更新不断追加流式片段,并使用 answerIndex 记录思考过程和正式答案的分界。生成结束后,complete 统一恢复按钮状态。
可中断生成也已经贯通:停止按钮发送 interrupt,Worker 修改 InterruptableStoppingCriteria,模型在下一轮检查时结束,而不销毁 Worker 和已经加载的模型。KV Cache 在这一版完成了结果接收与保存,但跨轮传入仍被注释,不能与完整聊天历史输入直接混用。
至此,项目从 WebGPU 检查、模型下载、进度展示和预热,一直推进到消息提交、聊天模板、流式生成、性能统计与中断控制。生成数据已经进入 React State,下一篇只需围绕 messages、answerIndex、tps 和 numTokens 完成聊天页面呈现与交互优化。
更多推荐


所有评论(0)