大模型服务做到多租户微调之后LoRA基本是绕不开的配置。同一个底座模型不同业务方各训一套低秩适配器推理阶段放在同一批请求里跑谁也不干扰谁。vLLM里这套机制通常叫动态LoRA落到昇腾NPU上事情就没有框架层看起来那么顺滑LoRA的核心计算就是两个小矩阵乘加残差但在AI Core上直接拆成四个普通矩阵乘去跑性能和显存都扛不住必须单独做融合算子。我最近把vLLM-Ascend里LoRA这条链路从头到尾梳理了一遍读完代码、调过算子、也踩过一些坑这篇文章就把“一个带LoRA的batch请求从调度到算子返回”的完整过程以及核心算子的内部设计逻辑一次讲清楚。1. 为什么LoRA推理在昇腾上需要单独设计算子1.1 LoRA的计算本质与两条实现路径先回到问题源头。LoRA微调得到的权重形态上就两个矩阵A形状是[r, K]B形状是[N, r]。r是秩通常16到64K是输入维度也就是hidden sizeN是输出维度比如4096。一个有LoRA的线性层计算等价于y_base x W^Ty_lora (x A^T) B^Ty y_base y_loraW是基座模型已经训练好的大矩阵。问题是这个式子怎么落实。有一条路是权重合并推理前把W整体变成W W BA的等价形式然后这个层就退化成普通线性层完全不感知LoRA。这一招对小规模部署或者单适配器场景很好用但代价也很直接每个适配器都要额外存一份合并之后的完整权重而且一旦切换到另一个LoRA要么整体换权重要么先把W恢复成原始值再合并新的I/O开销非常可观。如果服务里同时放着20个适配器切来切去会发现时间全耗在权重变换上。另一条路就是“动态展开”也是vLLM-Ascend实际采用的路径。不合并保留A和B推理的时候临时算x然后(x A)再B最后加到主路径输出上。好处是权重只在适配器加载的时候搬一次之后常驻显存batch里不同token可以用各自的LoRA互不影响。代价则是计算量上每次都要多跑两个小矩阵乘。当然这个“代价”正是整篇文章要解决的问题怎么让这两个小矩阵乘在昇腾上跑得快。1.2 动态LoRA场景下的性能难点动态LoRA真正麻烦的地方在于批量推断时的不规则性。举个例子batch里32个请求可能分布在20个不同的LoRA上。如果朴素做法每个请求的LoRA部分都单独调一次GEMM调度开销、算子启动开销、显存搬运开销都会迅速累积。而且还有一个token粒度的问题一个decode step中每个序列只有一个token32个序列就是32个token每个token可能映射到不同的adapter。如果按请求逐个计算每个GEMM的M维度经常是1或者很小的值这对昇腾这种强计算阵列来说是灾难AI Core利用率可能连1%都不到。另一个难点是秩不一致。适配器A的秩是16适配器B的秩是32还有可能是64。不同秩导致中间矩阵S在M和r维度上都不规整无法简单拼接成一个标准GEMM让算子一次算完。在GPU上CUDA可以用大量线程做灵活的分散计算来处理不规则性昇腾的cube阵列则更适合规则稠密分块这种动态性直接把实现复杂度拉高了。再加上注意力层和MLP层合起来一个token往后传播要经过多个带LoRA的线性层每层都要查一次权重、做两次小GEMM。这些如果不在一个算子内融合Host侧每层下发任务的次数会翻好几倍层数一多主机侧吞吐就成了新的瓶颈。1.3 昇腾AI Core的计算特征昇腾NPU的基本计算单元叫AI Core它有几个关键特性直接决定了LoRA算子的设计方向。第一cube单元负责矩阵乘但偏好规整的大矩阵对M、N、K维度有对齐要求常见是16的倍数。LoRA的A矩阵只有[r, K]r甚至只有16直接让cube算这种瘦长矩阵有效计算只占指令周期的一小部分。第二数据要从GM搬运到L2再到本地内存UB搬运成本极高算子间反复读写中间结果是非常奢侈的。第三算子内多核并行需要软件流水和同步机制每多一个算子同步和数据搬运的开销就多一层。所以答案就很明确了必须把LoRA计算写成一个自定义融合算子在算子内部通过合理切块发挥cube能力一次性完成“取权重、算A、算B、加回主输出”整条链路。这也是LoRA核心算子存在的根本原因。2. LoRA核心算子内部逻辑拆解2.1 算子输入输出与数据布局一个LoRA融合算子的输入通常包括这几部分hidden_states形状[M, K]M是本层累计参与计算的token数lora_id_map长度M每个token对应的LoRA id没有LoRA的记为-1权重表指针指向显存中预分配的LoRA权重表主路径GEMM的bias可空输出就是更新后的hidden_states形状[M, N]。算子内部要做的事情可以拆成三步第一根据每个token的lora_id去权重表里索引对应适配器的A、B矩阵第二执行两步低秩乘法第三把结果累加到主路径输出上。权重表在显存中的排布是我建议第一个关注的点。因为每个适配器秩不同如果直接按[max_rank, K]这种统一大小分配显存浪费很大如果按实际秩紧凑排列算子内部取数就麻烦。实际工程里常用做法是按秩分组连续存放r16的A全部放在一段r32的放另一段B矩阵也对应排布。这样算子内部读权重时段起始地址可以由[id]直接查表得到不需要遍历。权重加载时一次性把layout整理好后续每个step算子执行时都是零额外整理开销这一点对性能影响非常直接。2.2 两步乘法的结构选择与中间结果处理数学上LoRA路径是固定的先算S x A^T再算y_lora S B^T。这个串行依赖关系决定了不能用单次大GEMM替代。如果合成C A^T B^T再x C就回到权重合并了失去了动态适配的意义。中间矩阵S的形状是[M, r]如果r64、M2048那么S大小是2048×64×2字节约256KB而部分昇腾型号的UB总量也就256KB左右。这就意味着S不能整体驻留tiling时必须把M切成合适的块让每个核处理一个token块时产生的中间结果能塞进本地存储。另外A矩阵在内存里通常按[r, K]存放但计算需要的是A^T即[K, r]视角。在设备侧做显式转置是很亏的正确做法是利用cube接口里的转置标志位让硬件在读矩阵时自动按转置方式取数避免额外拷贝。同理B矩阵也可以预先整理成便于第二步直接计算的布局把转置成本转移到加载阶段的一次性处理上。2.3 按LoRA分桶与段式GEMM执行流一个算子内同时处理多个适配器的关键是把M按lora_id聚簇。例如token序列里id分布是[0, 2, 0, 1]那就做一次轻量下标重排让相同id的token连续排列。这一步得到段表每段对应一个lora_id、一组连续token下标、对应秩大小。实际上host侧每次forward时都会拿到当前step的lora_id_map基于它生成段表非常轻量。段式GEMM的执行流如下对每个LoRA段取出X_seg [m_i, K]与A_i^T [K, r_i]做矩阵乘得到[m_i, r_i]再与B_i^T [r_i, N]乘得到[m_i, N]最后加到y_base对应行上。每个段虽然是一个独立的低秩计算但关键点在于所有这些段都在同一个算子实例内一次性下发到NPU由tiling后的多个AI Core并行消化而不是每次启动一个kernel。这里容易有一个误解既然有多个段为什么不直接用循环顺序执行实际上循环顺序执行只是最简单朴素的写法真正的性能来源是让段与段之间在时空上重叠。两个不同段的token块可以由不同的AI Core并行处理每个核针对自己分到的token块连续执行它的“低秩两步乘”而核与核之间不需要互相等待。这就是“一次调度段级并行”的收益所在。2.4 Tiling策略与AI Core并行分配具体到多核并行通常按M维度切大块权重矩阵在L2缓存中复用。以一个实际场景估算总token数1024假设可用AI Core数为40按token切成40份每个核处理约25个token。25个token对一个核来说块很小但每个核要覆盖的r维度是完整的A矩阵[r, K]一次载入UB。因为r最多64K常见4096A矩阵64×4096个元素fp16下是512KB这又超过了UB容量所以还需要在秩方向二次切块比如一次处理16行。实际tiling的优先级通常是先切M让尽可能多的核并行再切r把A矩阵装进UB。切r时每个r_block对应一次cube计算部分和就地累加最终在一次算子生命期内完成所有r_block的累积。伪代码逻辑如下# 算子内核逻辑 for seg in segments: for m_block in m_blocks: load_x_block(seg, m_block) # 从GM取X块到UB for r_block in r_blocks: load_a_block(seg, r_block) # 取A矩阵的秩分块 cube_matmul_accumulate(x_block, a_block) # 累加得到部分S load_b_block(seg) # 取B矩阵 cube_matmul(x_block_s, b_block) # S与B相乘得到LoRA输出 add_to_base_output(seg, m_block) # 与主路径输出累加这只是一个示意结构真实实现里还会把第二步的B矩阵也分块让S块在UB里流转避免一次驻留。工程编码上的关键是“当前核处理哪个token块、该块属于哪个段、有多少有效秩”这些信息要提前算好在kernel启动前作为参数传进去不要在device侧做动态分支判断否则流水线会充满气泡。3. 从调度到算子的完整流程串联3.1 LoRA管理器与适配器加载vLLM-Ascend服务的代码路径里LoRA适配器不是普通权重那样直接塞进模型参数。每个适配器有自己的safetensors文件加载后由LoRA管理器统一管理。管理器做的事情包括把A/B矩阵读出来做dtype转换通常是fp16按我们前面说的秩分组布局整理好在主存中的排布然后拷贝到NPU显存里预分配的权重buffer最后注册一个adapter id返回给调度层。这个“预分配权重buffer”值得多说一句。NPU显存管理不比GPU温和频繁申请释放适配器权重会产生碎片碎片严重时一个64MB的权重都放不进去。实际项目里我建议按秩分组维护固定大小的显存池适配器加载时从池里取块卸载时整块回收而不是依赖底层分配器频繁malloc。这个经验来自线上服务适配器数量从10个涨到80个之后显存碎片急剧恶化的真实场景。3.2 调度器如何绑定LoRA与token每个序列进入推理循环时调度器会给它的seq_data打上lora_request字段。进入模型执行阶段后模型根据当前活跃序列表生成一个重要的中间数据结构lora_id_map它的长度等于当前step的token总数。decode阶段一个序列一个token所以lora_id_map天然是token粒度的prefill阶段一个序列可能包含大量token这些token共享同一个adapter idmap里会有连续一大段相同id。有了这个map算子就可以直接从权重表里取数。这里隐含了一个关键关联同一序列的所有token必然绑定同一个LoRA适配器所以序列级调度和token级计算之间的桥就是lora_id_map。host侧每次forward构造lora_id_map本质上就是把调度器已经决定的序列级信息展开成token级信息为算子提供索引依据。3.3 算子在模型中的实际注入位置LoRA不是加在模型的任意层只加在线性层上。具体来说Transformer里qkv投影、o_proj、gate_up_proj、down_proj这几个线性层都有对应的LoRA分支。昇腾版本的计算图里每个带LoRA的线性层被替换成一个自定义线性层它接收完整输入内部同时完成“主GEMM LoRA融合算子”两件事。主GEMM使用标准MatMul算子LoRA融合算子我们前面已经拆过了两者在host侧可能并行下发也可以合并成一个更大的融合算子具体取决于版本实现。对上层模型代码而言这种设计改动最小。模型仍然调用一个统一的线性层接口内部是走普通GEMM还是走LoRA融合路径由当前batch是否存在活跃LoRA决定。如果某个step完全没有LoRA请求就跳过融合算子的调用退回标准主GEMM不浪费任何额外计算。3.4 一个decode step的完整生命周期用一个具体例子把全过程串起来。假设当前batch有3个请求分别对应LoRA id 0、2和-1-1表示无LoRA。M3个tokenK4096N4096适配器0的秩是16适配器2的秩是32。第一步调度器给出序列列表同时给出lora_id数组[0, 2, -1]。第二步模型forward到第一个attention层的qkv投影时host侧基于这个数组构造kernel参数权重表基址、lora_id_map、每个adapter的A/B偏移。第三步主GEMM执行X [3, 4096]与Wqkv [12288, 4096]相乘得到Y_base [3, 12288]。同一时间LoRA融合算子并行执行id-1的token被跳过id0的token取A0 [16, 4096]id2的token取A2 [32, 4096]各自完成两步小矩阵乘再与主输出对应行累加。第四步算子返回更新后的hidden_states交给后续注意力计算。这个流程里最需要在意的点是id为-1的token不能进LoRA计算分支但它的主路径结果必须原样通过。如果host侧构造段表时把-1的token单独挑出来不走融合分支那它在输出位置上要对齐好如果采用“整个batch一起进融合算子内部跳过-1”那主GEMM的结果和LoRA算子的结果相加时要小心索引对应关系。不同实现会有不同取舍但核心逻辑都一样每个token最终拿到的是它自己的主路径结果加它自己的LoRA结果不能串。4. 实操踩坑与排查实录4.1 常见问题速查表现象可能原因解决方向输出结果乱序或张冠李戴段表索引与token原始顺序没有映射好检查lora_id_map和输出行号之间的重排逻辑多卡推理结果不一致权重表复制到多卡后缺少同步点在每卡权重加载后加同步确保可见性推理过程中显存持续增长权重buffer碎片化或段表反复动态分配设置显存池、按秩分桶预分配加LoRA后性能下降特别多每个step都重新整理权重或段过多过碎合并连续同id段权重加载与计算分离输出精度有细微偏差中间矩阵S用fp16累加导致累积误差S的累加改为fp32偶发kernel崩溃r维度未对齐到cube要求的边界对A矩阵秩方向做pad到16的倍数第一类问题我最常见到。段表按LoRA聚簇之后输出行的顺序与原始token顺序不同如果累加时没有还原索引位置后面整个tensor就乱了。排查方法很朴素写一个单层单token的单测用固定权重跑一遍逐行对比预期输出。这种问题靠肉眼盯日志很难发现构造最小复现用例是最快的路。第二类问题在多卡环境尤其隐蔽。权重表从主卡广播到其他卡之后如果某个算子过早执行而其他卡还没收到数据结果就会不一致而且不是必现往往压测跑一阵才冒出来。我自己的排查习惯是每次加载新adapter后在第一次forward前加一次全局同步代价很小但能排除大量偶发不一致问题。4.2 性能调优的几个实测心得第一个心得中间结果S一定要用高精度累加。LoRA的两个矩阵乘都是瘦长形状单步cube计算的有效数据比例低累加次数多如果全用fp16累加误差会随着层数和step数量累积最终在长文本生成时表现为输出质量下降。S在UB里用fp32做累加寄存器只在写回GM时转成fp16既稳又损失不了多少性能。第二个心得不要迷信“所有LoRA权重常驻显存”这一条。适配器很多但活跃度差异大时可以考虑对权重buffer做两级缓存热权重留在显存冷权重放在主存需要时异步预取。实测中如果活跃适配器只占20%冷热分离可以减少约50%的显存占用。当然这是服务级设计和算子本身无关但整体吞吐提升非常明显。第三个心得尽量让一个step里的段数量少。比如一批请求里6个适配器的秩都是16完全可以在host侧把它们归入同一组只生成一个段A矩阵按[6 * 16, K]这个大块一次装载效果比分成6个独立段好得多。段数减少tiling的切分更规整AI Core的流水线也更满。4.3 边界细节与易错点空batch的情况一定要处理。线上服务总会有某些step没有任何LoRA请求此时lora_id_map全部是-1如果不加判断直接进入融合算子至少白白浪费一次kernel launch最坏情况还会因为段表为空触发异常路径。host侧一行if判断就能避免。还有并发更新适配器的问题。如果一个推理循环还在进行后台另一个线程加载并替换了某个adapter的权重那当前step的算子可能读到一半新权重一半旧权重产生不可复现的结果。解决方向是给版本号加锁或对权重buffer做双缓冲。这个坑在单线程推测场景不会碰到但一旦接入动态热加载早晚会遇到。最后是秩方向的维度对齐。昇腾cube对rank维有对齐要求实际工程里通常把A矩阵行数从实际秩pad到16的倍数。比如秩16不用pad秩20要pad到32pad出来的行用0填充。B矩阵的对应维度也要同步pad否则两步乘法的维度就对不上。pad确实带来一些无效计算但比起牺牲对齐带来的正确性风险这个代价完全值得。我个人在实际调试中的体会是LoRA核心算子真正难的地方反而不是那两个矩阵乘本身而是在“动态”二字上。适配器集合会变、每个batch的分布会变、秩大小会变一个设计良好的算子必须把这些变化放进数据布局和索引逻辑里让计算路径本身尽量保持规整。如果你要在一个新硬件平台上做类似实现建议先画清权重布局和索引映射代码实现倒可以往后放一放这个顺序能帮你省掉大量返工的时间。