Coding Agent 推理性能优化:Decode、KV Cache 与多 Session 调度

Modal 公开了他们的 Coding Agent 推理服务优化实践,以 Kimi K2.6、SGLang 和 B200 为基础,将单用户 Decode 性能提升约 2.8 倍,Replica(单个模型服务实例)的吞吐提升约 5.6 倍。

这次推理性能优化主要集中在三部分:Decode 性能、长 Session 下的 KV Cache,以及多 Session 场景中的路由与负载均衡。本文就沿着这三部分展开,并补充一些同方向的工程实践。

图注:Coding Agent Session 中上下文逐轮累积示意图。图源 @Modal

Decode 性能优化

对于 Coding Agent,Decode TPS 会直接影响任务的推进速度。工具调用参数、代码修改和中间分析都要模型逐 Token 生成,Decode 越慢,后续工具执行和下一轮推理开始得也越晚。

在原始配置下,并发一旦超过 6 个用户,单用户的 Decode TPS 和整体吞吐就会明显下降。即使在较低并发下,生成速度也没有达到 Modal 预期的水平。

因此,Decode 优化这块,他们先把目标放在单个用户实际获得的 Decode TPS 上。

在 Decode 阶段,模型每生成一个 Token,都要读取大量的模型权重和状态。对于这种场景,HBM 带宽往往会先于计算能力成为瓶颈。这里可以采用两种方式进行优化:Tensor Parallelism(下文用缩写 TP)和 Speculative Decoding。前者把一次模型计算拆到多张 GPU 上,后者尝试在目标模型的一次前向计算中确认多个 Token。

TP4 还是 TP8

Kimi K2.6 的模型权重大约占 595 GB,而单张 B200 提供 180 GB HBM,因此 TP1 和 TP2 都无法在不卸载权重的情况下运行完整的模型,可选配置主要落在 TP4 和 TP8。

图注:TP4 / TP8 的 HBM 与 KV Cache 容量对比。图源 @Modal

TP4 使用 4 张 B200,总 HBM 为 720 GB,扣除模型权重后还剩约 125 GB;TP8 使用 8 张 B200,总 HBM 能达到 1.44 TB,剩余空间约 845 GB。按照当时的 BF16 KV 配置估算,TP4 大约能容纳 45 万 Token 的 KV,TP8 则能达到约 300 万 Token。

单看显存容量的话,TP8 优势十分明显。但在 Modal 的实际负载下,即使让 TP8 只负责 Prefill,每张 GPU 的吞吐也只是大致追平同时处理 Prefill 和 Decode 的 TP4。GPU 数量翻倍,并没有带来相应的单位 GPU 收益。

最终 Modal 选择了 TP4,并把更小的 Replica 作为扩展单位。虽然这个选择提高了单位 GPU 的使用效率,但也压缩了留给 KV Cache 的空间:扣除模型权重后,TP4 只剩约 125 GB HBM,后续继续提升并发时,KV Cache 很快就会成为新的瓶颈。

Speculative Decoding 的负载适配

TP 配置确定后,可以继续优化 Speculative Decoding。

Modal 采用 DFlash 作为推测模型。推测模型会先预测接下来的多个 Token,再由目标模型统一验证。只要一次能够接受多个预测结果,目标模型每轮读取大量权重和 KV 后,就能生成多个有效 Token,从而减少单 Token 的推理开销。

Coding Agent 的输出比较适合这类优化。代码中存在大量重复结构,工具调用通常使用固定格式,聊天模板里也包含不少控制 Token;Agent 还会频繁复用刚刚出现过的路径、函数名和代码片段。这些内容的可预测性较高,更容易提高 Speculative Decoding 的接受率。

Modal 也没有直接沿用通用的 DFlash 推测模型。他们先用通用数据训练,再使用目标模型真实产生的 Coding Trace 继续微调,使推测模型更贴近实际 Coding Agent 的输出分布。

微调后,代表性 Trace 上的平均接受长度从 5.00 提高到 5.84 Token / step ,并额外带来了约 20% 的速度提升。为了继续利用线上数据训练推测模型,Modal 还修改了 SGLang,让服务直接记录原始 Token ID,避免文本重新分词后无法完整还原原始 Token 序列。

相比单纯接入 Speculative Decoding,这部分更有价值的是对真实 Coding Agent 输出的适配。推测模型能带来多少收益,很大程度上取决于训练数据与实际请求分布是否一致。

Together AI 在生产推理中的实践也采用了类似思路。他们使用专门训练的推测模型,并加入自适应接受机制。他们公开的数据显示,代码和结构化输出这类可预测性较高的负载,通常更容易从 Speculative Decoding 中获得收益。2026 年推出的 Aurora 还进一步利用线上推理轨迹持续调整推测模型,以适应生产流量的变化。

因此,在 Coding Agent 场景下,Speculative Decoding 的效果不能只看固定推测模型的基准测试结果。Agent 输出中工具调用、代码和自然语言的占比不同,都会影响最终的平均接受长度和实际加速效果。

长 Session 下的 KV Cache 扩容

TP4 和 DFlash 提升了单用户 Decode 性能后,新的瓶颈转向了 HBM 容量:留给 Session 状态和 KV Cache 的空间不够了。

在最大输入长度约 100k Token 的实际负载下,TP4 的 KV 空间大约只能容纳 4 个活跃会话。继续提高并发后,KV Cache 命中率会快速下降,已经计算过的长前缀需要重新 Prefill,吞吐和交互速度也会随之下降。

接下来的优化重点,基本都围绕一个问题展开:怎么在有限的 HBM 里留出更多空间给 KV Cache。

DFlash 的显存占用优化

接下来先处理的是 DFlash 在 SGLang 中的显存占用。

原来的实现会先保留目标模型中间结果的指针,等一次前向计算结束后,再把这些数据复制到连续内存中。这个过程中,原始数据和用于存放结果的缓冲区会同时占用显存。

优化后的流程改为提前分配连续缓冲区,并在前向计算过程中直接写入中间结果,最终将这部分的 HBM 峰值占用降低了一半。相关修改后来也提交到了 SGLang。

这类改动没有量化或新算子那么显眼,但释放出来的 HBM 可以直接留给 KV Cache,为后续提高并发腾出空间。

KV Cache 量化:BF16 KV 换成 FP8

相比前面的内存优化,量化带来的 KV Cache 容量提升更明显。

在实际评测通过后,KV Cache 从 BF16 改成 FP8。单个 KV 元素的存储开销减半,同样的 HBM 可以容纳大约两倍 Token。

与此同时,原本以 FP8 保存、每个 Token 都会激活的共享专家也改成了 NVFP4,继续释放一部分 HBM。相关评测中,没有观察到这些修改带来超出运行波动范围的明显质量下降。

对于长 Session,KV Cache 的单位存储开销会直接影响一个 Replica 能保留多少历史状态,也会进一步影响缓存命中和并发能力。

这一方向还在继续往更低精度发展。NVIDIA 后续在 Qwen3-Coder-480B-A35B 上测试了 NVFP4 KV Cache。相比 FP8,KV Cache 占用可以进一步减半;在其测试环境中,缓存命中率最高提升约 20% ,延迟最高改善约 3 倍。随着缓存容量逐渐充足,两种精度之间的差距会缩小,最终收益仍然取决于模型、上下文长度和实际负载。NVIDIA 在 LiveCodeBench、MBPP、Ruler 64K 等测试中报告的精度损失低于 1%。

KV 量化带来的收益最终不只体现在显存占用上,它还会继续影响缓存命中率、并发能力和 TTFT。

HiCache 分层缓存

即使完成显存清理和量化,HBM 容量仍然有限。接下来启用了 SGLang 的 HiCache,把 KV Cache 从单层 HBM 扩展为分层缓存:GPU HBM 作为第一层,CPU 内存作为第二层,HiCache 还支持继续向分布式存储扩展第三层。

这次实践主要使用 CPU 内存作为第二层缓存,用来承接短时间内超出 HBM 容量的 KV 数据。CPU 内存容量更大,但访问速度低于 HBM,如果大量请求频繁从 CPU 内存读取 KV,同样会增加延迟。

因此,这一层缓存主要用于缓冲瞬时高并发:当某个 Replica 短时间超过理想并发时,尽量避免 KV Cache 命中率快速下降,导致大量长前缀重新 Prefill。

图注:HiCache 加入前后,随着并发增加的 KV Cache 命中率和性能变化。图源 @Modal

类似的思路也可以扩展到更大的集群。Together AI 在 Cache-aware Prefill--Decode Disaggregation(CPD)中使用了三层 KV Cache:GPU 显存、主机内存,以及通过 RDMA 访问的集群级分布式缓存。

同时,请求会按照 KV 复用程度区分为低复用和高复用两类,让需要完整 Prefill 的长请求与已经能够复用缓存的请求分开处理。在其测试中,这套方案在满足尾延迟要求的情况下,将可持续 QPS 最高提高约 35%~40%。

HiCache 主要缓解单个 Replica 短时间超出 HBM 容量的问题,而 Together AI 的 CPD 将 KV Cache 的复用范围继续扩展到了集群级。对于长上下文请求,已经计算过的 KV 能否继续复用,会直接影响后续 Prefill 开销和整体吞吐。

多 Session 路由与负载均衡

有了 KV Cache 之后,请求路由就不能只看哪个 Replica 当前更空闲。

对于已经运行几十轮的 Coding Agent Session,如果下一轮请求继续进入原来的 Replica,前面大量上下文就可以直接复用本地 KV;如果被分配到一个从未处理过这个 Session 的 Replica,几十万 Token 的前缀就可能需要重新 Prefill。

最开始采用的是 Session Affinity:客户端给请求附上 Modal-Session-Id,路由器根据 Session ID,尽量把同一个 Session 的后续请求固定到同一个 Replica。

这样可以提高缓存复用率,但在真实 Agent 流量下,也会带来新的调度问题。

Session Affinity 的负载偏斜

第一个问题来自 Hash 分配本身。

假设每个 Replica 平均承载 5 个 Session,即使总体 Session 数量足够多,随机分配后依然可能出现部分 Replica 很闲、部分 Replica 很忙的情况。以 250 个 Session 分配到 50 个 Replica 为例,平均值虽然正好是 5,但实际负载仍会出现明显偏差。再加上不同 Session 的计算量并不完全一致,只按照 Session ID 做映射,很难保证各个 Replica 获得接近的工作量。

另一个问题来自同一 Session 内部的并发。

传统对话大多是一问一答:一次请求完成后,再进入下一轮。但 Coding Agent 可能启动 Subagent,也可能从同一上下文分叉出多个并发任务。此时多个请求共享同一个 Session ID,如果仍然通过 Session Affinity 固定到同一个 Replica,很容易形成热点。

针对这种情况,同一 Session 的并发请求不再强制落到同一个 Replica,部分请求可以被分配到其他实例。这样会损失一部分 KV Cache 复用,但能换来更均匀的负载。

因此,在长 Session 场景下,路由需要在缓存复用和负载均衡之间做取舍:什么时候继续保持 Session Affinity,什么时候允许请求分散到其他 Replica。

扩容过程中的缓存复用

自动扩容还会带来另一类缓存复用问题。

传统一致性 Hash 或 Rendezvous Hash 在新增节点后,会把一部分已有 Session 重新映射到新的 Replica。对于无状态服务,这通常只是一次正常的流量再平衡;但对于已经积累大量 KV 的 Agent Session,重新分配意味着原 Replica 上的缓存无法继续复用,新 Replica 需要重新 Prefill。

因此,扩容时会优先让新 Session 进入新 Replica,尽量让已经建立 KV 状态的旧 Session 留在原来的 Replica。同时,路由还会结合当前 Replica 的负载、正在运行的请求等状态,决定新请求的落点。

调整后,Replica 之间的负载分布明显更加均匀,TTFT 的尾部波动也有所下降,多 Replica 部署的实际吞吐更接近单 Replica 基准测试推算出的水平。

图注:优化前后的 Replica 请求分布。图源 @Modal

类似的做法也出现在其他 Agent 推理服务中。Cloudflare 在 Workers AI 上使用 x-session-affinity,让同一 Session 尽量回到已经保存其输入状态的位置。OpenCode 等 Agent Harness 也加入了这一请求头。

Cloudflare 公开的数据中,在高流量场景启用 Session Affinity 后,高峰期输入 Token 的缓存命中率从约 60% 提升到 80%。与此同时,其内部路由还会参考 KV Cache 状态,因为只知道 Session ID,并不足以完成最终调度。

Session Affinity 能提高缓存复用率,但进入多 Replica 场景后,路由还需要同时考虑 Session 当前所在位置、KV Cache 的复用价值、Replica 的实时负载,以及同一 Session 是否产生了并发请求。

小结

从这些实践来看,Decode、KV Cache 和多 Session 调度很难完全拆开处理。TP 配置会影响显存余量,KV Cache 容量会影响并发和路由,Session Affinity 又会反过来影响缓存复用和负载分布。

因此,在真实 Coding Agent 负载下,很多优化都需要放回具体的模型、Session 规模和流量分布中重新验证。固定配置可以参考,但最终效果还是要靠实际测试确认。

相关推荐
GISer_Jing2 小时前
SuperAgent
ai·langchain
程序员柒叔2 小时前
Diagram Design 解构:45 种图型如何用语义路由与 SVG 质量门交付
人工智能·github·agent·skills
YYYing.2 小时前
【Agent系列 (十一) 】上下文工程
人工智能·python·ai·agent·上下文工程
Java后端的Ai之路4 小时前
【AIAgent开发】-Agent工具封装成CLI(纯干货!!!)
java·javascript·人工智能·agent·cli
最强小杰4 小时前
DeepSeek V4 Flash + Codex CLI 接入教程:原生 Responses API 适配,从配置到跑通收藏这篇就够了(2026)
java·前端·python·ai
lyai0075 小时前
读完 DeepSeek Harness 源码,我给自己画了一张「最小 CLI Agent」施工图
llm·agent
七夜zippoe5 小时前
Agent 错误处理体系:分类、重试、降级与熔断设计
ai·agent·降级·错误处理·熔断设计
武子康5 小时前
固定时间步已经开了,为什么游戏回放还会走出不同位置?
人工智能·后端·agent
VIP_CQCRE5 小时前
内容自动化越跑越乱?先补上状态、去重和失败恢复
python·ai·工作流·幂等设计·内容自动化