共享状态隔离问题：KV cache与批量调度中的状态污染实验

1. 从一次口令实验说起共享状态到底共享了什么第一次看到“共享状态隔离问题”这个说法是在一次内部技术分享的PPT标题上。当时我盯着这八个字看了很久脑子里冒出来的第一个疑问是共享和隔离这两个词放在一起不矛盾吗后来动手做了一组口令实验才慢慢理解这背后其实是在讲一个非常具体的工程问题——多个请求在同一个推理服务里并发跑的时候哪些东西可以共用哪些东西必须各管各的。这个实验的起点很简单我搭了一个小规模的语言模型推理服务用同一套权重同时处理多条输入。最开始我图省事把所有请求的状态都塞进一个全局字典里key是请求IDvalue是各种中间结果。跑单条请求的时候一切正常输出稳定延迟也低。但当我用脚本同时打进去五条请求问题立刻暴露了——有的请求输出串了有的请求直接报错说张量维度对不上还有的请求明明输入的是A输出却带着B的痕迹。这就是典型的共享状态污染。在推理服务里模型权重是只读的共享没问题但每个请求在计算过程中产生的中间状态比如KV cache、注意力掩码、位置编码的偏移量这些都是跟具体请求绑定的一旦混在一起就像把五个人的草稿纸叠在一起写不出错才怪。所以“共享状态隔离问题”这句话我的理解是共享的是那些只读的、与请求无关的资源隔离的是那些随请求变化的、有状态的中间结果。这个原则听起来像废话但真正落到代码里尤其是落到KV cache和注意力计算这种细节上坑非常多。接下来我会把这组实验的完整过程拆开讲包括我一开始怎么想的、中间踩了哪些坑、最后怎么改的以及从里面总结出来的几条实操经验。2. 核心概念拆解KV cache、注意力与批量调度到底在共享什么2.1 KV cache为什么不能随便共享KV cache是自回归生成里的一个关键优化。简单说模型在生成第t个token的时候需要用到前面t-1个token的Key和Value矩阵。如果每次都重新算一遍计算量会随序列长度平方增长。所以常规做法是把前面算过的K和V缓存下来生成新token的时候直接复用。问题在于这个缓存是跟具体请求的序列内容强绑定的。请求A的KV cache里存的是“今天天气怎么样”对应的K和V请求B的KV cache里存的是“帮我写一段代码”对应的K和V。这两个缓存如果混在一起注意力计算的时候就会把A的上下文当成B的上下文来用输出必然错乱。我一开始犯的错误就是用一个全局的past_key_values变量来存所有请求的缓存。单请求跑的时候没问题因为只有一份数据。多请求并发的时候后进来的请求会覆盖掉先进来的缓存或者更隐蔽的情况是两个请求的缓存被拼接在一起导致注意力分数矩阵的维度对不上。正确的做法是每个请求维护独立的KV cache实例。在批量调度的时候调度器需要知道每个请求当前生成到第几个token它的缓存长度是多少然后分别取出来做注意力计算。这就引出了下一个问题批量调度怎么处理不同长度的缓存。2.2 多头注意力里的QKV隔离逻辑多头自注意力机制的核心是把输入投影成Query、Key、Value三组矩阵然后分头做注意力计算。这里有一个很容易被忽略的细节Q、K、V的隔离级别是不一样的。Query是当前token产生的它只跟当前这一步有关所以天然是请求私有的。Key和Value则涉及历史信息需要跟KV cache配合。在批量场景下每个请求的Q长度可能不同比如有的请求刚进来Q长度是1有的请求已经生成了很多tokenQ长度还是1但KV长度很长所以Q、K、V三个矩阵在批量维度上的形状往往是不对齐的。我试过一种偷懒的做法把所有请求的Q拼成一个大的Q矩阵把所有请求的KV cache也拼成一个大的KV矩阵然后直接做一次批量矩阵乘法。结果发现注意力分数矩阵里出现了跨请求的注意力——请求A的Q attend到了请求B的K上。这是因为拼接之后注意力计算是在整个拼接维度上做的softmax没有做请求级别的掩码隔离。修正方案是引入请求级别的注意力掩码。具体来说在计算注意力分数之后、做softmax之前把不属于当前请求的位置的分数设成一个极大的负数比如-1e9这样softmax之后这些位置的权重就趋近于零。这个掩码的构造需要知道每个请求在拼接维度上的起止位置所以调度器必须把每个请求的序列长度信息传进来。2.3 批量调度中的状态隔离策略批量调度的目标是在一次前向计算里处理多个请求提高GPU利用率。但不同请求的输入长度、生成长度、缓存长度都不一样所以调度器需要做几件事第一按长度分桶。把长度相近的请求放在同一个批次里减少padding带来的浪费。我试过把长度差异很大的请求硬塞进一个批次结果padding比例超过70%GPU大部分时间都在算无效的零。第二维护每个请求的独立状态。包括当前生成了多少token、KV cache的当前长度、是否已经遇到结束符。这些状态不能放在全局变量里而要放在每个请求自己的上下文对象里。调度器在组批次的时候从各个上下文对象里读取状态拼成批次级别的输入。第三处理请求的动态进出。有的请求生成完了要退出批次有的新请求要加入。这时候不能直接把新请求的KV cache拼到已有批次后面因为已有批次的缓存位置是连续的新请求需要从自己的位置0开始。我当时的做法是给每个请求分配一个逻辑上的缓存槽位物理上可以用一个大的缓存池来管理但逻辑上每个请求只能访问自己的槽位。3. 口令实验的完整实操过程3.1 实验环境与基础模型搭建实验环境很简单一台带单张消费级显卡的机器Python 3.10PyTorch 2.1transformers库。模型选了一个小规模的自回归语言模型参数量在1亿左右主要是为了方便快速迭代不用等太久。基础推理代码我一开始写得很直白加载模型tokenizer然后写一个generate函数里面用model.generate或者手写自回归循环。手写循环的好处是能看清楚每一步的KV cache是怎么传递的。我选择手写因为这次实验的目的就是搞清楚状态共享的边界。手写循环的核心逻辑是这样的第一步输入完整的prompt模型返回logits和past_key_values第二步取最后一个token的logits做采样得到新token第三步把新token作为输入同时把上一步的past_key_values传进去模型返回新的logits和更新后的past_key_values。如此循环直到遇到结束符或达到最大长度。单请求跑的时候这个逻辑没有任何问题。past_key_values是一个元组里面每一层都有一对K和V张量形状是[batch_size, num_heads, seq_len, head_dim]。batch_size是1seq_len随着生成过程逐渐增加。3.2 多请求并发时的状态污染复现为了复现问题我写了一个简单的并发脚本用threading起五个线程每个线程调用同一个generate函数但传入不同的prompt。模型和tokenizer是全局共享的。第一次跑输出就开始乱了。具体表现是有的线程拿到的输出明显是另一个线程的prompt的续写有的线程直接抛异常说张量形状不匹配还有的线程虽然输出看起来正常但重复跑几次结果不一致。我加了一些日志打印每个线程在每一步拿到的past_key_values的形状和部分数值。发现问题是past_key_values被多个线程同时读写。线程A在生成第3个token的时候线程B刚好在生成第1个tokenB把它的缓存写进了同一个变量A读到的就是B的缓存。更隐蔽的问题是即使我用锁把past_key_values的读写保护起来输出还是会串。因为锁只能保证同一时刻只有一个线程在写但A的缓存和B的缓存还是混在同一个变量里只是写入顺序变了而已。A读的时候可能读到的是B上一次写入的内容。3.3 隔离改造从全局变量到请求级上下文改造的核心思路是把跟请求相关的所有状态都收进一个上下文对象里。这个对象包含request_id请求标识input_ids当前已经生成的token序列past_key_values该请求专属的KV cacheattention_mask该请求的注意力掩码position_ids位置编码的偏移finished是否已经生成结束每个请求创建自己的上下文对象generate函数接收这个对象作为参数所有中间状态都从对象里读、往对象里写。模型权重和tokenizer仍然是全局共享的因为它们是无状态的。改造之后单请求的逻辑不变多请求并发的时候每个请求操作自己的上下文互不干扰。但这时候又出现了新问题每个请求单独跑一次前向GPU利用率很低因为batch_size始终是1。这就引出了批量调度的需求。3.4 批量调度的实现与参数选择批量调度的目标是把多个请求的输入拼成一个批次一次前向计算处理多个请求。但不同请求的序列长度不同所以需要padding。padding的策略直接影响计算效率。我试了两种padding方式一种是左padding一种是右padding。对于自回归生成右padding更自然因为生成是从左往右的右边补零不影响已经生成的部分。但右padding会导致不同请求的有效长度不同注意力掩码需要精确构造。具体实现上我维护一个请求队列调度器每次从队列里取出若干个请求按当前序列长度排序取长度相近的组成一个批次。批次的最大长度设为这批请求里最长的那个。然后构造一个形状为[batch_size, max_len]的输入张量短的地方补pad token。同时构造注意力掩码有效位置为1pad位置为0。KV cache的拼接更麻烦一些。每个请求有自己的缓存长度不同。拼接的时候需要把每个请求的缓存放到批次维度的对应位置短的地方补零。然后在注意力计算的时候用掩码把补零的位置屏蔽掉。我实测下来当批次内请求长度差异不超过20%的时候批量调度的吞吐量比单请求串行高出3到4倍。但如果长度差异太大padding浪费严重吞吐量反而可能下降。所以调度器里加了一个长度分桶的逻辑把长度相近的请求放在一起。4. 常见问题与排查技巧实录4.1 输出串扰的三种典型表现与定位方法输出串扰是共享状态污染最直接的表现。我遇到过的有三种第一种是内容串扰请求A的输出里出现了请求B的prompt内容。这种通常是因为KV cache被覆盖或者拼接错误。定位方法是打印每个请求在每一步的KV cache形状和部分数值看是否有异常的长度或数值突变。第二种是维度报错比如RuntimeError: The size of tensor a (5) must match the size of tensor b (3)。这种通常是因为不同请求的缓存长度不一致拼接的时候没有对齐。定位方法是检查拼接后的张量形状是否符合预期特别是序列长度维度。第三种是结果不稳定同一个请求跑两次输出不一样。这种最隐蔽通常是因为多个线程同时读写共享状态导致每次执行的时序不同。定位方法是加锁之后看问题是否消失如果消失就说明是并发读写的问题。4.2 KV cache长度不一致导致的报错排查KV cache长度不一致是批量调度里最常见的问题。具体表现是批次里有的请求已经生成了10个token缓存长度是10有的请求刚进来缓存长度是0。拼接的时候如果直接按最大长度补零补零的位置在注意力计算时会被当成有效位置导致注意力分数异常。排查步骤是这样的首先打印每个请求的缓存长度确认是否一致然后检查拼接后的缓存张量形状看补零的位置是否正确最后检查注意力掩码确认补零位置是否被正确屏蔽。我踩过的一个坑是注意力掩码只屏蔽了输入端的pad位置没有屏蔽KV cache里的pad位置。结果模型在计算注意力的时候把缓存里的零向量也当成了有效信息导致输出出现无意义的重复。修正方法是在构造注意力掩码的时候同时考虑输入长度和缓存长度把缓存里的pad位置也屏蔽掉。4.3 批量调度中请求动态进出的处理请求动态进出是批量调度里比较复杂的场景。一个请求生成完了要退出批次同时可能有新请求要加入。如果处理不当会导致缓存错位或者内存泄漏。我的做法是每个请求分配一个固定的缓存槽位槽位的生命周期跟请求绑定。请求退出的时候槽位标记为空闲但不立即释放内存而是留给下一个新请求复用。新请求加入的时候从空闲槽位里分配一个初始化缓存为零。这样做的代价是需要维护一个槽位分配表但好处是避免了频繁的内存分配和释放性能更稳定。我实测下来槽位复用的方式比每次重新分配内存的吞吐量高出大概15%到20%。4.4 常见问题速查表问题现象可能原因排查方法修正方案输出内容串扰KV cache被覆盖或拼接错误打印每步缓存形状和数值每个请求独立缓存拼接时对齐张量维度报错缓存长度不一致拼接未对齐检查拼接后张量形状按最大长度补零构造精确掩码结果不稳定多线程并发读写共享状态加锁后看问题是否消失状态收进请求级上下文对象注意力异常掩码未屏蔽缓存pad位置检查掩码构造逻辑同时屏蔽输入和缓存的pad位置吞吐量低批次内长度差异大padding浪费统计padding比例按长度分桶相近长度组批内存泄漏请求退出后缓存未释放监控内存占用槽位复用标记空闲而非释放5. 从实验里总结出的几条硬核经验5.1 状态隔离的边界怎么划状态隔离的边界不是拍脑袋定的而是根据状态的生命周期和访问模式来划的。模型权重是只读的生命周期跟服务进程一样长访问模式是所有人共享读所以不需要隔离。KV cache是跟请求绑定的生命周期跟请求一样访问模式是每个请求只读写自己的所以必须隔离。注意力掩码和位置编码也是同理。我总结了一个简单的判断标准如果一个状态在请求A和请求B之间可以互换而不影响输出那它就可以共享如果互换之后输出会变那就必须隔离。模型权重互换当然不影响因为只有一份。KV cache互换肯定影响因为内容不同。5.2 批量调度的性能拐点在哪里批量调度不是批次越大越好。我实测下来批次大小从1增加到8的时候吞吐量提升最明显大概有3到4倍。从8增加到16提升放缓大概只有1.5倍。超过16之后提升非常有限甚至可能因为padding浪费和内存压力导致下降。拐点出现的原因有两个一是GPU的计算单元在批次较小时利用率不足批次增大能填满二是批次增大后padding比例上升无效计算增多。所以实际部署的时候批次大小要根据请求长度分布来调不能一刀切。5.3 注意力掩码的构造细节注意力掩码的构造是批量调度里最容易出错的地方。我踩过的坑包括掩码形状不对、掩码类型不对bool和float混用、掩码没有覆盖缓存维度。正确的做法是掩码的形状应该是[batch_size, 1, query_len, key_len]其中key_len包括输入长度和缓存长度。掩码的值在有效位置为0无效位置为负无穷或者一个极大的负数。这样在softmax之后无效位置的权重趋近于零。还有一个细节是如果用了因果自注意力causal attention掩码还要加上一个下三角矩阵保证每个位置只能看到自己及之前的位置。这个下三角矩阵的构造需要根据query_len和key_len的相对关系来定不能简单套用。5.4 实操心得先单请求跑通再逐步加并发我的建议是不要一上来就搞批量调度。先把单请求的逻辑跑通确认KV cache的传递、注意力掩码的构造、位置编码的偏移都正确。然后加两个请求手动构造批次看输出是否跟单请求一致。确认一致之后再加到四个、八个逐步增加。每一步都要做输出对比同样的prompt单请求跑出来的结果和批量跑出来的结果应该完全一致在贪心解码下。如果出现不一致就说明批量调度里有状态污染或者掩码错误。这个对比测试是最有效的排查手段比看日志快得多。另外KV cache的数值精度也值得注意。我试过用float16存缓存结果在长序列生成的时候出现了数值溢出输出变成乱码。后来改成float32问题消失。如果显存紧张可以用bfloat16它的动态范围比float16大不容易溢出。6. 这个实验后续还能怎么扩展这个口令实验虽然规模小但揭示的问题在真实推理服务里普遍存在。后续可以往几个方向扩展一是引入更复杂的调度策略比如连续批处理continuous batching让请求可以在生成过程中动态加入和退出而不是等一个批次全部生成完再处理下一批。这需要更精细的缓存管理和掩码构造。二是测试不同注意力机制的隔离需求比如多头自注意力、交叉注意力、因果自注意力它们在Q、K、V的共享和隔离上有不同的要求。交叉注意力的K和V来自编码器输出是只读的可以共享Q来自解码器是请求私有的必须隔离。三是把实验扩展到多模态场景比如同时处理文本和图像输入。图像特征通常是只读的可以共享文本的KV cache是请求私有的必须隔离。这种混合场景下的状态管理会更复杂但也更有意思。我在实际使用中发现状态隔离做得好不好直接决定了推理服务的稳定性和吞吐量。一开始图省事用全局变量后面排查问题花的时间远超预期。所以如果让我给建议我会说从第一行代码开始就把状态边界划清楚比后面出问题再改要省事得多。