[v1] refactor NPU kernel matching by model type (#10643)

src/llamafactory/model/loader.py

@@ -217,9 +217,9 @@ def load_model(
             "You are try to using future feature about kernels, please note that this feature "
             "is not supported for all models. If get any error, please disable this feature, or report the issue."
         )
-        from ..v1.plugins.model_plugins.kernels.interface import apply_default_kernels
+        from ..v1.plugins.model_plugins.kernels.interface import apply_v1_kernels
 
-        model = apply_default_kernels(model, include_kernels=model_args.use_v1_kernels)
+        model = apply_v1_kernels(model, use_v1_kernels=model_args.use_v1_kernels)
 
     trainable_params, all_param = count_parameters(model)
     if is_trainable:

src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py

@@ -29,15 +29,21 @@ import torch.nn.functional as F
 
 try:
     import torch_npu
-except ImportError:
+except ImportError as exc:
-    pass
+    _TORCH_NPU_IMPORT_ERROR = exc
+else:
+    _TORCH_NPU_IMPORT_ERROR = None
 
 from ......accelerator.helper import DeviceType, get_current_accelerator
+from ......utils.logging import get_logger
 from ......utils.packages import is_transformers_version_greater_than
 from ......utils.types import HFModel
 from ...base import BaseKernel, KernelPlugin
 
 
+logger = get_logger(__name__)
+
+
 class GmmFunction(torch.autograd.Function):
     """Custom autograd function for NPU Grouped Matrix Multiplication (GMM)."""
 
@@ -49,7 +55,7 @@ class GmmFunction(torch.autograd.Function):
             ctx: Context object to save tensors for backward pass.
             x (Tensor): Input tensor.
             weight (Tensor): Weight tensor.
-            group_list (list): List of group sizes.
+            group_list (Tensor): Number of tokens assigned to each expert.
 
         Returns:
             Tensor: The result of the grouped matrix multiplication.
@@ -174,14 +180,14 @@ class HybridGmmFunction(torch.autograd.Function):
         return (None, *grad_x_list, *grad_w_list)
 
 
-class NpuMoeFused:
+class NpuMoeFusedV4:
-    """Container for NPU fused MoE forward functions."""
+    """Container for Transformers v4 NPU fused MoE forward functions."""
 
     @staticmethod
-    def npu_moe_experts_forward(
+    def stacked_experts_forward(
         self, hidden_states: torch.Tensor, routing_weights: torch.Tensor, router_indices: torch.Tensor
     ) -> torch.Tensor:
-        """Forward pass for MoE experts using NPU fused operations.
+        """Forward pass for Transformers v4 MoE experts using NPU fused operations.
 
         Args:
             self: The MoE layer instance.
@@ -197,7 +203,9 @@ class NpuMoeFused:
         permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute(
             hidden_states, router_indices.to(torch.int32)
         )
-        tokens_per_expert = torch.histc(router_indices, bins=self.num_experts, min=0, max=self.num_experts)
+        tokens_per_expert = torch.histc(
+            router_indices.float(), bins=self.num_experts, min=0, max=self.num_experts
+        ).long()
         intermediate_hidden_states = GmmFunction.apply(permuted_hidden_states, self.gate_up_proj, tokens_per_expert)
         intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1)
         output = GmmFunction.apply(intermediate_activations, self.down_proj, tokens_per_expert)
@@ -206,61 +214,33 @@ class NpuMoeFused:
         return next_states
 
     @staticmethod
-    def npu_moe_sparse_block_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def stacked_sparse_block_forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        r"""Forward pass for sparse MoE block using NPU optimization.
+        r"""Forward pass for Transformers v4 sparse MoE block using NPU optimization.
 
         Args:
             self: The MoE sparse block instance.
             hidden_states (Tensor): Input hidden states.
 
         Returns:
-            Tensor: The routed output.
+            tuple: A tuple containing the routed output and router logits.
         """
         batch_size = hidden_states.shape[0]
         hidden_states = hidden_states.reshape(-1, self.hidden_size)
         router_logits = self.gate(hidden_states)
-        routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
         routing_weights, router_indices = torch.topk(routing_weights, self.top_k, dim=-1)
         routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
         routing_weights = routing_weights.to(hidden_states.dtype)
         hidden_states = hidden_states.reshape(batch_size, -1, self.hidden_size)
         routed_out = self.experts(hidden_states, routing_weights, router_indices)
-        return routed_out
+        return routed_out, router_logits
 
     @staticmethod
-    def npu_moe_experts_v5_forward(
+    def sparse_block_forward(self, hidden_states: torch.Tensor):
-        self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
+        """Forward pass for a Transformers v4 list-backed sparse MoE block using NPU fused operations.
-    ) -> torch.Tensor:
-        """Forward pass for Transformers v5+ MoE experts using NPU fused operations.
-
-        Transformers v5 stores expert weights in F.linear layout:
-        gate_up_proj: [num_experts, 2 * intermediate_dim, hidden_dim]
-        down_proj: [num_experts, hidden_dim, intermediate_dim]
-        The NPU grouped matmul path expects matmul layout, so both weights are transposed.
-        """
-        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
-        permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute(
-            hidden_states, top_k_index.to(torch.int32)
-        )
-        tokens_per_expert = torch.histc(top_k_index.float(), bins=self.num_experts, min=0, max=self.num_experts).long()
-
-        gate_up_proj = self.gate_up_proj.transpose(1, 2)
-        down_proj = self.down_proj.transpose(1, 2)
-        intermediate_hidden_states = GmmFunction.apply(permuted_hidden_states, gate_up_proj, tokens_per_expert)
-        intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1)
-        output = GmmFunction.apply(intermediate_activations, down_proj, tokens_per_expert)
-        return torch_npu.npu_moe_token_unpermute(output, row_ids_map, probs=top_k_weights)
-
-
-class Qwen3NpuMoeFused:
-    """Container for Qwen3 NPU fused MoE forward functions."""
-
-    @staticmethod
-    def qwen3moe_sparse_moe_block_forward(self, hidden_states: torch.Tensor):
-        """Forward pass for Qwen3 sparse MoE block using NPU fused operations.
 
         Args:
-            self: The Qwen3 MoE block instance.
+            self: The sparse MoE block instance.
             hidden_states (Tensor): Input hidden states.
 
         Returns:
@@ -304,33 +284,90 @@ class Qwen3NpuMoeFused:
         next_states = next_states.view(batch_size, sequence_length, -1)
         return next_states, router_logits
 
+    @staticmethod
+    def shared_sparse_block_forward(self, hidden_states: torch.Tensor):
+        """Forward pass for a Transformers v4 sparse MoE block with a shared expert."""
+        next_states, router_logits = NpuMoeFusedV4.sparse_block_forward(self, hidden_states)
 
-# moe patch config mapping
+        shared_expert_output = self.shared_expert(hidden_states)
-if is_transformers_version_greater_than("5.0.0"):
+        shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output
-    kernel_moe_mapping = {
+        next_states = next_states + shared_expert_output
-        "Qwen3MoeForCausalLM": {
+        return next_states, router_logits
-            "Qwen3MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward,
+
-        },
+
-        "Qwen3VLMoeForConditionalGeneration": {
+class NpuMoeFusedV5:
-            "Qwen3VLMoeTextExperts": NpuMoeFused.npu_moe_experts_v5_forward,
+    """Container for Transformers v5 NPU fused MoE forward functions."""
-        },
+
-        "Qwen3_5MoeForCausalLM": {
+    @staticmethod
-            "Qwen3_5MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward,
+    def experts_forward(
-        },
+        self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
-        "Qwen3_5MoeForConditionalGeneration": {
+    ) -> torch.Tensor:
-            "Qwen3_5MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward,
+        """Forward pass for Transformers v5+ MoE experts using NPU fused operations.
-        },
+
-    }
+        Transformers v5 stores expert weights in F.linear layout:
-else:
+        gate_up_proj: [num_experts, 2 * intermediate_dim, hidden_dim]
-    kernel_moe_mapping = {
+        down_proj: [num_experts, hidden_dim, intermediate_dim]
-        "Qwen3MoeForCausalLM": {
+        The NPU grouped matmul path expects matmul layout, so both weights are transposed.
-            "Qwen3MoeSparseMoeBlock": Qwen3NpuMoeFused.qwen3moe_sparse_moe_block_forward,
+        """
-        },
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
-        "Qwen3VLMoeForConditionalGeneration": {
+        permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute(
-            "Qwen3VLMoeTextExperts": NpuMoeFused.npu_moe_experts_forward,
+            hidden_states, top_k_index.to(torch.int32)
-            "Qwen3VLMoeTextSparseMoeBlock": NpuMoeFused.npu_moe_sparse_block_forward,
+        )
-        },
+        tokens_per_expert = torch.histc(top_k_index.float(), bins=self.num_experts, min=0, max=self.num_experts).long()
-    }
+
+        gate_up_proj = self.gate_up_proj.transpose(1, 2)
+        down_proj = self.down_proj.transpose(1, 2)
+        intermediate_hidden_states = GmmFunction.apply(permuted_hidden_states, gate_up_proj, tokens_per_expert)
+        intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1)
+        output = GmmFunction.apply(intermediate_activations, down_proj, tokens_per_expert)
+        return torch_npu.npu_moe_token_unpermute(output, row_ids_map, probs=top_k_weights)
+
+
+_V4_MODEL_TYPE_TO_PATCHES = {
+    "qwen3_moe": {
+        "Qwen3MoeSparseMoeBlock": NpuMoeFusedV4.sparse_block_forward,
+    },
+    "qwen3_next": {
+        "Qwen3NextSparseMoeBlock": NpuMoeFusedV4.shared_sparse_block_forward,
+    },
+    "qwen3_omni_moe": {
+        "Qwen3OmniMoeThinkerTextSparseMoeBlock": NpuMoeFusedV4.sparse_block_forward,
+        "Qwen3OmniMoeTalkerTextSparseMoeBlock": NpuMoeFusedV4.shared_sparse_block_forward,
+    },
+    "qwen3_omni_moe_thinker": {
+        "Qwen3OmniMoeThinkerTextSparseMoeBlock": NpuMoeFusedV4.sparse_block_forward,
+    },
+    "qwen3_vl_moe": {
+        "Qwen3VLMoeTextExperts": NpuMoeFusedV4.stacked_experts_forward,
+        "Qwen3VLMoeTextSparseMoeBlock": NpuMoeFusedV4.stacked_sparse_block_forward,
+    },
+}
+
+_V5_MODEL_TYPE_TO_PATCHES = {
+    "qwen3_moe": {
+        "Qwen3MoeExperts": NpuMoeFusedV5.experts_forward,
+    },
+    "qwen3_next": {
+        "Qwen3NextExperts": NpuMoeFusedV5.experts_forward,
+    },
+    "qwen3_omni_moe": {
+        "Qwen3OmniMoeThinkerTextExperts": NpuMoeFusedV5.experts_forward,
+        "Qwen3OmniMoeTalkerTextExperts": NpuMoeFusedV5.experts_forward,
+    },
+    "qwen3_omni_moe_thinker": {
+        "Qwen3OmniMoeThinkerTextExperts": NpuMoeFusedV5.experts_forward,
+    },
+    "qwen3_vl_moe": {
+        "Qwen3VLMoeTextExperts": NpuMoeFusedV5.experts_forward,
+    },
+    "qwen3_5_moe": {
+        "Qwen3_5MoeExperts": NpuMoeFusedV5.experts_forward,
+    },
+}
+
+_MODEL_TYPE_TO_PATCHES = (
+    _V5_MODEL_TYPE_TO_PATCHES if is_transformers_version_greater_than("5.0.0") else _V4_MODEL_TYPE_TO_PATCHES
+)
 
 
 @KernelPlugin("npu_fused_moe").register()
@@ -343,6 +380,17 @@ class NpuFusedMoEKernel(BaseKernel):
         if current != DeviceType.NPU:
             raise RuntimeError(f"NpuFusedMoEKernel requires NPU, current accelerator is {current}.")
 
+    @staticmethod
+    def check_deps() -> None:
+        if _TORCH_NPU_IMPORT_ERROR is not None:
+            raise RuntimeError("NpuFusedMoEKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR
+
+    @staticmethod
+    def _get_patch_forward(model_type: str, module: torch.nn.Module):
+        """Return the version-specific NPU forward function for a matched MoE module."""
+        model_patches = _MODEL_TYPE_TO_PATCHES.get(model_type, {})
+        return model_patches.get(module.__class__.__name__)
+
     @staticmethod
     def _apply(**kwargs) -> HFModel:
         """Applies the NPU fused MoE kernel to the model.
@@ -352,27 +400,21 @@ class NpuFusedMoEKernel(BaseKernel):
 
         Returns:
             HFModel: The model with patched MoE forward functions.
-
-        Raises:
-            ValueError: If the model is not provided.
-            RuntimeError: If dependencies are not met.
         """
-        model = kwargs.get("model", None)
+        model = kwargs["model"]
 
-        archs = getattr(model.config, "architectures", None) or []
+        model_type = getattr(model.config, "model_type", None)
-        target_moe_mapping = None
+        if model_type not in _MODEL_TYPE_TO_PATCHES:
-        for arch in archs:
-            if arch in kernel_moe_mapping:
-                target_moe_mapping = kernel_moe_mapping[arch]
-                break
-
-        if target_moe_mapping is None:
             return model
 
+        patched_count = 0
         for module in model.modules():
-            class_name = module.__class__.__name__
+            patch_forward = NpuFusedMoEKernel._get_patch_forward(model_type, module)
-            if class_name in target_moe_mapping:
+            if patch_forward is not None:
-                new_forward_func = target_moe_mapping[class_name]
+                module.forward = types.MethodType(patch_forward, module)
-                module.forward = types.MethodType(new_forward_func, module)
+                patched_count += 1
+
+        if patched_count:
+            logger.info_rank0(f"Applied NPU fused MoE kernel to {patched_count} modules for model type: {model_type}.")
 
         return model

src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_swiglu.py

@@ -20,20 +20,24 @@ Init Phase:
 
 """
 
-import re
 import types
 
 import torch
 
 from ......accelerator.helper import DeviceType, get_current_accelerator
+from ......utils.logging import get_logger
 from ......utils.types import HFModel
 from ...base import BaseKernel, KernelPlugin
 
 
+logger = get_logger(__name__)
+
 try:
     import torch_npu
-except ImportError:
+except ImportError as exc:
-    pass
+    _TORCH_NPU_IMPORT_ERROR = exc
+else:
+    _TORCH_NPU_IMPORT_ERROR = None
 
 
 def npu_swiglu_forward(self, hidden_state):
@@ -51,81 +55,69 @@ def npu_swiglu_forward(self, hidden_state):
     )
 
 
-def _npu_swiglu_glm4_forward(self, hidden_states):
+_MODEL_TYPE_TO_PATCHES = {
-    """SwiGLU forward pass for GLM4 on NPU.
+    "qwen3": {
-
+        "Qwen3MLP": npu_swiglu_forward,
-    Args:
+    },
-        self: The GLM4 MLP layer instance.
+    "qwen3_moe": {
-        hidden_states (Tensor): Input hidden states.
+        "Qwen3MoeMLP": npu_swiglu_forward,
-
+    },
-    Returns:
+    "qwen3_next": {
-        Tensor: Output of SwiGLU.
+        "Qwen3NextMLP": npu_swiglu_forward,
-    """
+    },
-    up_states = self.gate_up_proj(hidden_states)
+    "qwen3_omni_moe": {
-    gate, up_states = up_states.chunk(2, dim=-1)
+        "Qwen3OmniMoeThinkerTextMLP": npu_swiglu_forward,
-    return self.down_proj(torch_npu.npu_swiglu(torch.cat((gate, up_states), dim=-1), dim=-1))
+        "Qwen3OmniMoeMLP": npu_swiglu_forward,
-
+        "Qwen3OmniMoeTalkerTextMLP": npu_swiglu_forward,
-
+        "Qwen3OmniMoeCode2WavMlp": npu_swiglu_forward,
-def _npu_swiglu_gemma3ntext_forward(self, hidden_states):
+    },
-    """SwiGLU forward pass for Gemma3nText on NPU.
+    "qwen3_omni_moe_thinker": {
-
+        "Qwen3OmniMoeThinkerTextMLP": npu_swiglu_forward,
-    Args:
+    },
-        self: The Gemma3nText MLP layer instance.
+    "qwen3_vl": {
-        hidden_states (Tensor): Input hidden states.
+        "Qwen3VLTextMLP": npu_swiglu_forward,
-
+    },
-    Returns:
+    "qwen3_vl_moe": {
-        Tensor: Output of SwiGLU.
+        "Qwen3VLMoeTextMLP": npu_swiglu_forward,
-    """
+    },
-    gate_proj = self.gate_proj(hidden_states)
+    "qwen3_5": {
-    if self.activation_sparsity > 0.0:
+        "Qwen3_5MLP": npu_swiglu_forward,
-        gate_proj = self._gaussian_topk(gate_proj)
+    },
-    down_proj = self.down_proj(
+    "qwen3_5_moe": {
-        torch_npu.npu_swiglu(torch.cat((gate_proj, self.up_proj(hidden_states)), dim=-1), dim=-1)
+        "Qwen3_5MoeMLP": npu_swiglu_forward,
-    )
+    },
-    return down_proj
+}
 
 
 @KernelPlugin("npu_fused_swiglu").register()
 class NpuSwiGluKernel(BaseKernel):
     """NPU Kernel for fused SwiGLU activation."""
 
-    # just support apply to the following module layers
-    expect_modules = frozenset(
-        {
-            "Qwen3VLMoeTextMLP",
-            "Qwen3VLTextMLP",
-            "Qwen3OmniMoeThinkerTextMLP",
-            "Qwen3OmniMoeMLP",
-            "Qwen3OmniMoeTalkerTextMLP",
-            "Qwen3OmniMoeCode2WavMlp",
-            "Qwen3NextMLP",
-            "Qwen3MoeMLP",
-            "Qwen3MLP",
-            "Qwen2MLP",
-            "Qwen2MoeMLP",
-            "Qwen2_5_VLMLP",
-            "Qwen2_5OmniMLP",
-            "Llama4TextMLP",
-            "LlamaMLP",
-            "Glm4MLP",
-            "Glm4MoeMLP",
-            "Glm4vMoeTextMLP",
-            "Gemma3MLP",
-            "Gemma2MLP",
-            "Gemma3nTextMLP",
-            "Phi3MLP",
-            "DeepseekV2MLP",
-            "DeepseekV3MLP",
-            "SeedOssMLP",
-        }
-    )
-
     @staticmethod
     def check_device() -> None:
         current = get_current_accelerator().type
         if current != DeviceType.NPU:
             raise RuntimeError(f"NpuSwiGluKernel requires NPU, current accelerator is {current}.")
 
+    @staticmethod
+    def check_deps() -> None:
+        if _TORCH_NPU_IMPORT_ERROR is not None:
+            raise RuntimeError("NpuSwiGluKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR
+
+    @staticmethod
+    def _get_patch_forward(model_type: str, module: torch.nn.Module):
+        """Return the NPU forward function for a matched SwiGLU MLP module."""
+        model_patches = _MODEL_TYPE_TO_PATCHES.get(model_type, {})
+        patch_forward = model_patches.get(module.__class__.__name__)
+        if patch_forward is None:
+            return None
+
+        config = getattr(module, "config", None)
+        if getattr(config, "hidden_act", None) != "silu":
+            return None
+
+        return patch_forward
+
     @staticmethod
     def _apply(**kwargs) -> "HFModel":
         """Applies the NPU fused SwiGLU kernel to the model.
@@ -135,31 +127,21 @@ class NpuSwiGluKernel(BaseKernel):
 
         Returns:
             HFModel: The model with patched SwiGLU forward functions.
-
-        Raises:
-            ValueError: If the model is not provided.
-            RuntimeError: If dependencies are not met.
         """
-        model = kwargs.get("model", None)
+        model = kwargs["model"]
 
-        # Mapping of specific mlp modules to their corresponding kernel implementations
+        model_type = getattr(model.config, "model_type", None)
-        kernel_mapping = {
+        if model_type not in _MODEL_TYPE_TO_PATCHES:
-            "Glm4MLP": _npu_swiglu_glm4_forward,
+            return model
-            "Glm4vTextMLP": _npu_swiglu_glm4_forward,
-            "Phi3MLP": _npu_swiglu_glm4_forward,
-            "Gemma3nTextMLP": _npu_swiglu_gemma3ntext_forward,
-        }
 
-        swiglu_pattern = re.compile("MLP", re.IGNORECASE)
+        patched_count = 0
-        for name, module in model.named_modules():
+        for module in model.modules():
-            # Match any module whose class name contains "MLP"
+            patch_forward = NpuSwiGluKernel._get_patch_forward(model_type, module)
-            if (
+            if patch_forward is not None:
-                re.search(swiglu_pattern, module.__class__.__name__)
+                module.forward = types.MethodType(patch_forward, module)
-                and module.__class__.__name__ in NpuSwiGluKernel.expect_modules
+                patched_count += 1
-            ):
+
-                # Bind function as an instance method to preserve `self` semantics
+        if patched_count:
-                # and replace the original forward
+            logger.info_rank0(f"Applied NPU SwiGLU kernel to {patched_count} modules for model type: {model_type}.")
-                kernel_func = kernel_mapping.get(module.__class__.__name__, npu_swiglu_forward)
-                module.forward = types.MethodType(kernel_func, module)
 
         return model

src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py

@@ -20,54 +20,32 @@ Init Phase:
 
 """
 
-import re
 import types
 
 import torch
 import torch.nn.functional as F
 
 from ......accelerator.helper import DeviceType, get_current_accelerator
+from ......utils.logging import get_logger
 from ......utils.types import HFModel
 from ...base import BaseKernel, KernelPlugin
 
 
+logger = get_logger(__name__)
+
 try:
     import torch_npu
-except ImportError:
+except ImportError as exc:
-    pass
+    _TORCH_NPU_IMPORT_ERROR = exc
-
+else:
-
+    _TORCH_NPU_IMPORT_ERROR = None
-def _should_use_residual_rmsnorm(module):
-    """Detect whether the module uses residual RMSNorm parameterization.
-
-    Residual RMSNorm uses ``scale = 1.0 + weight`` where weight is initialized to 0,
-    while standard RMSNorm uses ``scale = weight`` where weight is initialized to 1.
-
-    Args:
-        module (nn.Module): The RMSNorm module to check.
-
-    Returns:
-        bool: ``True`` if the module uses residual parameterization, ``False`` otherwise.
-
-    .. note::
-        This must follow the module's forward semantics. Do not infer it from trained
-        weight values because standard RMSNorm weights can also be close to zero.
-    """
-    residual_rmsnorm_classes = {
-        "Qwen3_5RMSNorm",
-        "Qwen3_5MoeRMSNorm",
-        "Qwen3NextRMSNorm",
-    }
-
-    class_name = module.__class__.__name__
-    return class_name in residual_rmsnorm_classes
 
 
 def npu_rms_norm_forward(self, hidden_states):
     """NPU forward implementation for standard RMSNorm.
 
     Args:
-        self (nn.Module): The RMSNorm module instance with ``weight`` and ``variance_epsilon``.
+        self (nn.Module): The RMSNorm module instance with ``weight`` and either ``variance_epsilon`` or ``eps``.
         hidden_states (Tensor): Input hidden states tensor.
 
     Returns:
@@ -75,48 +53,108 @@ def npu_rms_norm_forward(self, hidden_states):
     """
     _eps = getattr(self, "variance_epsilon", None) or getattr(self, "eps", 1e-6)
 
-    if hasattr(self, "weight") and self.weight is not None:
+    weight = getattr(self, "weight", None)
-        if getattr(self, "_npu_use_residual_rmsnorm", False):
+    if weight is None:
-            effective_weight = 1.0 + self.weight.float()
+        raise RuntimeError(f"{self.__class__.__name__} has no RMSNorm weight for NPU RMSNorm kernel.")
-        else:
-            effective_weight = self.weight.float()
-    else:
-        effective_weight = None
 
-    if effective_weight is not None:
+    effective_weight = weight.float()
-        return torch_npu.npu_rms_norm(hidden_states, effective_weight.to(hidden_states.dtype), epsilon=_eps)[0]
+
-    else:
+    return torch_npu.npu_rms_norm(hidden_states, effective_weight.to(hidden_states.dtype), epsilon=_eps)[0]
-        return torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=_eps)[0]
+
+
+def npu_residual_rms_norm_forward(self, hidden_states):
+    """NPU forward implementation for residual RMSNorm.
+
+    Residual RMSNorm uses ``scale = 1.0 + weight`` where ``weight`` is initialized
+    to 0 in the original transformers implementation.
+
+    Args:
+        self (nn.Module): The residual RMSNorm module with ``weight`` and either ``variance_epsilon`` or ``eps``.
+        hidden_states (Tensor): Input hidden states tensor.
+
+    Returns:
+        Tensor: Normalized tensor consistent with residual RMSNorm behavior.
+    """
+    _eps = getattr(self, "variance_epsilon", None) or getattr(self, "eps", 1e-6)
+
+    weight = getattr(self, "weight", None)
+    if weight is None:
+        raise RuntimeError(f"{self.__class__.__name__} has no RMSNorm weight for NPU RMSNorm kernel.")
+
+    effective_weight = 1.0 + weight.float()
+
+    return torch_npu.npu_rms_norm(hidden_states, effective_weight.to(hidden_states.dtype), epsilon=_eps)[0]
 
 
 def npu_gated_rms_norm_forward(self, hidden_states, gate=None):
     """NPU forward implementation for Gated RMSNorm with high-precision FP32 computation.
 
     This function performs RMSNorm and gated SiLU multiplication in FP32 for numerical
-    stability. Unlike standard RMSNorm, Gated RMSNorm in Qwen3.5 uses standard
+    stability. The supported gated RMSNorm modules use ``scale = weight`` with weight
-    parameterization (``scale = weight`` where weight is initialized to 1), so the
+    initialized to 1, unlike the residual RMSNorm variants that use ``1.0 + weight``.
-    residual weight adjustment (``1.0 + weight``) is not applied here.
 
     Args:
         self (nn.Module): The Gated RMSNorm module instance.
         hidden_states (Tensor): Input hidden states tensor.
-        gate (Tensor, optional): Gate tensor for SiLU activation. Defaults to ``None``.
+        gate (Tensor): Gate tensor for SiLU activation.
 
     Returns:
         Tensor: Output tensor cast back to the original input dtype.
+
+    Raises:
+        ValueError: If the gate tensor is not provided.
     """
+    if gate is None:
+        raise ValueError(f"{self.__class__.__name__} requires a gate tensor for NPU Gated RMSNorm.")
+
     input_dtype = hidden_states.dtype
     hidden_states = hidden_states.to(torch.float32)
     _eps = getattr(self, "variance_epsilon", None) or getattr(self, "eps", 1e-6)
 
     hidden_states = torch_npu.npu_rms_norm(hidden_states, self.weight.float(), epsilon=_eps)[0]
-
+    hidden_states = hidden_states * F.silu(gate.to(torch.float32))
-    if gate is not None:
-        hidden_states = hidden_states * F.silu(gate.to(torch.float32))
 
     return hidden_states.to(input_dtype)
 
 
+_MODEL_TYPE_TO_PATCHES = {
+    "qwen3": {
+        "Qwen3RMSNorm": npu_rms_norm_forward,
+    },
+    "qwen3_moe": {
+        "Qwen3MoeRMSNorm": npu_rms_norm_forward,
+    },
+    "qwen3_next": {
+        "Qwen3NextRMSNorm": npu_residual_rms_norm_forward,
+        "Qwen3NextRMSNormGated": npu_gated_rms_norm_forward,
+    },
+    "qwen3_omni_moe": {
+        "Qwen3OmniMoeThinkerTextRMSNorm": npu_rms_norm_forward,
+        "Qwen3OmniMoeTextRMSNorm": npu_rms_norm_forward,
+        "Qwen3OmniMoeRMSNorm": npu_rms_norm_forward,
+        "Qwen3OmniMoeCode2WavRMSNorm": npu_rms_norm_forward,
+    },
+    "qwen3_omni_moe_thinker": {
+        "Qwen3OmniMoeThinkerTextRMSNorm": npu_rms_norm_forward,
+        "Qwen3OmniMoeTextRMSNorm": npu_rms_norm_forward,
+    },
+    "qwen3_vl": {
+        "Qwen3VLTextRMSNorm": npu_rms_norm_forward,
+    },
+    "qwen3_vl_moe": {
+        "Qwen3VLMoeTextRMSNorm": npu_rms_norm_forward,
+    },
+    "qwen3_5": {
+        "Qwen3_5RMSNorm": npu_residual_rms_norm_forward,
+        "Qwen3_5RMSNormGated": npu_gated_rms_norm_forward,
+    },
+    "qwen3_5_moe": {
+        "Qwen3_5MoeRMSNorm": npu_residual_rms_norm_forward,
+        "Qwen3_5MoeRMSNormGated": npu_gated_rms_norm_forward,
+    },
+}
+
+
 @KernelPlugin("npu_fused_rmsnorm").register()
 class NpuRMSNormKernel(BaseKernel):
     """NPU kernel wrapper for RMSNorm that applies the replacement within a model."""
@@ -127,34 +165,45 @@ class NpuRMSNormKernel(BaseKernel):
         if current != DeviceType.NPU:
             raise RuntimeError(f"NpuRMSNormKernel requires NPU, current accelerator is {current}.")
 
+    @staticmethod
+    def check_deps() -> None:
+        if _TORCH_NPU_IMPORT_ERROR is not None:
+            raise RuntimeError("NpuRMSNormKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR
+
+    @staticmethod
+    def _get_patch_forward(model_type: str, module: torch.nn.Module):
+        """Return the NPU forward function for a matched RMSNorm module."""
+        model_patches = _MODEL_TYPE_TO_PATCHES.get(model_type, {})
+        return model_patches.get(module.__class__.__name__)
+
     @staticmethod
     def _apply(**kwargs) -> "HFModel":
         """Iterate the model and apply NPU-optimized forward to matched RMSNorm modules.
 
-        Matches modules whose class name contains "RMSNorm" (case-insensitive) and binds
+        Matches modules configured for the current model type, then binds the corresponding
-        the appropriate NPU-optimized forward function as an instance method via
+        NPU-optimized forward function as an instance method via ``types.MethodType`` to
-        ``types.MethodType`` to replace the original ``forward``.
+        replace the original ``forward``.
 
         Args:
             **kwargs: Keyword arguments containing the model.
 
         Returns:
             HFModel: The model with NPU fused RMSNorm.
-
-        Raises:
-            RuntimeError: If ``torch_npu`` is not available.
-            ValueError: If the model is not provided.
         """
-        model = kwargs.get("model")
+        model = kwargs["model"]
 
-        rms_norm_pattern = re.compile("RMSNorm", re.IGNORECASE)
+        model_type = getattr(model.config, "model_type", None)
+        if model_type not in _MODEL_TYPE_TO_PATCHES:
+            return model
 
-        for _, module in model.named_modules():
+        patched_count = 0
-            if re.search(rms_norm_pattern, module.__class__.__name__):
+        for module in model.modules():
-                if "Gated" in module.__class__.__name__:
+            patch_forward = NpuRMSNormKernel._get_patch_forward(model_type, module)
-                    module.forward = types.MethodType(npu_gated_rms_norm_forward, module)
+            if patch_forward is not None:
-                else:
+                module.forward = types.MethodType(patch_forward, module)
-                    module._npu_use_residual_rmsnorm = _should_use_residual_rmsnorm(module)
+                patched_count += 1
-                    module.forward = types.MethodType(npu_rms_norm_forward, module)
+
+        if patched_count:
+            logger.info_rank0(f"Applied NPU RMSNorm kernel to {patched_count} modules for model type: {model_type}.")
 
         return model

src/llamafactory/v1/plugins/model_plugins/kernels/ops/rope/npu_rope.py

@@ -20,7 +20,7 @@ Init Phase:
 
 """
 
-import sys
+import importlib
 
 import torch
 
@@ -34,16 +34,18 @@ logger = get_logger(__name__)
 
 try:
     import torch_npu
-except ImportError:
+except ImportError as exc:
-    pass
+    _TORCH_NPU_IMPORT_ERROR = exc
+else:
+    _TORCH_NPU_IMPORT_ERROR = None
 
 
 def _apply_npu_rotary_emb(q, k, cos, sin):
     """Apply NPU-accelerated rotary embedding with automatic Partial RoPE detection.
 
-    This function automatically detects whether to use Partial RoPE or Full RoPE
+    Partial RoPE is detected when the ``cos/sin`` width is smaller than the ``q/k``
-    based on the dimension ratio between ``cos/sin`` and ``q/k`` tensors, ensuring
+    head dimension. The leading rotary dimensions are transformed and any trailing
-    compatibility with future model versions without hardcoding.
+    dimensions are passed through unchanged.
 
     Args:
         q (Tensor): Query tensor.
@@ -61,14 +63,14 @@ def _apply_npu_rotary_emb(q, k, cos, sin):
         q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
         k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
 
-        q_embed = torch_npu.npu_rotary_mul(q_rot, cos, sin).to(q.dtype)
+        q_embed = torch_npu.npu_rotary_mul(q_rot, cos, sin, "half").to(q.dtype)
-        k_embed = torch_npu.npu_rotary_mul(k_rot, cos, sin).to(k.dtype)
+        k_embed = torch_npu.npu_rotary_mul(k_rot, cos, sin, "half").to(k.dtype)
 
         q_embed = torch.cat([q_embed, q_pass], dim=-1)
         k_embed = torch.cat([k_embed, k_pass], dim=-1)
     else:
-        q_embed = torch_npu.npu_rotary_mul(q, cos, sin).to(q.dtype)
+        q_embed = torch_npu.npu_rotary_mul(q, cos, sin, "half").to(q.dtype)
-        k_embed = torch_npu.npu_rotary_mul(k, cos, sin).to(k.dtype)
+        k_embed = torch_npu.npu_rotary_mul(k, cos, sin, "half").to(k.dtype)
 
     return q_embed, k_embed
 
@@ -84,44 +86,43 @@ def _apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
         k (Tensor): Key tensor.
         cos (Tensor): Cosine part of embedding.
         sin (Tensor): Sine part of embedding.
-        position_ids (Tensor, optional): Position IDs. Defaults to ``None``.
+        position_ids (Tensor | int, optional): Ignored Transformers v4 position IDs, or the Transformers v5
+            ``unsqueeze_dim`` when supplied as the fifth positional argument.
         unsqueeze_dim (int): Dimension to unsqueeze cos and sin. Defaults to 1.
 
     Returns:
         tuple[Tensor, Tensor]: The embedded query and key tensors ``(q_embed, k_embed)``.
     """
+    # In transformers v5, the fifth positional argument is ``unsqueeze_dim``.
+    if isinstance(position_ids, int):
+        unsqueeze_dim = position_ids
+
     cos = cos.unsqueeze(unsqueeze_dim)
     sin = sin.unsqueeze(unsqueeze_dim)
 
     return _apply_npu_rotary_emb(q, k, cos, sin)
 
 
-def _apply_multimodal_rotary_pos_emb_qwen25_vl(q, k, cos, sin, mrope_section, unsqueeze_dim=1):
+def _default_rope_patch(module_type: str):
-    """Apply Rotary Position Embedding with multimodal sections (Qwen2-VL) on NPU.
+    return (
-
+        (
-    This function supports Partial RoPE for multimodal inputs with automatic dimension
+            f"transformers.models.{module_type}.modeling_{module_type}",
-    detection, ensuring compatibility with future model versions.
+            (("apply_rotary_pos_emb", _apply_rotary_pos_emb),),
-
+        ),
-    Args:
-        q (Tensor): Query tensor.
-        k (Tensor): Key tensor.
-        cos (Tensor): Cosine part of embedding.
-        sin (Tensor): Sine part of embedding.
-        mrope_section (list[int]): Multimodal RoPE section sizes.
-        unsqueeze_dim (int): Dimension to unsqueeze cos and sin. Defaults to 1.
-
-    Returns:
-        tuple[Tensor, Tensor]: The embedded query and key tensors ``(q_embed, k_embed)``.
-    """
-    mrope_section = mrope_section * 2
-    cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
-        unsqueeze_dim
-    )
-    sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
-        unsqueeze_dim
     )
 
-    return _apply_npu_rotary_emb(q, k, cos, sin)
+
+_MODEL_TYPE_TO_PATCHES = {
+    "qwen3": _default_rope_patch("qwen3"),
+    "qwen3_moe": _default_rope_patch("qwen3_moe"),
+    "qwen3_next": _default_rope_patch("qwen3_next"),
+    "qwen3_omni_moe": _default_rope_patch("qwen3_omni_moe"),
+    "qwen3_omni_moe_thinker": _default_rope_patch("qwen3_omni_moe"),
+    "qwen3_vl": _default_rope_patch("qwen3_vl"),
+    "qwen3_vl_moe": _default_rope_patch("qwen3_vl_moe"),
+    "qwen3_5": _default_rope_patch("qwen3_5"),
+    "qwen3_5_moe": _default_rope_patch("qwen3_5_moe"),
+}
 
 
 @KernelPlugin("npu_fused_rope").register()
@@ -135,50 +136,58 @@ class NpuRoPEKernel(BaseKernel):
             raise RuntimeError(f"NpuRoPEKernel requires NPU, current accelerator is {current}.")
 
     @staticmethod
-    def _apply(**kwargs) -> "HFModel":
+    def check_deps() -> None:
-        """Apply RoPE acceleration by monkey-patching ``apply_rotary_pos_emb``.
+        if _TORCH_NPU_IMPORT_ERROR is not None:
+            raise RuntimeError("NpuRoPEKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR
 
-        Iterates through the model's modules to find attention layers, identifies
+    @staticmethod
-        the module where they are defined, and replaces the original
+    def _apply_model_patches(model_type: str) -> int:
-        ``apply_rotary_pos_emb`` function in that module's namespace with the
+        patches = _MODEL_TYPE_TO_PATCHES.get(model_type)
-        NPU-accelerated version.
+        if patches is None:
+            return 0
+
+        patched_count = 0
+        for module_name, replacements in patches:
+            try:
+                target_module = importlib.import_module(module_name)
+            except Exception as e:
+                logger.warning_rank0_once(f"Failed to import {module_name} for NPU RoPE kernel: {e}")
+                continue
+
+            for target_function_name, replacement in replacements:
+                if not hasattr(target_module, target_function_name):
+                    logger.warning_rank0_once(f"{module_name} has no {target_function_name}, skip NPU RoPE patch.")
+                    continue
+
+                if getattr(target_module, target_function_name) is replacement:
+                    continue
+
+                setattr(target_module, target_function_name, replacement)
+                patched_count += 1
+
+        return patched_count
+
+    @staticmethod
+    def _apply(**kwargs) -> "HFModel":
+        """Apply RoPE acceleration by monkey-patching rotary embedding functions.
+
+        Selects the target transformers modeling module from ``model.config.model_type``
+        and replaces its rotary embedding helper with the NPU-accelerated version.
 
         Args:
             **kwargs: Keyword arguments containing the model.
 
         Returns:
             HFModel: The model with patched RoPE functions.
-
-        Raises:
-            RuntimeError: If ``torch_npu`` is not available.
-            ValueError: If the model is not provided.
         """
-        model = kwargs.get("model", None)
+        model = kwargs["model"]
 
-        _modules = set()
+        model_type = getattr(model.config, "model_type", None)
-        for module in model.modules():
+        if model_type not in _MODEL_TYPE_TO_PATCHES:
-            if "Attention" in module.__class__.__name__:
+            return model
-                module_name = module.__class__.__module__
+
-                if module_name in _modules:
+        patched_count = NpuRoPEKernel._apply_model_patches(model_type)
-                    continue
+        if patched_count:
-                try:
+            logger.info_rank0(f"Applied NPU RoPE kernel to {patched_count} functions for model type: {model_type}.")
-                    target_module = sys.modules[module_name]
-                    if hasattr(target_module, "apply_rotary_pos_emb"):
-                        if getattr(target_module, "apply_rotary_pos_emb") is not _apply_rotary_pos_emb:
-                            setattr(target_module, "apply_rotary_pos_emb", _apply_rotary_pos_emb)
-                            _modules.add(module_name)
-                    if hasattr(target_module, "apply_multimodal_rotary_pos_emb"):
-                        if (
-                            getattr(target_module, "apply_multimodal_rotary_pos_emb")
-                            is not _apply_multimodal_rotary_pos_emb_qwen25_vl
-                        ):
-                            setattr(
-                                target_module,
-                                "apply_multimodal_rotary_pos_emb",
-                                _apply_multimodal_rotary_pos_emb_qwen25_vl,
-                            )
-                            _modules.add(module_name)
-                except Exception as e:
-                    logger.warning_rank0_once(f"Failed to apply RoPE kernel to module {module_name}: {e}")
 
         return model

tests_v1/plugins/model_plugins/test_kernel_plugin.py

@@ -25,18 +25,19 @@ def _apply_kernel(rank) -> None:
         setattr(mock_device, "type", "npu")
         mock_get_accelerator.return_value = mock_device
 
-        # reload kernel modules to respect mocked accelerator
-        for k in list(sys.modules.keys()):
-            if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
-                del sys.modules[k]
-
-        from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels
-
         model = AutoModelForCausalLM.from_pretrained("llamafactory/tiny-random-qwen3")
         original_rmsnorm_forward = model.model.layers[0].input_layernorm.forward
         original_swiglu_forward = model.model.layers[0].mlp.forward
 
-        model = apply_kernels(model=model, config={"name": "npu_fused_rmsnorm"})
+        with patch.dict(sys.modules, {"torch_npu": MagicMock()}):
+            # Reload kernel modules so dependency checks use the mocked NPU environment.
+            for k in list(sys.modules.keys()):
+                if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
+                    del sys.modules[k]
+
+            from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels
+
+            model = apply_kernels(model=model, config={"name": "npu_fused_rmsnorm"})
 
         assert model.model.layers[0].input_layernorm.forward.__func__ is not original_rmsnorm_forward.__func__
         assert model.model.layers[0].mlp.forward.__func__ is original_swiglu_forward.__func__
@@ -48,18 +49,19 @@ def _apply_all_kernels(rank) -> None:
         setattr(mock_device, "type", "npu")
         mock_get_accelerator.return_value = mock_device
 
-        # reload kernel modules to respect mocked accelerator
-        for k in list(sys.modules.keys()):
-            if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
-                del sys.modules[k]
-
-        from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels
-
         model = AutoModelForCausalLM.from_pretrained("llamafactory/tiny-random-qwen3")
         original_rmsnorm_forward = model.model.layers[0].input_layernorm.forward
         original_swiglu_forward = model.model.layers[0].mlp.forward
 
-        model = apply_kernels(model=model, config={"name": "auto"})
+        with patch.dict(sys.modules, {"torch_npu": MagicMock()}):
+            # Reload kernel modules so dependency checks use the mocked NPU environment.
+            for k in list(sys.modules.keys()):
+                if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
+                    del sys.modules[k]
+
+            from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels
+
+            model = apply_kernels(model=model, config={"name": "auto"})
 
         assert model.model.layers[0].input_layernorm.forward.__func__ is not original_rmsnorm_forward.__func__
         assert model.model.layers[0].mlp.forward.__func__ is not original_swiglu_forward.__func__