ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

TensorRT-LLM 如何为 trtllm-serve 注入自定义 Post-Processing Hook 改写生成输出

TensorRT-LLM 如何为 trtllm-serve 注入自定义 Post-Processing Hook 改写生成输出 TensorRT-LLM 如何为 trtllm-serve 注入自定义 Post-Processing Hook 改写生成输出【免费下载链接】TensorRT-LLMTensorRT LLM provides users with an easy-to-use Python API to define Large Language Models (LLMs) and supports state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs. TensorRT LLM also contains components to create Python and C runtimes that orchestrate the inference execution in a performant way.项目地址: https://gitcode.com/GitHub_Trending/te/TensorRT-LLM如果你的部署需要在模型输出返回给客户端之前做改写、脱敏、拦截或提前终止例如检测到违禁词就停止流式输出TensorRT-LLM 的trtllm-serve提供了一个用户可注入的post-processing hook一个原生的、按请求执行的切面它在每个生成输出detokenization 之后、各端点 response formatter 之前运行无需修改 TensorRT LLM 源码即可改变客户端看到的输出。功能文档位于 docs/source/features/post-processor-hook.md。需要注意的前提该功能目前是prototype接口在未来的 release 中可能变化文档中的 note 原文。1. Hook 的运行时机与接口hook 是一个普通的 Python 可调用类通过 import path 提供机制类似--custom_tokenizer。它由LLM实例持有启用 post-processing worker 池时在每个 worker 进程内构建一次每个输出、每个流式 chunk 调用一次外加一次 final 调用因此可以自行持有按请求划分状态。hook 只需实现一个方法__call__(self, chunk) - verdict。接口定义见 tensorrt_llm/executor/postprocessor_hook.pyfrom tensorrt_llm.executor.postprocessor_hook import ( PostProcessorHookChunk, PostProcessorHookVerdict, emit, suppress, terminate, ) class MyPostProcessorHook: def __call__(self, chunk: PostProcessorHookChunk) - PostProcessorHookVerdict: return emit(chunk.text_diff) # pass through unchangedPostProcessorHookChunk是交给 hook 的单 chunk 载荷字段包括request_id同一响应的每个 chunk 传相同值可用来做 per-request 状态 key、output_index、text_diff本 chunk 新增的 detokenized 文本非流式请求下等于text、text到目前为止的累计文本、token_ids_diff、is_final、aborted、streaming。逐字段的权威说明见 postprocessor_hook.py 中的PostProcessorHookChunkdataclass。__call__返回三种 verdict 之一Helper效果emit(text)为本 chunk 发出text。可原样透传emit(chunk.text_diff)也可改写/脱敏。只影响 text 通道不会生成对应的token_ids。suppress()在所有客户端可见通道上扣住本 chunk——text、token_ids、logprobsdetokenizefalse返回的 token 输出同样被扣住。terminate(reason)停止该请求的流在全部通道上扣住终止 chunk。reason会作为响应的stop_reason返回引擎请求被取消。verdict 是per chunk的suppress()只扣住当前 chunkterminate()停止生成但保留此前已发出的 chunk。流式与非流式行为一致——非流式响应恰好包含 hook 在第一次suppress/terminate之前发出的内容与流式客户端看到的一致。如果一个 hook 要扣住整个输出all-or-nothing应从第一个 chunk 开始suppress()因为它能看到每个 chunk而不是先 emit 再 terminate。2. 编写一个可复用的 Hook文档给出三个可直接参考的示例。无状态改写——把所有 chunk 转大写from tensorrt_llm.executor.postprocessor_hook import PostProcessorHookChunk, PostProcessorHookVerdict, emit class UpperCaseHook: def __call__(self, chunk: PostProcessorHookChunk) - PostProcessorHookVerdict: return emit(chunk.text_diff.upper())有状态 guardrail——按请求累计文本命中违禁词即终止流并在请求结束时释放状态from tensorrt_llm.executor.postprocessor_hook import ( PostProcessorHookChunk, PostProcessorHookVerdict, emit, terminate, ) class BannedPhraseGuard: _BANNED (forbidden phrase,) def __init__(self): # Per-request accumulators owned entirely by the hook. self._buffers: dict[int, str] {} def __call__(self, chunk: PostProcessorHookChunk) - PostProcessorHookVerdict: buffer self._buffers.get(chunk.request_id, ) chunk.text_diff.lower() self._buffers[chunk.request_id] buffer if any(phrase in buffer for phrase in self._BANNED): self._buffers.pop(chunk.request_id, None) return terminate(banned_phrase) if chunk.is_final: self._buffers.pop(chunk.request_id, None) return emit(chunk.text_diff)抑制输出——扣住所有客户端可见文本from tensorrt_llm.executor.postprocessor_hook import PostProcessorHookChunk, PostProcessorHookVerdict, suppress class SuppressHook: def __call__(self, chunk: PostProcessorHookChunk) - PostProcessorHookVerdict: return suppress()关于 per-request 状态文档明确了三条约束hook 实例被所有它处理的请求共享所以任何 per-request 状态必须以chunk.request_id为 key并在看到chunk.is_final之后或terminate之后释放状态不跨进程、不跨独立LLM实例共享启用 worker 池时单个请求的所有 chunk 仍路由到同一 worker该请求的状态保持一致即使引擎层批处理了多个请求hook 仍是每请求每 chunk 调用一次没有 batch 调用形式。3. 启用 Hook把你的 hook 类的点分 import path 传给--post_processor_hooktrtllm-serve model --post_processor_hook my_pkg.guardrail.MyPostProcessorHook或在通过--extra_llm_api_options传入的 YAML 配置中等价地设置post_processor_hook: my_pkg.guardrail.MyPostProcessorHook其中model换成你的模型路径my_pkg.guardrail.MyPostProcessorHook换成你的模块路径与类名。类必须满足两个条件可导入——服务器及其 post-processing worker 进程如果有启动时已pip install或在PYTHONPATH上可 pickle 且无参可构造——hook 在各自进程内按引用重建__init__不接受必需参数它是唯一的一次性初始化点。hook 在有无 out-of-process post-processing worker 池--num_postprocess_workers时行为一致。4. 验证 Hook 是否生效仓库自带了两级测试可以直接运行来核对 hook 契约与端到端行为# hook 契约的单测rewrite / suppress / terminate、per-request state、loader pytest tests/unittest/executor/test_postprocessor_hook.py -v端到端测试位于 tests/unittest/llmapi/apps/_test_openai_post_processor.py它启动真实的trtllm-serve--post_processor_hook指向同目录下的示例 hook 模块 _postproc_hook_samples.py在 chat/completions 两个端点、流式与非流式、worker 池关闭in-proxy detok与开启两种配置下断言客户端可见的效果。该测试使用的判定标准可作为你自己验证时的参照UppercaseHook输出文本满足output output.upper()SuppressHook输出为空字符串TerminateHook输出为空且finish_reason stop而非跑满max_tokens。测试还特别验证了detokenizefalse无法绕过 hookterminate 场景下请求提前停止finish_reason stop而不是length且被扣住的内容不会从 token-id 通道泄漏。如果你自己接入 hook可以用同样的请求completions带extra_body{detokenize: False}检查token_ids是否为空来确认 suppress/terminate 覆盖了全部通道。5. 限制与边界以下限制直接来自功能文档接入前需要确认你的部署形态支持的端点chat/completions与completions流式与非流式。responses端点非 harmony 模型共享 detokenization 路径预期会应用 hook但当前端到端测试未覆盖该端点。不可被客户端绕过hook 是服务端 guardrail即使completions请求设置detokenizefalse也会执行——该 flag 只控制返回的通道text 还是 token ids。必须有 tokenizer--post_processor_hook与skip_tokenizer_init组合会在启动时被拒绝llm_args.py 中的校验会抛出post_processor_hook is not supported together with skip_tokenizer_init而不是静默失效。harmony / gpt-oss 模型不支持harmony 输出路径从原始 token ids 重建会绕过基于文本的 hook因此 openai_server.py 在启动时直接 fail fast--post_processor_hook is not supported with harmony/gpt-oss models。Disaggregated servingcontext 与 generation 是独立进程各自在自己的阶段、以不同的request_id运行 hookper-request 状态无法跨两个进程关联一个阶段的terminate不会传播到另一个阶段。text 与 token ids 会发散emit只改写 text 通道不改写底层token_ids/logprobs同时读取两个通道的客户端应预期二者不一致suppress/terminate扣住全部通道不会发散。Reasoning / tool 解析hook 在 reasoning 与 tool-call parser 之前运行改写或抑制文本的 hook 可能使这些 parser 失步文档建议这类场景优先用terminate或把 hook 只用于纯文本请求。Hook 异常 fail closedhook 抛异常时该请求以错误失败而不是发出未经审查的 chunk服务器与其他请求不受影响返回未知 action 的 verdict 同样被拒绝。n 1 / beam searchemit和suppress作用于单个输出序列但terminate取消整个请求所有序列因为引擎请求是取消的基本单位——需要对每条序列持有状态的 hook 应以(request_id, output_index)为 key。6. 进一步阅读功能文档docs/source/features/post-processor-hook.md接口与 verdict 实现tensorrt_llm/executor/postprocessor_hook.py可运行的示例 hook测试用、无状态、确定性tests/unittest/llmapi/apps/_postproc_hook_samples.py【免费下载链接】TensorRT-LLMTensorRT LLM provides users with an easy-to-use Python API to define Large Language Models (LLMs) and supports state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs. TensorRT LLM also contains components to create Python and C runtimes that orchestrate the inference execution in a performant way.项目地址: https://gitcode.com/GitHub_Trending/te/TensorRT-LLM创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表