LLaMA-Factory NPU 融合算子指南：RMSNorm、SwiGLU、RoPE 与 MoE 的硬件加速实现

本篇技术指南围绕 LLaMA-Factory（LlamaFactory）v1 插件体系中的融合算子（Fused Operators）展开：先讲清融合算子的注册与使能机制，再逐一深入 Ascend NPU 上四个核心融合算子——NpuFusedRMSNorm、NpuFusedSwiGLU、NpuFusedRoPE、NpuFusedMoE——的源码实现、适配模型范围与 apply_kernel 调用方式。读完本文，你可以在 NPU 训练环境中正确启用这些硬件加速算子，并了解如何为未适配的模型自行扩展新的融合算子。

一、融合算子体系总览：算子在哪里、如何注册

LLaMA-Factory 提供了一系列针对特定硬件优化的融合算子，核心实现位于 ops 目录 下，按算子类别分子目录组织：

• rms_norm/：NPU 融合 RMSNorm；
• mlp/：SwiGLU、融合 MoE（含 NPU 与 CUDA/Triton 实现）；
• rope/：NPU 融合 RoPE；
• linear_attention/：线性注意力相关算子。

文档中提到系统启动时会自动扫描该目录并注册所有可用算子。从当前仓库源码结构看，注册机制的具体实现是导入期自动注册：每个算子在模块底部使用装饰器 @KernelPlugin("<kernel_id>").register() 将实现类登记进注册表，而 interface.py 在导入时显式导入所有内置实现，从而让装饰器在模块加载阶段即完成注册：

# Import built-in implementations so their class decorators populate the registry.
from .liger_kernel_ops import LigerKernel  # noqa: F401
from .ops.linear_attention.fla import FlashLinearAttentionKernel  # noqa: F401
from .ops.mlp.cuda_fused_moe import CudaFusedMoEKernel  # noqa: F401
from .ops.mlp.npu_fused_moe import NpuFusedMoEKernel  # noqa: F401
from .ops.mlp.npu_swiglu import NpuSwiGluKernel  # noqa: F401
from .ops.rms_norm.npu_rms_norm import NpuRMSNormKernel  # noqa: F401
from .ops.rope.npu_rope import NpuRoPEKernel  # noqa: F401

所有融合算子统一继承 BaseKernel 抽象基类。该基类通过 __init_subclass__ + ensure_methods_implemented 强制子类实现 check_device() 与 _apply() 两个静态方法，并提供统一的 apply() 入口，其执行顺序为：

@classmethod
def apply(cls, **kwargs) -> HFModel:
    cls.check_device()   # 校验当前加速器类型（如必须是 NPU）
    cls.check_deps()     # 校验依赖（如 torch_npu 是否可导入）
    if kwargs.get("model") is None:
        raise ValueError(f"HFModel instance is required for {cls.__name__}.")
    return cls._apply(**kwargs)

这个“先校验、后打补丁”的两段式设计保证了：在非 NPU 环境或未安装 torch_npu 时，apply_kernel("npu_fused_...") 会抛出明确报错，而不是静默失败。

二、使能入口：apply_kernel 与 auto 模式

文档给出的使能方式是调用 apply_kernel 单独启用某个算子，或使用 auto 一键启用。对照当前源码 interface.py，仓库中实际提供的入口函数为：

def apply_kernel(kernel_id: str, **kwargs) -> HFModel:
    if kernel_id == "auto":
        return _apply_auto_kernels(**kwargs)
    return KernelPlugin(kernel_id).apply(**kwargs)

其中 auto 模式的设备映射关系如下：

_AUTO_KERNELS = {
    DeviceType.NPU: ("npu_fused_moe", "npu_fused_rmsnorm", "npu_fused_rope", "npu_fused_swiglu"),
}

即：当当前加速器为 Ascend NPU 时，apply_kernel("auto", model=model) 会依次应用全部四个 NPU 融合算子；其他设备类型没有内置 auto 列表。

在训练流水线中，算子由 kernel_config 参数驱动。从 model_engine.py 可以看到，模型加载完成后、进入训练前会执行：

if self.args.kernel_config is not None:
    from ..plugins.model_plugins.kernels.interface import apply_kernels
    model = apply_kernels(model, self.args.kernel_config, require_logits=self.is_train)

apply_kernels 接受 kernel_config.name 字段（逗号分隔的字符串，支持 auto），因此实践中可以通过配置一次性组合多个算子，例如 "npu_fused_rmsnorm,npu_fused_rope"，也可以直接用 Python API 针对单个模型调用：

from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernel

model = apply_kernel("npu_fused_rmsnorm", model=model)

启用后，算子会在日志中报告实际打补丁的模块数量，例如 Applied NPU RMSNorm kernel to N modules for model type: qwen3，可用于确认融合算子是否真正生效于当前模型。

三、NpuFusedRMSNorm：融合残差与归一化

RMSNorm（Root Mean Square Layer Normalization）是 Qwen 等大模型中最常见的归一化方式。原生 PyTorch 实现中，RMSNorm 需要多次独立的显存读写（求均方、归一化、乘权重）；NPU 融合算子将 bias、residual 等操作折叠为一次硬件指令，减少显存访问次数，从而加速计算。Ascend NPU 通过 torch_npu.npu_rms_norm 接口提供该融合算子，支持 float16、bfloat16、float 等数据格式。

由于 PyTorch 侧没有内置 RMSNorm 算子接口，模型中通常以自定义模块类出现，因此只需替换该类的 forward 方法即可使能。文档给出的核心替换示例：

def _npu_rms_forward(self, hidden_states):
    """NPU forward implementation for RMSNorm.

    Args:
        self: RMSNorm module instance with `weight` and `variance_epsilon`.
        hidden_states: Input hidden states tensor, same shape as the baseline.

    Returns:
        Normalized tensor consistent with the baseline RMSNorm behavior.
    """
    return torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=self.variance_epsilon)[0]

仓库中的完整实现位于 npu_rms_norm.py，相比文档示例做了更细致的工程化处理，值得注意的实现细节有：

1. eps 的兼容取值：_eps = getattr(self, "variance_epsilon", None) or getattr(self, "eps", 1e-6)，兼容不同模型类使用 variance_epsilon 或 eps 两种属性命名的情况；
2. 权重精度对齐：effective_weight = weight.float().to(hidden_states.dtype)，先将权重转 float32 再对齐到输入 dtype，保证数值稳定；
3. 三种变体。除标准 RMSNorm 外，还实现了：
   • npu_residual_rms_norm_forward：残差 RMSNorm，使用 scale = 1.0 + weight（原始实现中 weight 初始化为 0）；
   • npu_gated_rms_norm_forward：门控 RMSNorm，全程 FP32 计算归一化与 F.silu(gate) 相乘后转回输入 dtype，用于高数值精度场景。

_MODEL_TYPE_TO_PATCHES 表按 model_type → {模块类名: 替换函数} 映射，目前覆盖 qwen3、qwen3_moe、qwen3_next（含残差/门控两种变体）、qwen3_omni_moe（含 thinker/text/code2wav 多个文本塔）、qwen3_vl、qwen3_vl_moe、qwen3_5、qwen3_5_moe 等模型族。_apply 通过 types.MethodType(patch_forward, module) 把 NPU 前向函数以实例方法形式绑定到匹配模块上，替换其 forward。

使能方式：apply_kernel("npu_fused_rmsnorm", model=model)。对不在映射表中的 model_type，函数会原样返回模型，不产生任何副作用。

四、NpuFusedSwiGLU：门控激活的算子融合

SwiGLU（Swish-Gated Linear Unit）结合了 Swish 激活与门控线性单元，近年来被广泛应用于 LLM 的 MLP 层。原生实现中，一个 SwiGLU MLP 前向涉及 gate_proj、up_proj 两次矩阵乘、一次 Swish 激活、一次逐元素乘、一次 down_proj，产生多次内核启动与中间张量的显存往返。SwiGLU 融合算子将分割、激活等步骤合并为单一硬件指令，避免多次内核启动开销。

Ascend NPU 通过 torch_npu.npu_swiglu 接口提供该算子，支持 float16、bfloat16、float 等格式。文档给出的替换示例如下：

# 原始 MLP forward 方法：
def forward(self, x):
    down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
    return down_proj

# 替换后的 forward 方法：
def _npu_swiglu_forward(self, hidden_state):
    return self.down_proj(
        torch_npu.npu_swiglu(torch.cat((self.gate_proj(hidden_state), self.up_proj(hidden_state)), dim=-1), dim=-1)
    )

仓库实现见 npu_swiglu.py，替换函数与文档示例一致。额外值得注意的一点是 _get_patch_forward 中存在一个激活函数校验：

config = getattr(module, "config", None)
if getattr(config, "hidden_act", None) != "silu":
    return None

也就是说，只有当 MLP 的 hidden_act 为 silu 时才会打补丁——因为 npu_swiglu 的硬件语义即 Swish-Gated 激活，若模型的激活函数并非 silu，融合结果将不再与原实现等价，此处直接跳过以保证正确性。适配的模型类名与 RMSNorm 一节相同，覆盖 Qwen3MLP、Qwen3MoeMLP、Qwen3NextMLP、Qwen3OmniMoe*MLP、Qwen3VLTextMLP、Qwen3_5MLP 等。

使能方式：apply_kernel("npu_fused_swiglu", model=model)。对于未适配的模型，可以参照 NPU 算子开发者文档 与上述示例自行扩展 _MODEL_TYPE_TO_PATCHES。

五、NpuFusedRoPE：把位置编码下沉到硬件

RoPE（Rotary Positional Embedding，旋转式位置嵌入）广泛应用于 Qwen 等 LLM，用于编码序列位置信息，兼具绝对位置编码的稳定性与相对位置编码的灵活性，并有优秀的长度泛化能力。传统实现中 RoPE 通常以自定义函数形式存在，包含 unsqueeze、rotate_half、多次逐元素乘加，是注意力路径上的一个可融合热点。Ascend NPU 通过 torch_npu.npu_rotary_mul 提供融合接口，支持 float16、bfloat16、float32。

文档给出的以 Qwen3 为例的替换示例：

# 原始 apply_rotary_pos_emb：
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

# 替换 RoPE 融合算子后：
def _apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = torch_npu.npu_rotary_mul(q, cos, sin)
    k_embed = torch_npu.npu_rotary_mul(k, cos, sin)
    return q_embed, k_embed

仓库实现见 npu_rope.py，与前两个算子相比，它的打补丁方式有本质不同：RMSNorm 和 SwiGLU 替换的是模块实例的 forward，而 RoPE 的 apply_rotary_pos_emb 是 transformers modeling 模块里的自由函数。因此 NpuRoPEKernel._apply_model_patches 采用模块级 monkey-patch：

def _default_rope_patch(module_type: str):
    return (
        (
            f"transformers.models.{module_type}.modeling_{module_type}",
            (("apply_rotary_pos_emb", _apply_rotary_pos_emb),),
        ),
    )

即按 model_type 动态导入 transformers.models.{type}.modeling_{type} 模块，用 setattr 替换其 apply_rotary_pos_emb 函数；若目标模块导入失败或不含该函数，则只记录 warning 并跳过，不会中断训练。源码在文档示例之外还处理了两个兼容性细节：

1. Partial RoPE 自动检测：当 cos/sin 的宽度小于 q/k 的 head 维度时，只对前 rotary_dim 维执行 npu_rotary_mul，尾部维度直接拼接透传（见 _apply_npu_rotary_emb）；
2. transformers v4/v5 签名兼容：v5 中第五个位置参数变为 unsqueeze_dim，函数通过 isinstance(position_ids, int) 判断并相应调整，保证两个大版本下都能被正确调用。

使能方式：apply_kernel("npu_fused_rope", model=model)。未适配模型可按 NPU 算子开发者文档 中的规范，向 _MODEL_TYPE_TO_PATCHES 添加新的 (模块路径, 函数替换) 映射即可。

六、NpuFusedMoE：基于 GMM 的专家并行计算

MoE（Mixture of Experts）模型通过稀疏激活扩展参数容量。原生 Transformers 实现中，专家计算常用串行循环逐个遍历，内核启动开销大、硬件利用率低。NPU MoE 融合算子利用 GMM（Grouped Matrix Multiplication，分组矩阵乘） 技术，在单个硬件指令内并行处理多组行数不一的矩阵乘法，消除 Python 层循环开销，且无需额外显存复制，显著提升训练性能。Ascend NPU 通过 torch_npu.npu_grouped_matmul、torch_npu.npu_moe_token_permute / npu_moe_token_unpermute 等接口提供底层支持。

文档给出的核心逻辑替换示意（简化版）：

def _npu_moe_forward(self, hidden_states, routing_weights, router_indices):
    # 1. 排序：将乱序的 Token 按指派的专家归类，并生成索引映射
    permuted_states, row_map = torch_npu.npu_moe_token_permute(hidden_states, router_indices)

    # 2. 统计：计算每个专家需要处理的 Token 数量
    tokens_per_expert = torch.histc(router_indices, bins=self.num_experts, min=0, max=self.num_experts)

    # 3. 计算 (GMM)：一次性并行计算所有专家的权重，自动适配不同专家的输入长度
    inter_states = torch_npu.npu_grouped_matmul(permuted_states, self.gate_up_proj_weights, split_sizes=tokens_per_expert, ...)
    inter_states = torch_npu.npu_swiglu(inter_states)
    output = torch_npu.npu_grouped_matmul(inter_states, self.down_proj_weights, split_sizes=tokens_per_expert, ...)

    # 4. 还原：将结果恢复成原始 Token 顺序并应用路由权重
    return torch_npu.npu_moe_token_unpermute(output, row_map, routing_weights)

仓库中的真实实现见 npu_fused_moe.py，其工程复杂度明显高于文档示意，主要包括：

1. GMM 的自动微分封装：GmmFunction(torch.autograd.Function) 在 forward 中调用 npu_grouped_matmul(..., split_item=2, group_type=0, group_list_type=1)，在 backward 中分别计算 grad_input（split_item=2, group_type=0）与 grad_weight（split_item=3, group_type=2），保证训练时的反向传播完全在融合算子内完成；
2. HybridGmmFunction：针对以“每专家独立输入列表”形式调用的场景（split_item=0, group_type=-1），forward 输出各专家的独立结果元组，backward 中先逐专家计算 grad_x，再拼接后以 group_list 方式统一计算 grad_w；
3. 多结构前向变体：NpuMoeFusedV4 提供 stacked_experts_forward（权重堆叠为单张量的 experts 结构）、sparse_block_forward（列表式 experts，配合 HybridGmmFunction 与 F.silu(gate) * up）、shared_sparse_block_forward（带共享专家与 sigmoid 门控的 Qwen3-Next 结构）等实现；
4. transformers v4/v5 双轨适配：_MODEL_TYPE_TO_PATCHES 根据 is_transformers_version_greater_than("5.0.0") 选择版本表。v4 替换 Qwen3MoeSparseMoeBlock、Qwen3NextSparseMoeBlock、Qwen3OmniMoe*SparseMoeBlock、Qwen3VLMoeTextExperts 等模块；v5 中专家权重改为堆叠布局（gate_up_proj: [num_experts, 2*intermediate, hidden]），NpuMoeFusedV5.experts_forward 会先将权重 transpose(1, 2) 到 GMM 期望的 matmul 布局再调用 GmmFunction，覆盖 Qwen3MoeExperts、Qwen3NextExperts、Qwen3OmniMoe*Experts、Qwen3VLMoeTextExperts、Qwen3_5MoeExperts 等。

使能方式：apply_kernel("npu_fused_moe", model=model)。未适配模型同样可参考上述示例与 NPU 算子开发者文档 自行扩展。

七、适配模型矩阵与自定义扩展

综合四个算子的 _MODEL_TYPE_TO_PATCHES 映射表，当前仓库内各算子的模型覆盖情况如下（model_type 取自 model.config.model_type）：

算子	kernel_id	覆盖的 model_type
NPU 融合 RMSNorm	npu_fused_rmsnorm	qwen3, qwen3_moe, qwen3_next, qwen3_omni_moe, qwen3_omni_moe_thinker, qwen3_vl, qwen3_vl_moe, qwen3_5, qwen3_5_moe
NPU 融合 SwiGLU	npu_fused_swiglu	同上（要求 hidden_act == "silu"）
NPU 融合 RoPE	npu_fused_rope	同上（monkey-patch transformers modeling 模块）
NPU 融合 MoE	npu_fused_moe	qwen3_moe, qwen3_next, qwen3_omni_moe, qwen3_omni_moe_thinker, qwen3_vl_moe（v5 另含 qwen3_5_moe）

对任意 model_type 不在映射表中的模型，算子会静默跳过（返回原模型），因此可以放心在混合环境中开启 auto。

若需要为新模型或新硬件适配融合算子，扩展路径是明确的：

1. 在 ops 目录 对应子目录中新建实现文件，参照现有文件的模式定义替换用的 forward 函数与 _MODEL_TYPE_TO_PATCHES 映射；
2. 继承 BaseKernel 并实现 check_device / check_deps / _apply，通过 @KernelPlugin("<你的kernel_id>").register() 注册；
3. 在 interface.py 中导入该模块使注册生效，如需纳入 auto 模式则将其加入 _AUTO_KERNELS 的设备映射。

完整的插件开发规范（参数解析、注册约定等）可参见 NPU 算子开发者文档 与 内核插件单元测试，后者可用于验证自定义算子的注册与调用行为。

八、小结

LLaMA-Factory 的融合算子系统把“替换模型 forward 函数 / modeling 模块函数”这一 NPU 性能优化的通用模式，封装成了基于插件注册表的标准化机制：BaseKernel 统一设备与依赖校验，KernelPlugin 按名称路由，apply_kernel / apply_kernels（含 auto 模式与 kernel_config 配置驱动）提供从 Python API 到训练流水线的两条使能路径。四个 NPU 算子（RMSNorm、SwiGLU、RoPE、MoE）分别覆盖归一化、激活、位置编码与稀疏专家计算这些 LLM 前向路径上的主要热点，并以模型级映射表精确控制适配范围，在保证数值等价性的前提下将相关计算下沉为硬件融合指令。