1. 从一次真实的卡顿说起为什么你的AI代码助手越用越慢去年下半年我接手了一个内部代码辅助工具的调优工作。这个工具最初版本跑得挺欢补全一个函数也就一两百毫秒的事。但随着团队把代码库喂进去、把上下文从2K拉到8K、再拉到32K情况急转直下——补全延迟从200ms飙到3秒以上显存占用直接顶满服务端排队严重用户体验断崖式下跌。这不是个例。任何把大模型接入代码补全场景的人迟早都会撞上这堵墙上下文窗口越长推断越慢显存越吃紧。而代码场景偏偏又极度依赖长上下文——你得让模型看到当前文件的完整定义、相关引用、类型声明、甚至跨文件的接口约定它才能给出靠谱的补全建议。所以这篇东西我想把当时做性能重构的完整思路和落地细节摊开讲。核心就两件事上下文窗口怎么优化以及推断过程怎么加速。适合正在做AI代码助手、或者任何需要低延迟长上下文推断场景的开发者参考。不管你是刚接触这块还是已经被延迟问题折磨了一阵子下面这些实操经验应该都能直接拿去用。2. 上下文窗口优化不是塞得越多越好2.1 先搞清楚上下文到底被什么吃掉了很多人一上来就想着“怎么压缩上下文”但连上下文里装了什么都不知道。我当时的做法是先做一轮上下文成分审计把每次请求实际送入模型的token按来源分类统计。结果挺反直觉的真正用于代码补全的核心上下文当前光标附近的代码、函数签名、类型定义只占了不到30%。剩下的大头是整个文件的全文哪怕光标在文件末尾前面的代码也全塞进去了无关的import和注释块历史对话记录用户之前问过的所有问题检索回来的相似代码片段RAG结果往往冗余度极高实操心得先别急着优化算法花半天时间把上下文成分打出来看看。我见过太多团队在优化一个根本不存在的瓶颈。2.2 分层裁剪策略按优先级动态分配token预算搞清楚成分之后我设计了一套分层裁剪方案。核心思路是不是所有上下文都同等重要按优先级分配token预算高优先级的先占低优先级的能省则省。具体分层如下优先级上下文类型预算占比裁剪策略P0光标前后各50行代码40%不裁剪P1当前函数/类的完整定义20%超长时保留签名关键逻辑P2相关类型声明和接口15%只保留被引用的部分P3跨文件引用15%按引用距离衰减P4历史对话10%滑动窗口只保留最近3轮这套策略落地后平均上下文长度从28K降到了12K左右而补全质量用人工评估接受率衡量几乎没有下降。原因很简单模型真正需要的上下文远比我们以为的少。2.3 语义压缩用摘要替代原文对于P2和P3层直接裁剪可能会丢掉关键信息。我的做法是引入一个轻量级的语义压缩步骤对长类型定义和跨文件引用先用一个小模型或者规则引擎生成结构化摘要再送入主模型。举个例子一个500行的接口定义文件压缩后可能变成这样# 压缩后的类型摘要 class UserService: methods: [get_user(id)-User, update_user(id, data)-bool, delete_user(id)-bool] deps: [UserRepo, CacheClient, Logger] key_types: User{id:int, name:str, email:str, role:Role}这种压缩比原文少了80%的token但保留了模型做补全决策所需的核心结构信息。实测下来补全准确率只掉了1.2个百分点但延迟降低了35%。2.4 动态窗口根据任务类型调整上下文长度不是所有补全请求都需要长上下文。补一个变量名和补一个复杂业务逻辑需要的上下文量差了一个数量级。所以我加了一层任务分类器根据光标位置的语法结构和语义特征动态决定上下文窗口大小。简单补全变量名、简单表达式窗口压到2K中等补全函数体、控制流窗口8K复杂补全跨文件重构、架构级建议窗口16K-32K这个分类器本身很轻用一个小的BERT变体跑推断耗时不到5ms但带来的收益是平均延迟下降了40%以上。3. 推断加速从模型到工程的全链路优化3.1 KV Cache管理长上下文推断的命门长上下文推断慢很大一部分慢在KV Cache上。每生成一个token都要和之前所有token的Key/Value做注意力计算。上下文越长这个计算量越大而且是平方级增长。我做了几件事第一KV Cache分页管理。参考了操作系统的虚拟内存思路把KV Cache按页切分不活跃的页换出到CPU内存需要时再换入。这样显存占用从固定的32G降到了动态的8-16G服务端能同时处理的请求数翻了3倍。第二滑动窗口注意力。对于超长上下文不是所有历史token都同等重要。我用了滑动窗口全局token的混合注意力最近512个token用全注意力更早的token只保留少量全局关键token比如函数签名、类型定义。这样注意力计算量从O(n²)降到了O(n·w)w是窗口大小。第三KV Cache量化。把KV Cache从FP16量化到INT8显存直接减半精度损失在可接受范围内补全接受率下降不到0.5%。注意KV Cache量化要小心不同模型对量化的敏感度差异很大。建议先在小流量上做A/B测试确认精度损失可接受再全量。3.2 推测解码用小模型加速大模型推测解码是我认为性价比最高的加速手段之一。思路很简单用一个小的草稿模型先快速生成几个候选token然后用大模型一次性验证这些token是否正确。如果草稿模型猜得准大模型就能一次验证多个token相当于并行生成了。我当时的配置是草稿模型1.5B参数的小模型专门在代码数据上微调过主模型13B参数的代码大模型草稿长度4-8个token实测下来在代码补全场景下草稿模型的接受率能达到65%-75%整体推断速度提升了1.8-2.2倍。而且这个方案对输出质量几乎没有影响因为最终token都是主模型验证过的。3.3 批处理与连续批处理把GPU喂饱单请求优化做得再好如果GPU利用率上不去整体吞吐还是不行。我引入了连续批处理机制不等一个batch里所有请求都完成而是动态地把新请求插入到正在运行的batch中完成的请求立即释放资源。这个改动让GPU利用率从40%左右提升到了75%以上单位时间内处理的请求数翻了近一倍。配合前面的KV Cache分页管理服务端的整体吞吐提升了3倍多。3.4 算子融合与内核优化这一层偏底层但收益很直接。我用了一些现成的推断框架比如TensorRT-LLM、vLLM它们已经做了大量算子融合和内核优化。如果你在用原生PyTorch跑推断强烈建议换到这些专用框架上光这一项就能带来30%-50%的性能提升。具体来说FlashAttention是必开的。它把注意力计算的内存访问模式优化了长上下文下提速非常明显。另外RMSNorm融合、SiLU激活融合这些小算子融合累积起来也能省不少时间。4. 落地实操从零搭建一套优化后的推断服务4.1 环境准备与依赖选型先列一下我最终用的技术栈# 核心推断框架 vllm0.4.2 # 注意力优化 flash-attn2.5.8 # 量化支持 bitsandbytes0.43.0 # 服务框架 fastapi0.110.0 uvicorn0.29.0选vLLM而不是TensorRT-LLM主要考虑是vLLM对连续批处理和PagedAttention的支持更成熟而且社区活跃遇到问题好查。FlashAttention是必装的长上下文下没有它基本没法用。4.2 模型加载与KV Cache配置模型加载这块有几个关键参数from vllm import LLM, SamplingParams llm LLM( modelyour-code-model, tensor_parallel_size2, # 双卡并行 max_model_len32768, # 最大上下文长度 gpu_memory_utilization0.90, # GPU显存利用率 enable_prefix_cachingTrue, # 开启前缀缓存 block_size16, # KV Cache页大小 swap_space16, # CPU交换空间(GB) quantizationawq, # 权重量化 )enable_prefix_caching这个参数很关键。代码补全场景下很多请求的前缀是相同的比如同一个文件的不同位置开启前缀缓存后这些公共前缀的KV Cache可以复用省掉大量重复计算。block_size设成16是权衡的结果。太小了页表管理开销大太大了内部碎片多。16在大多数场景下是个比较稳的选择。4.3 上下文构建流水线上下文构建我单独抽了一个模块核心逻辑如下def build_context(cursor_pos, file_content, repo_index, history): # P0: 光标附近代码 p0 extract_nearby_code(file_content, cursor_pos, lines50) # P1: 当前函数/类定义 p1 extract_enclosing_scope(file_content, cursor_pos) # P2: 相关类型声明 p2 extract_relevant_types(repo_index, p0 p1) # P3: 跨文件引用 p3 extract_cross_file_refs(repo_index, p0 p1, max_refs5) # P4: 历史对话 p4 history[-3:] # 最近3轮 # 按预算裁剪 context allocate_budget([p0, p1, p2, p3, p4], total_budget12000) return contextallocate_budget函数按前面说的分层策略分配token预算。实际实现时token计数用tiktoken或者模型自带的tokenizer别用字符数估算误差太大。4.4 推测解码配置推测解码的配置在vLLM里比较简单llm LLM( modelyour-code-model, speculative_modelyour-draft-model, num_speculative_tokens5, speculative_draft_tensor_parallel_size1, )num_speculative_tokens设成5是我实测下来比较平衡的值。设太小了加速效果不明显设太大了草稿模型猜不准反而浪费计算。草稿模型的选择上我用了一个1.5B的代码模型在相同代码数据上做过微调。关键是草稿模型要和主模型用同一个tokenizer否则没法做token级别的验证。4.5 服务端部署与压测服务端用FastAPI包一层核心接口就一个补全接口app.post(/completion) async def completion(request: CompletionRequest): context build_context( request.cursor_pos, request.file_content, request.repo_index, request.history ) outputs llm.generate( context, SamplingParams( temperature0.2, max_tokens128, stop[\n\n, ] ) ) return {completion: outputs[0].outputs[0].text}压测我用的是locust模拟了50并发、100并发、200并发三种场景。优化前后的对比如下指标优化前优化后提升P50延迟1.8s380ms4.7xP99延迟5.2s1.1s4.7x吞吐(QPS)12484x显存占用32G14G2.3x最大并发301204x这个提升幅度是多项优化叠加的结果单看某一项可能只有20%-50%的提升但组合起来就是数量级的变化。5. 踩过的坑与排查实录5.1 常见问题速查表问题现象可能原因排查方法解决方案补全延迟突然飙升KV Cache换出频繁监控swap_space使用率增大显存或降低并发输出质量下降上下文裁剪过度对比裁剪前后补全接受率调整分层预算比例显存OOM批处理过大监控batch内请求数限制最大batch size推测解码加速不明显草稿模型接受率低统计草稿token接受率微调草稿模型或减小草稿长度前缀缓存命中率低请求前缀差异大统计前缀哈希分布规范化上下文构建顺序5.2 几个印象深刻的坑坑一上下文顺序影响缓存命中率。一开始我的上下文构建顺序是随机的导致前缀缓存基本命中不了。后来把上下文按固定顺序排列P0→P1→P2→P3→P4并且把最稳定的部分类型定义、跨文件引用放在最前面前缀缓存命中率从不到10%提升到了60%以上。坑二量化后的模型对温度参数更敏感。权重量化到INT8后模型输出的分布会有些微变化。同样的temperature0.2量化后的模型输出多样性明显增加导致补全结果不稳定。后来把temperature降到0.1并加了top_p0.9的截断才恢复稳定。坑三连续批处理下的请求优先级。连续批处理虽然吞吐高但如果不做优先级管理一个长上下文请求可能会阻塞后面一堆短请求。我的做法是给每个请求打上优先级标签短请求和高优先级请求优先调度长请求放到单独的队列里慢慢跑。坑四KV Cache量化的精度损失是累积的。单次推断看不出问题但连续多轮对话后量化误差会累积导致模型“忘记”前面的内容。后来对历史对话的KV Cache不做量化只对当前请求的KV Cache做量化问题才解决。实操心得性能优化不是一锤子买卖。上线后要持续监控P50/P99延迟、显存占用、缓存命中率、补全接受率这几个核心指标任何一个出现异常波动都要及时排查。5.3 一个容易被忽略的优化点请求合并代码补全场景下用户打字很快可能连续触发多次补全请求。如果每个请求都单独处理浪费很大。我加了一层请求合并在50ms窗口内的多个补全请求如果光标位置相近就合并成一个请求处理结果分发给所有等待的客户端。这个改动让实际QPS降低了30%左右因为合并了但用户感知的延迟反而更低了因为合并后的请求上下文更完整补全质量更高用户接受率提升了后续请求自然就少了。6. 一些关于扩展性的思考这套方案目前支撑了团队内部几百人的日常使用但我知道它还有扩展空间。比如上下文压缩这块现在用的是规则小模型摘要未来可以训练一个专门的上下文压缩模型把压缩率再提一个档次。推测解码的草稿模型也可以做得更智能根据当前代码上下文动态选择草稿策略而不是固定长度。另外多模型路由是个值得探索的方向。简单补全用1.5B的小模型就够了复杂重构才需要13B的大模型。根据任务复杂度动态路由到不同规模的模型能在保证质量的前提下进一步降低平均延迟。不过这些都是后话了。眼下这套方案已经能把P99延迟压到1秒出头显存占用控制在单卡16G以内对于大多数团队来说够用了。如果你也在做类似的事情建议先从上下文成分审计开始搞清楚瓶颈到底在哪再针对性优化。别一上来就想着换模型或者加卡很多时候问题出在工程层面而不是模型本身。最后分享一个我常用的排查手法当延迟出现异常时先把请求的上下文长度、KV Cache命中率、批处理大小这三个指标拉出来看一遍90%的问题都能定位到。剩下的10%多半是某个参数配错了或者某个依赖版本不兼容。