LoRA 服务化推理核心算子拆解：从 vLLM 调度到 Ascend NPU 实现

做 LLM 服务化的同学，这两年应该没少跟 LoRA 打交道。一个基座模型挂几十上百个 LoRA adapter，是低成本做垂直场景外挂最常见的玩法；而 vLLM-Ascend，就是把主流高性能推理框架 vLLM 迁移到 Ascend NPU 上跑通的一套工程实现。我在这条链路里啃得最久的一块，就是 LoRA 的核心算子逻辑和整体数据流程。这篇文章想把这部分彻底讲清楚：从带 LoRA 的请求进入调度器开始，到每个 Linear 层里增量是怎么算出来的，再到 prefill 和 decode 两种场景下算子为什么长成不一样的样子，最后落到 Ascend NPU 上实现这些小矩阵乘时容易踩的坑。适合正在做推理引擎适配、自定义算子开发，以及想弄懂多 LoRA 批处理原理的算法和工程同学参考。

1. LoRA 这条通路，为什么值得单独拆开看

1.1 不合并权重，也得把差值算出来

先回到一个基础问题。LoRA 的改动其实很小：把原始权重 W 冻结，在旁边加一个低秩分解 B·A，前向变成 Wx + (α/r)·B·A·x。微调完成后，推理侧最省事的做法是直接合并，算一个 W' = W + (α/r)·B·A 出来，之后整个模型跟没做过 LoRA 一样跑，额外开销为零。

合并方案在“一个模型只服务一个 adapter”的时候是绝对最优。可一旦进入服务化多租户场景，问题就来了。假设基座是一个 70 亿参数的通用大模型，上线时同时挂 100 个 LoRA adapter，每个 adapter 合并后都相当于一份新的全量权重。100 个 adapter 就是 100 份基座权重，显存直接爆炸；而且请求 A 用完想切到请求 B，还得重新加载一套合并后的权重，切换耗时完全无法接受。所以 vLLM 系的服务化方案选了另一条路：不合并，在线把 LoRA 增量算出来。这也意味着每一个注入 LoRA 的 Linear 层都要多算两个小矩阵乘和一个加法：y = Wx + (α/r)·(x·Aᵀ)·Bᵀ。共享一份基座权重，adapter 之间只差两个小矩阵，这才是多租户服务能跑起来的关键。

1.2 vLLM-Ascend 和 GPU 版的差异在哪

vLLM 最早在 GPU 侧把 LoRA 支持做到了生产可用，里面有一批针对多 adapter 场景设计的高性能核函数。vLLM-Ascend 的整体架构沿用这套设计，但不可能直接搬 GPU 核函数。Ascend NPU 有自己的调度模型和底层存储层次：矩阵乘法走专用矩阵计算单元，逐元素运算走向量单元，数据从全局显存到片上缓存的搬运需要显式安排。也就是说，上层流程（LoRA Manager、调度器、layer 结构）大体一致，真正要重写的是算子层：哪些逻辑 GPU 上有现成 kernel 而 NPU 上没有，哪些算子能复用底层库，哪些必须写成自定义算子补齐。这篇文章后面拆的，基本就是这个“必须补齐”的部分。理解这个背景，才能理解为什么 vLLM-Ascend 里 LoRA 的代码路径比 GPU 版多出一大截手工优化。

2. 整体流程：一个带 LoRA 的请求是怎么走完的

2.1 请求进来到权重就位：LoRA Manager 的职责

用户请求到达推理服务后，会带上自己的 adapter 标识。这个标识经过协议层解析后，交给全局唯一的 LoRA Manager。Manager 维护一个 adapter_id 到权重对象的注册表，第一次碰到某个 adapter 时，需要从权重文件里解析出每个目标模块对应的 A 矩阵和 B 矩阵。所谓目标模块，通常是注意力里的 q/k/v/o 投影和 MLP 里的 gate/up/down 投影，不同基座模型的名字略有差别，但基本都集中在这几类 Linear 上。

解析时还会读到两个关键参数：rank r 和缩放系数 alpha。工程实现对 alpha 的处理有个容易忽略的点：微调工具保存的权重文件通常不会把 α/r 这个缩放因子合并进 B 矩阵，而是作为独立配置存下来。推理引擎加载时必须自己在算子链路里乘上这个系数，或者加载时直接把它写进 B 矩阵。两种做法都对，但必须全链路保持一致。另外，vLLM-Ascend 会限制最大 rank（max_lora_rank），一方面控制权重缓存大小，另一方面也影响算子的 N 维度切分——rank 越大，中间结果越宽，片上 buffer 的压力也越大。

2.2 调度器视角：一个 batch 里可以有多个 adapter

调度器把不同请求凑成一个 batch 时，每个序列都携带着自己的 adapter id。同一批序列可以来自完全不同的 adapter，这给层计算提出了一个硬约束：同一个 Linear 算子，同一个 W 权重，但每个 token 行可能要用不同的 A/B 矩阵。注意这里不是简单的多权重副本问题，而是“在一条计算流里并行处理多种权重”的问题。

vLLM 的做法是给每个序列维护 LoRA 信息，每次 step 前汇总出本 batch 需要的 adapter 集合，并限制一个 batch 最多同时出现 max_loras 个 adapter。这个限制不是没道理：adapter 数量越多，SGMV 的分段就越多，权重搬运和索引整理的开销越大，而且会挤占为 LoRA 增量计算准备的临时空间。调度器还会做权重“就绪检查”，保证本 step 用到的 adapter 权重都已经加载完成、锁在显存里，不会在计算中途被淘汰掉。

2.3 层内数据流：prefill 和 decode 分两条路

进入某一个 Transformer 层后，对每一个带 LoRA 的 Linear，计算可以概括成四步：先算基座输出 Wx；再根据本 batch 各行的 adapter 信息，把对应行的激活收集起来；用 LoRA 权重算出增量 (α/r)·(x·Aᵀ)·Bᵀ；最后把增量加回基座输出。用伪代码表达大概是这样的：

def lora_linear_forward(x, base_w, lora_table, lora_ids, scale, is_prompt): # 1. 基座 GEMM，走正常的全量权重 out = gemm(x, base_w) # 2. 按 adapter 分组，取出本 batch 对应的 A/B # prefill: 走分组 GEMM，一个 adapter 对应一个小矩阵乘 # decode: 走 SGMV，一个 adapter 对应一组“单行向量” delta = compute_lora_delta(x, lora_table, lora_ids, scale, is_prompt) # 3. 累加回主输出 out += delta return out

prefill 阶段每个序列有一整段 token，属于同一 adapter 的 token 行可以拼成一个小 GEMM；decode 阶段每个序列只剩一个 token，整个 batch 变成若干条 1×d 的向量，必须换一种算子来承接。理解这两条路径的差异，是理解后面算子设计的钥匙。很多人看代码时觉得 LoRA 部分绕，绕就绕在同一个“增量”概念，在不同阶段有不同的算子形态。

3. 核心算子逻辑拆解

3.1 基座算子：还是一笔 GEMM 的事

基座部分没什么玄学。x·Wᵀ 在大模型里就是一个标准 GEMM：prefill 时 M 是当前 batch 的总 token 数，N 是输出维度，K 是输入维度；decode 时 M 通常等于 batch 大小（每个序列一个 token），实际退化成一个矩阵向量乘。NPU 上这类大 GEMM 明显要走矩阵计算单元，方块切分和流水调度都比较成熟，不需要额外折腾。LoRA 真正要解决的，是在这个大 GEMM 之外多出来的那两条细长的小矩阵乘。

这里有个值得强调的观点：LoRA 增量的计算量本来就不大，但它引入的“算子形态复杂度”很大。大 GEMM 是整齐的一块，切片、并行、数据复用都好做；而 LoRA 的两个小矩阵乘，M 和 N 都随 batch 内 adapter 分布剧烈变化，排列还不整齐。所以 LoRA 算子优化的核心矛盾，从来不是 FLOPs 不够小，而是怎么在数据排布和 kernel 调度上把这点计算“喂”给硬件。

3.2 两段式投影：为什么必须拆成 A 再 B

LoRA 增量的计算是 (x·Aᵀ)·Bᵀ。A 的维度是 d_in×r，B 的维度是 d_out×r，第一步把 d 维压缩到 r 维，第二步再涨回 d 维。以隐藏维度 4096、rank 16 为例，中间张量只有 16 维，两段小矩阵乘的 FLOPs 大约是 2×d×r×M，非常小。

有一个很多人问过的问题：能不能直接先算 B·A，合成一个 d_out×d_in 的稠密矩阵，再跟 x 乘一次？答案是：在“不合并、adapter 动态切换”的前提下不行。B·A 的计算量是 d_out×r×d_in，比两段式大了一个 d_out 倍，而且算出来的稠密矩阵还得临时存一份，等于把 LoRA 省下的显存又吃回去。只有在单 adapter 服务且允许 merge 的场景，合并才是对的。所以推理引擎里 LoRA 增量路径，天然就是两个小矩阵乘，这一步没有捷径。

scale 因子的位置也讲究。常见做法是把 α/r 乘在 A 的输出上，或者乘在 B 的输入上，两种方式数值上差别不大，但要注意和权重加载时的预缩放保持一致。我习惯在加载阶段就把 B 矩阵预乘好 scale，这样算子内部少一步乘法，kernel 更干净，排查 scale 问题也只用查一个地方。

3.3 decode 场景的 SGMV 算子：把所有小 GEMV 收进一个 kernel

decode 时每个序列贡献一行激活 x_i（1×d），每个序列有自己的 adapter。最笨的做法是给每个序列单独调两次矩阵向量乘：一次 (1×d) 乘 (d×r)，一次 (1×r) 乘 (r×d)。一个 batch 如果有 256 个序列，就是 512 次小 kernel 启动，光启动开销就比计算本身还大。

SGMV（分段聚合矩阵向量乘）的思路是：先按 adapter id 把序列分组，属于同一 adapter 的若干行激活在逻辑上拼成一个小的分块；kernel 内部按分块循环，把该 adapter 的 A/B 权重搬上片上，处理完这个分块的所有行，再切换到下一个 adapter。这样一次 kernel launch 就把整个 batch 的 LoRA 增量都算了。内部循环的伪代码类似：

for seg in segments: # 每个 seg 对应一个 adapter rows = gather(x, seg.ids) # (seg_len, d)，同一 adapter 的行 tmp = rows @ A[seg.lora_id].T # (seg_len, r) delta = (tmp * scale) @ B[seg.lora_id].T # (seg_len, d) scatter_add(out, seg.ids, delta)

看似只是“合并调用”，实际上内部有一个关键的数据复用逻辑：同一 adapter 的 A/B 权重只搬一次，在片上反复使用，而不是每个序列都搬一遍。这正是重算子设计的核心思想：把小而散的向量运算聚合起来，让权重产生复用。在 Ascend NPU 上实现这个 kernel 时，分组信息需要提前整理成紧凑的描述符：每个分块的起始行、行数、adapter 索引。索引整理本身也有开销，处理不好反而拖慢整体。我的做法是把索引整理放在每次 step 的 host 侧一次性完成，不要在 kernel 内部做复杂分支判断。

3.4 prefill 场景的分组 GEMM：从分段向量变成分段矩阵

prefill 时每个序列有若干 token，同一个 adapter 的数据自然形成一个 (tokens_i × d) 的矩阵，乘上 Aᵀ 得到 (tokens_i × r)，再乘 Bᵀ 回到 (tokens_i × d)。这里的关键问题是不同 adapter 的 tokens_i 不一样，甚至 rank 也可能不一样——如果允许多种 rank 的 adapter 混用。

实现上有三种选择。第一种是把所有 adapter 的权重按最大 rank padding 对齐，拼成一个大的 batch GEMM，多算的部分用 mask 忽略；第二种是真正的分组 GEMM，kernel 自适应地处理不同的 M 和 N；第三种是退回到 per-adapter 的小 GEMM 循环。padding 方案简单直接，但 rank 差异大时浪费明显，比如 rank 8 和 rank 64 混用，所有低 rank 的适配器都被迫多算 8 倍。分组方案最省计算，但 kernel 复杂度高，还要考虑不同分块在矩阵单元上的切分粒度。工程上通常先用 padding 或循环方案跑通正确性，再逐步迭代成分组实现。

这里还要提一个和 decode 路径的关系：prefill 和 decode 的 LoRA 算子虽然都叫“增量计算”，但一个是矩阵乘形态、一个是矩阵向量乘形态，没法自然共用一套实现。工程里一般靠 is_prompt 标志选择入口，权重本身是同一份，只是 kernel 的切分逻辑不一样。改造时要特别注意两条路径对权重排布的要求要保持一致，否则会出现 prefill 正常、decode 偶发错的诡异问题。

4. Ascend NPU 上的工程实现与优化要点

4.1 硬件单元怎么分配

Ascend NPU 上做矩阵运算有两类用户可感知的计算单元：处理大矩阵乘的矩阵计算单元，以及处理逐元素运算的向量单元。LoRA 两个小矩阵乘的常规选择是走矩阵单元，尤其当 batch 里同一 adapter 的 token 行数足够多时，矩阵单元的效率优势明显。但如果某个 adapter 在 decode 阶段只有一两行，矩阵单元的利用率会非常低，切块和流水调度的开销远大于计算收益。

在这个选择上我踩过坑：早期实现一律走矩阵单元，结果 decode 阶段大量分块只有 1 行，N 维度又只有 rank 那么大，整体时延比预期高不少。后来改成对分块粒度做判断：行数小于某个阈值的分块走向量单元的点积累加路径，行数够大的才走矩阵单元。实测下混合策略能明显改善小 batch 波动的场景，虽然阈值本身需要根据具体硬件和 hidden size 微调。

4.2 切分与数据排布：SGMV 的搬运问题

SGMV kernel 的性能瓶颈往往不在计算，而在数据搬运。A/B 权重在显存里的排布，直接决定了能不能把一整块权重连续搬进片上。按我的经验，每个 layer 的 LoRA 权重应该按“adapter 主键、rank 次键”连续存放，同时对 K/N 维度做硬件要求的对齐，一般是 16 或 32 的整数倍。权重不连续会让数据搬运变成大量小请求，性能立刻垮掉。这个约束在权重加载阶段就要处理，而不是等到 kernel 里再做 gather。

激活侧的搬运同样值得优化。分组之前需要知道每个 adapter 拿哪些行，这要求一个 index 数组。如果每次 step 都现排序，host 侧和设备侧就要多一层同步，而且这个同步恰好落在关键路径上。实测下来更稳的做法是利用“每序列 adapter id 固定”这个事实，维护一个稳定的分组缓存，只有新序列加入时才增量更新，而不是每步全量重建。这个优化看起来不起眼，但对 decode 这种每步只有一次小计算的场景，节省的 host 时间可能占整体时延的相当比例。

4.3 中间结果尽量留在片上

两段式投影天然有两个计算节点：A 的中间结果 (M×r)，然后是 B 的二次矩阵乘。如果实现里把中间结果写回全局显存再读出来，一个小矩阵乘就被拆成四次显存访问，时延翻倍都不止。正确做法是让 (M×r) 的中间结果留在片上 buffer，A 的算完直接接 B 的计算，形成流水。这一点在 decode 场景更重要，因为 M×r 往往非常小，来回搬运纯粹是浪费带宽。

具体到 NPU 实现，就是要把第一步矩阵乘的输出 buffer 地址直接作为第二步矩阵乘的输入地址，中间不经过全局内存。听起来是理所当然的事，但很多第一版实现因为图省事，直接复用了框架默认的算子调度，中间张量被隐式落到全局显存。排查这种问题有个窍门：对比中间张量在显存上的读写总量，如果明显大于理论值，基本可以断定落盘了。

4.4 什么时候该用“不合并”，什么时候该 merge

多 adapter 服务默认走 unmerged 加 SGMV，但这不是绝对的。实测下来，如果一个 batch 里实际只出现一两个 adapter，而且这两个 adapter 会持续服务比较长时间，那么把它们 merge 进基座权重再走普通 GEMM，往往更快。原因很简单：merge 之后就没有分段聚合的索引开销，主 GEMM 也能吃到最大的矩阵单元利用率。

更精细的做法是动态策略：批处理开始时统计 adapter 命中情况，命中数超过阈值就临时 merge，否则走 unmerged。这种混合策略在工程上会引入权重管理复杂度，比如 merge 后显存占用翻倍、切换时要回滚，所以一般作为后续优化项。正常的演进路径是：先把 unmerged 路径做稳做对，再用 profile 数据决定要不要做动态 merge。

5. 常见问题与排查实录

5.1 多个 adapter 结果串了怎么办

典型现象：batch 里同时有 adapter A 和 adapter B，输出里 A 序列混进了 B 的特征。绝大多数情况是索引错误，而不是数学错误。排查先从 host 侧的 adapter id 到权重表映射查起，确认权重缓存 key 唯一；再查 kernel 的分组描述符里行偏移和 adapter 索引是否对齐；最后检查 A/B 权重加载时有没有按层错位。用单个序列、单 adapter 的回归用例先跑通，再逐步增加混合 batch，是定位这类问题最快的方式。我见过一个案例，问题出在权重缓存的 key 只用了 adapter 名，没带 rank 维度，两个同名但不同 rank 的 adapter 互相覆盖，结果看起来像“串 adapter”，实际是权重被换掉了。

5.2 scale 因子不对导致推理结果偏差

α/r 的缩放如果漏乘或者乘了两次，正确性测试基本必挂。这类问题难查在于偏差可能很小，单条样本看不太出来，只有对比输出分布时才暴露。建议在代码里把 scale 计算收敛到一个常量，并且加一个单 adapter 的“merge 结果 vs unmerged 结果”一致性用例，用全量权重复算做基准，任何一步字节级偏差都能被及时捕获。这个用例应该集成到日常回归里，因为权重加载逻辑、算子实现、调度逻辑任何一处改动都可能影响它。

5.3 加了 LoRA 反而变慢

最常被忽略的三种原因：一是空 adapter 序列没有跳过，白白走了完整的 SGMV 流程；二是小 kernel 的启动开销没控制住，per-adapter 调度而不是整批调度；三是中间结果落显存，导致带宽浪费。遇到性能问题先做分层 profile：先看 kernel 总时长，再看 kernel 内部的时间分配，最后才看单算子细节。很多看似算子实现的问题，其实出在上层调度和索引整理。比如我优化过的一个案例，算子本身已经很快，但每次 step 的索引整理花了总时长的三成，最后靠增量更新分组缓存解决。

5.4 显存占用超标

LoRA 权重虽然比全量微调小很多，但几十上百个 adapter 叠起来也是可观的量。max_lora_rank、adapter 缓存淘汰策略、以及每步计算的临时 buffer 都要纳入统计。排查时建议单独记录“LoRA 权重缓存占用”和“LoRA 临时激活占用”两块，分别看增长曲线，避免和主模型的 KV cache 占用混在一起。有时候显存涨不是因为权重本身，而是淘汰策略触发不了：adapter 权重被引用计数卡住，看起来缓存满了，实际没有可淘汰对象。

下面用一个速查表总结最常遇到的问题：

我自己做这块最大的体会是：LoRA 的数学不难，难的是在多 adapter 批处理的约束下把数据流理清楚。Ascend NPU 上做算子适配时，先把整体流程里的每一处索引、缩放、搬运想明白，再动手写 kernel，比一上来就优化单个小 GEMM 靠谱得多。代码里看到“分段聚合”这类名字先别绕开，它背后其实就是“按 adapter 分组、让权重复用”八个字。后续如果再把动态 merge 策略和更细粒度的片上缓存调度做进去，LoRA 这条路径的收益空间还能再挖一截。

最后再分享一个小技巧：调试 LoRA 算子时，不要只看端到端输出对不对，先把某一个 layer 的基座输出、LoRA 增量、最终输出单独拉出来对比。这样能把问题精确锁定在“基座算错了”还是“增量算错了”，省掉一大半排查时间。