Kimi K3 三平台架构对比解读

同一个 Kimi K3 模型，在三个硬件平台上需要不同的底层算子实现。vLLM 的设计是：

• NVIDIA 和 AMD 是上游 vLLM 的内建实现，各自独立完整地实现了整个模型
• Ascend 是外部硬件插件，通过继承上游类、替换 NPU 特定组件来适配

上游入口：vllm/models/kimi_k3/__init__.py

上游通过检测当前平台，自动路由到对应的实现：

# vllm/models/kimi_k3/__init__.py
from vllm.platforms import current_platform

if TYPE_CHECKING:
    from .nvidia.model import KimiK3ForConditionalGeneration
elif current_platform.is_rocm():
    from .amd.model import KimiK3ForConditionalGeneration  # AMD
elif current_platform.device_type == "tpu":
    pass  # TPU插件自己注册，不在此加载
else:
    from .nvidia.model import KimiK3ForConditionalGeneration  # NVIDIA

Ascend 入口：vllm_ascend/models/kimi_k3.py

Ascend 插件绕过上述路由，直接继承上游 AMD 版本的类，并替换关键组件：

# vllm_ascend/models/kimi_k3.py 的导入链
from vllm.models.kimi_k3.amd.linear import KimiDecoderLayer as UpstreamKimiDecoderLayer
from vllm.models.kimi_k3.amd.linear import KimiMLAAttention as UpstreamKimiMLAAttention
from vllm.models.kimi_k3.amd.model import KimiK3ForConditionalGeneration as UpstreamKimiK3ForConditionalGeneration

from vllm_ascend.ops.kimi_kda import AscendKimiK3DeltaAttention  # NPU KDA后端
from vllm_ascend.ops.triton.kimi_k3.attention_residual import apply_attn_res  # NPU attn_res

Attention Residual 是 Kimi K3 独有的机制：将传统的加法残差替换为学习式 softmax 加权混合。三平台的数学逻辑完全相同，但算子实现不同。

关键代码对比：残差计算核心

# ─── AMD 版 (amd/linear.py:141) ─── 调用 ROCm Triton kernel
def _apply_attn_res(prefix_sum, block_residual, proj, norm, num_valid_blocks, ...):
    return attn_res(prefix_sum, delta, block_residual,
                    norm.weight, proj.weight.squeeze(0), ...)

# ─── Ascend 版 (kimi_k3.py:92) ─── 优先 NPU Triton，回退纯 PyTorch
def _apply_ascend_attn_res(prefix_sum, block_residual, proj, norm, num_valid_blocks):
    if apply_attn_res is not None and prefix_sum.device.type == "npu":
        return apply_attn_res(...)                    # ← NPU 加速路径
    # 纯 PyTorch 回退（数学等价）
    values = torch.cat((block_residual[:,:num_valid_blocks,:], prefix_sum.unsqueeze(1)), dim=1)
    values_fp32 = values.float()
    inverse_rms = torch.rsqrt(values_fp32.square().mean(-1, keepdim=True) + norm.variance_epsilon)
    normalized_without_gamma = values_fp32 * inverse_rms
    score_weight = norm.weight.float() * proj.weight.squeeze(0).float()
    scores = (normalized_without_gamma * score_weight).sum(-1)
    probabilities = scores.softmax(-1).unsqueeze(1)
    return torch.matmul(probabilities, values_fp32).squeeze(1).to(values.dtype)

Decoder Layer 是差异最大的组件。三平台的核心分歧在于 Attention 后端的选择 和 残差路径的处理。

继承策略图例

规律：越靠近底层（算子级），越需要替换；越靠近上层（模型结构、权重加载），越可以继承复用。

一图总结：代码复用 vs 硬件替换

最终结论

不是每个平台"都写一遍"同一份代码，而是：

• NVIDIA 和 AMD 各自完整独立实现（因为它们的算子生态完全不同，且是上游内建）
• Ascend 通过继承 AMD 版本，只替换 NPU 特有的底层算子和内存管理策略，上层的模型结构、权重加载、前向框架全部复用
• 越靠近算子层，差异越大（必须替换）；越靠近模型层，共性越多（可以继承）
• 如果未来 Ascend 被合并进上游，它也会变成 vllm/models/kimi_k3/ascend/ 下的一套完整实现，届时代码量会大幅增加

图9-1 把一次推理分成了六步，其中第④步「模型逐层计算」占 95% 以上的计算量。 接下来两节只聚焦它：9.2 放大看 30 层里一层的实际调用时序（顺便看清哪些代码来自 AMD 基类、 哪些是昇腾新写的）；9.3 顺着时序钻进源码。

先快速过一遍全流程（上一版已有，紫色是本章主角）：

第④步 = 30 次循环，每次调 AscendKimiDecoderLayer.forward_attn_residual （models/kimi_k3.py:478，这个类继承自 AMD 的 KimiDecoderLayer，同名方法被覆盖重写）。 下面是一次层计算的完整时序——注意三列参与者的来源：中间是昇腾写的，右列是 NPU 底层算子， 方法名旁标 ↳ 的行为/组件直接来自 AMD 基类。

那"依赖 AMD"到底是什么意思？——三种机制，对应到代码

# 机制一：整类继承 —— 类体是空的或只加参数，方法全部来自 AMD 基类
class AscendKimiMLAAttention(UpstreamKimiMLAAttention):   # Upstream = vllm/models/kimi_k3/amd/linear.py
    def forward(self, positions, hidden_states):
        return self.mla_attn(positions, hidden_states)     # 只改一个委托点

# 机制二：继承 + 覆盖 —— 方法签名沿用 AMD 基类，实现重写
class AscendKimiDecoderLayer(UpstreamKimiDecoderLayer):
    def forward_attn_residual(self, ...): ...              # 同名同序，换掉残差函数/注意力分派

# 机制三：整组件替换 —— 上游类彻底不用，接口对齐
self.self_attn = AscendKimiK3DeltaAttention(...)           # AMD 的 KimiK3DeltaAttention 被丢弃

一句话：AMD 提供的是"层的组织方式"（方法顺序、组件接线、权重加载），昇腾提供的是"每个动作怎么算"。 时序图里的编号 ①→⑤ 就是 AMD 基类定下的顺序；每个编号算什么，红色算子说了算。

时序图 alt 分支里 KDA 是最特殊的一条（上游没有、AMD 用 ROCm Triton、昇腾用自研 AscendC 算子）。 下面三段源码就是时序图 消息 4 → 6 → 7 的真身（截自 ops/kimi_kda.py， ⚡ 昇腾新逻辑 / ↳ 与 AMD·上游的关系 / →NPU 底层算子，见行内注释）。

def forward(  # ↳ 接口与上游 KimiK3DeltaAttention 一致（整组件替换，机制三）
    self,
    hidden_states: torch.Tensor,
    positions: torch.Tensor,
) -> torch.Tensor:
    if self.uses_mixed_projection:
        num_tokens = hidden_states.size(0)
        mixed_qkv = self.in_proj_qkvgfab(hidden_states)  # ⚡ 融合投影：一次出 q/k/v/门控多路[0]
        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)

def _forward(
    self,
    mixed_qkv: torch.Tensor,
    g1: torch.Tensor,
    g2: torch.Tensor,
    beta: torch.Tensor,
    core_attn_out: torch.Tensor,
) -> None:
    """Dispatch speculative, prefill, and decode tokens through KDA kernels."""
    forward_context = get_forward_context()
    attn_metadata_raw = forward_context.attn_metadata
    if attn_metadata_raw is None:
        core_attn_out.zero_()
        return

    assert isinstance(attn_metadata_raw, dict)
    attn_metadata = attn_metadata_raw[self.prefix]
    assert isinstance(attn_metadata, GDNAttentionMetadata)

    num_actual_tokens = attn_metadata.num_actual_tokens
    mixed_qkv = mixed_qkv[:num_actual_tokens]
    g1 = g1[:, :num_actual_tokens]

def _run_causal_conv1d(
    mixed_qkv: torch.Tensor,
    conv_weights_t: torch.Tensor,
    conv_state: torch.Tensor,
    query_start_loc: torch.Tensor,
    cache_indices: torch.Tensor,
    initial_state_mode: torch.Tensor | None,
    *,
    run_mode: int,
    num_accepted_tokens: torch.Tensor | None = None,
) -> torch.Tensor:
    output = torch.empty_like(mixed_qkv)
    # Consume the operator's declared output alias. Returning ``output``
    # independently would let graph functionalization treat the custom-op
    # result as dead and expose the uninitialized allocation instead.
    return torch.ops._C_ascend.npu_causal_conv1d_custom(  # →NPU：自研 AscendC 算子，CANN 调度到 AI Core
        output,
        mixed_qkv,
        conv_weights_t,
        conv_state=conv_state,
        bias_opt=None,
        query_start_loc_opt=query_start_loc,
        cache_indices_opt=cache_indices,
        initial_state_mode_opt=initial_state_mode,

读完这三段，时序图右边一列（NPU 底层算子）就有了实体： torch.ops._C_ascend.* 是 vllm-ascend 自己编译的 C++ 算子库（AscendC 语言编写、随包发布）， 调用之后编译与调度全部交给 CANN 落到 AI Core。MLA 分支同理，走的是第二个通道 torch_npu（第 10 章）。

第 9 章调用链里出现了三种"红色"底层调用：torch.ops._C_ascend.*、torch_npu.npu_*、 @triton.jit。它们就是 vllm-ascend 与"Ascend 底层"的三个接口。这一章把它们与 CANN、硬件串成一条线。

三个算子通道怎么区分（对应真实代码）

"和 Ascend 底层啥关系"——一次 KDA 计算的完整下沉过程

# ① Python 层：ops/kimi_kda.py（文档里能打开的源码）
def _run_prefill(self, ...):
    ...准备 q/k/v/beta 并按 AscendC 契约整理（_prepare_beta）...
    return torch.ops._C_ascend.chunk_kda_fwd(q, k, v, g, beta, ...)   # ← 跨过 Python/USA 边界

# ② C++ 层：vllm_ascend_C 扩展（随包发布的 .so，源码在仓库 vllm_ascend/C 下）
#    chunk_kda_fwd 的 C++ 实现负责：参数校验 → 取 NPU Stream → launch AscendC kernel

# ③ CANN 层：AscendC kernel（设备侧 C++/类 CUDA 语法）编译产物在 AI Core 上执行
#    Host 侧由 libascendcl（acl）管理设备上下文、内存、Stream 任务下发

# ④ 硬件层：AI Core 的 Cube 单元做矩阵乘（q·k、状态更新），Vector 单元做激活/门控

把两个问题串起来（本章结论）

• "看不到 Ascend 实现全景"：因为 70% 的实现是"继承来的"（引擎/runner/模型骨架），真正的 Ascend 代码集中在 worker/（执行器）、attention/（注意力后端）、ops/（算子）、models/（适配层）四处——第 9 章调用链把它们挂回了同一条链上。
• "和 Ascend 底层串不起来"：中间缺的是"算子通道"这一环——vllm-ascend 从不直接摸硬件， 它只会走 torch.ops._C_ascend（自研 AscendC）、torch_npu（华为 aclnn）、 triton-ascend 三个入口，再由 CANN 完成编译与调度落到 AI Core（本章图 10）。