KDA 门控投影融合与双流重叠

Kimi K3 性能优化 PR #15416 已合入 main 作者 Dawn952 · 2026-09-04 · 改动 4 文件 +597/-52

一、需求背景

二、相关流程

def forward(  # ↳ 接口与上游 KimiK3DeltaAttention 完全一致（整组件替换，机制三）
    self,
    hidden_states: torch.Tensor,
    positions: torch.Tensor,
) -> torch.Tensor:
    if self.uses_mixed_projection:  # ↳ mixed-precision 全秩门控布局才走本路径；普通上游路径不受影响
        num_tokens = hidden_states.size(0)
        mixed_qkv = self.in_proj_qkvgfab(hidden_states)[0]  # 主路：QKV 融合投影（大 GEMM）
        projected_gfab = self.in_proj_gfab(hidden_states)[0]  # ⚡优化前问题：门控侧投影排在其后串行执行
        split_sizes = [
            self.local_projection_size,
            self.head_dim,
            self.local_num_heads,
        ]
        if self.in_proj_padding:
            split_sizes.append(self.in_proj_padding)
        g_proj_states, f_a, beta = projected_gfab.split(split_sizes, dim=-1)[:3]  # 拆出输出门控 g / F 中间 / beta 三路
        beta = beta.unsqueeze(0)

        g1 = self.f_b_proj(f_a)[0]  # ⚡f_b/f_a 两个独立小投影（本需求将其离线组合掉）
        g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)
        g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
        core_attn_out = torch.empty(  # 输出缓冲：形状 (1, tokens, heads, head_dim)
            (1, num_tokens, self.local_num_heads, self.head_dim),
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
        self._forward(  # ⚡分派 spec/prefill/decode 三路到 AscendC 算子（→NPU）
            mixed_qkv=mixed_qkv,
            g1=g1,
            g2=g2,
            beta=beta,
            core_attn_out=core_attn_out,
        )
        core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
        return self.o_proj(core_attn_out)[0]  # 输出投影回 hidden 维度
    return super().forward(hidden_states, positions)

三、需求改动

    def _maybe_fuse_f_proj(self) -> None:  # ⚡改动①：f_a/f_b 权重加载完成后离线组合成 f_proj
        if not self._f_a_loaded or not self._f_b_loaded:  # 两份权重都到位才组合（各自加载器回调里触发）
            return
        output_dim = getattr(self.weight, "output_dim", None)
        if output_dim is None:
            raise ValueError("KDA fused f_proj requires an output-sharded parameter")
        shard_offset = sum(self.output_sizes[:_F_PROJ_SHARD_ID]) // self.tp_size
        shard_size = self.output_sizes[_F_PROJ_SHARD_ID] // self.tp_size
        param_shard = self.weight.narrow(output_dim, shard_offset, shard_size)
        fused_weight = torch.matmul(  # 离线 f_b @ f_a：运行期两个小投影 GEMM 合成一个
            self.f_b_weight.float(),
            self.f_a_weight.float(),
        ).to(dtype=param_shard.dtype)
        if fused_weight.shape != param_shard.shape:
            raise ValueError(
                "KDA composed f_proj shape mismatch: "
                f"expected {tuple(param_shard.shape)}, got {tuple(fused_weight.shape)}"
            )
        param_shard.copy_(fused_weight)  # 写回 fused_bfg_proj 的 TP 分片槽位


def _zero_padded_output(


def _prepare_beta(
    beta: torch.Tensor,

def _run_overlapped_qkv_bfg(  # ⚡改动③：门控侧整体搬到辅助流，与主流两阶段重叠
    self,
    hidden_states: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Fork BFG work to an auxiliary stream and join it back to main."""
    main_stream = torch.npu.current_stream()  # 主流：承载 DynamicQuant 与 QKV GEMM
    bfg_stream = _kda_bfg_stream()  # 辅助流：门控侧专用

    hidden_states_ready = main_stream.record_event()  # 事件①：从主流分叉
    hidden_states.record_stream(bfg_stream)
    with npu_stream_switch(bfg_stream):
        bfg_stream.wait_event(hidden_states_ready)
        fused_bfg = self._project_bfg(hidden_states)  # 辅流：BFG 融合投影（一次 GEMM）
        bfg_projection_ready = bfg_stream.record_event()  # 事件②：BFG 与 QKV 两块 Cube 乘保持串行

    quantized_qkv = self._quantize_fused_qkv(hidden_states)  # 主流：MXFP8 DynamicQuant（与辅流 BFG GEMM 重叠）
    quant_ready = main_stream.record_event()

    # Stage 1 join: DynamicQuant on main overlaps the BFG GEMM on the
    # auxiliary stream, but the two Cube matmuls remain serialized.
    main_stream.wait_event(bfg_projection_ready)  # Stage1 join：DynamicQuant 完成后再排 QKV GEMM
    mixed_qkv = self._matmul_fused_qkv(quantized_qkv)  # 主流：QKV GEMM（辅流小算子在旁重叠执行）

    with npu_stream_switch(bfg_stream):
        # Stage 2: after both first-stage branches complete, overlap the
        # QKV Cube matmul with beta's FP32 conversion and sigmoid vector
        # work. Split and reshape the F/output gates here as well so all
        # BFG output handling occurs after the QKV matmul is enqueued.
        bfg_stream.wait_event(quant_ready)  # 事件③：辅流 beta 向量活等主流量化完成
        beta, g1, g2 = self._postprocess_bfg(fused_bfg)  # 辅流重叠期：B/F/G 拆分 + beta sigmoid + gate reshape
        bfg_ready = bfg_stream.record_event()

    for tensor in (beta, g1, g2):
        tensor.record_stream(main_stream)
    # bfg_ready is the auxiliary stream tail. Joining that exact event is
    # required for multi-stream ACL graph capture as well as eager reuse.
    main_stream.wait_event(bfg_ready)  # 事件④：KDA 消费前辅流尾部全部到齐
    return mixed_qkv, beta, g1, g2

def _project_bfg(
    self,
    hidden_states: torch.Tensor,
) -> torch.Tensor:
    return self.fused_bfg_proj(hidden_states)[0]

def _postprocess_bfg(
    self,
    fused_bfg: torch.Tensor,

四、相关链接