SGLang 源码完全拆解（二十三）：权重装载管线——文件名如何变成 GPU 上切好的参数

4.1 前缀与模型级映射

_get_weights_iterator 在返回前做了一次前缀拼接：

1
return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)

Source.prefix 用于组合模型（例如视觉塔、草稿模型作为第二个权重源），通过给不同来源加前缀，避免两个模型的同名参数互相覆盖。_get_all_weights 先产出主模型，再遍历 model.secondary_weights 产出额外的源。

更复杂的是量化层的命名映射。某些模型（例如 DeepSeek 系列的 FP4 checkpoint）在 checkpoint 与 SGLang 的参数名之间有一套固定的改写规则，_get_quantization_config 会把它交给量化配置：

1
get_hf_to_sglang_mapper = getattr(model_class, "get_hf_to_sglang_mapper", None)
2
if get_hf_to_sglang_mapper is not None:
3
    hf_to_sglang_mapper = get_hf_to_sglang_mapper(model_config.hf_config)
4
else:
5
    hf_to_sglang_mapper = getattr(model_class, "hf_to_sglang_mapper", None)
6
if hf_to_sglang_mapper is not None and quant_config is not None:
7
    quant_config.apply_weight_name_mapper(hf_to_sglang_mapper)

为什么是量化配置来持有这个映射，而不是通用装载器？因为这类改名的动因几乎都来自量化格式：checkpoint 把 weight_scale_inv 写在某个位置，量化方法的 create_weights 却把 scale 注册在另一个名字下；映射规则只有量化方法自己最清楚。把它放在 quant_config 上，create_weights 与加载时的名字查找才能用同一张表。

4.2 融合层的「堆叠映射」

SGLang 的 Linear 层会把多个逻辑投影融合成一个参数，以减少 kernel 启动次数。最典型的是 MLP 的 gate_up_proj：checkpoint 里是两个独立的 gate_proj 与 up_proj，模型里是一个 gate_up_proj。DeepSeek 的加载器用一张三列表来表达这件事：

1
stacked_params_mapping = [
2
    # (param_name, shard_name, shard_id)
3
    ("gate_up_proj", "gate_proj", 0),
4
    ("gate_up_proj", "up_proj", 1),
5
]

加载循环对每个 checkpoint 名字依次尝试这张表。命中后先把名字改写成参数名，再调用参数的 weight_loader，并把 shard_id 传进去：

1
for param_name, weight_name, shard_id in stacked_params_mapping:
2
    if weight_name not in name:
3
        continue
4
    if ("mlp.experts." in name) and name not in params_dict:
5
        continue          # 专家由下一张表处理
6
    name = name.replace(weight_name, param_name)
7
    if name.endswith(".bias") and name not in params_dict:
8
        continue          # GPTQ 模型多出来的 bias
9
    param = params_dict[name]
10
    weight_loader = param.weight_loader
11
    maybe_executor_submit(..., func=weight_loader,
12
                          func_args=(param, loaded_weight, shard_id))
13
    break
14
else:
15
    ...  # 进入 expert_params_mapping

这里的 for ... else 结构很关键：只有当堆叠表没有任何一条命中时才进入专家表。如果顺序反过来，mlp.experts.0.gate_proj 会先被堆叠表命中，名字变成 mlp.experts.0.gate_up_proj，而专家表期待的是 experts.w13_weight，于是永远匹配不上。源码注释专门写了这一点，说明这是一个实际踩过的坑。

shard_id 不是「第几个张量」，而是融合参数内部的逻辑槽位。gate_up_proj 的权重沿输出维拼接，shard_id=0 表示写到前半段（gate），shard_id=1 表示写到后半段（up）。写入由参数自己的 weight_loader 完成，它同时处理了张量并行切分——这是一个「一次调用完成两件事」的合并：先按 shard_id 选槽，再按 TP rank 在自己的槽内取出局部切片。

4.3 专家映射：三列变四列

专家权重的映射更复杂，因为多了「哪个专家」这一维。FusedMoE.make_expert_params_mapping 返回的是四元组：

1
return [
2
    # (param_name, weight_name, expert_id, shard_id)
3
    (
4
        ("experts.w13_" if weight_name in [ckpt_gate_proj_name, ckpt_up_proj_name]
5
         else "experts.w2_"),
6
        f"experts.{expert_id}.{weight_name}.",
7
        expert_id,
8
        shard_id,
9
    )
10
    for expert_id in range(num_experts)
11
    for shard_id, weight_name in [("w1", ckpt_gate_proj_name),
12
                                  ("w2", ckpt_down_proj_name),
13
                                  ("w3", ckpt_up_proj_name)]
14
]

把这段展开读：对每个专家 e，checkpoint 里的 experts.e.gate_proj 映射到 experts.w13_weight 的 w1 槽，experts.e.up_proj 映射到同一参数的 w3 槽，experts.e.down_proj 映射到 experts.w2_weight 的 w2 槽。也就是说，FusedMoE 用两个大参数装下所有专家：w13_weight 形状是 [num_local_experts, 2 * intermediate, hidden]，w2_weight 是 [num_local_experts, hidden, intermediate]。这样一次分组 GEMM 就能处理整个批的专家计算，而不是每个专家一次 GEMM。

代价是加载时要多做三件事：

1. 名字匹配变成 O(num_experts) 的表扫描。num_experts 可能是 256，映射表有 768 条，每条都做一次子串匹配。这就是为什么循环里先查 stacked_params_mapping（只有两条）再查专家表——把便宜的表放前面。
2. 本地专家过滤。checkpoint 包含全局所有专家，本 rank 只拥有其中一部分。FusedMoE.weight_loader 在写之前调用 _map_global_expert_id_to_local_expert_id，返回 -1 就直接丢弃：

1
if not getattr(param, "_sglang_require_global_experts", False):
2
    expert_id = self._map_global_expert_id_to_local_expert_id(expert_id)
3
    if expert_id == -1:
4
        return

这一步决定了一件重要的事：被丢弃的张量仍然会被迭代器读出来。SGLang 没法在 safetensors 迭代器层面跳过「本 rank 不需要的专家」，因为那需要先解析名字再决定是否解包，而迭代器只负责产出。这让 EP 场景下的读盘量偏大：每个 rank 都要把全部专家读一遍，再丢掉大部分。generate_weight_name_filter 这类按 EP 布点生成的名字过滤器，以及 expert_pack 格式，就是为压缩这部分浪费而存在的。

3. 专家 id 可能被物理重映射。EPLB 在线重排后，同一个逻辑专家可能在多个物理槽位上有副本。weight_loader 里 logical_to_all_physical 会返回一组物理 id，循环把同一份权重写进所有副本。这保证了扩容或换点之后，副本与主本的数据一致，而不需要在重排时单独搬运权重。

4.4 加载之前先「剔骨」：哪些张量根本不用进显存

名字映射之前还有一组丢弃规则，它们不改变任何数值，但直接决定启动时的峰值内存：

1
if "rotary_emb.inv_freq" in name:
2
    continue
3
if ".indexer." in name and (name.rsplit(".indexer.", 1)[0] not in indexer_present_prefixes):
4
    continue

• rotary_emb.inv_freq 是旧版 HF checkpoint 留下的缓存张量，SGLang 的 RoPE 在构造时按公式算出频率表，不需要从文件读。
• 没有索引器的层，checkpoint 里却带着索引器权重（常见于把稀疏注意力的 checkpoint 复用到其他层数的模型），装载时直接丢弃，否则会报「参数找不到」。

还有两条与并行相关：

1
layer_id = get_layer_id(name)
2
if (layer_id is not None and hasattr(self.model, "start_layer")
3
    and (layer_id < self.model.start_layer or layer_id >= self.model.end_layer)):
4
    continue

流水线并行下，每个 PP stage 只拥有自己那几层。get_layer_id 从名字里正则提取 layers.(\d+).，落在区间外的层直接跳过——注意这是在读盘之后才发生的，因为名字来自迭代器。如果要进一步省读盘，就得像 EP 那样在迭代器层做过滤。

1
if ".norm." in name and not self.pp_group.is_last_rank:
2
    continue

最终 norm 只存在于最后一个 PP stage，其他 stage 连读都不该读。这几条规则共同构成一个原则：加载器的过滤越靠前，省下的 I/O 和内存越多；但过滤需要的元信息（层号、专家号、是否量化）也越靠前。SGLang 的现状是「名字级过滤在模型层，文件级过滤在装载器层」，中间那一档（按名字跳过读盘）只在专门的格式里做。

5.1 默认 loader 与它的两个前提

先把三种并行落位方式放在一起，后面每种都只在这张形状上做局部调整：

三种方式可以叠加：MoE 的专家权重既是 EP 的（本地专家子集），又是 TP 的（每个专家的中间维按 tp_rank 切）。叠加顺序在源码里是固定的——EP 先选专家，TP 再在选中的专家里切中间维，这也是下面 _load_w13 里同时出现 expert_data 与 moe_tp_rank 的原因。

最基础的 default_weight_loader 只有十几行，但它隐含了两个前提：

1
def default_weight_loader(param, loaded_weight):
2
    if param.numel() == 1 and loaded_weight.numel() == 1:
3
        param.data.fill_(loaded_weight.item())
4
    else:
5
        assert param.size() == loaded_weight.size(), (
6
            f"Attempted to load weight ({loaded_weight.size()}) "
7
            f"into parameter ({param.size()})"
8
        )
9
        param.data.copy_(loaded_weight)

第一个前提是 形状必须完全一致。当 TP size 大于 1 或 EP 打开时，参数形状是「本 rank 的那一份」，而 checkpoint 是全局形状，默认 loader 必然触发断言。这就是为什么带切分的参数必须换成 sharded_weight_loader。

第二个前提是 单元素标量按广播处理。k_scale、v_scale、某些 per-tensor 的量化 scale 在很多 checkpoint 里被存为零维标量，而参数可能是一维的 [1]。param.size() == loaded_weight.size() 对 () 与 (1,) 不成立，广播分支让这种常见的格式差异不至于报错。

张量并行的切分规则来自 Megatron-LM 的那套定义，下图是它的经典示意：MLP 的第一层按输出维切（列并行），第二层按输入维切（行并行），前向只在行并行层之后做一次 all-reduce：

SGLang 的 sharded_weight_loader(shard_axis) 里的 shard_axis 就是这张图里的「切哪一维」：列并行切输出维（权重张量 [out, in] 的第 0 维），行并行切输入维（第 1 维）。checkpoint 里的张量是全局形状，每个 rank 只保留自己那一段，截取起点由 tp_rank * shard_size 决定。

5.2 sharded_weight_loader：在构造时冻结 rank

SGLang 的切分 loader 是闭包工厂，rank 在一次构造中确定：

1
def sharded_weight_loader(shard_axis, tp_rank_getter=None, *, parallel_group=None):
2
    if parallel_group is not None:
3
        tp_rank, tp_size = resolve_linear_parallel_group(parallel_group)
4
    else:
5
        tp_rank = (tp_rank_getter() if tp_rank_getter is not None
6
                   else get_parallel().attn_tp_rank)
7
        tp_size = get_parallel().tp_size
8

9
    def loader(param, loaded_weight):
10
        shard_size = param.data.shape[shard_axis]
11
        start_idx = tp_rank * shard_size
12
        ...
13
        loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
14
        return default_weight_loader(param, loaded_weight)
15

16
    return loader

三个值得展开的点：

起点与宽度怎么来。 shard_size 取的是参数自己的维度长度而不是「全局长度除以 tp_size」。这在均匀切分下等价，但在不均匀场景下（例如 TP size 与 head 数不整除，或某些维度的输出被 padding）更稳健：参数形状是建模阶段按本 rank 真实容量申请出来的，用它反推切片是最直接的做法。start_idx = tp_rank * shard_size 也因此天然对齐。

为什么 rank 要冻结在闭包里。 注释写的是「a weight loader with placement frozen in its construction scope」。冻结的原因与上下文管理器有关：SGLang 有草稿模型（EAGLE/NextN）构造时临时切换并行拓扑的场景，如果 loader 在调用时再去读 get_parallel()，可能读到草稿模型的拓扑而不是目标模型的。构造时定死，等于把「这个参数的切法属于谁」记在 loader 上。SGLang 源码完全拆解（十七）：DP Attention 与专家并行里 DP attention 下 attention 层用的是 attn_tp_rank 而不是全局 tp_rank，默认的 tp_rank_getter 正好对应这一语义。

CPU 上的不均匀 padding。 CPU 路径多了一段分支：

1
if (is_cpu()
2
    and (loaded_weight.size(0) % tp_size != 0
3
         or loaded_weight.size(0) < tp_size * shard_size)
4
    and loaded_weight.dim() == 1):
5
    param_data = param.data
6
    param_data, loaded_weight = narrow_padded_param_and_loaded_weight(
7
        param_data, loaded_weight, 0, start_idx, shard_axis, shard_size)
8
    return default_weight_loader(param_data, loaded_weight)

loaded_weight.size(0) < tp_size * shard_size 对应「参数被 pad 得比 checkpoint 更宽」的情况（CPU 上的 AMX 要求某些维度对齐）。此时用 narrow_padded_param_and_loaded_weight 同时收缩参数视图与权重，再拷贝。

融合层的切分比普通 Linear 更绕一层。下图是 Megatron-LM 原论文里的自注意力张量并行布局：

head 维度被切到不同 rank，意味着融合的 qkv_proj 参数在每个 rank 上既承载了「Q/K/V 三个槽位」又承载了「本 rank 的 head 子集」。SGLang 用一个 shard_id 编码前者，用参数形状与 tp_rank 编码后者；加载时两者必须同时满足，否则切出来的切片会跨槽位。

5.3 融合层：一个 shard_id 同时表达「槽位」与「切分」

融合的 gate_up_proj 与 QKV 这类参数，加载器既要区分逻辑槽位，又要按 TP 切。以 DeepSeek 的 MoE 为例，FusedMoE._load_w13 把两件事写在一起：

1
if shard_id in {"w1", "w3"} and self.moe_runner_config.is_gated:
2
    shard_size = expert_data.shape[shard_dim] // 2
3
...
4
switch_w13 = getattr(_weight_layout_method(self), "load_up_proj_weight_first", False)
5
if ((switch_w13 and shard_id == "w1") or (not switch_w13 and shard_id == "w3")):
6
    start = shard_size
7
else:
8
    start = 0

先看 shard_size = expert_data.shape[shard_dim] // 2：参数把 gate 和 up 沿输出维拼在一段长度为 2 × intermediate 的缓冲区里，gate 占前一半、up 占后一半，所以每片宽度是一半。start 决定写到哪一半。

switch_w13 这一段是内核布局差异的直接体现。Triton 组的 MoE 内核与 TensorRT-LLM/FlashInfer 的 cutlass 内核期待 gate/up 的排列顺序不同，源码注释写的就是「trtllm cutlass kernel assumes differently」。同一个 checkpoint、同一份权重，落到不同内核后权重布局不同——如果加载器写死 up 永远在后半段，换内核后算出来的 SwiGLU 会静默错位（形状对、数值错，不会报异常）。

再往下是 TP 与 padding：

1
if self.use_padded_loading:
2
    weight_start = shard_size * self.moe_tp_rank
3
    if _is_cpu and self.use_presharded_weights:
4
        weight_start = 0
5
    expert_data, loaded_weight = narrow_padded_param_and_loaded_weight(
6
        expert_data, loaded_weight, start, weight_start, shard_dim, shard_size,
7
        not self.use_presharded_weights)
8
else:
9
    if not self.use_presharded_weights:
10
        if not is_bias and self.use_triton_kernels:
11
            loaded_weight = loaded_weight.transpose(-2, -1)
12
        loaded_shard_size = loaded_weight.shape[shard_dim] // self.moe_tp_size
13
        loaded_weight = loaded_weight.narrow(
14
            shard_dim, loaded_shard_size * self.moe_tp_rank, loaded_shard_size)
15
    expert_data = expert_data.narrow(shard_dim, start, shard_size)

use_padded_loading 的判定在 FusedMoE 里是：

1
return _is_cpu or self.use_flashinfer_trtllm_moe or aiter_padded

三种情况需要 padding：CPU 路径（AMX 等实现）要求中间维对齐；FlashInfer 的 TRT-LLM MoE 路径要求中间维 padding 到 128（gated NVFP4 是 64）；AITER 的权重也带 padding。因此参数缓冲区可能比 checkpoint 的 TP 切片更宽，加载时 expert_data.narrow(...).zero_() 先把多出来的部分清零，再拷数据：

1
if (loaded_weight.dim() == expert_data.dim()
2
    and loaded_weight.shape[shard_dim] < expert_data.shape[shard_dim]):
3
    loaded_size = loaded_weight.shape[shard_dim]
4
    expert_data.narrow(shard_dim, loaded_size,
5
                       expert_data.shape[shard_dim] - loaded_size).zero_()
6
    expert_data.narrow(shard_dim, 0, loaded_size).copy_(loaded_weight)

为什么要清零而不是留着 torch.empty 的垃圾值？ padding 区域虽然不参与真实计算，但会参与分组 GEMM：对齐表把每个专家的 token 行 pad 到块边界，GEMM 计算完再丢弃。如果 padding 权重里是 NaN 或随机大值，GEMM 的结果可能溢出到合法区域（某些内核的累加顺序与 padding 有关）或者触发 NaN 传播。清零让 padding 对结果严格无贡献。

transpose(-2, -1) 则是 Triton kernel 的布局要求：权重按 [N, K] 存、输入按 [K, N] 乘，与 checkpoint 的 [out, in] 差一次转置。同样的思路在 llama.cpp 源码完全拆解（十五）：量化 GEMM 的矩阵编队里也有对应，只是 llama.cpp 用 ggml_mul_mat 的转置语义表达。

5.4 不整除：TP 与专家数对不上时

TP size 不整除中间维或 head 数是常见的部署选择，例如 8 卡上的 96 个 attention head、或 6 卡上的 MoE。SGLang 用 pad_loaded_weight 解决「全局输出维不能被 tp_size 整除」的情况：

1
def pad_loaded_weight(loaded_weight, output_dim, output_sizes):
2
    total_output_size = sum(output_sizes)
3
    raw_output_size = loaded_weight.size(output_dim)
4
    if total_output_size > raw_output_size:
5
        weight_split_size = [int(o / total_output_size * raw) for o in output_sizes]
6
        assert sum(weight_split_size) == raw_output_size, ...
7
        split_weight = loaded_weight.split_with_sizes(weight_split_size, dim=output_dim)
8
        ...  # 每片末尾补零到对应的 output_sizes[i]

它是「按比例分配再补零」的策略：全局 96 个 head 分给 6 个 rank 时，有的 rank 拿 16 个、有的拿 17 个；参数按最大值 17 申请，少的一份补零。补零的 head 在注意力里被 mask 掉或参与计算但输出丢弃，所以数值无害；关键在于几何上必须补齐，否则 Q/K/V 的 head 数不同，注意力 kernel 的 stride 计算会出错。

7.1 共享专家融合：把共享专家伪装成第 N 个路由专家

DeepSeek 的 MoE 每个层有一个共享专家和 256 个路由专家。部分部署选择把共享专家与路由专家放进同一个 FusedMoE 参数，以便一次分组 GEMM 算完：

1
if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name:
2
    name = name.replace("mlp.shared_experts",
3
                        f"mlp.experts.{self.config.n_routed_experts}")

于是专家映射表的 num_experts 也变成 n_routed_experts + num_fused_shared_experts。共享专家占据索引 n_routed_experts（DeepSeek-V3 上是 256）这个槽位。FusedMoE.weight_loader 里的 _has_fused_shared 分支就是为它准备的：

1
shared_expert_id = (expert_id - global_expert_location_metadata.num_logical_experts
2
                    if self._has_fused_shared and expert_id is not None else -1)
3
if 0 <= shared_expert_id < self.num_fused_shared_experts:
4
    ...
5
    physical_expert_ids = [self._num_global_routed + shared_expert_id]

注释说得清楚：「Checkpoint shared experts start after logical routed experts, while local fused MoE weights store them after physical routed experts.」checkpoint 的共享专家排在逻辑专家之后，而本 rank 的参数里共享专家排在物理专家（包含 EPLB 复制出来的冗余专家）之后。两个「之后」不是同一个位置，因为 EP 下物理专家数可能多于逻辑专家数。加载器必须按物理布局重新计算目标槽位。

7.2 NextN：一份 checkpoint 两种角色

DeepSeek-V3 的 MTP（Multi-Token Prediction）层编号紧跟在主干层之后：主干有 61 层（编号 0–60），MTP 层编号 61。同一个 checkpoint 既可以被加载成「主干模型」，也可以被加载成「草稿模型」：

• NextN 关闭（主干模型）：所有 layers.{i} 且 i >= num_hidden_layers 的张量都会被丢弃。
• NextN 开启（草稿模型）：只保留 MTP 层，并且按张量类别改前缀——spec 权重（shared_head、enorm、hnorm、eh_proj）改成 model.*，其余改成 model.decoder.*。

1
if any(s in name for s in spec_weight_names):
2
    name = name.replace(layer_prefix, "model")
3
else:
4
    name = name.replace(layer_prefix, "model.decoder")

为什么同一层要拆成两个前缀？因为草稿模型的结构是「一个解码器层 + 一套 spec 模块」，model.decoder 对应那个解码器层，model.* 对应挂在模型根上的 spec 头。如果都用同一个前缀，参数树里要么多出找不到的槽位，要么把解码器层的名字挤掉。SGLang 源码完全拆解（十一）：投机解码讲了草稿前向如何使用这些参数，这里只负责把它们放进正确的槽。

还有一个细节："shared_head.head" in name or "embed_tokens" in name: continue。草稿模型复用主干的 lm_head 与 embed_tokens，checkpoint 里这份副本被丢弃以避免重复占显存。

7.3 融合专家大张量：一次读入、多次写入

FP4/NVFP4 的 checkpoint 往往把每个专家的 gate 与 up 预先拼接成 experts.gate_up_proj，形状是 [num_experts, 2 × intermediate, hidden]。名字里已经没有 experts.{id}. 这一层，专家表一条都匹配不上，于是走 _load_fused_expert_tensor：

1
mappings = (
2
    (".experts.gate_up_proj_scale", ".experts.w13_weight_scale_inv", "w13"),
3
    (".experts.down_proj_scale", ".experts.w2_weight_scale_inv", "w2"),
4
    (".experts.gate_up_proj", ".experts.w13_weight", "w13"),
5
    (".experts.down_proj", ".experts.w2_weight", "w2"),
6
)
7
for source, target, shard_id in mappings:
8
    if not name.endswith(source):
9
        continue
10
    param = params_dict.get(name[: -len(source)] + target)
11
    if param is None:
12
        return False
13
    for expert_id, expert_weight in enumerate(loaded_weight):
14
        if shard_id == "w13":
15
            gate, up = expert_weight.chunk(2, dim=0)
16
            weight_loader(param, gate, param_name, shard_id="w1", expert_id=expert_id)
17
            weight_loader(param, up, param_name, shard_id="w3", expert_id=expert_id)
18
        else:
19
            weight_loader(param, expert_weight, param_name, shard_id="w2",
20
                          expert_id=expert_id)
21
    return True

这一段是加载器的「二次分发」：一个大张量在 Python 层枚举出各专家的切片，再分别交给 weight_loader，后者继续做本地专家过滤、TP 切分、padding。优点是不必修改内核或参数布局；代价是 enumerate 会把整个大张量驻留在内存中（迭代器只产出一次），一个 256 专家、2 × 2048 × 7168 的 BF16 张量约 15 GB（256 × 2 × 2048 × 7168 个元素 × 2 字节）。这也是 enable_multithread_load 配合滑动窗口重要的原因——同一时刻只有一个大张量在内存里。

7.4 融合索引器与「等 scale 才能反量化」

DeepSeek-V3.2/V4 的稀疏索引器有一个 wk 与一个 weights_proj。CUDA 路径把两者融合成一个 bf16 参数 wk_weights_proj：

1
fused_name = name.rsplit(".indexer.", 1)[0] + ".indexer.wk_weights_proj.weight"
2
fused_param = params_dict.get(fused_name)
3
if fused_param is None or fused_param.dtype != torch.bfloat16:
4
    return False

融合参数是 bf16，但 checkpoint 里可能是块状 FP8：wk 分两片（weight 与 weight_scale_inv），weights_proj 也分两片。于是加载器要跨两个名字攒齐再反量化：

1
entry = pending.setdefault(fused_name, {})
2
entry["scale" if is_scale else "weight"] = _clone_if_runai_streamed_tensor(loaded_weight)
3
if "weight" in entry and "scale" in entry:
4
    pending.pop(fused_name)
5
    wk_bf16 = block_quant_dequant(entry["weight"], entry["scale"], block_size, torch.bfloat16)
6
    fused_param.data[: wk_bf16.shape[0]].copy_(wk_bf16)

pending 字典的生命周期是整个加载循环，攒齐后立刻清除对应的键。这种「按名字配对」的写法在加载器里很常见，常见错误是跨层的同名键互相覆盖——这里用完整前缀 fused_name 作键，按层区分。

_clone_if_runai_streamed_tensor 不是多余的拷贝。RunAI streamer 模式下张量是复用 CPU 缓冲的零拷贝视图，如果你把它存进 pending 而不拷贝，streamer 填下一批时视图内容会被覆盖。源码在 should_async_load 的注释里也强调这一点。

7.5 线程池与 CUDA device 的坑

所有 weight_loader 调用都经过 maybe_executor_submit：

1
if use_async:
2
    current_device = torch.cuda.current_device() if torch.cuda.is_available() else None
3
    def device_aware_wrapper(*args, **kwargs):
4
        if current_device is not None:
5
            torch.cuda.set_device(current_device)
6
        return func(*args, **kwargs)
7
    futures.append(executor.submit(device_aware_wrapper, *func_args, **func_kwargs))
8
else:
9
    func(*func_args, **func_kwargs)

use_async 由 should_async_load 决定，判断依据有三条：张量带 RunAI streamer 标记就同步（零拷贝视图必须当场消费），不在任何 device 上就同步，其余只在 CPU 张量上异步。注释解释了原因：CPU 张量的 copy_ 是 H2D 拷贝，把多个专家的拷贝分散到线程池可以重叠 I/O 与拷贝；GPU 张量的拷贝已经在 GPU 侧异步执行，Python 线程只会带来 GIL 争用与全局解释器锁下的上下文切换。

device_aware_wrapper 处理的是一个 PyTorch 的老问题：torch.cuda.current_device() 是线程本地状态，新线程的默认 device 是 0。在多 GPU 机器上，如果 rank i 的加载线程没有 set_device(i)，copy_ 会落到 GPU 0 上；运气好是报「device mismatch」，运气差是把数据写进另一个 rank 的显存。源码里链接的 PyTorch issue 就是这个问题的出处。

线程池在 with concurrent.futures.ThreadPoolExecutor() as executor: 里创建，没有传 max_workers，按 Python 默认规则取 min(32, os.cpu_count() + 4)。提交完 futures 之后，循环末尾统一 as_completed 收集，任何一个 future.result() 抛异常都会中止加载；这里没有 try/except，注释写的是「Wait for all tasks to complete and raise any exceptions」——加载出错就应该启动失败，而不是带着残缺权重进入服务。

8.1 量化：先变成一个统一的 dtype 再拆

post_load_weights 之所以长，是因为它要处理 kv_b_proj 的量化输入。它依次判断 AWQ 解包、FP8 块状、FP8 通道、INT8 块状、INT8 通道：

• 块状 FP8 且块大小是 128×128：如果 DeepGEMM 的 BMM 可用（JIT DeepGEMM 启用、非 Blackwell、且 SGLANG_USE_DEEPGEMM_BMM 打开），就保留块状 FP8 与块状 scale，直接交给 DeepGEMM；否则先 block_quant_dequant 反量化成 BF16，再拆。
• 块状 FP8 但块大小不是 128×128：block_quant_to_tensor_quant 把所有块的 scale 压成一个 per-tensor scale，存到 self_attn.w_scale。注释写明原因：「bmm_fp8 only supports per-tensor scale」。这是精度与通用内核之间的取舍：块状量化会损失一部分精度，但能跑在更广泛的硬件上。
• 通道状 FP8：scale 是一维 [out]，先 view(-1, 1) 让它在 [out, in] 上正确广播，再 channel_quant_to_tensor_quant 合并成 per-tensor。

三个分支的共同终点是：拆分之前，w 必须是一份布局为 [heads × (d_nope + d_v), r] 的张量。量化格式不同只决定它怎么从量化形式变回这个布局，拆分本身完全不变。

这段代码还留了 AMD 特有处理：gfx95 上如果架构是 GlmMoeDsaForCausalLM 且权重是 BF16，会量化成 per-tensor 的 float8_e4m3fn（不是 fnuz），以匹配 forward_mla_rocm 的 dtype 门槛；w_scale 在 HIP 上要乘 2.0，因为 fnuz 格式的指数偏置不同。这些是「加载期按平台调整布局」的典型例子——同一个 checkpoint 在不同加速器上需要不同的物理表示，全部集中在加载时做一次，前向路径就不必分叉。

8.2 为什么不在前向里拆

一个自然的问题是：既然 w_kc、w_vc 只是 kv_b_proj 的视图，为什么不保留 kv_b_proj，在前向里按需切？理由有三条：

1. 内存。保留原张量再切出两个视图意味着两份参数同时驻留（视图与原张量共享存储时还好，但 transpose 后需要新布局的副本）。对 MLA 来说 kv_b_proj 每层是 [128 × (128 + 128), 512]，BF16 下约 33.5 MB，61 层合计约 2 GB；拆开之后原张量可以释放，显存峰值更可控。
2. 前向路径的干净程度。注意力内核只认 w_kc 与 w_vc，前者在吸收形式的 Q 投影里当 [heads, d, r] 用、后者在输出投影里当 [heads, r, d] 用。让内核每次前向都从一个大张量里 narrow、transpose，等于每次都在重复同一个只做一次的转换。
3. 量化后处理。FP8 反量化、块状 scale 合并、MUSA 上「bmm 不支持 FP8 就转 BF16 并把 scale 乘进去」这些操作都只与加载时的一次性转换有关，放在前向里会随 batch 重复执行。

TensorRT-LLM 与 vLLM 的思路一致：前者在 Linear 的加载后处理里重排权重，后者在 process_weights_after_loading 里做，区别只在拆分发生在哪个钩子里。可参考 TensorRT-LLM 源码完全拆解（十二）：注意力层的张量并行，那一篇从 MLA 的张量并行视角讲了 w_kc/w_vc 的切分方式。

9.1 capture-safe model：用占位值先把 CUDA Graph 抓完

GPU 上启动耗时的很大一部分不是加载权重，而是 CUDA Graph 捕获与 warmup。可先读 SGLang 源码完全拆解（九）：CUDA Graph 与编译路径，那一篇讲过整图捕获的形状约束；抓图需要模型参数已经存在（形状、stride、device 都对），但不需要参数是真值。于是 SGLang 把「建壳 + 抓图」和「读真权重」拆开并允许重叠：

1
def prepare_model_for_capture(self, *, model, model_config):
2
    with set_default_torch_dtype(model_config.dtype):
3
        initialize_capture_safe_weights(model)   # 全部浮点参数填 1e-3
4
        post_load_weights(model)
5
        for _, module in model.named_modules():
6
            quant_method = getattr(module, "quant_method", None)
7
            if quant_method is None or (hasattr(module, "is_weights_quantized")
8
                                        and module.is_weights_quantized()):
9
                continue
10
            quant_method.process_weights_after_loading(module)
11
    return model.eval()

initialize_capture_safe_weights 把所有浮点参数填成一个可检测的哨兵值 1e-3（CAPTURE_SAFE_WEIGHT_SENTINEL），持久化 buffer 不动。为什么是 1e-3 而不是 0？两个原因：0 会让所有输出为 0，某些内核在图捕获时对全零输入可能有特化分支；取一个小正值能让前向不出现 NaN，又能被 ModelStorageManifest.unchanged_parameter_names 精确识别出「哪些参数没被真权重覆盖」。

抓完图之后，commit_model_weights 加载真权重，并靠 ModelStorageManifest 校验：

1
def commit_model_weights(self, *, model, model_config, resolved_sources,
2
                         target_device, startup_prefetch_active):
3
    def weights_iterator():
4
        for resolved_source in resolved_sources:
5
            yield from self._get_weights_iterator(
6
                resolved_source.source,
7
                resolved_source=resolved_source,
8
                startup_prefetch_started=True,
9
                startup_prefetch_active=startup_prefetch_active)
10
    with set_default_torch_dtype(model_config.dtype):
11
        self.load_weights_and_postprocess(model, weights_iterator(), target_device)

ResolvedSource 这个类型就是为这次拆分准备的：resolve_model_weights 在抓图前把「读哪些文件」一次性定下来（hf_folder、weight_files、use_safetensors），抓图期间后台线程去预热这些文件的页缓存，提交阶段复用同一份结果，不再重复 glob 与过滤。

ModelStorageManifest 记录每个参数与 buffer 的 data_ptr、形状、stride、dtype、device、storage_offset，提交后对比一遍，凡是存储发生变化的名字都会被找出来：

1
changed_names = manifest.changed_names(self._model)
2
if changed_names:
3
    raise RuntimeError(f"Startup weight commit changed graph-visible tensor storage: {preview}")
4
unchanged_names = manifest.unchanged_parameter_names(CAPTURE_SAFE_WEIGHT_SENTINEL)
5
if unchanged_names:
6
    raise RuntimeError("Startup weight commit did not replace capture-safe dummy values: ...")

两行断言正好覆盖两个方向的错误：前者抓「图里引用的地址被换了」，后者抓「哨兵值没被真权重覆盖」。这个设计解决了一个隐蔽问题：抓图时 process_weights_after_loading 已经把哨兵权重做了量化后处理（例如打包成 4-bit），如果加载真权重时换成了一块新的 Parameter 对象，图里引用的地址就失效了。Manifest 用「同一对象 + 同一指针 + 同一 stride」严格判定，任何替换都会被当作错误。

StartupWeightLoadManager.finalize 里还藏着一条顺序信息：它在提交前自己 capture 一次 manifest，而不是用 prepare 阶段的旧快照。原因是 finalize 之前可能已经插入了 CUDA Graph 捕获，捕获过程本身会分配临时 buffer、可能改动部分 param 的存储；以「提交前的状态」为基线，changed_names 才能只报告提交动作造成的差异。

源码的 docstring 也点出了一个已知代价：prepare_model_for_capture 与 commit_model_weights 会分别调用一次 process_weights_after_loading，比串行路径多一次。对当前支持的矩阵（CUDA 上未量化的 process_weights_after_loading 是空操作）是安全的；但对 在原位修改权重 的量化方法，这个多出来的调用可能把已经打包好的权重再打一次。所以启动重叠路径对支持的架构与量化方法做了准入检查（StartupWeightLoadManager._get_unsupported_reason），不满足就直接走串行加载。

9.2 checkpoint prefetch：先把页缓存喂热

第二个改造是页缓存预热。_prefetch_checkpoint_file 的行为非常简单：

1
with open(file_path, "rb") as f:
2
    while cancel_event is None or not cancel_event.is_set():
3
        if not f.read(_get_prefetch_block_size()):   # 默认按 16 MB 块读
4
            break

顺序读完整文件，什么都不保留。目的不是拿到数据，而是让内核把文件页缓存在内存里；之后所有 rank 的 mmap 访问都命中缓存，而不是各自去争磁盘带宽。CheckpointFilePrefetchHandle 提供 wait / cancel / stop 生命周期，failed 会把「后台线程报错」与「线程结束但没有成功事件」都算作失败。

为什么 prefetch 在多线程路径上会主动降级为单线程？第四段已经引过那段条件，这里补上原因：多线程读盘同样会填页缓存，但每条线程读的是不同文件，页缓存的竞争让预热线程与加载线程互相挤掉对方的页；单线程 + prefetch 的组合让读盘交给一个线程、预热交给另一个线程，两者不在同一批文件上争抢。注释里给出的判断线索是「用户是否显式指定多线程」——本地 NVMe 上 prefetch 没用，用户显式点名多线程说明场景是本地盘。

9.3 IPC 权重缓存：多个进程共享同一份权重

第三个改造更彻底：权重只加载一次，所有进程共享。maybe_enable_ipc_weight_cache 只改两行：

1
if load_config.load_format != LoadFormat.IPC_CACHE:
2
    load_config.fallback_load_format = load_config.load_format
3
    load_config.load_format = LoadFormat.IPC_CACHE

原来的 load_format 被保存成 fallback，主格式换成 IPC_CACHE。之后 get_model_loader 会返回 IpcModelLoader，从权重缓存守护进程（Weight Cache Daemon）通过 CUDA IPC 拿到已经建好的权重。IpcModelLoader 的注释说明了两种模式的区别：daemon 模式下引擎与守护进程共享同一份权重，如果守护进程不可用就报错，而不是退回磁盘——因为在同一张 GPU 上再建一份权重会直接吃满显存；client 模式下允许退回磁盘，但只有 Unix socket 文件不存在这一种情况才退回；连接被拒（守护进程绑定后崩溃）、CacheConfig 不匹配、协议或传输错误全部直接抛出，防止「IPC 路径坏了却静默变慢」或者「在同一张卡上按另一份配置的指纹读了权重」。这个「失败即报错、不降级」的契约就写在类的 docstring 里，对应的判定集中在 _fetch_from_cache。多进程场景（例如每个进程一个 TP rank 同机部署，或者 P/D 分离中同一台机器上跑多个实例）下，这样能省掉 tp_size 倍的读盘与 tp_size 倍的 CPU 内存。

这套 IPC 路径对加载管线的其余部分有一个硬约束：init_weights 阶段的参数必须允许被重新绑定到 IPC 内存上，所以像 bind_or_assign、_rebind_parameter_aliases 这类「改写指针而不替换对象」的工具要贯穿整条管线。前面第八段 MLA 分裂时用 bind_or_assign 的原因也在这一层——它同时服务 IPC 共享与 RL 热更新两种场景。

9.4 LayeredModelLoader：逐层加载，压低峰值

LayeredModelLoader 继承 DefaultModelLoader，docstring 一句话说清了取舍：「loads weights layer by layer so that one can quantize a layer before loading another to make the peak memory envelope smaller」。在线量化（BF16 checkpoint → FP8/NVFP4）时，如果先把整体 BF16 读进来再量化，峰值是「全量 BF16 + 全量量化后权重」；逐层加载则每一层读完、量化完、释放原 BF16，峰值只多一层。代价是读盘与量化无法跨层并行，启动会变慢。这也是在线量化部署里常见的选择：显存够就开多线程整体加载，显存紧就切逐层。