SGLang 源码完全拆解（二十八）：线性层的并行契约——一个 Linear 同时管着怎么切、怎么装、怎么算、怎么归约

这一篇要回答的问题

第二十五篇站在模型层的视角，看过一个 LLaMA 层如何被装配起来；第二十六篇从 MoE 的 MLP 出发，看一个算子在模型层、算子层与后端之间如何被分派。两篇都碰到同一个对象，但都只是擦肩而过：layers/linear.py 里的那几个 xxxParallelLinear。

它们是整个框架里最不起眼、也最容易被误解的一类代码。表面上，它们只是「带切分的 nn.Linear」：拿一个浮点张量进来，做一次矩阵乘，吐一个浮点张量出去。但一个推理框架里的 Linear 远不止做这一次乘法。它同时承担四份互相独立的契约：

1. 并行契约：这条权重在一组 rank 上怎么切——切输出维还是切输入维，切到「头」的粒度还是「列」的粒度，K/V 头数少于切分份数时又该怎么办。
2. 装载契约：磁盘上的一个张量（或者三个分开的张量）如何被切成本 rank 的那一片，scale 和打包格式在这条路径上被如何换算。
3. 计算契约：真正发出 GEMM 的不是 Linear 自己，而是挂在上面的量化方法；Linear 只负责「把形状和参数交给它」。
4. 归约契约：行并行的输出是部分和，谁来做那次 all-reduce——是这一段自己，还是上层的层边界，还是更外层的 DP Attention 组。

这四份契约在代码里是分开定义、分开演化的。一个刚读代码的人最容易犯的错，是把它们当成同一件事：看到 tp_size 就以为切分、装载、归约用的是同一个组；看到 reduce_results=False 就以为「这里不做归约」；看到 weight_loader 就以为只是「切片再 copy」。

这一篇要做的，就是沿着 sglang/srt/layers/linear.py（约两千行）和 sglang/srt/layers/parameter.py，把这四份契约逐份拆开，说明每一份在哪些类里、以什么参数形式出现、为什么必须这样设计。读完应当能做到：给任意一条模型层的 Linear...(...) 构造语句，都能在脑海里复原出「每个 rank 上这块权重是什么形状、它从磁盘的哪个位置来、前向发出去的是哪个 kernel、结果要不要跨卡求和」。

和前三篇的关系是这样的：第二十五篇讲的是模型层怎么摆这些 Linear，第二十六篇讲的是MoE 算子怎么分派后端，本篇讲的是Linear 自己——它是两者共用的最底层积木。KV 池、注意力后端、层边界这些相邻话题在前文都已拆过，本篇只在必需时引用结论。

先说清楚本文的语境。SGLang 的 linear.py 文件头第一行就写着「Adapted from vLLM」，核心的四个类——不切、列并行、行并行、QKV 并行——沿用的是 vLLM 的设计。SGLang 在其上做了两类改造：一是把并行组显式化（parallel_group 参数），二是在前向里接入了层边界与 DP Attention 的归约归属。这两类改造正是本篇与「只想了解 TP 基本原理」的文章之间的差别。

张量并行的基本形态：列并行与行并行

在拆 SGLang 的代码之前，先把张量并行（Tensor Parallelism，TP）的基本算术摆清楚。Megatron-LM 在 2019 年的论文《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》（Shoeybi 等，arXiv:1909.08053）给出了 Transformer 的两种切法，这套切法至今在所有推理框架里都能看到影子。

一个全连接层写成 $Y = XA + b$，其中 $X \in \mathbb{R}^{M \times K}$（$M$ 是 token 数、$K$ 是输入维），$A \in \mathbb{R}^{K \times N}$，$Y \in \mathbb{R}^{M \times N}$。要在 $P$ 个 rank 上分摊，只有两种不引发中间通信的切法：

列并行（Column Parallel）：把 $A$ 按输出维切成 $P$ 段，$A = [A_1, \dots, A_P]$，每段 $A_r \in \mathbb{R}^{K \times N/P}$。每个 rank 拿完整的 $X$，算 $Y_r = X A_r$，尺寸 $M \times N/P$。切分沿「列」方向进行，输入不需要任何重排，前向无通信；代价是输出被切开了，如果下游需要完整的 $Y$，就得一次 all-gather。

行并行（Row Parallel）：把 $A$ 按输入维切成 $P$ 段，叠起来 $A = [A_1; \dots; A_P]$，每段 $A_r \in \mathbb{R}^{K/P \times N}$。此时每个 rank 只能算一部分内积：

其中 $X_r$ 是 $X$ 沿输入维的第 $r$ 段。每段的乘积 $X_r A_r$ 是一张完整的 $M \times N$ 矩阵，把 $P$ 份加起来才得到真正的 $Y$。所以行并行的输出是部分和，必须做一次跨 rank 的 all-reduce（或 reduce-scatter）。

图：Megatron-LM（Shoeybi 等，2019）给出的 MLP 并行切分。第一段 `f` 是列并行，各 rank 只持有本 rank 的中间激活；第二段 `g` 是行并行，各 rank 的结果先求和（`f` 箭头处的 $\oplus$）再进入 dropout。这正是 SGLang 的 `gate_up_proj` + `down_proj` 这一对的结构。来源：arXiv:1909.08053。

图：Megatron-LM 的多头注意力切分。Q、K、V 三个投影都按「头」切分，每个 rank 只算自己那几个头的注意力；输出投影按输入维（也就是头维的拼接维）切成行并行，最后在 $Z$ 处求和。SGLang 的 `QKVParallelLinear` + `RowParallelLinear`（`o_proj`）与此一一对应，区别只在「头数不能被整除时怎么办」。来源：arXiv:1909.08053。

把两张图放在一起看，Transformer 层的 TP 形态就出来了：注意力的 QKV 投影是列并行、输出投影是行并行；MLP 的 gate/up 投影是列并行、down 投影是行并行。也就是说，一次前向里每段行并行之后恰好跟着一段可以吸收归约的逐元素操作（dropout、残差加法、下一层的 norm），Megatron 的原始设计就是把归约夹在两段列并行之间，让通信次数等于「层数 × 2」而不是「层数 × 每层 4 个线性层」。

SGLang 的四个类——ReplicatedLinear、ColumnParallelLinear、RowParallelLinear、QKVParallelLinear——与这套形态一一对应。其中 QKVParallelLinear 是列并行的一个特化，专门处理「按头切、且 GQA 下 K/V 头可能不够切」的情形。MergedColumnParallelLinear 则是把若干段列并行在输出维上拼接成一个参数，目的是把 gate 与 up 两次投影合并成一次 GEMM。

四份契约的对照表

在展开细节前，先把四份契约在代码里的落点列出来。这张表是本文的纲，后面的每一节都是在它的一行里往下挖。

最容易出 bug 的地方恰恰在这四份契约的交叉点上：并行契约选了 attn_tp 组，装载契约却按 tp 组去切；归约契约把 reduce_results 交给上层，上层却没接住；量化方法换了一套参数名，装载器却没换 v2 版本。这些错的共同形态是「每一份契约单独看都自洽，拼在一起就错位」。

并行契约（一）：并行组是显式传入的

打开 ColumnParallelLinear.__init__，第一个不寻常的地方是它不读全局的 tp_size，而是先解析一个 parallel_group：

1
class ColumnParallelLinear(LinearBase):
2
    def __init__(self, input_size, output_size, ..., parallel_group=None, ...):
3
        self._initialize_partition(..., parallel_group=parallel_group,
4
                                   tp_group=_resolve_linear_group(parallel_group))
5

6
    def _initialize_partition(self, ..., tp_group, parallel_group=None, ...):
7
        self.tp_group = tp_group
8
        self.parallel_group = parallel_group
9
        _, tp_size = get_group_rank_size(tp_group)
10
        super().__init__(input_size, output_size, skip_bias_add, params_dtype,
11
                         quant_config, prefix)
12
        ...
13
        self.output_size_per_partition = divide(self.output_size, tp_size)

parallel_group 的取值是一个联合类型：

1
LinearParallelGroup = (
2
    Literal["tp", "attn_tp", "replicated", "shared_experts_tp"]
3
    | ReplicatedParallelGroup
4
)

四个字符串，加上一个结构体。它们分别对应「普通 TP」「注意力 TP」「完全不切」「共享专家 TP」，以及「在某个组内再按相邻 rank 复制份数合并权重」。_resolve_linear_group 负责把名字翻译成一个组对象，翻译过程分三档：

1
def _resolve_linear_group(parallel_group):
2
    if isinstance(parallel_group, ReplicatedParallelGroup):
3
        group = _resolve_linear_group(parallel_group.group)
4
        if group is None:
5
            if parallel_group.replica_size != 1:
6
                raise ValueError("Cannot replicate a size-one group")
7
            return None
8
        return _ReplicatedGroup(group, parallel_group.replica_size)
9
    if parallel_group == "replicated":
10
        return None
11
    name = "tp" if parallel_group is None else parallel_group
12
    if name not in ("tp", "attn_tp", "shared_experts_tp"):
13
        raise ValueError(f"Unknown linear parallel_group: {parallel_group!r}")
14
    parallel = get_parallel()
15
    try:
16
        group = getattr(parallel, f"{name}_group")
17
    except RuntimeError:
18
        group = None
19
    if group is None:
20
        if name == "shared_experts_tp":
21
            raise ValueError(f"The {name} group must exist before construction")
22
        size = getattr(parallel, f"{name}_size")
23
        if size > 1:
24
            return _LogicalGroup(getattr(parallel, f"{name}_rank"), size)
25
    return group

这里有几个设计点值得停一下。

第一，None 表示「不切」，而不是「默认 TP」。 parallel_group == "replicated" 返回 None；而默认值 None（没传参数）反而被翻译成 "tp"。这个反直觉的约定必须记住，因为 get_group_rank_size(None) 返回 (0, 1)——不切时 rank 是 0、规模是 1，divide(output_size, 1) 就是原尺寸。所以「传 None」和「不传」是两件完全不同的事，读代码时看到 parallel_group=parallel_group 这种透传，要回溯到调用方才知道真正的值。

第二，构造期可能拿不到真正的通信组。 模型层在构造 Linear 时会建一堆 M 维大小为 0 的空参数，这时 NCCL 通信组可能还没建好。代码的处理是退回 _LogicalGroup——一个只记了 rank_in_group 和 world_size 的冻结结构体：

1
class _LogicalGroup(msgspec.Struct, frozen=True):
2
    """Own a weight partition before distributed communication is initialized."""
3
    rank_in_group: int
4
    world_size: int

也就是说，切分算术只需要「我是第几个、一共有几个」两条信息，根本不需要通信对象。这个抽象让「构造期定形状」与「运行期做集合通信」彻底解耦：形状可以早早定下来，通信等到前向时才发生。shared_experts_tp 是唯一的例外，它要求组必须已经存在——因为共享专家的切分份额要与其他组对齐，不能等到运行期再补。

第三，组的选择不是越多越好，而是有严格语义的。 attn_tp 是 DP Attention 拓扑下「注意力线性层实际使用的 TP 组」。SGLang 的 derive_attn_tp_size 写得极简：

1
def derive_attn_tp_size(*, tp_size: int, attn_cp_size: int, attn_dp_size: int) -> int:
2
    return tp_size // attn_dp_size // attn_cp_size

TP 组被切成 attention-DP 与 attention-CP 两个维度之后，剩下的才是 attention-TP。rank 的排列顺序是 (dp, cp, tp)、tp 变化最快：

代码里对应的两行是 tp_rank % attn_tp_size 与 tp_rank // (attn_tp_size * attn_cp_size)。想通这个公式的办法是：最低若干位是 tp、其上是 cp、再上是 dp——一个标准的「多维坐标按行主序展平」问题。DP Attention 的动机（注意力走数据并行以扩 KV 容量、MoE 走专家并行）在 SGLang 源码完全拆解（十七）：DP Attention 与专家并行 里已详细拆过，这里只取它的结论：注意力线性层必须用 attn_tp 组，不能用全局 tp 组，否则在 DP 副本之间存在多余的切分，前向会算错形状。

模型层的调用点印证了这一点：Qwen3Attention 的 qkv 与 o_proj 都写死 parallel_group="attn_tp"（models/qwen3.py），而 LlamaMLP 的默认值是 "tp"。同一个注意力里，o_proj 还额外把 reduce_results 设成 False——归约交给层边界。这就把「用哪个组」与「在哪个组上归约」分成了两件事，后面的第四节专门讲。

1
# Bind the selected group before quantization creates weight loaders.
2
self.tp_group = tp_group
3
self.parallel_group = parallel_group
4
_, tp_size = get_group_rank_size(tp_group)
5
super().__init__(input_size, output_size, skip_bias_add, params_dtype,
6
                 quant_config, prefix)

1
class ReplicatedParallelGroup(msgspec.Struct, frozen=True):
2
    """Share weight partitions across consecutive ranks; collectives stay separate."""
3
    group: Literal["tp", "attn_tp"]
4
    replica_size: int

1
class _ReplicatedGroup(msgspec.Struct, frozen=True):
2
    group: GroupCoordinator | _LogicalGroup
3
    replica_size: int
4

5
    def __post_init__(self):
6
        divide(self.group.world_size, self.replica_size)
7

8
    @property
9
    def rank_in_group(self):
10
        return self.group.rank_in_group // self.replica_size
11

12
    @property
13
    def world_size(self):
14
        return self.group.world_size // self.replica_size

1
kv_tp_group = (_resolve_linear_group(kv_parallel_group)
2
               if kv_parallel_group is not None else tp_group)
3
_, kv_tp_size = get_group_rank_size(kv_tp_group)
4
self.num_heads = divide(self.total_num_heads, tp_size)
5
if kv_tp_size >= self.total_num_kv_heads:
6
    self.num_kv_heads = 1
7
    self.num_kv_head_replicas = divide(kv_tp_size, self.total_num_kv_heads)
8
else:
9
    self.num_kv_heads = divide(self.total_num_kv_heads, kv_tp_size)
10
    self.num_kv_head_replicas = 1

计算契约：权重不是 Linear 自己造的

LinearBase 的构造函数里有一行是整个设计的关键：

1
if quant_config is None:
2
    from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
3
    self.quant_method = UnquantizedLinearMethod()
4
else:
5
    self.quant_method = quant_config.get_quant_method(self, prefix=prefix)

Linear 从来不 torch.empty 自己的权重。它把「创建权重」这件事托付给 quant_method.create_weights(...)。对应的抽象在 layers/quantization/base_config.py：

1
class LinearMethodBase(QuantizeMethodBase):
2
    def create_weights(self, layer, input_size_per_partition, output_partition_sizes,
3
                       input_size, output_size, params_dtype, **extra_weight_attrs):
4
        ...
5
    def apply(self, layer, x, bias=None):
6
        ...
7
    def process_weights_after_loading(self, layer) -> None:
8
        return

这个接口只有三件事：造参数、发计算、装载后整理。不同量化方法给出完全不同的参数集合：无量化时是一张 weight（加可选 bias）；GPTQ 是 qweight、qzeros、scales、g_idx；FP8 有 per-tensor 与 block 两套布局；MXFP8 还有按块打包的块 scale。Linear 的构造代码对这些名字一无所知——它只知道「参数创建好了，apply 能消费当前输入」。

第三个方法 process_weights_after_loading 常常被低估。它发生在「权重全部装载完、可以开始服务」之前，是唯一允许改形状、换布局、换数据类型的时机。以 SGLang 的 Fp8LinearMethod 为例（layers/quantization/fp8.py），它在 block 量化路径下会做三件不同的事：

• checkpoint 本来就是 MXFP8 时，convert_mxfp8_weight_to_block_fp8 把它就地转换成 128×128 的块 FP8 布局，并重新赋给 layer.weight——参数的形状与语义在装载后被整体替换。
• 某些加速器的小 batch 路径需要一份 row-wise FP8 的副本（由 dequant_mxfp8_2d_to_bf16 先反量化到 BF16、再逐行重新量化），这一份只服务 decode 时的小 M 场景。
• 如果模型声明了 deepgemm_mxfp8_input，它还会准备 DeepGEMM 需要的 MXFP8 输入路径，并把 Linear 上的对应标志置为可用。

这三件事有一个共同点：它们都依赖实际加载到的权重与目标硬件，只有装载后才可能知道。这就是为什么接口里要单独留出这一阶段，而不是塞进 create_weights。对应的 restore_weights_before_loading 则处理「需要重新装载」的场景（例如权重热更新或 QuantizedRL 重载）：撤销那些改变了参数形状的就地重排，让 checkpoint 格式的权重能再装一次。两个方法合起来，构成了「装载—整理—（可能）回滚」的闭环。

create_weights 接收的尺寸参数是每 rank 尺寸与全局尺寸同时传入：

1
self.quant_method.create_weights(
2
    layer=self,
3
    input_size_per_partition=self.input_size,          # 列并行：输入不切
4
    output_partition_sizes=self.output_partition_sizes, # 本 rank 每段输出宽
5
    input_size=self.input_size,                         # 全局输入维
6
    output_size=self.output_size,                       # 全局输出维
7
    params_dtype=self.params_dtype,
8
    weight_loader=(self.weight_loader_v2 if ... else self.weight_loader),
9
)

两套尺寸同时在场的理由，是要 create_weights 能同时服务造形状与造装载器：参数本身的形状取每 rank 尺寸；而装载器要知道全局尺寸才能把磁盘张量切成 P 份。这也就是第二十五篇提醒过的「同一个层里同时存在全局尺寸和每 rank 尺寸，脑子要分开放」。

前向一侧就更简单了：

1
def forward(self, input_):
2
    bias = self.bias if not self.skip_bias_add else None
3
    output_parallel = self.quant_method.apply(self, input_, bias)
4
    ...

注意 apply(self, ...) 传的是整个 layer 对象，而不是单独的权重张量。这让量化方法可以访问 layer.weight、layer.weight_scale、layer.input_size_per_partition 等一切它需要的东西，也可以在 apply 里按输入形状临时选择不同的内核（比如 decode 时用 GEMV、prefill 时用 GEMM）。核函数的选择权在量化方法手里，不在 Linear 手里——SGLang 里一个 Fp8LinearMethod 在不同 SM 架构、不同量化方案下可以落到不同的 GEMM 实现，Linear 不需要知道任何细节。

skip_bias_add 把 bias 从 GEMM 里剥出来返回给调用者，也是为融合服务的：调用者可以把「加 bias」和后续的激活、量化写回融进一个 kernel。这个参数在 vLLM 与 SGLang 里一直存在，是 epilogue 融合的入口——GEMM Epilogue 完全拆解 一文里讨论的「主循环只算乘积，真正决定性能的是尾声」，在 Linear 这一层就体现为「bias 不进主循环」。

1
scheme = None

1
weight_loader=(
2
    self.weight_loader_v2
3
    if (self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED
4
        or (self.scheme is not None and self.scheme.requires_weight_loader_v2))
5
    else self.weight_loader
6
)

装载契约：一个张量要同时满足四个约束

装载是四份契约里最琐碎、也最容易出错的一份。它要回答的问题看起来只有一句「把磁盘上的张量搬进本 rank 的参数」，但实际上要同时处理四类约束：

1. 并行约束：本 rank 该拿哪一段（列并行按输出维、行并行按输入维）。
2. 合并约束：磁盘上的三个张量（q、k、v）要拼进一个连续参数，或者反过来，一个合并张量要按 shard 切开。
3. 量化约束：打包格式（每 pack_factor 个值压进一个元素）、scale 的形状（标量、按通道、按块）。
4. 复制约束：GQA 下 K/V 头被复制时，每个 rank 拿的是哪一份。

除了这四条「从哪里来」的约束，还有一条方向相反的约束：磁盘上可能已经切好了。QKV 层的构造参数 load_presharded_attn=True 会把 use_presharded_weights 设成真，表示磁盘里的张量本身就是「每 rank 一份」的预切分格式。此时装载器不能再切一次：

1
if not use_presharded_weights:
2
    loaded_weight = loaded_weight.narrow(self.output_dim, tp_rank * shard_size, shard_size)
3
copy_with_check(self.data, loaded_weight)

同时，MergedColumnParallelLinear.weight_loader 里计算各段偏移时也要把 use_presharded_weights 算进去：

1
effective_size = (output_size // tp_size if self.use_presharded_weights else output_size)

否则「磁盘已是每 rank 尺寸」和「按全局尺寸切 P 份」会被重复应用一次，切出来的片只有期望的 1/P 那么大。预切分是一个全路径都要感知的开关，它出现在切片、偏移、断言三处；漏掉任何一处，结果都会是形状对不上或者静默读错段。

1
weight_loader=(
2
    self.weight_loader_v2
3
    if (self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED
4
        or (self.scheme is not None and self.scheme.requires_weight_loader_v2))
5
    else self.weight_loader
6
)

1
def load_column_parallel_weight(self, loaded_weight, tp_rank,
2
                                use_presharded_weights=False):
3
    if not use_presharded_weights:
4
        shard_size = self.data.shape[self.output_dim]
5
        loaded_weight = loaded_weight.narrow(
6
            self.output_dim, tp_rank * shard_size, shard_size)
7
    copy_with_check(self.data, loaded_weight)

1
param_data = param_data.narrow(self.output_dim, shard_offset, shard_size)
2
start_idx = tp_rank * shard_size
3
end_idx = start_idx + shard_size
4
if end_idx > loaded_weight.shape[self.output_dim]:
5
    loaded_weight = pad_or_narrow_weight(loaded_weight, self.output_dim,
6
                                         start_idx, shard_size)
7
else:
8
    loaded_weight = loaded_weight.narrow(self.output_dim, start_idx, shard_size)

1
if packed_dim == output_dim:
2
    shard_size = round(shard_size // param.pack_factor)
3
    shard_offset = round(shard_offset // param.pack_factor)
4
    shard_size, shard_offset = adjust_marlin_shard(param, shard_size, shard_offset)

1
if isinstance(param, BlockQuantScaleParameter):
2
    weight_block_size = self.quant_method.quant_config.weight_block_size
3
    raw_block_n, _ = weight_block_size[0], weight_block_size[1]
4
    block_n = 1 if getattr(param, "format_ue8m0", False) else raw_block_n
5
    shard_offset = (shard_offset + block_n - 1) // block_n
6
    shard_size = (shard_size + block_n - 1) // block_n

特殊的三种层

linear.py 里除了四个主类，还有三种值得单独讲的形态：QKVParallelLinear 的按头装载、RowParallelLinear 的「bias 只落一次」、以及 MergedColumnParallelRepeatedLinear 与 ColumnParallelBatchedLinear 这两类为特殊模型准备的变体。

1
kv_tp_rank, _ = get_group_rank_size(self.kv_tp_group)
2
tp_rank, _ = get_group_rank_size(self.tp_group)
3
param.load_qkv_weight(
4
    loaded_weight=loaded_weight,
5
    num_heads=self.num_kv_head_replicas,
6
    shard_id=loaded_shard_id,
7
    shard_offset=shard_offset,
8
    shard_size=shard_size,
9
    tp_rank=(tp_rank if loaded_shard_id == "q" else kv_tp_rank),
10
    use_presharded_weights=self.use_presharded_weights,
11
)

1
shard_id = tp_rank if shard_id == "q" else tp_rank // num_heads
2
param_data = param_data.narrow(self.output_dim, shard_offset, shard_size)
3
...
4
loaded_weight = loaded_weight.narrow(self.output_dim, shard_id * shard_size, shard_size)

1
shard_offset_mapping = {
2
    "q": 0,
3
    "k": self.num_heads * self.head_size,
4
    "v": (self.num_heads + self.num_kv_heads) * self.head_size,
5
    ...
6
}

1
# Only fuse bias add into GEMM for rank 0 (this ensures that bias will not
2
# get added more than once in TP>1 case)
3
bias_ = None if (tp_rank > 0 or self.skip_bias_add) else self.bias

1
if (input_dim is not None and not use_bitsandbytes_4bit
2
        and not self.use_presharded_weights):
3
    shard_size = param_data.shape[input_dim]
4
    start_idx = tp_rank * shard_size
5
    end_idx = start_idx + shard_size
6
    if end_idx > loaded_weight.shape[input_dim]:
7
        loaded_weight = pad_or_narrow_weight(loaded_weight, input_dim, start_idx, shard_size)
8
    else:
9
        loaded_weight = loaded_weight.narrow(input_dim, start_idx, shard_size)

1
def weight_loader(self, param, loaded_weight, loaded_shard_id):
2
    output_dim = param.output_dim
3
    shard_offset = sum(self.output_partition_sizes[:loaded_shard_id])
4
    shard_size = self.output_partition_sizes[loaded_shard_id]
5
    param_data = param.data.narrow(output_dim, shard_offset, shard_size)
6
    tp_rank, _ = get_group_rank_size(self.tp_group)
7
    if loaded_shard_id < self.num_column_parallel:
8
        start_idx = tp_rank * shard_size
9
        loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
10
    param_data.copy_(loaded_weight)

归约契约：部分和由谁求和

归约契约是 SGLang 相对 vLLM 改动最大的一份。RowParallelLinear.forward 的签名本身就透露出这不是一个简单的 Linear：

1
def forward(self, input_, skip_all_reduce=False, forward_batch=None, output_tensor=None):

四个参数，其中三个与归约和内存有关。拆开看：

1
if self.use_dp_attention_reduce:
2
    output = get_parallel().attn_tp_group.all_reduce(output_parallel)
3
else:
4
    if quantize_communications:
5
        output = tensor_model_parallel_quant_all_reduce(output_parallel)
6
    elif self.parallel_group == "tp":
7
        output = get_parallel().tp_group.all_reduce(output_parallel)
8
    else:
9
        output = tensor_model_parallel_all_reduce(output_parallel)

1
if (((self.reduce_results and tp_size > 1) or self.use_decode_attn_tp)
2
        and not skip_all_reduce
3
        and not should_skip_mlp_all_reduce()):
4
    ...do all_reduce...
5
else:
6
    output = output_parallel

1
if output_tensor is None:
2
    output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
3
else:
4
    apply_into = getattr(self.quant_method, "apply_into", None)
5
    if apply_into is None:
6
        raise RuntimeError(f"{type(self.quant_method).__name__} cannot write into "
7
                           f"caller-owned linear output")
8
    output_parallel = apply_into(self, input_parallel, output_tensor, bias=bias_)

1
if self.use_dp_attention_reduce:
2
    symm_ctx = use_symmetric_memory(get_parallel().attn_tp_group)
3
else:
4
    symm_ctx = use_symmetric_memory(
5
        get_parallel().tp_group,
6
        disabled=self.tp_group is None or not is_allocation_symmetric(),
7
    )
8
with symm_ctx:
9
    output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)

1
if (get_forward().sp_active and tp_size > 1 and not skip_all_reduce
2
        and output_tensor is None):
3
    output = layernorm_sp.row_parallel_gbar_matmul(self, input_parallel, bias_)
4
    ...

1
if get_forward().sp_active and tp_size > 1:
2
    output = layernorm_sp.column_parallel_g_matmul(self, input_, bias)

一次前向走查：以 LLaMA 的一个解码 batch 为例

把四份契约串起来，用 LLaMA 的注意力 + MLP 走一次解码的前向，看每一段 Linear 在每一步发生了什么。设 TP=8，Q 头 32、KV 头 8（不复制），hidden=4096，head_size=128，decode 时 batch 有 64 条请求、每请求 1 个 token，共 64 行输入。

对照表看两件事：

第一，每一段的「本 rank 形状」都可以从构造参数直接算出来，不需要看磁盘。 qkv_proj 的 Q 段是 num_heads = divide(32, 8) = 4，宽度 4 × 128 = 512；K/V 段是 num_kv_heads = divide(8, 8) = 1，宽度 1 × 128 = 128，合计 768。这里有一个很容易犯的错：读代码时把全局头数当成每 rank 头数，把 K 宽算成 8 × 128 = 1024——一旦这么算，qkv.split([q_size, kv_size, kv_size]) 就会对不上。q_size 与 kv_size 都是 LlamaAttention 里按「每 rank」算出来的，与 QKVParallelLinear 内部的算术是同一套。任何形状都从 divide(total, tp) 开始推，这是读 Linear 代码必须养成的习惯。

第二，归约的次数与位置由层边界决定，与 Linear 自己无关。 注意力里有一段行并行、MLP 里有一段行并行，但归约发生了两次而不是四次，而且都不在 Linear 的 forward 里。这就是第二十五篇讲的「层边界把残差和通信收走」在更深一层的表现：Linear 只是「可能」做归约，「一定」做归约的是边界。

同一段前向，换成 DP Attention 拓扑

上面的例子是纯 TP。把注意力切成两个 DP 副本、让 FFN 继续留在完整的 TP 组上，就能看到「线性层同时使用两个不同的组」到底意味着什么。设 TP=8、attn_dp_size=2、attn_cp_size=1，于是 attn_tp_size = 8 // 2 // 1 = 4，注意力 TP 组是 [0,1,2,3] 与 [4,5,6,7] 两个副本，FFN 仍用 [0..7]。

这张表和上一张表只差了几个数字，却对应两个完全不同的服务形态：

注意力侧，两个副本各自持有一份完整的 8 个 KV 头，各自服务自己的请求 batch。 每副本 4 个 rank、每 rank 2 个头，合起来 8 个头；另一个副本也同样有 8 个头。物理上 KV 头被存了两份，但两份服务的是不同的 token，所以聚合的 KV 容量翻倍、显存也翻倍——这不是浪费，而是用更多卡换取更多并发上下文。如果只看「每 rank 2 个头」就以为容量和纯 TP 一样，就漏掉了副本数这个乘数。

FFN 侧，权重仍按完整的 TP 组切分，没有复制。 注意力输出的 token 需要先跨两个副本收集成全局 batch，FFN 才能一次算完；算完后再散回各自的副本。注意力权重（qkv/o）虽然被复制了两份，但它们在总权重里占比很小；FFN 与 MoE 才是参数大头，必须共享。「注意力 DP、FFN TP/EP」这个组合的本质，是把复制代价留给小的那部分权重，把容量收益留给大的那部分 KV。

再看归约的通信量。o_proj 的 all-reduce 从 8 个 rank 缩到 4 个 rank：ring 算法的每 rank 搬运量约为 $\frac{2(P-1)}{P} S$（$S$ 是张量字节数），$P=8$ 时是 $1.75S$，$P=4$ 时是 $1.5S$。单看一条请求，归约字节数下降约 14%；更重要的是归约不再跨副本，两个副本的集合通信互不等待，慢 rank 的抖动只影响自己那一半请求。这与 AllReduce 完全拆解 里讨论的「通信域越小、同步面越窄」是同一件事。

use_dp_attention_reduce=True 的 RowParallelLinear 就是为这张表服务的：它让 o_proj 的归约发生在 attn_tp 组内，而不是全局 tp 组。对照上一张表，同样的 LLaMA 结构、同样 8 张卡，仅仅因为组的选择不同，o_proj 的归约域就从 8 变成 4。「Linear 用哪个组」并不是一个可以随便默认的细节，它直接决定了通信拓扑。

与 vLLM、TensorRT-LLM 的对照

SGLang 这套设计与另外两个框架对照着看，能更清楚它的取舍。

vLLM：vllm/model_executor/layers/linear.py 的四个主类与 SGLang 几乎同构（SGLang 文件头明确引用 vLLM v0.6.4）。vLLM 的 parallel_group 是一个后加的雏形，主要服务于 DP Attention 场景；SGLang 把它扩展成完整的类型系统（四个名字加复制结构体），并把 _LogicalGroup 引入构造期。vLLM 有 weight_loader_v2 与 parameter.py 的参数类体系，两边的设计基本一致。

TensorRT-LLM：新一代 PyTorch 后端里，线性层由 Linear 模块配合 QuantizeMethodBase，同时承担 TP 切分与 GEMM-AllReduce 融合。它的特殊之处是把「列并行 GEMM 后立刻 all-reduce」作为可选融合路径，归约的归属由模块自己决定；这比 SGLang 的「层边界接管」更集中，但灵活性低一些。TensorRT-LLM 源码完全拆解（二十三）：Linear 与层内张量并行 有完整拆解。

llama.cpp：走的是另一条路——ggml-backend-meta.cpp 用虚拟设备表达张量并行，线性层在计算图里是普通的 MUL_MAT，只不过权重被切分到多个后端上，最后用 all-reduce 节点合并。它没有「Linear 类」，切分与归约都是图节点。可参考 llama.cpp 源码完全拆解（二十九）：Meta device 与张量并行。

这三条路线放在一起，可以看到「四份契约」的另一种组织方式。vLLM 与 SGLang 把四份契约塞进一个类对象，用构造参数与挂载的量化方法表达；TensorRT-LLM 把归约契约部分上提到模块自身的融合选项；llama.cpp 则干脆把并行契约与归约契约外化成图，Linear 只剩计算契约。

前两列的共同点是「契约在对象上」，第三列是「契约在图上」。对象化写法的优点是能在一个类里读到全部规则，缺点是规则分散在构造参数、类属性、前向分支三处；图化写法的优点是规则一目了然，缺点是每种拓扑组合都要在图里表达一次，代码复用困难。SGLang 选择对象化，并用 parallel_group 把最易错的部分显式化，本质上是在「规则分散」与「形态单一」之间做了一个偏向可读性的取舍。

三个框架的共同点是：「怎么切」与「谁归约」必须分开表达，否则同一份切分逻辑无法同时服务纯 TP、DP Attention、序列并行、上下文并行这四种拓扑。

图：SGLang 官方文档给出的 DP Attention + EP 架构示意。左侧是各 DP 副本各自持有注意力与 KV、右侧是专家分片；中间的搬运与归约正是本文「归约契约」里 `use_dp_attention_reduce` 与 `attn_tp` 组所服务的路径。来源：SGLang 文档《DP Attention / DPA + SMG 使用指南》。

读这段代码的实用路径

最后给出一条可执行的阅读顺序，适合需要在自己的框架或模型里复用这套设计时对照。

1. 先看模型层的构造语句（models/llama.py 的 LlamaMLP、models/qwen3.py 的 Qwen3Attention），记下每条 Linear 用了哪个 parallel_group、reduce_results 是什么值。这一步回答「契约的输入是什么」。
2. 进 linear.py 看 _resolve_linear_group 与 _initialize_partition，确认本 rank 形状是怎么从名组、total_num_heads、head_size 算出来的。这一步建立「形状算术」。
3. 判断走 v1 还是 v2 装载器（查 WEIGHT_LOADER_V2_SUPPORTED 与 scheme.requires_weight_loader_v2），然后只在对应的那份代码里读切片。这一步避免在两代装载器之间来回跳。
4. 对量化参数，去 parameter.py 找对应的参数类（ModelWeightParameter、BlockQuantScaleParameter……），读它的 load_column_parallel_weight 与 load_row_parallel_weight。这一步补上打包与 scale 的坐标变换。
5. 回到 forward，只读归约分支：use_dp_attention_reduce、quantize_communications、reduce_results/skip_all_reduce/should_skip_mlp_all_reduce。这一步确认归约的实际归属。
6. 最后核对：构造期选的组、装载期用的 rank、归约时用的组，三者是否一致。不一致必然有一个是 bug 或特例，顺着它就能找到原因。

第 6 步是本篇最想留下的一点。读 Linear 代码时的多数困惑，不是因为某处代码写得晦涩，而是因为你在三份契约之间串错了线。把「组、装载 rank、归约组」三个名字并排写在纸上，逐一对照，大多数问题会立刻浮现。

小结

• 一个推理框架里的 Linear 同时承担并行、装载、计算、归约四份互相独立、分别演化的契约；它们共享同一个类，但不共享同一个参数。
• 并行契约的关键是并行组显式传入：tp、attn_tp、shared_experts_tp、replicated，外加一个可按相邻 rank 复制的结构体；构造期可以用 _LogicalGroup 在没有通信组时先定形状。
• Q/K/V 可有独立的 kv_parallel_group；GQA 下 num_kv_head_replicas 把「KV 头不够切」变成「相邻 rank 读同一段」。
• 计算契约把「造参数、发 GEMM」全部托付给 quant_method；Linear 只管形状与装载器的选择，核函数的选择权在量化方法手里。
• 装载契约有两代：v1 在 Linear 里切片，v2 下沉到参数类；打包因子、Marlin tile、块 scale 的坐标变换都在装载路径上发生。
• 归约契约的三个开关（reduce_results、skip_all_reduce、前向标志）和三条组选择（attn_tp、量化通信、按组兜底）决定了一次前向里部分和在哪求和；reduce_results=False 不等于不归约。
• 对照 vLLM、TensorRT-LLM、llama.cpp，三者的共同结论是：「怎么切」与「谁归约」必须分开表达，才能同时适配纯 TP、DP Attention、序列并行、上下文并行四种拓扑。

下一篇会接着往「层边界」走：残差、norm 与归约如何在 layers/layer_boundary/ 里被组织成一段可替换的执行段，以及它如何与 Linear 的归约契约对接、与 CUDA Graph 的形状约束对齐。

参考资料

1. SGLang 源码：layers/linear.py
2. SGLang 源码：layers/parameter.py
3. SGLang 文档：DP Attention 与 DPA + SMG 使用指南
4. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism（Shoeybi 等，2019）
5. NVIDIA Megatron-LM 仓库
6. vLLM 源码：model_executor/layers/linear.py
7. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints（Ainslie 等，2023）
8. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM（Narayanan 等，2021）