ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

Prefill阶段优化实战:分块Prefill与动态负载分配解决快慢卡问题

Prefill阶段优化实战:分块Prefill与动态负载分配解决快慢卡问题

1. 先说清楚:Prefill阶段到底是什么,为什么值得专门优化

做大模型推理部署的人,对“Prefill”这个词一定不陌生。但我在实际交流中发现,不少刚接触推理优化的开发者,对它的理解还停留在“就是算一次前向传播”的层面,这就会导致真正遇到性能问题时无从下手。

Prefill阶段,指的是模型在生成第一个token之前,对用户输入的全部历史token做一次完整的前向计算。这个过程会产生每个token对应的Key和Value向量,也就是我们常说的KV Cache,并缓存在显存里,供后续的Decode阶段逐token生成时反复读取。一次对话输入了512个token,那Prefill就要把512个token一次性算完;输入了4096个token,那就得把4096个token全部算完,才能开始吐第一个字。

这里有个关键的直觉差异:Prefill阶段的计算特征和Decode阶段完全不同。Decode阶段是每次只算1个token,是典型的访存密集型(memory-bound)操作,大量时间花在把KV Cache从显存搬到计算单元上,计算量反而不是大头;而Prefill阶段是一次性算完一整条输入序列,是典型的计算密集型(compute-bound)操作,每个token都要做完整的矩阵乘(GEMM),计算量非常饱满。

打个比方:Decode像是在流水线上一个一个地拧螺丝,瓶颈在于手能不能跟上;Prefill则像是提前把一整批螺丝全部冲压成型,瓶颈完全在于冲压机本身的功率有多大。

既然Prefill是算力密集型的,那优化的思路就非常明确了——尽可能提高算力利用率(MFU),让每张卡上的计算单元都饱和起来。听起来简单,但实际落地时,你会撞上一个非常普遍又非常磨人的问题,就是标题里提到的“快慢卡问题”。

2. 快慢卡问题:现象、根因与定位方法

2.1 快慢卡现象的三种典型表现

我一开始做多卡推理时,遇到过一个特别让人困惑的现象:8张卡一起跑,监控面板上画出来的利用率曲线像锯齿一样参差不齐。有的卡算力打到了90%以上,有的卡却只有30%多一点;有的卡显存几乎占满,有的卡还剩一大截空余。整个集群的上限被慢卡死死拽住,快卡干着急使不上劲。

这种现象通常有三种典型表现:

第一种是“计算利用率不均”。同一批请求进来,某些卡的SM(流式多处理器)占用率接近满载,还有些卡长期在低负载徘徊。这往往是因为请求序列长度差异太大,长序列请求被分到了某张卡上,短序列请求集中在另一张卡上,导致两张卡的计算压力完全不对称。

第二种是“显存水位不均”。这类问题更隐蔽,显存占用高不代表算力占用高。有些卡因为KV Cache预留空间大,或者承载的请求并发数多,几乎要把显存耗尽;而另一些卡显存还很宽裕,却没有请求分过来,资源白白浪费着。

第三种是“同步等待时间过长”。多卡推理几乎绕不开张量并行或流水线并行,卡与卡之间要靠通信来同步中间结果。慢卡没算完,快卡就得干等。通信耗时一旦占了整个step的30%以上,那基本可以确定你的并行策略或者负载分配出了问题。

2.2 根因拆解:绝不只是一句“负载不均衡”能概括的

快慢卡问题的根因,表面上看是负载不均衡,但深挖起来至少能拆出四层:

第一层是请求长度分布不均。这是最直接的原因。一次batch里输入了5个2048 token的长请求和35个64 token的短请求,如果按请求数均分到8张卡上,那长请求落在哪张卡上,哪张卡就天然要干更多的活。Prefill的计算量跟序列长度是线性关系,长度差10倍,计算量就差10倍,靠调度器是救不回来的。

第二层是KV Cache分配不均衡。现在的推理框架普遍采用分页KV Cache(PagedAttention)机制,显存按块(block)来管理。但有些实现里,KV Cache的预留空间是跟序列长度直接挂钩的,长请求预留的空间大,短请求预留的空间小,分配起来就很容易出现碎片,导致某张卡的可用KV Cache块明显少于其他卡,影响到后续batch的调度。

第三层是硬件拓扑带来的天然差异。多卡服务器上,不同的GPU插槽对应的PCIe通道或者NVLink带宽并不完全一致。相邻的两张卡之间通信带宽可能是900GB/s,跨节点通信可能只有几十GB/s。如果请求被调度到了通信拓扑较差的卡上,即使计算量相同,单次通信的耗时也更高。

第四层是动态显存预留和时钟降频。有些推理框架在初始化时会根据显存总量预留KV Cache空间,预留多了,留给计算图本身的空间就少;预留少了,遇到长序列又容易OOM。再加上如果某张卡温度过高触发了降频保护,算力会直接掉到标称值的一半以下,这也是快慢卡问题的隐性来源之一。

2.3 定位慢卡:从监控工具到时序分析

定位快慢卡问题,我建议按“先看系统层,再看框架层,最后看算法层”的顺序排查。

系统层用nvidia-smi和dcgmi配合看。nvidia-smi dmon可以实时查看每张卡的SM占用率、显存占用量、温度、功耗。功耗是个好抓手,如果某张卡的功耗明显低于其他卡,大概率它的计算负载不高或者处于等待通信的状态。dcgmi可以看到NVLink的收发吞吐和重传率,重传率异常升高说明通信链路可能有问题。

框架层要打日志。我习惯在每个step结束后输出每张卡的耗时分布,具体到“Prefill阶段耗时”和“Decode阶段耗时”分开统计。这样做的好处是能精准定位慢的是哪个环节——如果所有卡Prefill耗时都正常,只有某张卡Decode耗时异常高,那问题多半出在KV Cache读写的访存路径上。

算法层就要拿数据说话了。把一次请求的batch按序列长度排序,画一条长度分布图,再看每张卡实际分配到的请求长度之和,基本就能判断出调度策略是否合理。如果长请求全塞给了同一张卡,那调度器这块就得动刀了。

注意:不要一上来就怀疑硬件损坏。我遇到过好几次“慢卡”,查到最后都是调度策略或者内存碎片的问题,真正属于硬件故障的反而少见。先把软件层排查干净,再考虑走硬件售后流程。

3. 优化方案选型:为什么我最终选择了“Chunked Prefill + 动态负载分配”这套组合拳

3.1 备选方案扫描:连续批处理、分块Prefill与微批次调度

针对快慢卡问题,业界主流的优化方案大致有三条路,我逐一评估过它们的优劣:

第一条路是连续批处理(Continuous Batching)。核心思想是让不同请求的Prefill和Decode阶段在一个step内穿插执行,避免“先全部Prefill完再全部Decode”的瀑布式等待。这个方案的优点是能显著提升Decode阶段的设备利用率,因为Decode阶段是访存密集型,跟Prefill阶段的计算密集型正好互补。缺点是它主要解决的是算力“空闲”问题,对“分布不均”问题帮助有限——快卡慢卡的问题核心是请求分配不均,不是整体利用率低。

第二条路是分块Prefill(Chunked Prefill)。把长序列的Prefill计算切成多个小块,跟其他请求的Prefill或Decode混在一起调度。这块方案的优点在于能打散长序列的“独占效应”,让长请求不再长期霸占某张卡的计算资源,从而天然缓解快慢卡问题。缺点是实现复杂度高,需要框架层面配合调度器细粒度控制,而且切块本身会带来一些额外的拼接和调度开销,切太小了反而划算。

第三条路是微批次调度(Micro-batching)。把一个大的batch切分成N个微批次,依次经管线送入多卡并行执行。这个方案的优点是实现相对简单,能利用流水线并行隐藏通信延迟;缺点是对请求长度敏感,N个微批次的执行时间取决于其中最慢的那个批次,尾部效应明显。

综合评估之后,我最终选定了“分块Prefill + 动态负载分配”的组合路线。核心原因是这套方案能同时解决“算力空闲”和“分布不均”两个问题。分块Prefill先打散长序列的计算粒度,动态负载分配再根据每张卡的实际负载情况调整后续请求的分配权重,两手抓,快慢卡问题才能根治。

3.2 动态负载分配的“为什么”:从静态哈希到水位感知调度

多卡推理框架里最常见的请求分配方式是哈希取模:对请求ID取模,决定分配到哪张卡。举个例子,8卡推理,请求ID对8取模,余数为0就分卡0,余数为1就分卡1,以此类推。这种方案实现简单、确定性好,但它完全没有感知到请求的长度差异,所以快慢卡问题几乎必然会发生。

我最终落地的是水位感知调度(Watermark-based Scheduling)。核心逻辑是把每张卡预估的“当前耗时”当成一个水位线,新请求到了之后,优先分配给水位线最低的卡。

怎么估算“当前耗时”呢?我用的公式是:

[ T_{card} = \alpha \cdot \frac{total_input_tokens_{card}}{perf_{card}} + \beta \cdot queued_requests_{card} ]

其中total_input_tokens是当前卡上所有请求输入token的总和,perf是这张卡的实测Prefill吞吐(token/s),queued_requests是排队中的请求数,(\alpha)和(\beta)是两个可调权重。这个公式的本质是:把当前卡上“累积待处理的计算量”除以“卡的算力”,得到预估完成时间,同时考虑排队项,避免某张卡因为分配了太多短请求而锁死。

实际效果很直观:长请求来了,不会一股脑塞给同一张卡,而是会优先落到当前负载最低的卡上。快卡的水位线低,自然接到的活多;慢卡水位线高,自然会少接活。这比静态哈希的“一刀切”要灵活得多。

3.3 为什么“切块”能破局:从资源独占到时间片共享

分块Prefill的底层逻辑,用一句话就能说透——把大任务拆成小任务,让多张卡都能分到一杯羹。

举个例子,如果有一个8192 token的长请求,不做分块的话,Prefill阶段它会被完整地塞进一张卡里,在这段时间里,这张卡几乎没法处理其他请求的Prefill计算。而分块Prefill会把8192 token切分成4个2048 token的块,每个块可以在不同的卡上执行,也可以在同一张卡上跟其他短请求交错执行。

这样一来,长请求不再独占任何一张卡的计算资源,而是以时间片的形式与其他请求共享算力。短请求可以插在长请求的块与块之间执行,整体的设备利用率就上来了,快慢卡的“卡点”也就被拆解掉了。

当然,分块不是没有代价。切块之后要做额外的KV拼接——第一个块算出的KV Cache要缓存起来,第二块算的时候要同时读取第一块的KV作为上下文——这里就多了通信和内存拷贝的开销。切块粒度越大,开销越小但调度灵活性越差;切块粒度越小,调度越灵活但额外开销越大。我实测下来,把单块长度控制在512到1024 token之间比较合适,既能把长请求的独占时间控制在百毫秒级别,又不会让额外拼接开销高到反噬收益。

4. 实操过程:从参数调整到效果复现的完整记录

4.1 实验环境与基线数据

先说实验环境,方便你复现时参考。我用的是8卡GPU服务器,单卡显存80GB,CPU端64核,系统内存512GB,GPU间通过NVLink全互连。推理框架选的是当前社区里流行的开源推理引擎,模型是一个70B参数的稠密大模型,精度FP16,张量并行度TP=8,即整模型权重分布在8卡上。

基线配置用的是框架的默认参数:连续批处理开启,静态请求分配(哈希取模),不分块Prefill,KV Cache自动预留,最大batch数设为128。输入输出长度配置上,短请求(64-256 token)约占70%,长请求(1024-4096 token)约占30%,平均输入长度约1580 token。

用这套基线配置压测,首轮我用的是“固定并发数 + 无限请求流”的方式,发压工具模拟了32路并发持续请求。跑满10分钟之后取稳定数据,基线表现如下:

  • 整体吞吐率(输出token吞吐):约4800 token/s
  • Prefill阶段平均单请求耗时:约380ms
  • 8卡SM利用率:最高93%,最低41%,标准差约19%
  • 最大尾延迟(P99):约5.8s

注意看第四行,SM利用率标准差19%已经很能说明问题了,最慢的卡比最快的卡利用率低了一半还多。这就是典型的快慢卡问题形态。

4.2 第一刀:开启分块Prefill(Chunked Prefill)

第一步我先只改分块配置,其他参数保持基线不动。把enable_chunked_prefill打开,max_prefill_chunk_size设为1024。这个参数的含义是:单次Prefill最多算1024个token,如果一个请求的输入长度超过1024,就把多余的token拆成后续的块再算。

重启服务,重新压测。我预期的是长请求的尾部延迟能降下来,但没想到整体吞吐率先涨了。数据如下:

  • 整体吞吐率:约5600 token/s,比基线提升了16.7%
  • Prefill阶段平均单请求耗时:约220ms,降幅42%
  • 8卡SM利用率:最高89%,最低55%,标准差缩小到13%
  • 最大尾延迟(P99):约4.2s

这里的收益来源不难理解:长请求不再独占整张卡,短请求有更多机会穿插执行,整体设备利用率上去了,吞吐自然提高。但注意最低卡利用率还是只有55%,标准差13%意味着快慢卡问题还没根治。

4.3 第二刀:换成水位感知动态调度

保持分块Prefill打开,把请求分配方式从哈希取模改成水位感知调度。这里我没有改动框架源码,而是用框架暴露的自定义调度接口,在应用层写了一个调度插件,核心逻辑就是我上面那个水位公式。权重参数(\alpha)取0.6,(\beta)取0.4,这个比例是我试了几组值之后挑出来的,主要兼顾了长请求和排队项的敏感度。

重启压测,数据如下:

  • 整体吞吐率:约6400 token/s,比第二刀又提升了14.3%,比基线提升了33.3%
  • Prefill阶段平均单请求耗时:约190ms
  • 8卡SM利用率:最高91%,最低70%,标准差缩小到8%
  • 最大尾延迟(P99):约3.1s

卡的利用率分布明显收窄了,最低利用率从55%拉到了70%,高位卡和低位卡的差距被压缩到了接近20个百分点以内。吞吐率的提升主要来自慢卡不再拖后腿,每一张卡都更饱和地在干活。

4.4 第三刀:KV Cache显存水位优化与gc策略

分块Prefill和水位调度解决了算力分配不均,但显存利用率还不够理想。我观察了压测时的显存监控,发现有几张卡的KV Cache预留空间很大,空闲块很多,但实际利用率不到60%。

我做了两个调整。第一个是打开KV Cache的按需分配(demand-based allocation),让框架根据实际batch需求增长KV Cache,而不是启动时就圈死一大块。第二个是设置KV Cache清理阈值,当已用比例超过85%时,把超过空闲时间阈值的请求提前清出,释放KV Cache块。

这两刀下来的压测数据:

  • 整体吞吐率:约6780 token/s,比上一轮再提升5.9%,相比基线提升41.3%
  • 最大尾延迟(P99):约2.7s
  • 显存占用率:8卡更均衡,标准差从11%降到5%

到这里为止,快慢卡问题从“算力不均”“显存不均”两个维度都被明显缓解了,整体吞吐率从4800提到了6780左右,提升了约41%。这个幅度在不动模型结构、不换硬件的前提下,已经相当可观了。

4.5 一版可直接参考的参数速查表

我把这次优化中用到的关键参数整理成一个表格,方便你对照自己的环境快速预设:

参数项基线值优化后推荐值说明
分块Prefill开关关闭开启长请求切块,打散资源独占
单块最大长度无512-1024 token切太小费通信,切太大加剧不均
请求分配策略哈希取模水位感知调度快卡多接活,慢卡少接活
调度权重α无0.6权衡计算量占比
调度权重β无0.4权衡排队请求占比
KV Cache分配启动时全量预留按需分配减少显存空耗
KV清理阈值无85%避免高峰OOM

提示:这一版参数是基于TP=8、单卡80GB的环境调出来的。如果你的环境是单卡或者TP=4,调度权重(\alpha)和(\beta)需要重新标定。我一般是先跑一组短压测观察吞吐和延迟,再按结果微调。

5. 常见问题与排查技巧实录

5.1 优化后P99延迟反而变高了,怎么回事?

这是我踩过的一个很实在的坑。第一轮开分块Prefill的时候,整体吞吐确实上去了,但P99尾延迟从5.8s降到4.2s,看着挺好。结果后来有一组压测里,P99突然飙到了7s多,比基线还差。

排查思路:把P99对应的请求日志拉出来,逐个看它们的输入长度和调度记录。发现延迟飙升的全是超长输入(超过8k token)且跟其他长请求扎堆出现的请求。原因是分块Prefill虽然能把长请求拆开,但拆出来的块依然要按顺序执行,如果多个长请求同时进来,块与块之间会在同一张卡上排队,反而把尾延迟拉高了。

解决办法:给系统加了按长度分级的优先级队列。短请求(输入小于512 token)走最高优先级,长请求(大于2048 token)走低优先级。这样短请求能被及时穿插执行,长请求的块等待时间虽然略有增加,但P99整体的表现反而稳定了。我实测下来,加了这个队列之后P99从7s回落到3s以内。

5.2 水位调度之后,最慢的那张卡还是慢,怎么继续定位?

水位调度只能做到“请求按负载合理分配”,但如果某张卡本身的硬件或软件环境就有问题,调度器再聪明也补不回来。这时候建议做一次干净的基线测试:只启动一个很小的batch,只跑一份固定长度的输入,逐卡记录执行耗时。

如果单卡独立执行都慢,重点检查三件事:

  • 温度是不是偏高,有没有触发降频。nvidia-smi -q -d TEMPERATURE看看当前温度,如果超过80度,先改善散热再谈优化。
  • 不可屏蔽中断(NMI)和PCIe错误计数有没有异常。nvidia-smi -q -d ECC,如果有大量报错,建议走售后检测流程。
  • GPU驱动和推理框架的版本跟其他卡是不是一致,不一致的话,重新安装统一版本。

如果单卡独立执行不慢,只是合在一起跑的时候慢,那就要回到通信上排查——查NVLink的link利用率是否均衡。如果某张卡的NVLink吞吐明显低于其他卡,多半是这台机器的PCIe拓扑有异常,或者NVLink的链路有降速,这类物理问题调度器无解,但可以作为性能基线记录在案,后续做任务分配时手动避开这张卡。

5.3 分块Prefill的块大小设多少合适?

这是一个典型的“既要又要”问题——块太大,长请求的独占时间太长,分块的意义就折扣了;块太小,KV拼接和额外调度的开销会反噬性能。

我提供一个实测节奏:先在默认配置下跑一版,然后分别试512、768、1024、1280四个档位,每个档位跑10min压测,记录吞吐率和P99。以我的环境来说,最终定了1024,因为在这个档位上吞吐率最高且P99没有明显劣化。块大小跟模型的每层维度、注意力头数有关,不同类型的模型最优档位可能有差异。如果你用的模型规模比我小(比如7B或13B),可以从512开始试;如果模型更大(比如超过100B),可以从1536开始试。

5.4 读完这一篇,你动手时最容易忽略的3个细节

第一,压测之前先把环境变量和框架版本锁定。我见过不少团队的环境里框架版本不一致,同一套参数跑出来数据差异很大,导致后面所有优化结论都不可信。我习惯把框架版本、CUDA版本、GPU驱动版本、模型精度全部记录在实验文档里,和压测数据放一起存档。

第二,不要只看平均吞吐,一定要分段看Prefill和Decode的单独耗时。平均延迟会把两个阶段的特征混在一起,掩盖问题。我通常都会单独记录Prefill阶段的平均耗时和Decode阶段的每token耗时,分别关注。优化Prefill阶段时,重点看的目标就是“在吞吐不降的前提下,Prefill耗时降了多少”。

第三,切分点要放在解码框架的调度循环里,别放在模型前向里。这个是我踩过的坑,放在模型前向里切块,等于把注意力计算改了,KV拼接逻辑全乱掉。分块Prefill的执行顺序是“调度器拆块->逐个算块->块间KV衔接”,必须由框架的调度循环来控制,模型前向结构保持不变。如果发现某个实现方案需要改动模型内部的注意力代码,最好先停下来仔细想想这个方案是不是绕了远路。

6. 最后说点实在的

这一轮Prefill阶段优化做下来,我最大的体会是:性能问题的表象千奇百怪,但根子往往藏在“资源分配”和“执行粒度”这两个维度里。快慢卡问题看起来是负载不均衡,本质上却是调度策略太粗糙——哈希取模没有感知请求长度,不分块Prefill让长请求天然独占资源。你去把这两个维度补上,大部分问题都能缓解一大半。

我自己实际跑项目的顺序也供你参考:第一优先做分块Prefill,因为它改动量小、风险低、见效快;第二优先做水位感知调度,这需要一点代码量,但收益巨大;第三才是扣KV Cache显存的细活,属于锦上添花。如果你时间有限,就先把前两步做了,吐吞掉的性能都能捞回来一大截。

最后一个小技巧:优化做完之后,把基线配置的压测场景保存成脚本,以后每次升级框架版本或者换硬件,都跑一遍同一套脚本对比数据。性能优化这事,最怕的就是“感觉变快了”但没有数据。有数据对比在手里,你后面遇到任何性能回归,都能第一时间定位问题出在哪一层。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进