From 9ffb54964df8a542ff22a32691b69b72a0169743 Mon Sep 17 00:00:00 2001 From: jiaqiw09 Date: Thu, 5 Mar 2026 21:42:32 +0800 Subject: [PATCH] add omni kernels --- swift/model/npu_patcher.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/swift/model/npu_patcher.py b/swift/model/npu_patcher.py index 7f5b3045ea..f270914d21 100644 --- a/swift/model/npu_patcher.py +++ b/swift/model/npu_patcher.py @@ -11,6 +11,7 @@ from transformers.models.qwen2 import modeling_qwen2 from transformers.models.qwen3 import modeling_qwen3 from transformers.models.qwen3_moe import modeling_qwen3_moe +from transformers.models.qwen3_omni_moe import modeling_qwen3_omni_moe from transformers.models.qwen3_vl_moe import modeling_qwen3_vl_moe from typing import Any @@ -332,6 +333,15 @@ def _apply_patch_map(root: Any, patch_map: dict[str, Any]) -> None: 'apply_rotary_pos_emb': npu_apply_rotary_pos_emb, }, ), + ( + modeling_qwen3_omni_moe, + { + 'Qwen3OmniMoeTextRMSNorm': NpuRMSNorm, + 'Qwen3OmniMoeRMSNorm': NpuRMSNorm, + 'apply_rotary_pos_emb': npu_apply_rotary_pos_emb, + 'Qwen3OmniMoeThinkerTextSparseMoeBlock.forward': npu_moe_block_forward, + }, + ), ) for _module, _patch_map in _PATCH_TABLE: