RTPLLM推理优化：K3芯片上MLA注意力的DCP与通算融合实践

干推理优化这行久了你会发现一个怪现象显存省下来了时间却没省下来。这次折腾的事情一句话说清就是在自研推理引擎RTPLLM里针对K3加速芯片上的MLA注意力模块做了DCP动态计算划分和通算融合通信计算重叠把原本被跨卡通信拖住的那部分时延硬生生又抢回来一截。说起来背景也不复杂。我们手里的K3是一颗片上资源相当猛的加速芯片算力不差缓存也不小但跑长序列MLA模型时多卡之间的同步等待经常把计算单元饿着。一开始我以为是大模型推理的老问题——显存带宽不够、算子碎片化后来把profile数据拉出来一看AllGather潜在KV的等待时间能占到单步时延的35%上下这才意识到问题的重心不是“怎么算得慢”而是“怎么让通信别挡着计算”。这篇文章就把我们在RTPLLM里给K3 MLA做DCP和通算融合的过程、踩过的坑、实测数据都抖搂出来供同样折腾推理引擎的朋友参考。1. 为什么偏偏是MLA又为什么非要做通算融合1.1 从KV Cache到MLA省下显存却不省通信MLAMulti-head Latent Attention多头潜在注意力核心思路很简单把传统MHA那个又大又笨的KV Cache压缩成一个低维的潜在向量。推理的时候每来一个新token只需要往这个潜在向量里追加一小段压缩后的K和V而不是存完整的高维K/V矩阵。从显存角度讲长序列下MLA能比MHA省一个数量级的KV Cachebatch可以开得更大上下文可以拉得更长这也是当初RTPLLM选择适配MLA的主要原因。但MLA有个反直觉的地方它省的是单卡显存和单卡带宽跨卡通信并没有相应减少甚至因为通信块的形状变了调度反而更别扭。分布式推理里我们要做张量并行或序列并行每张卡手里的潜在KV都只是一部分必须通过集合通信把完整的潜在KV汇齐后续才能做attention。这在MHA时代就存在不过当时KV块又大又整一个AllGather下去虽然慢但同步关系清晰MLA的潜在向量维度小单次通信的数据量看着不大可一旦序列变长、batch变大通信次数和等待次数照样往上涨。实测最直观序列长度32K、8卡并行时AllGather潜在KV的耗时占到step时间的37%最差一次冲到45%。这个占比意味着计算单元大量时间在空转——卡里头的矩阵单元已经闲下来了但数据没到齐只能等。这场景下靠堆算力已经没意义必须让通信和计算叠起来跑。1.2 RTPLLM里K3的处境算力够用被同步拖住先交代一下K3在RTPLLM里的位置。RTPLLM是我们内部一直维护的推理引擎算子库、调度器都是自己写的目标是把新硬件上的性能榨干。K3这芯片有个特点片上SRAM和L2容量都挺阔绰矩阵单元的执行效率高但跨卡互联的带宽参数相对普通不是那种互联带宽特别夸张的卡。这种“计算强、通信弱”的组合在长序列MLA场景下会非常快地滑向通信瓶颈。传统做法是搞通信掩盖前一个batch计算的时候预取下一个batch的通信数据。但MLA解码阶段对通信的依赖不是一次性的而是每个token生成都要同步一次潜在KV。一旦batch内的序列长度涨落比较大固定的预取窗口根本跟不上实际需求。我们最初试过把通信提前一个step预取效果有但跟理想差距很大——因为decode阶段每个step的Q不一样KV在生成后立刻就要用能预取的空间本来就小。DCP和通算融合就是在这种压力下提出来的。DCP解决“算力分布不均”的问题通算融合解决“通信等待”的问题两者结合效果不是11而是互相成就DCP把计算切成细粒度的块之后通信块也随之变小天然就有了重叠的缝隙。2. RTPLLM中MLA推理的算子边界与瓶颈模型2.1 推理时MLA的计算流到底长什么样要把DCP做好首先得把MLA在推理时的计算流掰碎了看。RTPLLM里MLA单层前向大致走这几步输入隐藏状态h过输入投影分成两条路一条算query的潜在表示以及RoPE相关部分另一条算压缩后的KV潜在表示。KV潜在表示这步是关键它不再是完整的高维K和V而是一个低维潜在向量c_KV。这个向量就是KV Cache里存的东西也是通信时要AllGather的东西。真正做attention之前需要把潜在KV通过上投影拆回完整的K和V或者更准确地说在注意力计算时才逐块恢复不必先把完整K/V物化出来。之后就是标准的QK点积、softmax、加权求和最后再做输出投影。这里边有个细节值得注意MLA的QK点积不是直接在低维潜在空间做的通常要把潜在向量先升维或者是分开处理RoPE部分的key和latent部分的key具体取决于实现。RTPLLM为了省显存选择的是“边恢复边计算”的方式也就是attention计算时不提前把完整K/V铺满全局内存而是按块恢复、按块消费。这个设计天然适合DCP既然K/V都是动态恢复出来的那我按query的chunk来切分attention计算就非常顺——每个chunk只需要恢复该chunk需要的部分K/V算完就扔不需要等全量数据到齐。2.2 通信与计算的比值量化一下到底差多少为了说明问题严重性我们给一个简化但足够指导决策的模型。设并行卡数为Nbatch大小为B序列长度为LMLA的低维KV潜在向量维度为d_l单元素字节数为2FP16。单个解码step每张卡需要AllGather的数据量大致是C_comm B × L × d_l × N × 2 字节而同一step的实际计算量核心是整个attention部分q矩阵、pointwise操作、以及恢复K/V的矩阵乘。若只看QK^T这个热点其浮点量约等于 2 × B × L_q × L_k × head_dim其中L_q是query长度decode阶段通常为1或小chunkL_k是KV长度。把两者放在一起看通信量和KV长度成正比计算量也随KV长度线性增长。但问题在于通信带宽远低于计算对带宽的消耗速度。以K3为例矩阵单元每个周期能消费的数据量远大于互联链路能送过来的数据量。结果就是计算单元在等待KV数据的空窗期里浪费掉的算力非常可观。DCP要做的事情就是把这个“等全量KV”的线性节奏打散变成“来一块算一块”的流水线节奏。理论上只要通信切片足够细并且计算单元能及时消费到达的切片等待时间就能趋近于零。这也是我们在RTPLLM里选DCP而不是单纯加大预取深度的根本原因。3. DCP动态计算划分的具体设计与落地3.1 DCP决策三阶段动态划分DCP全称Dynamic Compute Partitioning动态计算划分。它不是什么新概念但在MLA场景下有特殊的实现角度。传统算子融合里的切分往往是静态的按照固定的block大小切谁先到谁算。DCP强调的是“动态”两个字——切分的粒度、切分的时机、分配到哪个计算单元都由运行时根据当前batch的状态来决定。RTPLLM里的DCP分三个阶段第一是探测阶段。每个step开始前调度器拿一组轻量信号当前batch的有效query长度、KV cache的实际长度、上一step的通信时延、计算队列深度。这些信号不需要很精确反映趋势就够。第二是决策阶段。根据这些信号用一个启发式函数决定本step的切分参数。核心产出是query chunk大小也就是把当前这些query切成多少个块。决策的目标函数是让“通信完成曲线”和“计算消耗曲线”尽量贴在一起。第三是执行阶段。把决策结果交给下层调度器生成对应的小任务下发到K3的计算队列。每个query chunk对应一个任务任务之间允许乱序完成但最终会按顺序合并结果。这个流程说起来简单真正难的是决策函数的设计。我们第一版直接用固定阈值比如seq_len大于16K就切128否则不切结果遇到混合batch时效果很渣——有的请求算完了没得可切有的请求还在等通信。后来改成实时预估时延差根据前一步的实际通信时延修正切分系数效果才稳定下来。3.2 chunk大小与K3片上缓存的匹配关系DCP切分粒度不能拍脑袋必须跟硬件缓存容量对齐。K3的片上SRAM和L2总容量是有限的attention中间结果最怕的就是爆缓存。我们推导了一个简单的约束关系一个query chunk的partial attention中间结果主要由partial QK矩阵决定。设batch为Bquery chunk长度为Q_chead数为Hhead维度为D则中间结果是 B × Q_c × H × D 个元素。如果Q_c选得太大这个矩阵就超出片内容量一部分就得写回L2甚至HBM反而得不偿失。在我们压测的配置里batch32、Q_c64、H32、D128时中间矩阵大小是 32 × 64 × 32 × 128 × 2字节算下来约16MB勉勉强强能在K3片内SRAM放下。如果Q_c提到128直接翻倍到32MB就会跟其他数据抢L2性能不升反降。这就是为什么我们后续多数实验Q_c固定取64宁可稍微牺牲少数计算单元的利用率也要保住缓存命中率。3.3 决策函数的具体形式这里给一段我们调过的伪代码帮助理解DCP的决策逻辑。它不是最终版本但具备代表性def decide_chunk_size(batch_size, max_q_len, kv_len, prev_comm_latency_us, calc_queue_depth, sr_available_bytes): # 估算通信时间通信量和上一拍实际通信带宽拟合 kw_est batch_size * kv_len * latent_dim * num_shards * 2 comm_us max(kw_est / comm_bw_est, prev_comm_latency_us * 0.8) # 估算每个候选chunk的计算时间 best_chunk None best_score float(inf) for chunk in [32, 64, 128]: partial_bytes batch_size * chunk * num_heads * head_dim * 2 if partial_bytes sr_available_bytes * 0.8: continue calc_us chunk * kv_len * batch_size * gflops_per_peak * 1e-6 score abs(comm_us - calc_us) 0.2 * calc_us if score best_score: best_score score best_chunk chunk return best_chunk or 128核心思想是让通信时延和计算时延互相咬合。score函数里那个0.2是用来惩罚计算时间过长的小权重防止选出来的chunk让计算成为新瓶颈。实际运行下来这个函数在绝大多数step都能收敛到合理值偶尔出现抖动时我们会额外加一个“惯性系数”让chunk大小不要每step都剧烈跳变减少调度器重排的代价。4. 通算融合让通信成为计算的填充4.1 三层重叠策略DCP把计算切碎之后真正发挥威力还得靠通算融合。通算融合或者叫通信计算融合就是把原本串行的“通信→计算”变成“通信→计算→通信→计算”的流水线。我们在RTPLLM里分三个层次做了融合第一层是通信切片级融合。通信的数据量再小如果必须等全量到达才开始算那还是白搭。我们不再一次性AllGather整个潜在KV块而是把它切成多个slice。每收到一个slice立刻触发一个计算任务对已到达的K/V做partial QK点积。第二层是算子内融合。把RoPE变换、潜在KV恢复、partial QK计算合并成一个kernel。这样KV数据到达后不需要在全局内存里呆一圈直接在片上SRAM里完成恢复和消费减少内存往返。第三层是引擎级融合。K3的通信DMA引擎和计算引擎是独立的我们利用这个特性把通信slice的数据搬移任务和计算任务分别放入两个队列计算任务不主动等待通信任务完成而是注册一个“slice就绪”的回调通信完成多少就算多少。4.2 流式AllGather与partial softmax的实现实现里比较关键的是不能直接用集合通信库的阻塞式AllGather得用流式传输原语。RTPLLM里我们封装了一个非常轻量的接口// 伪代码示意 StreamAllGather(kv_latent_slices, [](Slice s) { // 每个KV slice到达后触发partial计算 PartialQK(s); if (all_slices_ready) { MergeSoftmaxAndFinish(); } });PartialQK是这个方案的核心。传统attention的softmax需要全局的max和sum切片式partial QK显然没法一次性算出标准softmax。我们的做法是维护online softmax的状态每个slice到达时更新本地的running max和running sum所有slice完成后做一次最终归一化。这个坑必须说清楚如果用普通softmax的公式硬算在slice边界会出现数值错位生成质量直接崩。online softmax的rescale步骤不能省而且要注意并行slice之间做rescale的先后顺序否则可能引入不确定浮点误差。我们调试的时候有整整一个下午都耗在这个问题上。4.3 实测效果从0.63利用率拉到0.88把DCP和通算融合都启用之后我们跑了一轮8卡压测。模型是一个中等规模的MLA架构模型单层per-step数据如下非端到端但趋势有代表性方案query chunk单step平均时延吞吐(tok/s)计算单元利用率基线无DCP、无融合不切分1.30ms24500.63仅DCP1281.08ms29500.72DCP 通算融合640.87ms39000.85DCP 通算融合 深度流水640.71ms46000.88利用率从0.63提升到0.88主要功劳是把等通信的那段空窗填上了。数据也符合预期单靠DCP只能改善约17%再多就触及天花板一旦通信计算叠加起来改善幅度直接放大时延下降接近45%。这里有个值得注意的反直觉点启用深度流水后平均时延其实还能再压但计算单元利用率只涨了0.03。原因是深度流水借助多倍buffer让计算队列时刻有活干但K3的调度器在切换task时仍有开销。压到后面主要瓶颈已经变成片上任务切换的overhead了。5. 调参经验与避坑记录5.1 三个关键参数q_chunk、通信切片、流水深度基于这段时间的折腾我觉得最有必要讲清楚的是三个参数q_chunk_size决定计算切分的粒度。我们试过32、64、128、256结论是64和128最稳。32太小调度和kernel启动的固定开销占比太高128以上在长序列下会明显挤占片上缓存。实际项目里建议从64起步用profile数据决定要不要往128调。recv_slice_bytes是通信切片的字节数一般取16KB到64KB。这个值跟DMA描述符数量、片上缓存行大小都有关系。切得太小DMA中断和回调频率过高CPU/控制器反而成为瓶颈切得太大计算侧能开始干活的时间点就晚了重叠效果变差。pipeline_depth是同时驻留的通信slice数量。深一点能吸收网络抖动但每个slice都要占一块SRAM做staging buffer不能无限加深。我们最终常用值是4在内存占用和抖动容忍度之间平衡。5.2 踩坑实录五个让人头大的问题第一个坑是online softmax的rescale顺序问题前面已经提过。这个属于“不亲自动手看不出问题”的类型最终我们通过单元测试构造极端数值比如某些slice的max远超其他slice才暴露出来。第二个坑是通信buffer被覆盖。流式通信回调里用到的buffer如果在计算任务还没读完就被下一个slice复用会出现偶发的错误结果。这个问题的可怕之处在于它不是必然触发只在特定时序下出现排查起来特别费劲。解决办法是在slice descriptor里加引用计数计算任务持有引用时通信侧不能复用这块buffer。第三个坑是乱序问题。DCP切分后的多个chunk是乱序完成的如果直接把结果写回输出采样阶段会拿错token。这里没有捷径必须是chunk按序合并后才能进入下一步。我们一开始图省事想让调度器保证完成顺序结果牺牲了大量并行性后来干脆放弃保序改成结果缓冲加按序归并。第四个坑是prefill阶段开DCP反而变慢。prefill阶段query长度本身很大KV长度也大通信量巨大但计算量同样巨大切分后调度开销造成的影响比通信等待还明显。最终我们的策略是DCP只在decode阶段启用prefill阶段保持原有静态切分。第五个坑是“省了显存但没省时间”的假象。光看显存占用会以为MLA已经解决了长序列问题但实际推理时间没降多少。我们中间有一版只做了DCP没做通算融合显存降了40%端到端延迟只降了8%差点让我误判方案无效。后来把通算融合叠上去延迟才开始明显往下走。这也说明显存和延迟是两个独立的优化维度不能互相替代。5.3 一点取舍建议说到底DCP和通算融合不是万能的。batch比较小、序列比较短的时候通信量本身就小静态AllGather加上简单的预取已经够用上DCP反而因为调度开销让性能方差变大。我们的经验门槛大概是batch≥32、有效序列长度≥2048这个区间DCP的收益才会稳定超过其成本。另外MLA的潜在KV维度是固定的DCP切分的真正对象是query不是KV。理解了这一点你就明白为什么DCP对MLA这么友好——query长度在decode阶段往往很短切成小chunk后每个chunk的计算量刚好和通信slice到达的节奏匹配。如果你在别的架构上想复用这套方案先确认你的注意力模块是否存在某种低维瓶颈没有的话通信切片的计算颗粒度很可能对不上。6. 几个尚未解决的老问题这套方案整体跑通了但我心里清楚还有不少地方只能说是“可用”不是“优雅”。DCP的决策模型依赖时延估算而系统噪声一上来估算就容易失准。K3上偶尔会遇到某个slice的通信突然慢了30%决策函数还在按上一拍的平均值切分于是这一个step的重叠效果就明显打折。后续想把决策改成在线强化学习式的自适应直接根据实时观测去微调q_chunk但工作量不小。通信slice到片上缓存位置的映射还比较机械。现在我们是固定staging buffer区域能不能根据KV slice的访问局部性动态跳转缓存位置让后续的注意力计算直接从L2命中而不是再绕一次SRAM理论上能再挤出一点性能但工程复杂度高得吓人。另一个方向是把DCP扩展到跨节点场景。目前这套优化严格限定在单机多卡范围内跨节点时通信延迟跳变一个数量级slice到达节奏会变得很不规律现有的固定深度流水可能扛不住。这可能需要把pipeline_depth也变成动态可调的当前还没有想好怎么设计得足够稳。把这些未解问题写出来不是为了凑篇幅而是想说明一个事实推理优化这行没有终点你解决掉一个瓶颈下面永远藏着下一个瓶颈。但这就是有意思的地方。最后分享一个让我印象很深的体会一开始做DCP和通算融合的时候我以为最大的障碍是找不到合适的切分算法或者模型精度受影响。真正开干之后才发现最大的障碍是把“计算单元有空闲时调度器愿意立刻填上通信驱动的计算任务”这件事做到足够顺滑。K3本身不缺算力RTPLLM的算子库也不缺性能缺的恰恰是让流水线时刻饱满的那种协调感。DCP给了我们切分工具通算融合给了我们流水线两者配合起来才终于把MLA在长序列下的性能潜力真正放了出来。如果你也在折腾类似的东西建议从“通信切片触发计算回调”这个小环节入手先把重叠的手感找到再谈更花哨的动态调度。