TensorRT-LLM 源码完全拆解（二十一）：LoRA 服务——adapter 池、批内多适配器与前向替换的完整链路

从 CUDA Graph 继续：一个批次里的多个「模型」

第二十篇把动态批冻进了静态图：为了让一段前向可以被 torch.cuda.graph 捕获，权重指针、缓冲地址、批大小都得在捕获那一刻就固定下来，任何「每轮都会变」的量都必须事先搬进固定缓冲，靠原地改写而不是重新分配来更新。这一篇要拆的 LoRA 服务，正是那个约束最尖锐的场景——同一个 decode 批次里的不同请求，可能要用不同的权重。

LoRA（Low-Rank Adaptation）的出发点是一句很轻的假设：微调一个基座模型时，权重变化量 $\Delta W$ 是低秩的，可以写成两个小矩阵的乘积

原始权重 $W$ 不动，前向从 $h = Wx$ 变成 $h = Wx + BAx$。一个 adapter 需要落的盘只有 $(A, B)$ 两个小矩阵，rank 取 8、16 时参数量常常是基座的千分之几。对服务系统来说，这句话翻译过来是：同一份基座权重可以服务成百上千个业务适配器——每个客服场景、每个垂直领域、每个租户的私有微调，都只是一组小小的 $(A, B)$。

图源：LoRA 原始论文 Figure 2（arXiv:2106.09685）。论文用实测曲线说明了一件事：把 $\Delta W$ 投影到低秩子空间后，它的能量绝大部分落在前几个奇异方向上——这就是「低秩假设」不只是省参数、而且几乎不损精度的地方。

问题在于，朴素做法做不到「一个批次里混跑多个 adapter」。如果批次里一半请求用 adapter 0、一半用 adapter 1，逐请求拆成两次前向等于把批处理的好处全部丢掉；而如果每轮把权重换成当前请求的那一份，一个几十 MB 的 H2D 拷贝就要吃掉大半 decode 预算。TensorRT-LLM（下称 TRT-LLM）在 PyTorch 后端里的答案分成三步：

1. 权重常驻 + 固定槽位：所有要用的 adapter 权重一次性放进 GPU 的固定 buffer，槽位数量在启动时定死，前向期间不再搬运（AdapterSlotManager 与 C++ 侧 PeftCacheManager）；
2. 批内多适配器路由：把「哪些 token 用哪个 adapter」变成内核里的索引查表，而不是 Python 层的分支（CudaGraphLoraParams）；
3. 前向替换：用一个低秩算子在前向路径上「就地接上」修正项 $BAx$，并与基座计算重叠（LoraLayer）。

本文的四个问题：

1. LoraConfig 里哪些字段决定了「哪些矩阵可以被替换」，以及 LoraManager 从 HF 检查点到 GPU 权重的那条流水线做了什么变换；
2. AdapterSlotManager 怎样用固定槽位 + LRU 淘汰来管理 adapter 池，槽位与 C++ 缓存如何保持一致；
3. 一次 decode 批次如何从请求列表变成 slot_ids、sorted_ids 与按槽索引的权重指针表；
4. LoraLayer 的三条前向路径（eager、CUDA Graph、split-K 调优）分别做什么，以及 MoE、张量并行、FP8 三处需要额外对齐的口径。

本文沿用系列前文相同的固定提交 41913522c7f9f085c022421fb6183bd7592e076f。SGLang 侧的同题拆解可参考 SGLang 源码完全拆解（十三）：多 LoRA 服务——adapter 池、SGMV 内核与批内多适配器共存；两篇对照看，能看到同一套论文思想在两条不同运行时里的两种落地方式。

LoraConfig：哪些矩阵可以被替换

一切从配置开始。tensorrt_llm/_torch/peft/lora/config.py 里的 LoraConfig 决定了 LoRA 的能力边界，字段不多但每一个都对应后续的一处约束：

1
class LoraConfig(StrictBaseModel):
2
    lora_dir: List[str]              # adapter 检查点目录列表（HF 或 NeMo 格式）
3
    lora_ckpt_source: Literal["hf", "nemo"] = "hf"
4
    max_lora_rank: int = 64          # 所有 adapter 的最大秩，用于预分配 workspace
5
    lora_target_modules: List[str]   # 要在哪些 TRT-LLM 模块上挂 LoRA
6
    trtllm_modules_to_hf_modules: Dict[str, str] = {}
7
    max_loras: Optional[int] = None      # 设备侧同时驻留多少 adapter
8
    max_cpu_loras: Optional[int] = None  # 主机侧预取缓存多少 adapter
9
    swap_gate_up_proj_lora_b_weight: bool = True
10
    overlap_lora_and_base: bool = False
11
    cuda_graph_specialize_lora: bool = False

几个字段值得单独说。

max_lora_rank 默认 64，它的作用不是「限制 adapter 的秩」，而是预分配 workspace 的上界：内核要按最大的秩准备缓冲，所以启动时就得知道上界。这也是 LoRA 服务里最常被低估的一个工程约束——它事后无法放大，只能重启。

lora_target_modules 用的是 TRT-LLM 自己的模块名，与 HF 模块名之间靠 trtllm_modules_to_hf_modules 映射。默认映射表把两边的命名对齐了一遍：

missing_qkv_modules 是一个只读属性，背后是一个很实际的兼容处理：在当前设计里，q_lora_params、k_lora_params、v_lora_params 必须同时启用或同时禁用——因为内核一次算的就是 QKV 三路。但有些检查点（配置注释里点名了 BART）只带其中两个，于是 TRT-LLM 的策略是：检测到「三缺一」时，记下缺的那个模块，后面用零张量补齐。cross_attn_q/k/v 同理。这样做的好处是内核路径永远只看到「三路齐全」的形态，代价是一个恒等于零的 GEMM 白算一遍——用一点算力换掉一整类分支。

swap_gate_up_proj_lora_b_weight 默认打开，是因为不同来源的检查点对 gate_up_proj 的排列约定不同：HF 的 gate_up_proj 在权重里通常是 [gate; up] 顺序，而 TRT-LLM 的融合路径用的是相反顺序，所以要交换 lora_B 的上下两半。模型如 Phi-4-MM 排列一致，才需要把这个开关关掉。这类「同一份数据两种约定」的分歧在融合算子里非常常见，处理方式也很统一：在装载时一次性改写成自己内部的约定，前向路径上不再做任何判断。

从检查点到 GPU：LoraManager 的加载流水线

tensorrt_llm/_torch/peft/lora/manager.py 里的 LoraManager 负责把磁盘上的 adapter 变成内核对得上账的权重。load_from_hf() 的主循环对每个 adapter 走一遍下面这条流水线：

第一步：读检查点。 load_state_dict(get_model_path(model_dir, "adapter_model")) 会把 adapter_model.bin 或 adapter_model.safetensors 读成一张 state dict，键名是 HF 的 ...q_proj.lora_A.weight 这种形态。

第二步：gate/up 交换。 前面说的约定修正就发生在这里：

1
def preprocess_lora_weights(lora_model, model_config):
2
    if getattr(model_config, "swap_gate_up_proj_lora_b_weight", True):
3
        for key, value in lora_model.items():
4
            if "gate_up_proj.lora_B.weight" in key:
5
                half_split = value.shape[0] // 2
6
                first_half, second_half = value[:half_split], value[half_split:]
7
                lora_model[key] = torch.cat((second_half, first_half), dim=0)
8
    return lora_model

第三步：按 TP 交织。 如果 tp_size > 1，融合模块的 LoRA 权重必须先重排——HF 检查点里 attn_qkv 的 lora_B 是 cat([Wq, Wk, Wv]) 的顺序，而 TP 分片是按「每个 rank 拿一段 q、一段 k、一段 v」切的，所以要把顺序改成 cat([Wq_rank0, Wk_rank0, Wv_rank0, ..., Wq_rankN, Wk_rankN, Wv_rankN])：

1
assert weight.shape[rank_dim] == sum(part_sizes)
2
weight_parts = [weight.narrow(rank_dim, sum(part_sizes[:i]), part_sizes[i])
3
                for i in range(len(part_sizes))]
4
weight_parts_tp_weights = [
5
    torch.split(part, part.shape[rank_dim] // tp_size, dim=rank_dim)
6
    for part in weight_parts
7
]
8
return list(itertools.chain.from_iterable(zip(*weight_parts_tp_weights)))

其中 part_sizes 对 attn_qkv 是 [q_size, kv_size, kv_size]，q_size = head_size * num_heads * tp_size（乘回 tp_size 是因为 ModelConfig 里的 num_heads 已经被除过一次），对 mlp_gate_up 则是均分的两半。tp_size == 1 时整段直接跳过。交织完之后代码还会 torch.cat 拼回完整权重——注释写得很直白：C++ 实现期望拿到未切分的完整权重，切分这件事留给后续的指针表按 rank 计算偏移。

第四步：定标。 rank 与缩放系数来自 adapter 自己的配置：

1
rank = int(hf_config["r"])
2
rs_lora = bool(hf_config.get("use_rslora", False))
3
cache_dtypes = {...}
4
if rs_lora:
5
    scale = float(hf_config["lora_alpha"]) / np.sqrt(effective_rank)
6
else:
7
    scale = float(hf_config["lora_alpha"]) / effective_rank

这里有个容易踩的坑：use_rslora 打开时（Rank-Stabilized LoRA），分母是 $\sqrt{r}$ 而不是 $r$。两种口径下同一个 adapter 的有效缩放能差出数倍，如果服务端按普通 LoRA 的口径去读一个 rs-LoRA 检查点，模型不会报错，只是输出质量悄悄变差。

第五步：FP8 特判。 如果基座是量化模型、adapter 权重是原生 FP8，且当前设备支持原生 FP8 LoRA 内核（Hopper SM90 与 Blackwell B200 SM100 支持，SM103/SM107/SM120/SM121 不支持），权重就保持 FP8 不降精度。但 FP8 没有标量乘法，缩放只能在 BF16 里做完再夹回 FP8 值域：

1
fp8_max = torch.finfo(t_out.dtype).max
2
t_out = (t_out.to(torch.bfloat16) * scale).clamp(-fp8_max, fp8_max).to(t_out.dtype)

同时这里还有两条硬性校验：t_in 与 t_out 必须同 dtype；并且同一个 adapter 的所有模块必须落在同一个 PEFT 缓存 dtype 上——混用「原生 FP8 稠密模块」和「计算 dtype 的路由专家 MoE 模块」会直接抛错，而不是悄悄降级。

第六步：打包成 C++ 侧认识的结构。 每个 (layer_idx, lora_module) 的权重被拍平、拼接、pad 到统一的 max_weight_size，再 stack 成一整块：

1
cpp_lora_weights.append(torch.cat(weights_to_concat))   # [t_in, t_out(, t_mag)]
2
cpp_lora_config.append(torch.tensor(
3
    [self.LORA_MODULE_IDS[lora_module], layer_idx, effective_rank, is_dora],
4
    dtype=torch.int32))

LORA_MODULE_IDS 是一张 26 项的模块名到整数 id 的映射表（attn_qkv 为 0，moe_latent_fc2 为 25），这份整数 id 就是内核对「这一段权重属于哪个矩阵」的全部认知。is_dora 位也是在这里埋下的：DoRA 会多一个 magnitude 向量，跟着一起打包，但DoRA 与 MoE 的组合被显式拒绝（源码里那句注释「this might work, but I had no MoE DoRA models to test」），理由是 per-expert 布局下 magnitude 的语义没有定义。

装载的最后一步是把打包结果写进 _cpp_lora_weights / _cpp_lora_config，并把低秩表记进 _lora_uid_to_low_ranks[uid]。注意 _retain_device_tensors 这个开关：只有 cpp_peft_cache_manager is None（也就是老的 TRT 后端路径）时才需要 Python 侧常驻 GPU 张量并记录 data_ptr()——PyTorch 后端下，设备侧缓存由 C++ 的 PeftCacheManager 自己管，Python 只留 CPU 侧打包好的权重。

adapter 池：AdapterSlotManager 与两级缓存

权重装载完之后，真正决定「一个批次里能同时跑几个 adapter」的是槽位。tensorrt_llm/_torch/peft/lora/adapter_slot_manager.py 里只有 130 行左右，但它是整条 CUDA Graph 兼容链路的支点：

1
class AdapterSlotManager:
2
    def __init__(self, max_num_adapters: int):
3
        self.max_num_adapters = max_num_adapters
4
        self.slot2task: List[Optional[int]] = [None] * max_num_adapters
5
        self.task2slot: OrderedDict[int, int] = OrderedDict()  # LRU 顺序
6
        self.slots_changed = False

为什么是「固定槽位」而不是字典？ 因为 CUDA Graph。捕获一次图之后，内核读的权重指针必须落在捕获时确定的地址上；如果槽位是按需分配的，指针表就得每轮重写地址，图也就白捕获了。TRT-LLM 的做法是：槽位容器固定（slot2task 是一个定长列表），每个槽位的内容按需换入换出，指针表则始终按「槽号」索引而不是按「任务号」索引。这样一来，槽位归属变化时只改指针表的内容（图内允许的原地改写），槽位数量与地址布局都不变。

LRU 的记账方式很直观：task2slot 是一个 OrderedDict，队尾是最近使用的（MRU），队首是最久未用的（LRU）。get_or_assign_task() 返回一对值——分配到的槽号与被淘汰的任务号：

1
def get_or_assign_task(self, task_id: int) -> tuple[int, Optional[int]]:
2
    evicted_task = None
3
    if task_id in self.task2slot:
4
        self.task2slot.move_to_end(task_id)          # 命中：刷新 LRU
5
    else:
6
        self.slots_changed = True
7
        if len(self.task2slot) < self.max_num_adapters:
8
            free_slot = self.find_free_slot()        # slot2task.index(None)
9
            self.slot2task[free_slot] = task_id
10
            self.task2slot[task_id] = free_slot
11
        else:
12
            evicted_task, evicted_slot = self.task2slot.popitem(last=False)
13
            self.slot2task[evicted_slot] = task_id
14
            self.task2slot[task_id] = evicted_slot
15
    return self.task2slot[task_id], evicted_task

「基座模型」也占一个槽——但不在这个表里。 update_slots() 把没有 adapter 的请求统一映射成 max_num_adapters 这个哨兵值：

1
tasks = [request.lora_task_id for request in requests]
2
unique_tasks = {t for t in tasks if t is not None}
3
assert len(unique_tasks) <= self.max_num_adapters, (
4
    f"Batch with more unique LoRA adapters ({len(unique_tasks)}) than "
5
    f"max_num_adapters={self.max_num_adapters} is not supported")
6
for i, task in enumerate(tasks):
7
    tasks[i] = self.max_num_adapters if task is None else self.get_or_assign_task(task)[0]

也就是说，槽号的有效范围是 $[0, \text{max\_num\_adapters}]$，最后一个值是「不挂 LoRA」的专用槽。这个设计让前向路径不需要区分「这个请求有没有 adapter」——它只看见一个槽号，而基座槽对应的 rank 为 0，内核在 rank 为 0 时自然算出零修正项。一条判断从热路径挪到了参数表里。

两级缓存的一致性。 Python 侧的槽位表之外，还有 C++ 侧 PeftCacheManager 管理的设备/主机缓存。两边可能不同步：C++ 侧可能因为容量压力提前淘汰了某个 adapter，而 Python 侧的槽位表还以为它在。所以 update_slots() 的第一件事就是对齐：

1
def remove_evicted_slots_in_cpp(self, peft_cache_manager):
2
    for task_id in self.slot2task:
3
        if task_id is not None and not peft_cache_manager.is_task_cached_device(task_id):
4
            self.remove_task(task_id)

slots_changed 标志是整条链路的性能开关。 它只在「槽位归属真的变了」时置位——新任务进槽、任务出槽都会置位，而命中已有槽位（LRU 刷新）不会。下游 CudaGraphLoraManager 每轮检查这个标志，只有它为真时才重刷整张权重指针表。考虑一个稳态服务：批次里反复出现同样的几个 adapter，槽位归属长期不变，标志几乎一直是假，指针表在捕获期之后就再没被改写过。这正是「固定槽位」这个设计换来的收益。

一次 decode 批次的路由：从请求到 slot_ids

tensorrt_llm/_torch/pyexecutor/engine/lora.py 里的 LoraParamBuilder.build() 是每轮前向之前的入口。它的第一个决策是走哪条路：

1
use_cuda_graph_mode = cuda_graph_lora_manager is not None and maybe_graph
2
if use_cuda_graph_mode:
3
    if not use_lora_graph:
4
        cuda_graph_lora_manager.prepare_base_only_batch(peft_cache_manager)
5
        return None
6
    tokens_per_seq = 1
7
    if enable_spec_decode and runtime_draft_len > 0 and self._spec_config.is_linear_tree:
8
        tokens_per_seq = runtime_draft_len + 1
9
    return cuda_graph_lora_manager.prepare_cuda_graph_lora_params(
10
        scheduled_requests, attn_metadata, peft_cache_manager, tokens_per_seq)
11
else:
12
    if cuda_graph_lora_manager is not None:
13
        cuda_graph_lora_manager.adapter_slot_manager.remove_evicted_slots_in_cpp(
14
            peft_cache_manager)
15
    peft_table = peft_cache_manager.get_and_reset_batch_peft_table()
16
    lora_params = peft_table and self._build_eager(...)
17
    if lora_params:
18
        lora_params["data_type"] = peft_cache_manager.data_type
19
    return lora_params

两点值得注意。第一，prepare_base_only_batch() 存在的原因是：一个批次可能整体都不需要 LoRA（比如 cuda_graph_specialize_lora 打开时，框架会为「无 LoRA 的批次」单独捕获一张图）。这种批次虽然不跑 LoRA 内核，但仍要维护 PEFT 缓存状态——把 batch_peft_table 消费掉、淘汰失效槽位——否则下一轮真正需要 LoRA 时会读到过期的表。第二，走 CUDA Graph 路径时有一条硬性前提：

1
assert scheduled_requests.num_context_requests == 0, (
2
    "Context requests are not supported with LoRA CUDA Graph path. "
3
    f"Have {scheduled_requests.num_context_requests} context requests")

LoRA 的 CUDA Graph 只服务纯 decode 批次（或投机解码的验证批次）。prefill 的 token 数随请求变化，形状不固定，跑不进静态图，只能退回 eager 路径。

CUDA Graph 路径的核心在 CudaGraphLoraManager.prepare_cuda_graph_lora_params() 里，一共四步：

1
peft_table = peft_cache_manager.get_and_reset_batch_peft_table()
2
request_slot_ids = self.adapter_slot_manager.update_slots(request_list, peft_cache_manager)
3
self.cuda_graph_lora_params.update_sorted_indices(request_slot_ids, tokens_per_seq)
4
slot2task = self.adapter_slot_manager.get_slot_to_task_mapping()
5
if self.adapter_slot_manager.has_slots_changed():
6
    self.cuda_graph_lora_params.update_weight_pointers(peft_table, slot2task)
7
    self.adapter_slot_manager.reset_slots_changed()
8
self.cuda_graph_lora_params.update_slots_params(batch_slot_ids=request_slot_ids,
9
                                               tokens_per_seq=tokens_per_seq)

第一步：请求 → 槽号。 就是上一节 update_slots() 的产物，一个长度等于批内序列数的 slot_ids 列表。

第二步：按槽号重排。 update_sorted_indices() 把「槽号 → 该槽在本批次里占据的行」变成一张置换表。为什么要重排？因为分组 GEMM 的语义是「连续的一段行共用一个权重」，所以必须先把属于同一个 adapter 的行聚到一起：

1
sorted_indices = self.get_sorted_indices(slot_ids)
2
if tokens_per_seq > 1:
3
    token_sorted_ids = (sorted_indices.unsqueeze(1) * tokens_per_seq
4
                        + torch.arange(tokens_per_seq, dtype=sorted_indices.dtype)).flatten()
5
else:
6
    token_sorted_ids = sorted_indices
7
num_tokens = actual_batch_size * tokens_per_seq
8
self.sorted_ids = self.persistent_sorted_ids     # 捕获期固定地址
9
self.sorted_ids[:num_tokens].copy_(sorted_ids_host, non_blocking=True)

tokens_per_seq > 1 的分支服务于投机解码的线性树验证：每个序列在这一轮前向里贡献不止一个 token（草稿 token 加一个被验证的位置），所以「按槽重排」要先把序列级的置换展开成 token 级的置换。这里所有写入都落在 persistent_sorted_ids 与 pinned host 缓冲上，地址在捕获期就固定——图外的 Python 只做 copy_，不做分配。

同一个函数里还顺手生成了 token_to_slot，这是路由专家 MoE 路径专用的「token → 槽」表。它的 padding 区域在构造时清零一次之后就再不动了，注释解释得很清楚：C++ 算子里的地址算术会读到 num_tokens 之外的位置，那里必须永远是 0，否则会读到上一轮的陈旧槽号，把某个 token 的修正项算到错误的 adapter 上。

第三步：刷新权重指针表（只在必要时）。 update_weight_pointers() 遍历全部槽位，按「槽 → 任务 → PEFT 表」三级查出一组指针。它有三种分支：

1
for slot_id in range(self.max_lora_size):
2
    task_id = slot_to_task_mapping[slot_id]
3
    if task_id is None:                       # 空槽：rank 归零、指针清零
4
        self.slot_ranks_host[slot_id] = 0
5
        zero_out_weight_pointers(slot_id)
6
    elif peft_table is None or task_id not in peft_table:
7
        continue                              # 本轮没有为该任务准备新条目：沿用旧指针
8
    else:
9
        config = peft_table[task_id][0]       # 同槽内所有模块 rank 相同
10
        self.slot_ranks_host[slot_id] = config.adapter_size
11
        ...

第二条分支是这类「按需刷新」代码的典型形态：「暂时没有新信息」不等于「信息失效」。如果某个 adapter 上一轮已经装好、这一轮没有再出现在 peft_table 里（因为它没被重新准备），那么它现有的 rank 和指针依然有效，直接跳过——重刷一遍不仅浪费，还会因为拿不到新指针而把有效表项清掉。

指针表本身的布局也是为图捕获准备的：每个槽一行，每行三个 int64（A 指针、B 指针、DoRA magnitude 指针），打包成 [max_lora_size, 3] 的整型张量：

1
assert self.PTR_DTYPE == torch.int64, (
2
    f"PTR_DTYPE must hold a full device address, got {self.PTR_DTYPE}")

这条断言看着多余，其实是这类「把地址当数据存在张量里」的设计里最要紧的一道防线：任何比 int64 窄的类型都会静默截断设备地址，而截断后的指针不会报错，只会让 grouped GEMM 读到一片无关显存。同理 LD_DTYPE（leading dimension）也是 int64，SIZES_DTYPE（m,n,k）才是 int32。

MoE 路径还有一套按 (layer_idx, module_id) 展开的槽缓存，也就是每个 MoE 层、每个专家模块各有一份 [max_lora_size] 的 rank 表和 [max_lora_size, 3] 的指针表。之所以要按层展开，是因为一个 adapter 不必覆盖每一层：

1
# The ranks are the global per-slot ranks masked to the slots this
2
# layer/module actually has pointers for: an adapter need not cover every
3
# MoE layer, and the grouped-GEMM problem builder reads an unmasked
4
# rank > 0 on a null-pointer slot as an active rank-sized GEMM.

换句话说，如果某层没有该 adapter 的权重却带着一个非零 rank，问题构造器会认真地为它排一段 GEMM，然后在解引用空指针时崩掉——所以「没有指针」必须显式表达成「rank 为 0」。

第四步：按批次更新 GEMM 尺寸与偏移。 update_slots_params() 根据本轮的 batch_slot_ids 重算每个槽里有多少个 token，并把它变成前缀和形式的槽偏移：

1
slot_counts = self.get_slot_counts(batch_slot_ids, self.max_lora_size) * tokens_per_seq
2
self.slot_counts_host.copy_(slot_counts)
3
self.get_offset_from_counts(slot_counts, full=True, out=self.slot_offsets_full_host)
4
self.slot_counts.copy_(self.slot_counts_host, non_blocking=True)
5
self.slot_offsets_full.copy_(self.slot_offsets_full_host, non_blocking=True)

get_slot_counts() 本身只是一次 torch.bincount(..., minlength=max_lora_size) 再截断到 max_lora_size——统计每个槽在批次里出现几次，用的是计数排序的同一套工具；get_offset_from_counts() 则对计数做 cumsum_，得到每个槽在重排后张量里的起始行。* tokens_per_seq 再一次呼应投机解码：序列级的计数要放大成 token 级的计数，否则偏移会少算。这些写入同样只做 copy_，不重新分配——slot_counts、slot_offsets_full、slot_ranks 三个张量以及它们对应的 pinned host 副本都是构造期一次性分配的，这是「图内只改数据」原则的又一次应用。

slot_offsets_full 的长度是 max_lora_size + 1 而不是 max_lora_size（self.slot_offsets = self.slot_offsets_full[:-1]）：前缀和数组天然要比计数多一格，最后那一格是「所有槽的 token 总数」，也就是重排后张量的总行数。这类「多一格」的边界在写内核参数时最容易出错，所以在数据结构层面就把它固定下来，而不是让调用点各自去推断。

前向替换：LoraLayer 的三条路径

前面所有准备都是为了这一刻：模型前向走到某个线性层时，怎么把 LoRA 的修正项接上去。

1
@staticmethod
2
def forward_with_base(base_forward, lora_layers, x, lora_params, layer_idx):
3
    cuda_graph_params = lora_params.get('cuda_graph_params')
4
    has_lora_layer = bool(cuda_graph_params) and any(
5
        CudaGraphLoraParams.LoraLayerKey(layer_idx=layer_idx,
6
                                         module_ids=tuple(layer.module_types))
7
        in cuda_graph_params.layer_info for layer in lora_layers)
8

9
    lora_aux_stream = lora_params.get("lora_aux_stream")
10
    execute_in_parallel = (has_lora_layer and lora_aux_stream is not None
11
                           and do_multi_stream() and not torch.compiler.is_compiling())
12
    ...
13
    if execute_in_parallel:
14
        base_output, lora_outputs = maybe_execute_in_parallel(
15
            base_forward, lora_forward, lora_layers[0]._par_events[0],
16
            lora_layers[0]._par_events[1], lora_aux_stream, disable_on_compile=True)
17
    else:
18
        base_output, lora_outputs = base_forward(), lora_forward()
19

20
    for lora_output in lora_outputs:
21
        if not isinstance(lora_output, torch.Tensor):
22
            continue
23
        base_output = add_lora_result(base_output, lora_output)
24
    return base_output

1
def add_lora_result(output: torch.Tensor, lora_result: Optional[torch.Tensor]) -> torch.Tensor:
2
    if lora_result is not None:
3
        output.add_(lora_result.to(output.dtype))
4
    return output

1
input_dtype = x.dtype
2
data_type = lora_params.get("data_type")
3
if data_type is not None and input_dtype != data_type:
4
    if data_type == torch.float8_e4m3fn and input_dtype in (
5
            torch.float16, torch.bfloat16, torch.float32):
6
        fp8_max = torch.finfo(data_type).max
7
        x = x.clamp(min=-fp8_max, max=fp8_max).to(data_type)
8
    else:
9
        raise TypeError(f"LoRA input dtype {input_dtype} must match PEFT cache dtype {data_type}.")

1
layer_key = CudaGraphLoraParams.LoraLayerKey(layer_idx=layer_idx,
2
                                             module_ids=tuple(self.lora_module_types))
3
if (not cuda_graph_params or not cuda_graph_params.layer_info
4
        or layer_key not in cuda_graph_params.layer_info):
5
    return None
6
layer_params = cuda_graph_params.get_layer_params(layer_key)
7
if layer_params is None:
8
    return None  # 该层没有 LoRA 模块，直通

1
batch_size, hidden_size = x.shape[0], x.shape[-1]
2
num_layer_modules = len(self.lora_module_types)
3
total_output_size = sum(self.output_hidden_sizes)
4
if x.dtype == torch.float8_e4m3fn:
5
    min_kn = _validate_fp8_lora_cuda_graph_alignment(
6
        cuda_graph_params.slot_ranks_host, hidden_size,
7
        self.output_hidden_sizes, max_rank)
8
else:
9
    min_kn = min(hidden_size, 8, max_rank)   # rank 若小于 8 会有对齐问题
10

11
output_buffer = torch.empty(batch_size, total_output_size, dtype=x.dtype, device=x.device)
12
intermediate_buffer = torch.empty([num_layer_modules, batch_size, max_rank], ...)

1
torch.ops.trtllm.lora_group_gemm_param_fill_row_reorder_fusion(
2
    in_sizes, out_sizes, a_offset, d_offset, d_prime_offset, lda, ldd,
3
    ldb_prime, ldd_prime, splitk_offsets, reordered_input,      # 输出
4
    input.max_lora_size, input.max_rank, sum_out_sizes, input_hidden_size,
5
    bs, input.slot_counts, input.slot_ranks, input.slot_offsets,
6
    input.output_hidden_sizes, input.output_sizes_offset,
7
    input.b_ptrs, input.b_prime_ptrs, input.x, input.sorted_ids[:bs],
8
    input.intermediate_buffer, input.output_buffer, input.x.dtype)  # 输入

1
torch.ops.trtllm.lora_grouped_gemm_cuda_graph(
2
    grouped_gemm_params.in_sizes, grouped_gemm_params.out_sizes,
3
    grouped_gemm_params.a_offset, layer_params.d_b_ptrs,
4
    grouped_gemm_params.d_offset, layer_params.d_b_prime_ptrs,
5
    grouped_gemm_params.d_prime_offset,
6
    cuda_graph_params.get_problem_count(layer_key),
7
    grouped_gemm_params.lda, grouped_gemm_params.ldb, grouped_gemm_params.ldd,
8
    grouped_gemm_params.ldb_prime, grouped_gemm_params.ldd_prime,
9
    host_max_in_sizes, host_max_out_sizes,
10
    grouped_gemm_params.splitk_offsets, grouped_gemm_params.reordered_input.dtype,
11
    min_kn, split_k)

1
restored_output = torch.empty_like(output_buffer)
2
restored_output.index_copy_(0, cuda_graph_params.sorted_ids[:batch_size], output_buffer)

1
# PyTorch does not implement index_copy_ for FP8 tensors.
2
if output_buffer.dtype == torch.float8_e4m3fn:
3
    output_buffer = output_buffer.to(torch.bfloat16)

1
runner = self._split_k_runner     # 每层一个，惰性创建
2
runner.lora_params = runner.copy_lora_params(lora_params)
3
runner_inputs = [x]
4
_, split_k = AutoTuner.get().choose_one(
5
    "trtllm::lora_grouped_gemm_cuda_graph", [runner], runner.tuning_config, runner_inputs)
6
return runner(runner_inputs, tactic=split_k)

1
layer_params = lora_params.get(layer_idx, {})
2
for module_idx in self.lora_module_types:
3
    if module_idx in layer_params:
4
        active_lora_module_ids.append(module_idx)
5
        lora_ranks.append(layer_params[module_idx]['adapter_size'])
6
        lora_weight_pointers.append(layer_params[module_idx]['weight_pointers'])
7
if len(active_lora_module_ids) == 0:
8
    return None
9
lora_outputs = torch.ops.trtllm.lora_grouped_gemm(
10
    x, lora_params['host_request_types'][:num_seqs], lora_ranks,
11
    lora_weight_pointers, lora_params['prompt_lens_cpu'][:num_seqs],
12
    self.output_hidden_sizes, False, True, max([r.max() for r in lora_ranks]), 0, True)

MoE LoRA：per-expert 布局与 kernel slot 约定

路由专家（routed-expert）LoRA 是这条链路里最绕的一块。layer.py 里有一段「单一事实来源」的注释，把约定写死了：

1
MOE_LORA_MODULE_NAMES = ("moe_h_to_4h", "moe_4h_to_h", "moe_gate")
2
MOE_LORA_MODULE_TO_KERNEL_SLOT = {
3
    LoraModuleType.MOE_H_TO_4H: "fc1",
4
    LoraModuleType.MOE_GATE: "gated",
5
    LoraModuleType.MOE_4H_TO_H: "fc2",
6
}

它解释的是内核为什么需要「slot」而不是「模块名」：MoE 内核把 FC1 的两半打包在一起输出，其中 gate（接 SiLU）那一半与 up（线性）那一半是相邻的。内核约定：「fc1」这个 LoRA 槽作用在 gate 半上，「gated」这个槽作用在 up 半上，「fc2」作用在第二个投影上。而 TRT-LLM 规范的模块命名是 moe_h_to_4h = w1（gate/SiLU）、moe_gate = w3（up/linear）、moe_4h_to_h = w2（down）——三者的对应关系是交叉的（moe_h_to_4h → fc1，moe_gate → gated），不是字面上的顺序对齐。这张映射表被 MoE 工厂、校验器、布局工具与 CUTLASS MoE 后端共享，注释特意点名「keep this the single source of truth」——因为一旦某一处独立推导这套对应关系，就会得到「gate 上的 LoRA 被加到 up 半上」这种数值正确、语义全错的 bug。

布局上，每个专家各有自己的一对 $(A, B)$。 moe_layout.py 的注释把寻址方式讲清楚了：

1
# Each expert gets its own (A, B) low-rank pair; the underlying MoE kernel
2
# offsets both A and B pointers by `expert_index * dim * rank` to read the
3
# per-expert slice.

所以在装载时 MoE 模块的权重形状与稠密模块不同：t_in 是 [E, rank, in_dim]、t_out 是 [E, out_dim, rank]，rank_dim 从稠密的 0 变成 1（因为多了一个专家维），这也是 prepare_fused_lora_modules_for_tp 需要显式接收 rank_dim 的原因。而前向时，内核靠 token_to_slot 找到每个 token 属于哪个槽，再靠路由结果定位它命中的专家，把指针偏移到对应的专家切片上。

与稠密模块最大的不同是：MoE 的 GEMM 不由 LoraLayer 执行。 layer.py 里有一个 MoeLoraLayer 子类，它的 forward() 直接抛异常：

1
class MoeLoraLayer(LoraLayer):
2
    """Marker LoraLayer for routed-expert MoE modules.
3

4
    Routed-expert LoRA is *fused* into the MoE kernel
5
    (torch.ops.trtllm.fused_moe with LoRA kwargs); the actual GEMMs are not
6
    performed by LoraLayer.forward. This subclass exists so that
7
    CudaGraphLoraManager._initialize_from_model and the LoRA target-module
8
    validator can discover MoE LoRA layers via isinstance(module, LoraLayer)
9
    and inspect their lora_module_types / output_hidden_sizes ...
10
    """
11

12
    def forward(self, *args, **kwargs):
13
        raise NotImplementedError(
14
            "MoeLoraLayer is a discovery marker; routed-expert LoRA is applied "
15
            "inside torch.ops.trtllm.fused_moe via the MoE module. "
16
            "Do not call MoeLoraLayer.forward directly.")

这是一个「空壳子类」，存在的唯一目的是让发现逻辑工作：CudaGraphLoraManager._initialize_from_model() 与目标模块校验器都用 isinstance(module, LoraLayer) 来扫模型，如果 MoE 模块不继承 LoraLayer，它们就会在 layer_info 里漏掉整层 MoE 的 LoRA 参数——而漏掉的后果是图捕获时少了一层的关键字，运行到那层直接 KeyError。让 forward() 抛异常则是另一道防线：MoE 的修正项是在 fused_moe 内部与专家 GEMM 一起算的（同一遍读输入、同一遍写输出），如果谁误以为它和稠密模块一样、在外面单独调一次 LoraLayer.forward，等于把修正项白算两遍，第二遍还会因为缺少 MoE 上下文而算错。用「显式报错」替代「静默多算一遍」，在融合算子边界上是很有必要的。

MoE LoRA 有两条明确的限制，都在装载期就拦下来：其一，原生 FP8 adapter 权重的路由专家 LoRA 不支持（混合 dtype 会触发前面那条 cache_dtypes 校验）；其二，DoRA 与 MoE 不组合。

三处口径必须对齐：TP、FP8 与 gate/up

把前面散落的细节收一收，TRT-LLM 的 LoRA 实现里有三处「同一份数据、两套约定」的分歧点，处理方式各不相同但思路一致。

第一处：张量并行的分片口径。 前面讲过，融合模块的 LoRA 权重需要按 TP 交织，且交织完之后还要拼回完整权重交给 C++，由指针表按 rank 计算偏移。这与 S-LoRA 的「把 adapter 也按 TP 维度切开」是两种不同的工程路线：S-LoRA 的目标是让每个 rank 只存自己那一份、从而支持极多的 adapter；TRT-LLM 这里保留完整权重、由指针表做偏移，换来的好处是同一份打包权重在所有 rank 上是同一份数据，不需要为分片再做一次落盘一致性检查。

图源：S-LoRA 论文 Figure 3（arXiv:2311.03285）。两条路线对照看很有意思：S-LoRA 选择「切得更多、省显存」，TRT-LLM 选择「不切、靠偏移」，前者适合 adapter 数量极多的多租户场景，后者更适合与既有 fused kernel、CUDA Graph 路径无缝衔接。

第二处：FP8 的值域与缩放。 前面讲过三段：装载时用 BF16 完成缩放再夹回 FP8；前向入口把 BF16 激活夹到 FP8 值域；输出还原时因为 index_copy_ 不支持 FP8 而先转 BF16。三段的共同前提是「基座是 FP8 量化模型、adapter 也是原生 FP8」——如果设备不支持原生 FP8 LoRA 内核（比如 SM120/SM121），adapter 会被转成计算 dtype，链路自动退回普通的 BF16 路径，行为与未量化模型一致。

第三处：gate/up 的排列。 装载期的上下半交换（swap_gate_up_proj_lora_b_weight）与 MoE 的 kernel slot 交叉映射是同一类问题的两个实例：融合算子的内部排列是私有的，外部来源的约定必须在入口处完成一次性改写。 这两种处理都把「翻译」挤在装载阶段，前向路径上没有任何分支——这是能让内核保持扁平的通用手法。

工程取舍与陷阱

max_lora_rank 事后无法放大。 它决定的是 workspace 上界，一旦启动就冻结。上线前把 rank 估小、后面想加一个大秩 adapter，只能改配置重启。稳妥的做法是取「已知最大秩 × 1.5」再向上取整到 16 的倍数。

max_loras 与 max_cpu_loras 是两个容量。 前者是设备侧同时驻留的 adapter 数（也就是槽位数量），后者是主机侧预取缓存的容量。两个数字都偏小时，最常见的行为不是报错而是抖动：adapter 被频繁换入换出，每次都要过一次 H2D，性能掉得毫无规律。压测时如果看到 P99 抖动而 P50 正常，先怀疑这里。

adapter 没有出现在 peft_table 里不代表它失效。 上一节的第二条分支说明过这一点：只有「新的准备条目」才会触发指针刷新，「没有新条目」会沿用旧指针。排查「换了 adapter 但输出没变」这类问题时，要看的不是 peft_table 的内容，而是 slots_changed 是否被正确置位——如果某一轮槽位归属变化却被判成没变，指针表就会停留在旧槽位上。

预填充批次不走 CUDA Graph 路径。 前面那条 assert num_context_requests == 0 就是证据。这意味着长 prompt 的 prefill 阶段 LoRA 走的是 eager 算子；如果业务里 prefill 占比高，调优重点应该在 eager 路径的 lora_grouped_gemm 上，而不是 CUDA Graph 的 split-K。

稀疏的 per-expert 权重必须显式表达成 rank 0。 MoE 那条注释值得再念一遍：rank 大于 0 但指针为空，会被问题构造器当成「有一个活跃的 rank-sized GEMM」，然后在解引用时崩掉。做自定义 adapter 生成工具时，凡是没有权重的层/模块，都要显式写 0，不能留默认值。

小结

这一篇把 TRT-LLM PyTorch 后端的 LoRA 服务链路拆完了，核心结论有五条：

1. LoRA 服务的技术核心不是低秩数学，而是「把动态身份塞进静态数据结构」。 低秩分解只是前提；真正的工程量集中在 AdapterSlotManager 的固定槽位、CudaGraphLoraParams 的定长张量、persistent_sorted_ids 的地址稳定这几件事上。原因很简单：第二十篇所拆的 CUDA Graph 路径要求地址与形状在捕获期固定，而 LoRA 天生是「每轮都可能不同」的东西，两者要对齐，只能把「不同」编码成固定容器里的内容变化。

2. 槽位号是整条链路的通用货币。 从 update_slots() 产出槽号，到 sorted_ids 按槽重排，到 slot_counts/slot_offsets/slot_ranks 按槽索引，到指针表按槽排列——前向路径上的每一层都只认槽号，不认任务号、不认 adapter 名。基座模型也用槽号表达（max_num_adapters 哨兵 + rank 0），于是「有没有 LoRA」这件事在热路径上根本不需要判断。

3. 权重装载是一次性的翻译层。 gate/up 排列、TP 交织、rs-LoRA 缩放口径、FP8 值域处理，全部挤在 load_from_hf() 里完成；前向路径上没有任何「这一步该怎么解释」的分支。这也是 _cpp_lora_config 里那个 [module_id, layer_idx, rank, is_dora] 四元组能承载全部语义的原因——内核只需要整数。

4. 两级缓存的一致性靠一次显式对齐保证。 Python 侧的槽位表与 C++ 侧的 PeftCacheManager 之间可能出现「Python 以为还在、C++ 已经淘汰」的分歧，靠每轮开头的 remove_evicted_slots_in_cpp() 抹平。缓存不一致在这类系统里的默认表现不是报错，而是数值悄悄错位——这是把「以谁为准」写成一个独立步骤、而不是靠多处防御性检查的原因。

5. 代价落在三个地方： 槽位数量（max_loras）与秩上界（max_lora_rank）都在启动时定死，事后只能重启；prefill 批次退回 eager 路径，长 prompt 场景享受不到静态图；MoE + FP8、MoE + DoRA 两类组合被显式拒绝，adapter 侧要提前避开。

从系列视角看，这一篇补上了「请求身份」这条维度——前面二十篇反复出现的一个主题是「一个批次里有很多不同的请求」，而这一篇增加的是一层「这些请求可能属于不同的模型版本」。下一篇会继续停在服务面：不管选题落在哪一块，主线仍然是同一个问题——怎么把一个批次里并存的多种身份，拆成可以被静态捕获的数据结构。

参考资料

1. TensorRT-LLM GitHub 仓库
2. TensorRT-LLM LoRA 官方文档（docs/source/features/lora.md，固定提交）
3. TensorRT-LLM LoRA 管理器源码（peft/lora/manager.py，固定提交）
4. TensorRT-LLM LoRA 前向层源码（peft/lora/layer.py，固定提交）
5. TensorRT-LLM LoRA 配置源码（peft/lora/config.py，固定提交）
6. TensorRT-LLM adapter 槽位管理源码（peft/lora/adapter_slot_manager.py，固定提交）
7. TensorRT-LLM CUDA Graph LoRA 参数源码（peft/lora/cuda_graph_lora_params.py，固定提交）
8. TensorRT-LLM CUDA Graph LoRA 管理源码（peft/lora/cuda_graph_lora_manager.py，固定提交）
9. TensorRT-LLM adaptor 装载源码（peft/lora/loaders.py，固定提交）
10. TensorRT-LLM MoE LoRA 布局工具源码（peft/lora/moe_layout.py，固定提交）
11. TensorRT-LLM 每批次 LoRA 参数构造源码（pyexecutor/engine/lora.py，固定提交）
12. LoRA 原始论文：Low-Rank Adaptation of Large Language Models（arXiv:2106.09685）
13. S-LoRA 论文：Serving Thousands of Concurrent LoRA Adapters（arXiv:2311.03285）
14. Punica 论文：Efficient Multi-Tenant LoRA Serving（arXiv:2310.18547）
15. TensorRT-LLM 源码完全拆解（二十）：CUDA Graph 与编译路径
16. TensorRT-LLM 源码完全拆解（十三）：MoE 层前向执行——路由、Fused Kernel 与专家并行的三段拼装
17. TensorRT-LLM 源码完全拆解（十二）：注意力层的张量并行——head 怎么切、MLA 为什么不能照切
18. SGLang 源码完全拆解（十三）：多 LoRA 服务——adapter 池、SGMV 内核与批内多适配器共存