大模型服务做到多租户微调之后，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到64，K是输入维度也就是hidden size，N是输出维度，比如4096。一个有LoRA的线性层，计算等价于：
y_base = x @ W^T
y_lora = (x @ A^T) @ B^T
y = y_base + y_lora
W是基座模型已经训练好的大矩阵。问题是这个式子怎么落实。
有一条路是权重合并：推理前把W整体变成W' = W + B@A的等价形式，然后这个层就退化成普通线性层，完全不感知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中，每个序列只有一个token，32个序列就是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]这种统一大小分配，显存浪费很大；如果按实际秩紧凑排列，算子内部取数就麻烦。实际工程里常用做法是按秩分组连续存放：r=16的A全部放在一段，r=32的放另一段，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]，如果r=64、M=2048，那么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最多64，K常见4096，A矩阵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 id，map里会有连续一大段相同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）。M=3个token，K=4096，N=4096，适配器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被跳过；id=0的token取A0 [16, 4096]，id=2的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到32，pad出来的行用0填充。B矩阵的对应维度也要同步pad，否则两步乘法的维度就对不上。pad确实带来一些无效计算，但比起牺牲对齐带来的正确性风险，这个代价完全值得。
我个人在实际调试中的体会是，LoRA核心算子真正难的地方反而不是那两个矩阵乘本身，而是在“动态”二字上。适配器集合会变、每个batch的分布会变、秩大小会变，一个设计良好的算子必须把这些变化放进数据布局和索引逻辑里，让计算路径本身尽量保持规整。如果你要在一个新硬件平台上做类似实现，建议先画清权重布局和索引映射，代码实现倒可以往后放一放，这个顺序能帮你省掉大量返工的时间。
