Rust 流式推理输出:背压要从 token 生成处开始

📅 2026/7/6 4:51:54
Rust 流式推理输出:背压要从 token 生成处开始
Rust 流式推理输出背压要从 token 生成处开始一、流式输出不是把 token 往 channel 里塞大模型服务常用流式输出改善体验。模型每生成一个 token就通过 SSE、WebSocket 或 gRPC stream 发给客户端。实现上很容易写成生成线程不停往 channel 里写网络线程慢慢发送。客户端慢、网络慢或代理缓冲时channel 会堆积内存和延迟一起失控。流式推理的背压要从 token 生成处开始。发送端慢生成端就要减速、暂停或取消。否则系统只是把压力从网络层搬到了内存里。二、把生成、缓冲和网络发送串成可控链路流式链路至少包含模型解码、token 缓冲、协议编码和网络发送。每一层都要有容量上限。flowchart TD A[模型解码] -- B[有界 Token Channel] B -- C[SSE 或 gRPC 编码] C -- D[网络发送] D -- E{客户端是否慢} E --|是| F[背压或取消] F -- A无界 channel 是流式服务的常见坑。它能让测试很顺但生产中会把慢客户端变成内存炸弹。三、Tokio 里优先使用有界 channel下面示例用有界 channel 表达背压。send().await会在缓冲满时等待。use tokio::sync::mpsc; pub async fn stream_tokens(mut decoder: Decoder, tx: mpsc::SenderString) - anyhow::Result() { while let Some(token) decoder.next_token().await? { if tx.send(token).await.is_err() { decoder.cancel(); break; } } Ok(()) }如果接收端断开send返回错误生成端应取消推理。继续生成没有意义只会浪费算力。有界 channel 的容量选择需结合 token 生成速率和网络吞吐来定。容量太小生成端频繁阻塞导致 GPU 空转容量太大失去背压意义。实用经验公式是capacity max_tokens_per_second × network_rtt_p95_seconds × 3预留 3 个 RTT 窗口缓冲。以 50 tokens/s 生成速率、100ms RTT 为例容量约 15。短输出请求可用较大 channel 保证低延迟长输出请求用较小 channel 强制背压生效。另一个优化是批量化发送——不在每个 token 后立即 send而是攒到固定间隔如 50ms批量发送减少系统调用次数和小包开销。Tokio 的mpsc::Sender线程安全批量逻辑需在生成循环内维护 local buffer 和 tick 间隔。四、背压策略要和计费、超时一起设计客户端慢不一定是恶意也可能是网络环境差。系统可以设置最大缓冲、最大空闲时间和最大流式时长。超过阈值后明确取消并记录取消原因。还要区分模型计算和网络发送耗时。只看总请求耗时会误判瓶颈。Trace 中应记录首 token 时间、token 间隔、发送等待时间和取消点。最后流式输出要支持清理。请求取消后KV Cache、采样状态、channel buffer 都要释放。Rust 的所有权能帮忙但跨任务资源仍要显式设计。代理层也会影响背压。某些网关会缓冲 SSE导致服务端以为发送成功客户端却迟迟收不到。压测时要穿过真实网关链路观察首 token 到达客户端的时间而不是只测服务进程内的生成时间。还要处理半关闭连接。客户端断网时服务端未必立刻收到断开事件。可以设置心跳或发送超时超过阈值主动取消生成。否则长尾断连请求会占着 KV Cache 和模型 slot。计费系统也要接入取消点。生成了多少 token、成功发送多少 token、因为什么取消都要能对账。流式链路里的“生成”和“交付”不是一回事。五、总结Rust 流式推理输出要从有界缓冲和取消语义开始。客户端慢时压力必须反馈到 token 生成处而不是堆在内存里。实现上用有界 channel、断开检测和资源清理守住边界。流式体验好不好不只看首 token 快还要看慢客户端下系统是否稳定。