企业官网建设流程全解析

WebLLM LogitProcessor 与底层 forwardTokensAndSample 详解：在采样前自由改写 logits

【免费下载链接】web-llmHigh-performance In-browser LLM Inference Engine项目地址: https://gitcode.com/GitHub_Trending/we/web-llm

导读

本指南以 examples/logit-processor 示例为完整实战载体，讲解 WebLLM（@mlc-ai/web-llm）中LogitProcessor的核心用法：如何在 token 采样之前直接操纵原始 logits（例如把某个 token 的 logit 置为inf/-inf），以及如何使用底层 APIforwardTokensAndSample()获得比chat.completions.create()更细粒度的生成控制。读完本文，你将掌握自定义 logit 处理器、在 Web Worker 与无 Worker 两种模式下注册处理器、利用处理器维护有状态逻辑并清理状态，以及使用底层逐 token 前向采样接口完成可控生成的完整方案。

一、LogitProcessor：采样前的 logits 修改钩子

1.1 它在推理链路中的位置

在 WebLLM 的推理流程中，模型每次前向（forward）后会产出一份 logits，随后才执行采样得到下一个 token。LogitProcessor正是插入在这两个环节之间的钩子：在前向之后、采样之前，对 logits 进行后处理。接口定义位于 src/types.ts：

export interface LogitProcessor { /** * Process logits after forward() and before sampling implicitly, happens on the CPU. * @param logits The logits right after forward(). * Returns the processed logits. */ processLogits: (logits: Float32Array) => Float32Array; /** * Use the sampled token to update the LogitProcessor's internal state. Called implicitly * right after the next token is sampled/committed. * @param token Token sampled from the processed logits. */ processSampledToken: (token: number) => void; /** * Called when in `MLCEngine.resetChat()`. Can clear internal states. */ resetState: () => void; }

接口包含三个方法，构成了处理器的完整生命周期：

• processLogits(logits: Float32Array): Float32Array：核心改写入口。接收前向计算得到的 logits（Float32Array，每个元素对应词表中的一个 token），返回处理后的 logits。可用于约束解码，例如把禁止出现的 token 对应 logit 置为-inf强制屏蔽，或把必须出现的 token 置为inf强制选中。需要注意的是，该步骤发生在CPU 侧。
• processSampledToken(token: number): void：在下一个 token 被采样/提交后隐式调用，用于让处理器感知"上一步实际采到了哪个 token"，从而维护跨步的内部状态（如记录已生成的 token 序列）。
• resetState(): void：当引擎调用MLCEngine.resetChat()时被调用，用于清空内部状态，保证新一轮对话/生成从干净状态开始。

1.2 底层调用时序

从源码可以确认完整的执行顺序。在 src/llm_chat.ts 中，logits 处理发生在采样前：

1. 若已注册logitProcessor，将 GPU 上的 logits 拷贝到 CPU（updateLogitsOnCPU）；
2. 调用logitProcessor.processLogits(logitsOnCPUArray)，并用返回值回写 GPU/CPU 侧的 logits 副本；
3. 随后基于处理后的 logits 执行采样，并在采样提交后调用logitProcessor.processSampledToken(sampledToken)（见 src/llm_chat.ts）；
4. 而resetState()则在MLCEngine.resetChat()流程中被触发（见 src/llm_chat.ts）。

另外接口注释明确指出：如果同时使用了GenerationConfig.logit_bias，则logit_bias是在processLogits()之后应用的，因此LogitProcessor的改写结果会被 logit_bias 进一步调整——这一点在需要二者叠加控制时需特别注意。

二、示例工程结构速览

logit-processor示例位于 examples/logit-processor，核心源码文件如下：

示例默认使用phi-2-q4f32_1-MLC模型，package.json（examples/logit-processor/package.json）中声明了对@mlc-ai/web-llm: ^0.2.84的依赖，并通过 Parcel 提供开发服务器（默认端口8885）。

三、自定义 LogitProcessor：强制采样 token 0

examples/logit-processor/src/my_logit_processor.ts 给出了一个完整可运行的自定义实现：

import * as webllm from "@mlc-ai/web-llm"; // Define LogitProcessor export class MyLogitProcessor implements webllm.LogitProcessor { private tokenSequence: Array<number> = []; processLogits(logits: Float32Array): Float32Array { logits[0] = 100.0; // should be enough so that we always sample token 0 below return logits; } processSampledToken(token: number): void { this.tokenSequence.push(token); console.log("processSampledToken: " + this.tokenSequence.length); } resetState(): void { this.tokenSequence = []; console.log("resetState"); } }

3.1 核心改写逻辑

processLogits中仅做了一件事：将 token 0 的 logit 手动设置为100.0。由于该数值远大于词表中其他 token 的 logit，采样器几乎必然选中 token 0——正如 README 所述，观察控制台输出即可验证"始终采样到 token 0"这一预期结果。这演示了最直接的 logits 操纵手法：

• 将某 token 置为极大值（如100.0）→ 强制采样该 token；
• 反向操作，将某 token 置为-inf→ 完全禁止采样该 token，可用于敏感词过滤、格式约束等场景。

3.2 有状态的处理器

MyLogitProcessor还演示了处理器的有状态特性：tokenSequence数组在processSampledToken中不断累积每次实际采样到的 token，并在控制台打印累计计数。这意味着处理器可以跨 token 步长记住"已经生成了什么"，进而实现依赖于前文上下文的动态 logits 修改（例如防止某序列重复出现）。

3.3 状态清理

resetState()将tokenSequence重置为空数组，保证状态不会泄漏到新一轮生成中。它在引擎调用resetChat()时被自动触发，用户通常无需手动调用（示例入口代码中的engine.resetChat()即会联动触发，详见第七节）。

四、无 Worker 模式：通过logitProcessorRegistry注册

在不使用 Web Worker 时，直接在调用CreateMLCEngine时传入一个"模型 ID → 处理器实例"的映射表（logitProcessorRegistry），WebLLM 引擎会在加载模型时按模型 ID 自动查找并绑定对应的处理器。入口代码见 examples/logit-processor/src/logit_processor.ts：

// Instantiate myLogitProcessor, registering in the logitProcessorRegistry const myLogitProcessor = new MyLogitProcessor(); const logitProcessorRegistry = new Map<string, webllm.LogitProcessor>(); logitProcessorRegistry.set("phi-2-q4f32_1-MLC", myLogitProcessor); let engine: webllm.MLCEngineInterface; if (USE_WEB_WORKER) { // see worker.ts on how LogitProcessor plays a role there engine = await webllm.CreateWebWorkerMLCEngine( new Worker(new URL("./worker.ts", import.meta.url), { type: "module" }), "phi-2-q4f32_1-MLC", { initProgressCallback: initProgressCallback }, ); } else { engine = await webllm.CreateMLCEngine("phi-2-q4f32_1-MLC", { initProgressCallback: initProgressCallback, logitProcessorRegistry: logitProcessorRegistry, }); }

关键的注册链路在引擎源码中得到印证：MLCEngine构造时接收engineConfig.logitProcessorRegistry（src/engine.ts），并在reloadInternal()中通过this.logitProcessorRegistry?.get(modelId)依据模型 ID 取出对应处理器（src/engine.ts）。因此：

• 注册表是"模型 ID → 处理器"的多模型映射，同一页面中不同模型可各自绑定不同的处理器；
• 只有注册表包含当前加载模型 ID 的条目时，处理器才会被实际挂载；
• 若希望为所有会话统一注入处理器，可在每次创建引擎时传入同一个注册表。

五、Web Worker 模式：在 Worker 线程内注册

当USE_WEB_WORKER = true时，引擎运行在独立的 Web Worker 线程中，此时需要在 Worker 线程代码 examples/logit-processor/src/worker.ts 内完成注册：

// Serve the chat workload through web worker import * as webllm from "@mlc-ai/web-llm"; import { MyLogitProcessor } from "./my_logit_processor"; console.log("Use web worker for logit processor"); const myLogitProcessor = new MyLogitProcessor(); const logitProcessorRegistry = new Map<string, webllm.LogitProcessor>(); logitProcessorRegistry.set("phi-2-q4f32_1-MLC", myLogitProcessor); const handler = new webllm.WebWorkerMLCEngineHandler(); handler.setLogitProcessorRegistry(logitProcessorRegistry); self.onmessage = (msg: MessageEvent) => { handler.onmessage(msg); };

该模式下的注册链路是：主线程通过CreateWebWorkerMLCEngine(new Worker(new URL("./worker.ts", import.meta.url), { type: "module" }), ...)创建引擎（注意示例采用 ESM Worker），Worker 内部实例化WebWorkerMLCEngineHandler并调用handler.setLogitProcessorRegistry(logitProcessorRegistry)，把注册表注入到 Worker 侧引擎（对应实现见 src/web_worker.ts）。主线程与 Worker 之间通过消息协议通信，forwardTokensAndSample等调用同样经由 Worker 消息路由分发。

两种模式的核心区别：无 Worker 模式在创建引擎时传logitProcessorRegistry；Worker 模式则必须先在worker.ts中通过setLogitProcessorRegistry注册。二者二选一，由USE_WEB_WORKER布尔常量一键切换。

六、底层 API：forwardTokensAndSample()逐 token 细粒度生成

6.1 与 chat/completions API 的差异

chat.completions.create()面向自回归对话场景，帮你管理 prompt 拼装、历史上下文与终止条件；而forwardTokensAndSample()是低层级的逐 token 接口，直接接收 token ID 数组并返回下一个 token ID，把生成循环完全交到开发者手中，适用于需要精确控制每一步前向与采样的场景。

6.2 示例中的自回归循环

示例在 examples/logit-processor/src/logit_processor.ts 中演示了完整用法：

// Below we demonstrate the usage of a low-level API `forwardTokensAndSample()` const prompt: Array<number> = [42]; let nextToken = await engine.forwardTokensAndSample( prompt, /*isPrefill=*/ true, ); console.log(nextToken); let counter = prompt.length; while (counter < AUTOREGRESS_LIMIT) { counter += 1; nextToken = await engine.forwardTokensAndSample( [nextToken], /*isPrefill=*/ false, ); console.log(nextToken); }

要点拆解：

• 首步调用：传入初始 prompt token 序列（示例为[42]，即单个 token ID），并将isPrefill置为true，表示本次调用负责对输入做前向填充（prefill），返回第一个生成 token；
• 自回归循环：此后每次仅把上一个 token 作为输入（[nextToken]），isPrefill置为false，不断追加生成，直到达到AUTOREGRESS_LIMIT = 32的上限；
• token ID 是原始数值：forwardTokensAndSample的输入输出都是 token ID（number），不做文本解码，因此需要使用者自行处理 tokenizer/文本映射。

6.3 底层实现细节

从 src/llm_chat.ts 的实现可见其内部流程：

1. 将输入按prefillChunkSize分块（getChunkedPrefillInputData）；
2. 逐块执行embedAndForward完成前向，同时更新 KV Cache 长度并校验；
3. 对最后一块产出的 logits 调用sampleTokenFromLogits采样下一个 token；
4. 按isPrefill标记累加 prefill 时间与 token 数统计。

在 src/engine.ts 中，forwardTokensAndSample被定义为引擎级入口，会转发到当前所选模型的 pipeline 上执行；同时 src/engine.ts 的注释指出，截至目前runtimeStatsText()的主要使用场景正是forwardTokensAndSample()流程——示例结尾也调用了console.log(await engine.runtimeStatsText())打印运行统计。换言之，这套底层 API 与运行时性能统计天然兼容，适合做延迟/吞吐测量。

七、状态清理联动：resetChat()触发resetState()

有状态处理器的状态需要在会话重置时清除，WebLLM 已内置这一联动机制。示例代码在完成第一轮生成后调用：

// By calling `engine.resetChat()`, we triggers MyLogitProcessor.resetState() engine.resetChat(); counter = prompt.length; nextToken = await engine.forwardTokensAndSample(prompt, /*isPrefill=*/ true); // ... 再次进行第二轮自回归生成

调用engine.resetChat()后，MyLogitProcessor.resetState()会被隐式触发（控制台将打印resetState），随后示例重新以同一 prompt 启动第二轮生成，以验证"清理后从零开始"的行为。对应实现可见 src/engine.ts 中resetChat的转发逻辑，以及 src/llm_chat.ts 中this.logitProcessor?.resetState()的调用点。

这个设计的意义在于：处理器状态与对话会话的生命周期保持一致。若你实现了依赖前文做约束的处理器（例如禁止连续重复、强制特定格式），一轮对话结束后无需手动清状态，resetChat()会统一处理。

八、运行方式与本地开发

8.1 快速运行

在examples/logit-processor目录下执行：

npm install npm start

npm start实际执行parcel src/logit_processor.html --port 8885（见 examples/logit-processor/package.json），即用 Parcel 在8885 端口启动开发服务器。浏览器打开页面后，模型加载进度会显示在init-label中，而生成的 token 序列与processSampledToken计数等日志均输出在浏览器控制台，因此务必打开 DevTools Console 观察效果。

8.2 切换两种运行模式

通过修改 examples/logit-processor/src/logit_processor.ts 顶部的常量即可切换：

const USE_WEB_WORKER = true; // Toggle this to use Logit Processor without a web worker const AUTOREGRESS_LIMIT = 32; // How many tokens to generate for this test

• USE_WEB_WORKER = true：引擎运行在 Web Worker 中，处理器在worker.ts内注册（推荐用于不阻塞 UI 的生产场景）；
• USE_WEB_WORKER = false：引擎与处理器均运行在主线程，走CreateMLCEngine+logitProcessorRegistry路径。

8.3 本地构建 WebLLM 核心包（可选）

如果你希望直接改动 WebLLM 核心包本身，可以把package.json中的@mlc-ai/web-llm依赖改为本地路径形式"file:../.."，然后按项目根目录 README.md 中的 build from source 指引本地构建 WebLLM。此选项仅推荐给需要深入 hack WebLLM 核心源码的开发者，普通示例使用无需此操作。

九、总结

LogitProcessor与forwardTokensAndSample()构成了 WebLLM 中"采样前改写 logits + 逐 token 细粒度控制"的两大利器：

• LogitProcessor提供processLogits/processSampledToken/resetState三个生命周期钩子，可在 CPU 侧修改 logits、感知每次采样结果、随resetChat()清理状态；通过"模型 ID → 处理器"注册表，在无 Worker（CreateMLCEngine+logitProcessorRegistry）与 Web Worker（WebWorkerMLCEngineHandler.setLogitProcessorRegistry）两种架构下均可挂载。
• forwardTokensAndSample(inputIds, isPrefill)则绕过高层对话 API，直接以 token ID 为单位驱动 prefill 与自回归循环，并天然兼容runtimeStatsText()性能统计，适合需要精细控制与性能观测的场景。

参考实现：完整的可运行示例位于 examples/logit-processor，接口定义见 src/types.ts，引擎侧注册与转发逻辑见 src/engine.ts，底层处理与采样流程见 src/llm_chat.ts 与 src/llm_chat.ts。

【免费下载链接】web-llmHigh-performance In-browser LLM Inference Engine项目地址: https://gitcode.com/GitHub_Trending/we/web-llm