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

WebLLM LogitProcessor 与底层 forwardTokensAndSample 详解在采样前自由改写 logits【免费下载链接】web-llmHigh-performance In-browser LLM Inference Engine项目地址: https://gitcode.com/GitHub_Trending/we/web-llm导读本指南以 examples/logit-processor 示例为完整实战载体讲解 WebLLMmlc-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.tsexport 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 LogitProcessors 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核心改写入口。接收前向计算得到的 logitsFloat32Array每个元素对应词表中的一个 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 处理发生在采样前若已注册logitProcessor将 GPU 上的 logits 拷贝到 CPUupdateLogitsOnCPU调用logitProcessor.processLogits(logitsOnCPUArray)并用返回值回写 GPU/CPU 侧的 logits 副本随后基于处理后的 logits 执行采样并在采样提交后调用logitProcessor.processSampledToken(sampledToken)见 src/llm_chat.ts而resetState()则在MLCEngine.resetChat()流程中被触发见 src/llm_chat.ts。另外接口注释明确指出如果同时使用了GenerationConfig.logit_bias则logit_bias是在processLogits()之后应用的因此LogitProcessor的改写结果会被 logit_bias 进一步调整——这一点在需要二者叠加控制时需特别注意。二、示例工程结构速览logit-processor示例位于 examples/logit-processor核心源码文件如下文件作用examples/logit-processor/src/logit_processor.ts入口主逻辑注册处理器、按开关选择 Web Worker 模式、驱动forwardTokensAndSample()循环examples/logit-processor/src/worker.tsWeb Worker 线程代码在 Worker 内注册 LogitProcessor 并接管消息分发examples/logit-processor/src/my_logit_processor.ts自定义MyLogitProcessor实现展示了强制采样 token 0 与有状态处理examples/logit-processor/src/logit_processor.html测试页面打开浏览器控制台观察处理器效果示例默认使用phi-2-q4f32_1-MLC模型package.jsonexamples/logit-processor/package.json中声明了对mlc-ai/web-llm: ^0.2.84的依赖并通过 Parcel 提供开发服务器默认端口8885。三、自定义 LogitProcessor强制采样 token 0examples/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: Arraynumber []; 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 → 处理器实例的映射表logitProcessorRegistryWebLLM 引擎会在加载模型时按模型 ID 自动查找并绑定对应的处理器。入口代码见 examples/logit-processor/src/logit_processor.ts// Instantiate myLogitProcessor, registering in the logitProcessorRegistry const myLogitProcessor new MyLogitProcessor(); const logitProcessorRegistry new Mapstring, 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.logitProcessorRegistrysrc/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 Mapstring, 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 WorkerWorker 内部实例化WebWorkerMLCEngineHandler并调用handler.setLogitProcessorRegistry(logitProcessorRegistry)把注册表注入到 Worker 侧引擎对应实现见 src/web_worker.ts。主线程与 Worker 之间通过消息协议通信forwardTokensAndSample等调用同样经由 Worker 消息路由分发。两种模式的核心区别无 Worker 模式在创建引擎时传logitProcessorRegistryWorker 模式则必须先在worker.ts中通过setLogitProcessorRegistry注册。二者二选一由USE_WEB_WORKER布尔常量一键切换。六、底层 APIforwardTokensAndSample()逐 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: Arraynumber [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 IDnumber不做文本解码因此需要使用者自行处理 tokenizer/文本映射。6.3 底层实现细节从 src/llm_chat.ts 的实现可见其内部流程将输入按prefillChunkSize分块getChunkedPrefillInputData逐块执行embedAndForward完成前向同时更新 KV Cache 长度并校验对最后一块产出的 logits 调用sampleTokenFromLogits采样下一个 token按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 startnpm 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 testUSE_WEB_WORKER true引擎运行在 Web Worker 中处理器在worker.ts内注册推荐用于不阻塞 UI 的生产场景USE_WEB_WORKER false引擎与处理器均运行在主线程走CreateMLCEnginelogitProcessorRegistry路径。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 → 处理器注册表在无 WorkerCreateMLCEnginelogitProcessorRegistry与 Web WorkerWebWorkerMLCEngineHandler.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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考