423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2026-08-17 05:25:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
xvxuopop
2026-08-04 19:38:41 +08:00
committed by GitHub
parent 713b5a3f95
commit 84576b1408
6 changed files with 413 additions and 329 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/llamafactory/model/loader.py
View File

@@ -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:

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

@@ -29,15 +29,21 @@ import torch.nn.functional as F
try:
import torch_npu
except ImportError:
pass
except ImportError as exc:
_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:
"""Container for NPU fused MoE forward functions."""
class NpuMoeFusedV4:
"""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:
r"""Forward pass for sparse MoE block using NPU optimization.
def stacked_sparse_block_forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
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(
self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
) -> 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.
def sparse_block_forward(self, hidden_states: torch.Tensor):
"""Forward pass for a Transformers v4 list-backed 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
if is_transformers_version_greater_than("5.0.0"):
kernel_moe_mapping = {
"Qwen3MoeForCausalLM": {
"Qwen3MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward,
},
"Qwen3VLMoeForConditionalGeneration": {
"Qwen3VLMoeTextExperts": NpuMoeFused.npu_moe_experts_v5_forward,
},
"Qwen3_5MoeForCausalLM": {
"Qwen3_5MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward,
},
"Qwen3_5MoeForConditionalGeneration": {
"Qwen3_5MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward,
},
}
else:
kernel_moe_mapping = {
"Qwen3MoeForCausalLM": {
"Qwen3MoeSparseMoeBlock": Qwen3NpuMoeFused.qwen3moe_sparse_moe_block_forward,
},
"Qwen3VLMoeForConditionalGeneration": {
"Qwen3VLMoeTextExperts": NpuMoeFused.npu_moe_experts_forward,
"Qwen3VLMoeTextSparseMoeBlock": NpuMoeFused.npu_moe_sparse_block_forward,
},
}
shared_expert_output = self.shared_expert(hidden_states)
shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output
next_states = next_states + shared_expert_output
return next_states, router_logits
class NpuMoeFusedV5:
"""Container for Transformers v5 NPU fused MoE forward functions."""
@staticmethod
def experts_forward(
self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
) -> 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)
_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 []
target_moe_mapping = None
for arch in archs:
if arch in kernel_moe_mapping:
target_moe_mapping = kernel_moe_mapping[arch]
break
if target_moe_mapping is None:
model_type = getattr(model.config, "model_type", None)
if model_type not in _MODEL_TYPE_TO_PATCHES:
return model
patched_count = 0
for module in model.modules():
class_name = module.__class__.__name__
if class_name in target_moe_mapping:
new_forward_func = target_moe_mapping[class_name]
module.forward = types.MethodType(new_forward_func, module)
patch_forward = NpuFusedMoEKernel._get_patch_forward(model_type, module)
if patch_forward is not None:
module.forward = types.MethodType(patch_forward, 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

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

@@ -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:
pass
except ImportError as exc:
_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):
"""SwiGLU forward pass for GLM4 on NPU.
Args:
self: The GLM4 MLP layer instance.
hidden_states (Tensor): Input hidden states.
Returns:
Tensor: Output of SwiGLU.
"""
up_states = self.gate_up_proj(hidden_states)
gate, up_states = up_states.chunk(2, dim=-1)
return self.down_proj(torch_npu.npu_swiglu(torch.cat((gate, up_states), dim=-1), dim=-1))
def _npu_swiglu_gemma3ntext_forward(self, hidden_states):
"""SwiGLU forward pass for Gemma3nText on NPU.
Args:
self: The Gemma3nText MLP layer instance.
hidden_states (Tensor): Input hidden states.
Returns:
Tensor: Output of SwiGLU.
"""
gate_proj = self.gate_proj(hidden_states)
if self.activation_sparsity > 0.0:
gate_proj = self._gaussian_topk(gate_proj)
down_proj = self.down_proj(
torch_npu.npu_swiglu(torch.cat((gate_proj, self.up_proj(hidden_states)), dim=-1), dim=-1)
)
return down_proj
_MODEL_TYPE_TO_PATCHES = {
"qwen3": {
"Qwen3MLP": npu_swiglu_forward,
},
"qwen3_moe": {
"Qwen3MoeMLP": npu_swiglu_forward,
},
"qwen3_next": {
"Qwen3NextMLP": npu_swiglu_forward,
},
"qwen3_omni_moe": {
"Qwen3OmniMoeThinkerTextMLP": npu_swiglu_forward,
"Qwen3OmniMoeMLP": npu_swiglu_forward,
"Qwen3OmniMoeTalkerTextMLP": npu_swiglu_forward,
"Qwen3OmniMoeCode2WavMlp": npu_swiglu_forward,
},
"qwen3_omni_moe_thinker": {
"Qwen3OmniMoeThinkerTextMLP": npu_swiglu_forward,
},
"qwen3_vl": {
"Qwen3VLTextMLP": npu_swiglu_forward,
},
"qwen3_vl_moe": {
"Qwen3VLMoeTextMLP": npu_swiglu_forward,
},
"qwen3_5": {
"Qwen3_5MLP": npu_swiglu_forward,
},
"qwen3_5_moe": {
"Qwen3_5MoeMLP": npu_swiglu_forward,
},
}
@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
kernel_mapping = {
"Glm4MLP": _npu_swiglu_glm4_forward,
"Glm4vTextMLP": _npu_swiglu_glm4_forward,
"Phi3MLP": _npu_swiglu_glm4_forward,
"Gemma3nTextMLP": _npu_swiglu_gemma3ntext_forward,
}
model_type = getattr(model.config, "model_type", None)
if model_type not in _MODEL_TYPE_TO_PATCHES:
return model
swiglu_pattern = re.compile("MLP", re.IGNORECASE)
for name, module in model.named_modules():
# Match any module whose class name contains "MLP"
if (
re.search(swiglu_pattern, module.__class__.__name__)
and module.__class__.__name__ in NpuSwiGluKernel.expect_modules
):
# Bind function as an instance method to preserve `self` semantics
# and replace the original forward
kernel_func = kernel_mapping.get(module.__class__.__name__, npu_swiglu_forward)
module.forward = types.MethodType(kernel_func, module)
patched_count = 0
for module in model.modules():
patch_forward = NpuSwiGluKernel._get_patch_forward(model_type, module)
if patch_forward is not None:
module.forward = types.MethodType(patch_forward, module)
patched_count += 1
if patched_count:
logger.info_rank0(f"Applied NPU SwiGLU kernel to {patched_count} modules for model type: {model_type}.")
return model

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

@@ -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:
pass
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
except ImportError as exc:
_TORCH_NPU_IMPORT_ERROR = exc
else:
_TORCH_NPU_IMPORT_ERROR = None
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:
if getattr(self, "_npu_use_residual_rmsnorm", False):
effective_weight = 1.0 + self.weight.float()
else:
effective_weight = self.weight.float()
else:
effective_weight = None
weight = getattr(self, "weight", None)
if weight is None:
raise RuntimeError(f"{self.__class__.__name__} has no RMSNorm weight for NPU RMSNorm kernel.")
if effective_weight is not None:
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, self.weight, epsilon=_eps)[0]
effective_weight = weight.float()
return torch_npu.npu_rms_norm(hidden_states, effective_weight.to(hidden_states.dtype), 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
parameterization (``scale = weight`` where weight is initialized to 1), so the
residual weight adjustment (``1.0 + weight``) is not applied here.
stability. The supported gated RMSNorm modules use ``scale = weight`` with weight
initialized to 1, unlike the residual RMSNorm variants that use ``1.0 + weight``.
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]
if gate is not None:
hidden_states = hidden_states * F.silu(gate.to(torch.float32))
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
the appropriate NPU-optimized forward function as an instance method via
``types.MethodType`` to replace the original ``forward``.
Matches modules configured for the current model type, then binds the corresponding
NPU-optimized forward function as an instance method via ``types.MethodType`` to
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():
if re.search(rms_norm_pattern, module.__class__.__name__):
if "Gated" in module.__class__.__name__:
module.forward = types.MethodType(npu_gated_rms_norm_forward, module)
else:
module._npu_use_residual_rmsnorm = _should_use_residual_rmsnorm(module)
module.forward = types.MethodType(npu_rms_norm_forward, module)
patched_count = 0
for module in model.modules():
patch_forward = NpuRMSNormKernel._get_patch_forward(model_type, module)
if patch_forward is not None:
module.forward = types.MethodType(patch_forward, module)
patched_count += 1
if patched_count:
logger.info_rank0(f"Applied NPU RMSNorm kernel to {patched_count} modules for model type: {model_type}.")
return model

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

@@ -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:
pass
except ImportError as exc:
_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
based on the dimension ratio between ``cos/sin`` and ``q/k`` tensors, ensuring
compatibility with future model versions without hardcoding.
Partial RoPE is detected when the ``cos/sin`` width is smaller than the ``q/k``
head dimension. The leading rotary dimensions are transformed and any trailing
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)
k_embed = torch_npu.npu_rotary_mul(k_rot, cos, sin).to(k.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, "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)
k_embed = torch_npu.npu_rotary_mul(k, cos, sin).to(k.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, "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):
"""Apply Rotary Position Embedding with multimodal sections (Qwen2-VL) on NPU.
This function supports Partial RoPE for multimodal inputs with automatic dimension
detection, ensuring compatibility with future model versions.
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
def _default_rope_patch(module_type: str):
return (
(
f"transformers.models.{module_type}.modeling_{module_type}",
(("apply_rotary_pos_emb", _apply_rotary_pos_emb),),
),
)
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":
"""Apply RoPE acceleration by monkey-patching ``apply_rotary_pos_emb``.
def check_deps() -> None:
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
the module where they are defined, and replaces the original
``apply_rotary_pos_emb`` function in that module's namespace with the
NPU-accelerated version.
@staticmethod
def _apply_model_patches(model_type: str) -> int:
patches = _MODEL_TYPE_TO_PATCHES.get(model_type)
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()
for module in model.modules():
if "Attention" in module.__class__.__name__:
module_name = module.__class__.__module__
if module_name in _modules:
continue
try:
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}")
model_type = getattr(model.config, "model_type", None)
if model_type not in _MODEL_TYPE_TO_PATCHES:
return model
patched_count = NpuRoPEKernel._apply_model_patches(model_type)
if patched_count:
logger.info_rank0(f"Applied NPU RoPE kernel to {patched_count} functions for model type: {model_type}.")
return model

34
tests_v1/plugins/model_plugins/test_kernel_plugin.py
View File

@@ -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__