`npu_patch.py` — 华为 NPU 优化补丁¶

文件路径¶

verl/models/transformers/npu_patch.py

文件概述¶

为华为昇腾 (Ascend) NPU 提供优化的算子替换。当 verl 运行在 NPU 上时，通过导入此文件自动将 HuggingFace 模型中的标准 PyTorch 算子替换为 NPU 专用的高性能实现。覆盖了 RMSNorm、SiLU、RoPE、MoE 等核心计算。

关键代码讲解¶

NPU 优化的 RMSNorm¶

def rms_norm_forward_npu(self, x):
    if x.dtype != self.weight.dtype:
        x = x.to(self.weight.dtype)
    return torch_npu.npu_rms_norm(x, self.weight, epsilon=self.variance_epsilon)[0]

使用 torch_npu.npu_rms_norm 替代 PyTorch 原生实现，利用 NPU 硬件加速。

NPU 优化的 SiLU (SwiGLU MLP)¶

def silu_forward_npu(self, hidden_state):
    gate_up = torch.cat((self.gate_proj(hidden_state), self.up_proj(hidden_state)), dim=-1)
    return self.down_proj(torch_npu.npu_swiglu(gate_up, dim=-1))

将 gate_proj 和 up_proj 的结果拼接后，使用 torch_npu.npu_swiglu 融合计算 SiLU 激活和乘法。

NPU 优化的 RoPE¶

def apply_rotary_pos_emb_npu(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.to(q.dtype), k_embed.to(k.dtype)

使用 torch_npu.npu_rotary_mul 替代手工实现的旋转位置编码。

NPU 分组矩阵乘法 (GMM)¶

class NPUGmmFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, group_list, group_list_type=1):
        """
        Args:
            x: (tokens_num * top_k, hidden_size)
            weight: (n_experts, hidden_size, intermediate_size)
            group_list: 每个专家处理的 token 数量
        """
        output = torch_npu.npu_grouped_matmul(
            [x], [weight], bias=None, group_list=group_list,
            split_item=2, group_type=0, group_list_type=group_list_type
        )[0]
        return output

MoE (Mixture of Experts) 模型中，不同专家处理不同数量的 token。NPUGmmFunction 使用 torch_npu.npu_grouped_matmul 将所有专家的矩阵乘法一次性完成，避免逐专家循环。

NPU MoE 前向（以 Qwen3 为例）¶

def _qwen3_sparse_moe_routed_forward_npu(self, hidden_states):
    router_logits = self.gate(hidden_states)
    routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
    routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)

    # 收集所有专家权重
    w1 = torch.stack([e.up_proj.weight for e in self.experts]).transpose(1, 2)
    w2 = torch.stack([e.gate_proj.weight for e in self.experts]).transpose(1, 2)
    w3 = torch.stack([e.down_proj.weight for e in self.experts]).transpose(1, 2)

    # NPU 专用的 token 排列和反排列
    permuted_tokens, row_ids_map = torch_npu.npu_moe_token_permute(hidden_states, selected_experts)
    tokens_per_expert = torch.histc(selected_experts, bins=self.num_experts, ...)

    # 融合计算
    up_res = NPUGmmFunction.apply(permuted_tokens, w1, tokens_per_expert)
    gate_res = NPUGmmFunction.apply(permuted_tokens, w2, tokens_per_expert)
    act_res = torch_npu.npu_swiglu(torch.cat([gate_res, up_res], dim=-1))
    down_res = NPUGmmFunction.apply(act_res, w3, tokens_per_expert)

    routed_hidden_states = torch_npu.npu_moe_token_unpermute(down_res, row_ids_map, probs=routing_weights)

流程：token 排列 -> 分组矩阵乘法 -> SwiGLU 激活 -> 分组矩阵乘法 -> token 反排列。全部使用 NPU 专用算子。

模块级自动 Patch¶

# 文件底部直接执行替换，导入时自动生效
modeling_qwen2.Qwen2RMSNorm.forward = rms_norm_forward_npu
modeling_qwen2.Qwen2MLP.forward = silu_forward_npu
modeling_qwen2.apply_rotary_pos_emb = apply_rotary_pos_emb_npu

modeling_qwen3.Qwen3RMSNorm.forward = rms_norm_forward_npu
modeling_qwen3.Qwen3MLP.forward = silu_forward_npu
# ... 更多模型的 patch

这些替换在 import npu_patch 时自动生效，不需要显式调用。

核心类/函数列表¶

名称	作用
`rms_norm_forward_npu()`	NPU 优化的 RMSNorm
`silu_forward_npu()`	NPU 优化的 SwiGLU MLP
`apply_rotary_pos_emb_npu()`	NPU 优化的 RoPE
`NPUGmmFunction`	NPU 分组矩阵乘法自定义 Autograd 函数
`_qwen3_sparse_moe_routed_forward_npu()`	NPU 优化的 MoE 路由前向
`qwen3_moe_sparse_moe_block_forward_npu()`	Qwen3 MoE 块的 NPU 前向
`NPUQwen3VLMoeTextExperts`	NPU 优化的 Qwen3-VL MoE 专家类
`NPUQwen3VLMoeTextSparseMoeBlock`	NPU 优化的 Qwen3-VL MoE 块

支持的模型¶

模型	替换的算子
Qwen2	RMSNorm, MLP, RoPE
Qwen2.5-VL	RMSNorm, MLP
Qwen3	RMSNorm, MLP, RoPE
Qwen3 MoE	RMSNorm, MoE Block, RoPE
Qwen3-VL	RMSNorm, MLP
Qwen3-VL MoE	MoE Block, RMSNorm, RoPE
Qwen3 Next	MoE Block, RMSNorm (含 Gated), RoPE

与其他模块的关系¶

在 NPU 环境下被自动导入
替换的是 HuggingFace transformers 库中的原始实现
与 monkey_patch.py 独立，是硬件级别的优化

小结¶

这个文件是 verl 对华为 NPU 硬件的适配层。通过将标准 PyTorch 算子替换为 NPU 专用算子（torch_npu 库），可以充分利用 NPU 的硬件加速能力。特别是 MoE 模型的分组矩阵乘法和 token 排列/反排列操作，使用 NPU 专用算子可以获得显著的性能提升。

npu_patch.py — 华为 NPU 优化补丁¶