TensorRT-LLM 源码完全拆解（十二）：注意力层的张量并行——head 怎么切、MLA 为什么不能照切

第一步：确认并行度

1
tp_size = self.mapping.tp_size
2
pp_size = self.mapping.pp_size
3
cp_size = self.mapping.cp_size
4
dp_size = 1
5
if self.mapping.enable_attention_dp:
6
    dp_size = tp_size
7
    tp_size = 1

开头这段代码有个容易被忽略的转折：开了 attention DP，就把 tp_size 折叠成 1，把 dp_size 设成原来的 tp_size。

原因是这样。在纯 TP 模式下，所有 rank 持有同一个请求的不同 head 切片，算完要 all-reduce；在 attention DP 模式下，每个 rank 独立处理不同的请求，注意力层根本不并行，自然也不需要通信——原来那组 TP rank 被重新解释成一组数据并行副本。这也是为什么上一篇里 DP Attention 与专家并行必须”错开”：注意力侧用 DP，MoE 侧用 EP，两者在同一个进程组里的划分方式不同。

注意这里还构造了一个新的 Mapping 对象，而不是复用 self.mapping：

1
mapping = Mapping(
2
    world_size=dp_size * tp_size * pp_size * cp_size,
3
    tp_size=tp_size,
4
    pp_size=pp_size * dp_size,
5
    cp_size=cp_size,
6
    ...
7
)

world_size 不变，但 tp_size 已经是从 DP 里剥出来的那个，pp_size 则乘上了 dp_size。这是一个”重新解释坐标系”的操作：同一个物理进程组，用另一套数学约定去看它。后面所有 Linear 构造都传这个 mapping，它们才会认为”我的 TP 组是这么多卡”。这个技巧在后面讲 MLA 的 mapping_o 时会再用一次。

第二步：head 数的除法，以及一个不对齐的取整

1
assert self.num_heads % (tp_size * cp_size) == 0
2
self.num_heads = self.num_heads // tp_size
3
self.num_heads_tp_cp = self.num_heads // cp_size
4
self.num_key_value_heads = (self.num_key_value_heads + tp_size - 1) // tp_size

三行代码，三种处理方式，各有各的道理。

Q head 用整除，并且断言必须整除。 这是硬约束：如果 num_heads 不能被 tp_size 整除，注意力就无法均匀分配工作，框架直接拒绝启动。以 DeepSeek-V3 为例，num_attention_heads = 128，配 tp_size = 8 时每个 rank 拿 16 个 Q head。

num_heads_tp_cp 是 TP 切完再 CP 切。 HELIX 上下文并行会把每个 rank 的 head 再按 CP 组切一次——注意这个变量在非 CP 场景下就等于 num_heads，代码里那句注释说明了这一点：

1
# note: if we do not have CP, then num_heads_tp_cp == num_heads_tp

KV head 用向上取整，不整除也不报错。 这个差异值得展开。GQA（Grouped-Query Attention）下 num_key_value_heads 通常远小于 num_heads——比如 Llama-3-70B 是 64 个 Q head 配 8 个 KV head。当 tp_size 大于 KV head 数时（例如 tp_size=16、num_key_value_heads=8），整除是不可能的：8 个 KV head 分不到 16 张卡上。

向上取整的语义是”允许多个 rank 重复持有同一个 KV head”。8 个头分到 16 张卡，结果是每个 rank 拿到 1 个 KV head，但实际只有 8 个不同的 head，剩下 8 个 rank 持有的是重复的。这看起来浪费，但在注意力层里是可以接受的——KV head 的重复持有只影响权重加载（多读一点）和 KV Cache 分配（每个 rank 多存一份），不影响计算正确性，因为每个 rank 的 Q head 仍然只和它本地的那份 KV 算注意力，结果在 o_proj 的 all-reduce 中合并。

代价是 KV Cache 的显存有冗余。这就是为什么 resource_manager.py 里 KV 显存估算也用了同一个公式：

1
self.num_kv_heads_per_layer = [
2
    (num_kv_heads + tp_size - 1) // tp_size
3
    ...
4
]

估算公式和模块的切分公式必须严格一致，否则会算错池子大小——这是这类”两处独立实现同一个除法”的地方最容易出的 bug。

第三步：导出本 rank 的形状

1
self.q_size = self.num_heads * self.head_dim
2
self.kv_size = self.num_key_value_heads * self.head_dim

到这里，self.num_heads 已经是本 rank 的 head 数（__init__ 里被改写过了），所以 q_size 和 kv_size 都是”本 rank 的 Q/KV 投影输出宽度”。用 DeepSeek-V3（128 head、head_dim=128、tp_size=8）算：本 rank 的 q_size = 16 × 128 = 2048，kv_size 因为 num_key_value_heads 也是 128，按向上取整得 16，所以 kv_size = 16 × 128 = 2048。

第四步：构造 QKV 融合权重

1
qkv_shard_indices_mapping = {
2
    "q": (0, self.q_size * (2 if self.attn_output_gate else 1)),
3
    "k": (self.q_size * (2 if self.attn_output_gate else 1), self.kv_size),
4
    "v": (self.q_size * (2 if self.attn_output_gate else 1) + self.kv_size,
5
          self.kv_size),
6
}
7

8
self.qkv_proj = Linear(
9
    self.hidden_size,
10
    tp_size * self.q_size * (2 if self.attn_output_gate else 1) +
11
    2 * tp_size * self.kv_size,
12
    bias=bias,
13
    dtype=dtype,
14
    mapping=mapping,
15
    tensor_parallel_mode=TensorParallelMode.COLUMN,
16
    weights_loading_config=WeightsLoadingConfig(
17
        weight_mode=WeightMode.FUSED_QKV_LINEAR),
18
    ...
19
    fused_weight_shard_indices_mapping=qkv_shard_indices_mapping,
20
    ...)

这里有三个关键点。

输出维度是 tp_size × 全局尺寸，不是本 rank 尺寸。 Linear 的 out_features 参数描述的是全局形状（tp_size * q_size + 2 * tp_size * kv_size），Linear 内部再按 COLUMN 模式把它切成 tp_size 份。这一点如果读错，会以为每个 rank 分配了一大堆用不到的权重。q_size 是本 rank 的，乘上 tp_size 就还原成全局的了。

attn_output_gate 让 Q 的宽度翻倍。 带输出门的注意力（如 Qwen3）在 Q 投影后额外产出一个 gate，用于对注意力输出做 out * sigmoid(gate) 的调制。所以 Q 部分的宽度是 2 × q_size，qkv_shard_indices_mapping 里三处 2 if self.attn_output_gate else 1 就是在处理这个。

qkv_shard_indices_mapping 解决的是”融合权重的加载切分”问题。 这是整个 __init__ 里最容易被跳过、但实际最容易出错的部分。回忆一下上一篇讲过的：Mapping 决定进程组，而这里决定的是权重加载时从 checkpoint 的哪个片段取数据。

COLUMN 模式下 Linear 会沿输出维度（第 0 维）切片，切片边界由 _calc_shard 算出：

1
@staticmethod
2
def _calc_shard(total, tp_size, rank):
3
    return (total // tp_size) * rank + min(total % tp_size, rank)

但当权重模式是 FUSED_QKV_LINEAR 时，这个自动切分不能用。原因很直白：checkpoint 里存的是分开的 q_proj.weight、k_proj.weight、v_proj.weight 三个张量，而模块持有的是一个融合后的大权重。切”融合后的张量”和切”三个独立张量再拼接”得到的是不同的结果——除非每一段的宽度都恰好能被 tp_size 整除，且段边界与 rank 边界对齐。

fused_weight_shard_indices_mapping 就是来声明这件事的。load_shard 里对它的处理是：

1
if isinstance(self.tp_sharding, dict):
2
    slice_start, slice_end = self.tp_sharding[name]

name 就是 load_weights_fused_qkv_helper 里分别传入的 'q'、'k'、'v'：

1
q_weight = module.load_shard(weights[0], 'weight', device=device,
2
                             name='q', elm_packing=elm_packing)
3
k_weight = module.load_shard(weights[1], 'weight', device=device,
4
                             name='k', elm_packing=elm_packing)
5
v_weight = module.load_shard(weights[2], 'weight', device=device,
6
                             name='v', elm_packing=elm_packing)

于是三个独立张量各自被切到 qkv_shard_indices_mapping 声明的区间，再写进融合权重的对应位置。bias 的处理更直接，用 copy_weight_shard 做原地分段拷贝：

1
def copy_weight_shard(dst: Parameter, src: torch.Tensor, shard_offset: int,
2
                      shard_size: int):
3
    if dst.dtype != src.dtype:
4
        src = src.to(dst.dtype)
5
    dst[shard_offset:shard_offset + shard_size].data.copy_(src)

shard_offset/shard_size 直接来自同一张映射表。这段代码解释了为什么 GQA 下必须显式声明 shard 边界：q_size 和 kv_size 不相等，自动均分必然错位。

第五步：o_proj，以及一个”额外一套 mapping”

1
mapping_o = Mapping(
2
    world_size=dp_size * tp_size * pp_size * cp_size,
3
    tp_size=tp_size * cp_size,
4
    pp_size=pp_size * dp_size,
5
    cp_size=1,
6
    ...
7
)
8
self.mapping_o = mapping_o
9

10
self.o_proj = Linear(
11
    tp_size * self.q_size,
12
    self.hidden_size,
13
    bias=self.dense_bias,
14
    dtype=dtype,
15
    mapping=mapping_o,
16
    tensor_parallel_mode=TensorParallelMode.ROW,
17
    ...
18
)

o_proj 是全篇最值得停下来想的地方，因为它同时出现两种形状。

输入维度写的是 tp_size * self.q_size——这又是”全局”尺寸。因为 self.q_size = num_heads × head_dim，而 num_heads 是本 rank 的 head 数，所以 tp_size * self.q_size 等于全局的 Q 宽度，也就是 o_proj 真正的输入维度。ROW 模式下 Linear 会沿输入维度（第 1 维）把它切成 tp_size 份，每个 rank 正好拿到 q_size——和它输出的注意力张量宽度严丝合缝对上。

但传进去的 mapping 是 mapping_o 而不是 mapping。 mapping_o 的 tp_size 是 tp_size * cp_size，cp_size 是 1。也就是说：从 o_proj 的视角看，TP 组变大了，包含了原本的 CP 组。

这就是 HELIX 上下文并行的核心安排。回到 __init__ 里的注释：

1
# For Helix CP, combine TP and CP for the output projection so each
2
# rank's o_proj input is num_heads_tp_cp * head_dim.

在 HELIX 模式下，注意力侧每个 rank 只持有 num_heads_tp_cp = num_heads / tp_size / cp_size 个 head，算出的注意力张量宽度是 num_heads_tp_cp × head_dim。但对 o_proj 来说，“我应该切多少份”必须和新加入的 CP 维度一起算——否则会有 cp_size 个 rank 认为自己持有同一段权重，all-reduce 的结果就错了。把 tp_size * cp_size 作为 o_proj 的 TP 度，恰好让每个 rank 的输入宽度对上。

图源：TensorRT-LLM 官方技术博客 Helix Parallelism 的分片示意（docs/source/blogs/media/tech_blog22_attention_sharding.png）。图中展示的是 HELIX 对注意力的双重切分：head 维按 TP/CP 分给不同 rank，序列（token）维也按 CP 切开。这正是 mapping_o 要把 tp_size × cp_size 合并成一个 TP 度的原因——o_proj 必须同时容纳这两条切分线。

这里可以对照上一篇结尾提到的层内模式：“all-gather → 注意力 → o_proj → reduce-scatter”。Attention.forward 的两端就是这两个操作的落点：

1
def _helix_cp_allgather_input(hidden_states, attn_metadata, mapping, layer_idx):
2
    """AllGather hidden states from CP group for layers after the first.
3

4
    The first layer already has the full input from the embedding.
5
    Subsequent layers need to undo the previous layer's reduce-scatter.
6
    """
7
    if (mapping.has_cp_helix() and mapping.enable_attention_dp
8
            and layer_idx > 0):
9
        hidden_states = cp_allgather(hidden_states, mapping, dim=0)
10
        hidden_states = hidden_states[:attn_metadata.num_tokens]
11
    return hidden_states

注意 layer_idx > 0 这个条件：第一层的输入直接来自 embedding，本来就是完整的；只有第二层及以后，输入是上一层 reduce-scatter 之后的token 片段，才需要先 all-gather 拼回完整序列。

另一端：

1
def _helix_cp_output_projection(o_proj, attn_output, attn_metadata,
2
                                all_reduce_params, mapping, mapping_o,
3
                                layer_idx, lora_params=None):
4
    """Apply output projection with reduce-scatter when Helix CP+DP is active.
5

6
    Reduce-scatter sums partial sums across the CP group and scatters the
7
    result so each CP rank processes a distinct token chunk through the MLP.
8
    Falls back to the standard AllReduce path otherwise.
9
    """
10
    if mapping.has_cp_helix() and mapping.enable_attention_dp:
11
        attn_output = o_proj(
12
            attn_output,
13
            all_reduce_params=AllReduceParams(enable_allreduce=False),
14
            lora_params=lora_params, layer_idx=layer_idx)
15

16
        attn_output, _ = _helix_cp_pad(attn_output, attn_metadata.num_tokens,
17
                                       mapping.cp_size)
18
        attn_output = reducescatter(attn_output, mapping_o, dim=0)
19
    else:
20
        attn_output = o_proj(attn_output, all_reduce_params=all_reduce_params,
21
                             lora_params=lora_params, layer_idx=layer_idx)
22
    return attn_output

enable_allreduce=False 这行是上一篇里见过的技巧，现在它有了具体的语境：不是关掉通信，而是把 all-reduce 换成 reduce-scatter。all-reduce 让所有 rank 拿到相同的完整结果；reduce-scatter 求和之后把结果按 token 段切开，每个 rank 只拿一段——正好接上 MLP 需要”每个 rank 处理不同 token”的 DP 语义。

_helix_cp_pad 是补齐操作。reduce-scatter 要求 token 总数能被 cp_size 整除，否则切片边界不整齐。

先看传统 MHA/GQA 的 KV Cache 是怎么长的

普通注意力的 KV Cache 每一层存什么？对每个 head，存一份 K 和一份 V。形状是 [num_kv_heads, seq_len, head_dim]，再乘 2（K 和 V 各一份）。

GQA 已经省了一层：多个 Q head 共享一个 KV head，num_kv_heads 从 num_heads 降到 num_kv_heads。但每个 KV head 仍要存完整的 head_dim 维 K 和 V。

MLA 的做法完全不同：它不缓存 K 和 V，而是缓存一个低秩的 latent 向量。DeepSeek-V3 的配置里 kv_lora_rank = 512，而按传统方式算的话，每层每 token 要存 128 heads × 128 dim × 2 (K,V) = 32768 个元素；MLA 只存 512 + 64 = 576 个元素（512 维的 compressed KV latent，加上 64 维的 RoPE key）。压缩率超过 50 倍。

图源：DeepSeek-V2 论文（arXiv:2405.04434）中的注意力结构对比图。自左向右依次是 MHA、GQA 与 MLA。前两者的 KV Cache 随 head 数线性增长，MLA 则把 K/V 压进一个与 head 数无关的低秩 latent —— 这正是它改变张量并行布局的根本原因：缓存里的东西不再有 head 维度。

代价是：K 和 V 要从 latent 现场”展开”出来。这就是 MLA 的张量并行必须重新设计的原因。

MLA 的权重全景

MLA.__init__ 里构造的模块比 Attention 多得多：

1
if not self.is_lite:
2
    kv_a_out_features = self.kv_lora_rank + self.qk_rope_head_dim
3
    if self.fuse_qkv_a_proj:
4
        kv_a_out_features += self.q_lora_rank
5
    self.kv_a_proj_with_mqa = Linear(
6
        hidden_size, kv_a_out_features, ...)   # 注意：没有 mapping！
7
    ...
8
    self.q_a_layernorm = RMSNorm(hidden_size=self.q_lora_rank, ...)
9

10
    self.q_b_proj = Linear(
11
        self.q_lora_rank,
12
        self.num_heads * self.qk_head_dim,
13
        mapping=mapping,
14
        tensor_parallel_mode=TensorParallelMode.COLUMN,
15
        ...)

再往下：

1
self.kv_b_proj = Linear(
2
    self.kv_lora_rank,
3
    self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
4
    bias=bias, dtype=dtype,
5
    mapping=mapping,
6
    tensor_parallel_mode=TensorParallelMode.COLUMN,
7
    ...)
8
self.v_b_proj = nn.Parameter(
9
    torch.empty(
10
        (self.num_heads_tp_cp, self.v_head_dim, self.kv_lora_rank),
11
        dtype=dtype,
12
    ),
13
    requires_grad=False,
14
)

以及两个注意力后端：

1
self.mqa = create_attention(
2
    config.attn_backend, self.layer_idx,
3
    self.num_heads_tp,
4
    head_dim=self.kv_lora_rank + self.qk_rope_head_dim,
5
    num_kv_heads=1,
6
    ...)
7
...
8
self.mha = create_attention(
9
    config.attn_backend, self.layer_idx,
10
    self.num_heads_tp,
11
    head_dim=self.qk_head_dim,
12
    num_kv_heads=self.num_key_value_heads_tp,
13
    ...)

把这几段摆在一起，MLA 的张量并行全景就出来了：

第一行的”不切”是理解 MLA 张量并行的钥匙。

为什么 kv_a_proj_with_mqa 不能切

看这个 Linear 的构造参数：它没有 mapping= 参数。

1
self.kv_a_proj_with_mqa = Linear(
2
    hidden_size,
3
    kv_a_out_features,
4
    bias=bias, dtype=dtype,
5
    quant_config=quant_config,
6
    skip_create_weights_in_init=config.skip_create_weights_in_init,
7
    use_custom_cublas_mm=True,
8
    force_dynamic_quantization=config.force_dynamic_quantization,
9
    use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm,
10
    use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm,
11
)

不传 mapping，Linear 就走 tp_size = 1 的路径，权重在每个 rank 上完整复制一份。

这是必须的，不是偷懒。kv_a_proj_with_mqa 的输出是 kv_lora_rank + qk_rope_head_dim（+ q_lora_rank，如果融合）维的 latent 向量。这个 latent 向量没有 head 的概念——它是所有 head 共享的一份压缩表示。如果按 COLUMN 模式把它切开，每个 rank 只会拿到 latent 的一部分维度，而后续的 kv_b_proj 需要完整的 latent 才能展开出本 rank 的 head，直接从数学上就错了。

所以 MLA 的并行分界线画在这里：latent 空间之前的投影不切，latent 空间之后的展开切。同样的道理也适用于 q_a_proj 和 q_a_layernorm——它们都在 q_lora_rank 这个压缩空间里工作。

kv_b_proj 的切分：COLUMN 模式为什么是对的

kv_b_proj 从 kv_lora_rank 维的 latent 展开到 num_heads * (qk_nope_head_dim + v_head_dim)。COLUMN 模式沿输出维度切，切出来的是”本 rank 的 head 对应的那段权重”。

正确性论证：展开后的 K_nope 和 V 都是 per-head 的，第 $h$ 个 head 的展开只依赖权重矩阵的第 $h$ 段行和完整的 latent。每个 rank 都有完整的 latent（因为上一步没切），所以每个 rank 都能独立完成自己那段的展开。这正好复刻了普通注意力”head 独立、无需通信”的性质。

_get_kvbproj_head_weights 展示了这个”按 head 分段”的物理布局：

1
def _get_kvbproj_head_weights(self):
2
    """Per-head views into kv_b_proj.weight, layout [H, N, K] (contiguous).
3

4
    kv_b_proj.weight rows are [H*qk_nope (k_nope block) | H*v_head (v
5
    block)] with K = kv_lora_rank, so both blocks reshape to head-major
6
    [H, N, K] without any copy.
7
    """
8
    w = self.kv_b_proj.weight
9
    nope_rows = self.num_heads_tp * self.qk_nope_head_dim
10
    w_k = w[:nope_rows].view(self.num_heads_tp, self.qk_nope_head_dim,
11
                             self.kv_lora_rank)
12
    w_v = w[nope_rows:].view(self.num_heads_tp, self.v_head_dim,
13
                             self.kv_lora_rank)
14
    return w_k, w_v

权重矩阵的行被分成两段：前 num_heads_tp * qk_nope_head_dim 行是 K_nope 的展开权重，后 num_heads_tp * v_head_dim 行是 V 的展开权重。每段都能直接 reshape 成 head-major 的 [H, N, K]，不需要任何拷贝——这意味着这个 per-head 视图可以被真正当成一个 batched GEMM 来跑。

这种”零拷贝 reshape”能成立的前提是权重加载时按 head-major 顺序排列。回到 modeling_deepseekv3.py 里的反例，能看到另一种布局：

1
k_nope_weight.reshape(local_num_heads * local_qk_nope_head_dim,
2
                      local_kv_lora_rank),
3
v_weight.reshape(local_num_heads * local_v_head_dim,
4
                 local_kv_lora_rank)

这是另一种拆分方式——把 K_nope 和 V 各自 reshape 成 [H * N, K]，也就是保持 [head, dim] 的嵌套顺序。两种布局在数学上都对，但影响后续能否直接 view 成 batched GEMM 的输入。

吸收（Absorption）：把 kv_b_proj 从推理路径上删掉

MLA 最反直觉的优化是吸收。

标准 MLA 推理流程：缓存 latent → 展开成 K_nope/V → 算注意力。但展开这一步本身有计算量，而且是每步都要做的重复劳动。

吸收的想法是：把展开矩阵乘进 Q 和输出投影里，让注意力直接在 latent 上算。

图源：TensorRT-LLM 官方技术博客的 MLA 吸收图示（docs/source/blogs/media/tech_blog3_mla_absorb.png）。图中把吸收拆成两步：第一步把 kv_b_proj 的 K 半段权重转置后提前乘到 Q 上，注意力因此在 latent 空间做内积；第二步把 V 半段权重挪到注意力输出之后，一次性把 latent 维度的注意力结果抬回 v_head_dim。这两步对应的正是代码里的 k_b_proj_trans 和 v_b_proj。

对 Q 侧，注意力内积是 $q_h^\top k_h$，其中 $k_h = W_{kb}[h] \cdot c$（$c$ 是 latent）。于是：

右边把 $W_{kb}[h]^\top$ 乘到 $q_h$ 上，得到一个和 latent 同维的向量，再和 latent 做内积。这就是代码里的 k_b_proj_trans。

对 V 侧同理：输出是 $\sum_h o_h \cdot W_{vb}[h]$，把 $W_{vb}[h]$ 移到注意力输出之后，注意力输出保持 latent 维度，最后一次性做 per-head BMM 抬回 v_head_dim。

forward_absorption_generation 就是这条吸收路径：

1
def forward_absorption_generation(self, q, compressed_kv, k_pe, attn_metadata,
2
                                  output, position_ids=None,
3
                                  latent_cache=None, **kwargs):
4
    num_tokens = q.shape[0]
5
    q_nope, q_pe = q.view([-1, self.num_heads_tp, self.qk_head_dim]).split(
6
        [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
7
    ...
8
    if hasattr(self, "k_b_proj_trans"):
9
        fused_q = torch.empty(
10
            [num_tokens, self.num_heads_tp,
11
             (self.kv_lora_rank + self.qk_rope_head_dim)],
12
            dtype=q.dtype, device=q.device)

fused_q 的形状是 [num_tokens, num_heads_tp, kv_lora_rank + qk_rope_head_dim]——注意第二维仍是本 rank 的 head 数，num_heads_tp。head 并行的结构没有被动摇，被压缩的只是每个 head 的向量宽度：从 qk_head_dim 变成 latent 宽度。

吸收后的 Q 用 BMM 算出来：

1
if self.k_b_proj_trans.dtype == torch.bfloat16:
2
    # [num_heads, num_tokens, self.qk_nope_head_dim]
3
    q_nope_t = q_nope.transpose(0, 1)
4
    # [num_heads, num_tokens, self.kv_lora_rank]
5
    q_nope_out = fused_q[..., : self.kv_lora_rank].transpose(0, 1)
6

7
    # [num_heads, num_tokens, qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim]
8
    # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank]
9
    # The output of bmm is written directly into fused_q
10
    maybe_execute_in_parallel(
11
        lambda: self._bmm_bf16_out(
12
            q_nope_t, self.k_b_proj_trans,
13
            self.k_b_proj_trans.transpose(1, 2),
14
            q_nope_out, self._locality_domain_k_b_proj_trans_shards),
15
        _mla_gen_rope,
16
        self.ln_events[0], self.ln_events[1], rope_stream)

逐行读这段：q_nope 是 [tokens, heads, qk_nope_head_dim]，转置成 [heads, tokens, qk_nope_head_dim]。k_b_proj_trans 的形状是 [heads, kv_lora_rank, qk_nope_head_dim]——正是 $W_{kb}[h]^\top$ 的 head-major 堆叠。两者做 batched matmul，得到 [heads, tokens, kv_lora_rank]，再转置回 [tokens, heads, kv_lora_rank] 直接写进 fused_q 的前段。

bmm 的第一维是 head，而且 k_b_proj_trans 的第 0 维是 num_heads_tp——每个 rank 的 batched GEMM 只用自己那份 head 权重。这就是吸收路径下张量并行的具体形态：head 维同时是并行维和 batch 维。

_has_full_absorption_weights 里的形状断言把这一点钉死了：

1
def _has_full_absorption_weights(self):
2
    if getattr(self, "kv_lora_rank", None) != 512 or \
3
       getattr(self, "v_head_dim", None) != 128:
4
        return False
5
    # The Rubin BF16 BMM requires its K dimension to be 16-byte aligned.
6
    # qk_nope_head_dim is K for the absorbed K projection.
7
    if getattr(self, "qk_nope_head_dim", 0) % 8 != 0:
8
        return False
9
    ...
10
    expected_k_shape = (self.num_heads_tp, self.kv_lora_rank,
11
                        self.qk_nope_head_dim)
12
    expected_v_shape = (self.num_heads_tp_cp, self.v_head_dim,
13
                        self.kv_lora_rank)
14
    return (...)

k_b_proj_trans 的第 0 维是 num_heads_tp，v_b_proj 的第 0 维是 num_heads_tp_cp。这两个变量的差别不是笔误——它精确反映了 TP 和 CP 在两个吸收矩阵上的不同处理，和前面 Attention.__init__ 里 num_heads_tp / num_heads_tp_cp 的划分是同一套逻辑。

再看 v_b_proj 的构造位置：

1
self.v_b_proj = nn.Parameter(
2
    torch.empty(
3
        (self.num_heads_tp_cp, self.v_head_dim, self.kv_lora_rank),
4
        dtype=dtype,
5
    ),
6
    requires_grad=False,
7
)

它是从 kv_b_proj.weight 里视图出来的参数，不是独立权重——_bind_v_b_proj_weight 做这件事。这解释了为什么它会出现在 need_absorption 分支里：只有走吸收路径时，才需要这个 per-head 的 V 展开视图。

mqa 和 mha：两个后端，分别对应两条路径

最后看那两个注意力后端实例的构造参数差异：

1
self.mqa = create_attention(..., head_dim=self.kv_lora_rank + self.qk_rope_head_dim,
2
                            num_kv_heads=1, is_mla_enable=True, ...)
3
self.mha = create_attention(..., head_dim=self.qk_head_dim,
4
                            num_kv_heads=self.num_key_value_heads_tp, ...)

mqa（Multi-Query Attention 后端）的 head_dim 是 kv_lora_rank + qk_rope_head_dim，num_kv_heads=1。这两个参数一起说明：它处理的正是吸收后的表示——latent 维度的 head，所有 head 共享同一份 latent（所以是”1 个 KV head”，即 MQA 式的共享）。

mha 的 head_dim 是 qk_head_dim = qk_nope_head_dim + qk_rope_head_dim，num_kv_heads 是本 rank 的 KV head 数。这是展开后的标准注意力表示。

mqa 永远被创建（if self.mqa is None: raise RuntimeError(...)），mha 则可能为 None：

1
create_dense_mha = (self.sparse_attn_hooks is None or
2
                    self.sparse_attn_hooks.need_dense_mha)
3
if create_dense_mha:
4
    self.mha = create_attention(...)
5
else:
6
    self.mha = None

这说明吸收路径（走 mqa）是主路径。上下文阶段需要展开到 mha 走标准 FlashAttention，而生成阶段走 mqa 的吸收路径——这也符合”上下文阶段算一次、生成阶段每步都算”的收益结构：吸收省下的是每步都要付的展开开销。