ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

text-generation-inference 中的 Router(Webserver)深度解析:连续批处理、Token 流式输出与 gRPC 调度架构

text-generation-inference 中的 Router(Webserver)深度解析:连续批处理、Token 流式输出与 gRPC 调度架构 text-generation-inference 中的 RouterWebserver深度解析连续批处理、Token 流式输出与 gRPC 调度架构【免费下载链接】text-generation-inferenceLarge Language Model Text Generation Inference项目地址: https://gitcode.com/GitHub_Trending/te/text-generation-inferenceRouter在 docs 中也被称为webserver是 text-generation-inference 服务架构中的大脑它负责承接所有 HTTP 请求、对请求做校验与排队并把何时让新请求进入 prefill、何时暂停 decode、哪些请求该出队等核心调度逻辑集中在这一层处理。本文将以仓库中 router/README.md 为骨架结合 router/src 与 backends/v3/src 的源码实现拆解其连续批处理、prefill/decode 调度与 Token 流式输出的完整原理与落地细节帮助你理解如何用这一层去榨干 GPU 算力、同时兼顾用户可感知的延迟。Router 的职责为什么需要一个路由层Router 之所以存在核心动机是把复杂的调度逻辑与高效的模型前向计算解耦它处理了批处理的大部分逻辑判断何时向 GPU 提交新的prefill请求、何时暂停正在进行的decode请求、以及哪些请求应该被合并或移出批次它通过gRPC与后端 shard 通信如 backends/v3/src/client/grpc_client.rs因此 shard 本身可以保持极简专注于尽可能高效的 forward 计算客户端只需要发送单条请求由 Router 来负责 mix match混排/组合批次从而最充分地利用可用算力。从源码结构看Router 是一个独立的 Rust craterouter/Cargo.toml由 router/src/lib.rs 定义核心数据类型请求、参数、响应router/src/server.rs 实现 HTTP 服务器与各 API 端点router/src/validation.rs 负责请求校验与 tokenizationrouter/src/infer/mod.rs 定义Infer推理入口与Backend抽象router/src/config.rs 则承载模型配置解析含视觉模型的图像分辨率选择逻辑。此外还有 router/src/chat.rs、router/src/sagemaker.rs、router/src/vertex.rs、router/src/kserve.rs 等模块分别提供 Chat 会话、SageMaker/Vertex/KServe 兼容端点。连续批处理Continuous Batching吞吐与延迟的甜蜜点为什么 LLM 场景需要连续批处理text-generation-inference 的一个重要特性——连续批处理continuous batching正是由 Router 实现的。它的含义是在 LLM 同一次forward步骤即一个 batch中定期地加入新查询同时把已完成的查询移出。连续批处理要真正发挥作用前提是算力相对于模型的内存需求是过剩的。这一点对 LLM 而言几乎天然成立模型越大需要 pooling 的 GPU 越多你手上可用的算力总量也就越大——这也解释了为什么大模型场景下连续批处理的收益尤其显著。与它相对的是静态批处理static batching同样是把多个查询放进同一次计算但批的大小由客户端预先决定服务端只能被动接受。对内存受限memory bound的 LLM 文本生成来说静态批处理无法根据实时负载动态调整往往造成算力浪费。简单连续批处理一个非常通用的思路文本生成的过程是把 prompt 喂给模型然后反复调用forward一次只产生 1 个 token。最简单的连续批处理思路是查询到达后立即开始处理新查询到达时等当前forward结束把正在运行的 prompt 与新查询拼进同一个 batch再调用forward每当任一查询结束模型产出 EOS 结束符或达到允许的长度上限把它从 batch 中移除、释放所有已分配的内存继续处理剩余部分直到 batch 清空。这个朴素思路泛化能力很强理论上可以在同一个 batch 里堆叠非常多的请求。值得注意的是不同查询可能使用不同的采样参数是否采样、temperature、top_k 等这在上述方案里完全不是问题——只需对 batch 中每个成员独立进行采样即可。Prefill、Decode 与 past key values把两种 forward 区分开让 LLM 文本生成高效运转的关键技巧是缓存注意力矩阵past key values第一次过 prompt 时必须计算完整的注意力矩阵而后续的 forward 只需计算新 token 的注意力。因此第一次 pass 被称为prefill整个代码库中的统一叫法后续的 pass 被称为decode。由于prefill比decode昂贵得多我们不想频繁执行它但一个正在运行的查询大概率处于decode阶段如果要做连续批处理就必须在某个时机执行prefill为新查询生成所需的注意力矩阵使其能够加入decode组。text-generation-inference 使用多种策略和参数帮你找到压榨硬件利用率与可感知延迟之间的平衡点批处理模式延迟Latency吞吐Throughput说明无连续批处理极好极差本质上等于 1每个请求独立串行处理静态批处理极差可能达到最大需要等请求凑齐批大小最大总批大小受硬件限制连续批处理可接受的甜蜜点高通用原则2 倍延迟换取同一硬件上 10 倍的用户量是可接受的源码视角调度循环里到底发生了什么prefill/decode的区分在 Router 的调度循环中体现得非常直接。以 backends/v3/src/backend.rs 为例调度主循环约 L160-L290持续从队列取下一批请求根据当前 batch 的max_tokens与current_tokens计算剩余的token_budget与prefill_token_budgetL173-L200通过queue.next_batch(min_size, max_size, prefill_token_budget, token_budget)尝试取下一批L204-L206取到后调用prefill(mut client, new_batch, ...)为新请求执行首次前向L229 或 L253并随后对合并后的批次执行decode(mut client, batches, ...)L285。队列本身在 backends/v3/src/queue.rs 中实现State用VecDeque(u64, Entry)维护排队条目next_batch依据min_size/max_size/prefill_token_budget/token_budget决定放行哪些条目L86-L114并配合tgi_queue_size等 metrics 暴露队列状态L143、L158。此外还有两个直接影响调度行为的关键参数max_waiting_tokens如果等待新请求入队达到该阈值即使 batch 很小也会强行放行L186-L189避免请求饿死waiting_served_ratio用来计算最小批大小(batch_size * waiting_served_ratio)L194作为背压backpressure机制max_batch_prefill_tokens/max_batch_total_tokens分别约束一次 prefill 可容纳的 token 数与整个 batch 的 token 上限L131-L135。需要说明的是若模型支持 prefill chunkingprefill 分块则waiting_served_ratio和max_waiting_tokens会被忽略backends/v3/src/backend.rs L39、L175-L184此时新 batch 会与当前 batch 在服务端拼接后一起 prefill从而持续运行在算力上限。这些参数的完整列表与默认值由 launcher 端环境变量控制可参考 launcher/src/env_runtime.rs 与 docs/source/reference/launcher.md对应 batch 参数在 backends/v3/src/backend.rs 的BackendV3::new中接收并传入调度循环。并发上限信号量如何保护后端除了批次级别的调度Router 还用**信号量semaphore**限制同时进入推理的请求数。在 router/src/infer/mod.rs 中Infer::new创建Semaphore::new(max_concurrent_requests)L83-L84generate_stream在入口处通过try_acquire_owned()获取许可L111-L120获取失败会递增tgi_request_failure{erroverloaded}指标并返回过载错误请求完成后许可随OwnedSemaphorePermit释放从而严格保证并发不超过max_concurrent_requests。该上限正是/info端点中max_concurrent_requests字段的来源router/src/lib.rs L276。Token 流式输出Token Streaming让第一 token 延迟决定用户体验为什么流式输出至关重要Token 流式输出是客户端 UX 中非常重要的一环。正如上文所述延迟是 LLM API 用户最关心的感知质量指标。有了 token streaming服务端在完成第一次prefillpass 后就可以立即开始应答而无需等待整个生成过程结束。对于极长的查询客户端可以在工作真正完成之前几个数量级的时间就看到有内容在产出看到进展中的输出用户可以提前判断结果是否符合预期及时中断cut short从心理体验上讲边生成边看到远比长时间空白后一次性吐出要舒服。源码视角SSE 流式响应是如何实现的Router 的流式输出基于SSEServer-Sent Events实现。在 router/src/server.rs 中非流式请求走generateL273最终返回VecGenerateResponse流式请求走generate_streamL476构造Sseimpl StreamItem ResultEvent, Infallible并通过.keep_alive(KeepAlive::default())维持长连接L504、L984、L1247流中每个事件携带StreamResponserouter/src/lib.rs L1503-L1513index、token、top_tokens、generated_text仅在最后一条携带与details含finish_reason、generated_tokens、input_length等。generate_stream_internal会消费Infer返回的InferStreamResponse流Prefill/Intermediate/End等变体把解码产生的每个 token 即时转发给客户端Infer侧则在 router/src/infer/mod.rs 的generate_stream中通过self.backend.schedule(valid_request)拿到后端的生成流并包裹统计逻辑记录first_queued、first_start等时间点供tgi_request_*系列延迟指标使用。值得注意的细节CompatGenerateRequest兼容模式请求中stream: bool默认falserouter/src/lib.rs L1362当为true时compat_generate直接转调generate_streamrouter/src/server.rs L139-L150保证两条路径行为一致。从客户端到后端一次完整请求的生命周期综合上述源码一次文本生成请求在 Router 中的完整路径可以概括为HTTP 接入客户端请求到达 router/src/server.rs 中注册的路由POST /、POST /generate、POST /v1/chat/completions等路由表位于 L2238-L2251请求解析与参数归一化router/src/lib.rs 中的GenerateRequest/ChatRequest等结构体通过 serde 反序列化请求体ChatRequest::try_into_generateL941-L1036把 OpenAI 风格的 chat 请求转换为内部GenerateRequest温度 0 时切换为贪心解码do_samplefalse、presence_penalty映射为repetition_penalty 2.0、response_format与tools互斥校验等校验与 Tokenizerouter/src/validation.rs 的Validation维护一组后台 tokenization workerworkers个Python tokenizer 时强制为 1L56-L60通过 round-robin 分发请求L62-L91异步返回Encoding与Chunk列表同时校验max_best_of、max_stop_sequences、max_input_length、max_total_tokens等上限L31-L35限流排队router/src/infer/mod.rs 的generate_stream先取信号量许可再进入backend.schedule()gRPC 调度与执行后端如 v3在 backends/v3/src/backend.rs 的调度循环中通过 gRPC 客户端backends/v3/src/client/grpc_client.rs向 shard 发起prefill/decode调用并在 backends/v3/src/queue.rs 的队列状态机里做批次编排响应回传非流式请求聚合为GenerateResponse返回流式请求通过 SSE 逐 token 推送StreamResponse。该架构中 shard 只负责模型前向全部调度智能集中在 Router这与 router/README.md 开头shards can be kept much simpler的设计目标完全一致。总结与进一步阅读Routerwebserver是 text-generation-inference 的调度中枢它以 gRPC 与极简的 shard 通信集中实现了连续批处理prefill 入门、decode 续跑、完成后出队、并发信号量限流以及 SSE 的 token 流式输出从而在吞吐与可感知延迟之间取得平衡——正如原文档所说2 倍延迟换取 10 倍用户量通常是一笔划算的交易。若想深入理解 Router 的更多细节可以继续阅读router/README.mdRouter 的设计初衷与概念总览router/src/server.rsHTTP 端点、SSE 流式输出与路由表实现router/src/infer/mod.rsInfer推理入口、Backend抽象与并发信号量router/src/validation.rs请求校验与后台 tokenization 工作池router/src/lib.rs请求/响应数据类型与 OpenAI 兼容参数映射backends/v3/src/backend.rs 与 backends/v3/src/queue.rsprefill/decode 调度循环与队列状态机docs/source/reference/launcher.mdmax_batch_prefill_tokens、max_waiting_tokens等运行参数的完整说明。【免费下载链接】text-generation-inferenceLarge Language Model Text Generation Inference项目地址: https://gitcode.com/GitHub_Trending/te/text-generation-inference创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表