资讯动态

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

发布时间:2026/10/21 18:11:01 来源:尧图企业网站定制
做 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 会限制最大 rankmax_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 outprefill 阶段每个序列有一整段 token属于同一 adapter 的 token 行可以拼成一个小 GEMMdecode 阶段每个序列只剩一个 token整个 batch 变成若干条 1×d 的向量必须换一种算子来承接。理解这两条路径的差异是理解后面算子设计的钥匙。很多人看代码时觉得 LoRA 部分绕绕就绕在同一个“增量”概念在不同阶段有不同的算子形态。3. 核心算子逻辑拆解3.1 基座算子还是一笔 GEMM 的事基座部分没什么玄学。x·Wᵀ 在大模型里就是一个标准 GEMMprefill 时 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 再 BLoRA 增量的计算是 (x·Aᵀ)·Bᵀ。A 的维度是 d_in×rB 的维度是 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 收进一个 kerneldecode 时每个序列贡献一行激活 x_i1×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 忽略第二种是真正的分组 GEMMkernel 自适应地处理不同的 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) 的中间结果留在片上 bufferA 的算完直接接 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 权重被引用计数卡住看起来缓存满了实际没有可淘汰对象。下面用一个速查表总结最常遇到的问题现象可能原因排查入口输出混入其他 adapter 特征分组索引错位、权重缓存 key 冲突查 adapter id 映射与 kernel 段描述符数值偏差小但稳定存在scale 漏乘或重复乘对照 merge 基准做一致性用例加了 LoRA 时延反而升高小 kernel 启动频繁、空 adapter 未跳过分层 profile统计 kernel 次数显存持续增长权重缓存淘汰失效、临时 buffer 未复用分项统计权重与临时激活占用prefill 正常但 decode 偶发错两条路径权重排布不一致检查两条路径索引字段与对齐方式我自己做这块最大的体会是LoRA 的数学不难难的是在多 adapter 批处理的约束下把数据流理清楚。Ascend NPU 上做算子适配时先把整体流程里的每一处索引、缩放、搬运想明白再动手写 kernel比一上来就优化单个小 GEMM 靠谱得多。代码里看到“分段聚合”这类名字先别绕开它背后其实就是“按 adapter 分组、让权重复用”八个字。后续如果再把动态 merge 策略和更细粒度的片上缓存调度做进去LoRA 这条路径的收益空间还能再挖一截。最后再分享一个小技巧调试 LoRA 算子时不要只看端到端输出对不对先把某一个 layer 的基座输出、LoRA 增量、最终输出单独拉出来对比。这样能把问题精确锁定在“基座算错了”还是“增量算错了”省掉一大半排查时间。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑