From e500c00545b80446d58c4ade4cbc831a26932042 Mon Sep 17 00:00:00 2001
From: Weiming Liao <liaowm5@gmail.com>
Date: Thu, 25 Jun 2026 18:50:08 +0800
Subject: [PATCH] [Bugfix][NPU] Register vllm-ascend custom ops in
 NPUOmniPlatform.set_device (#4712)

Signed-off-by: Weiming Liao <liaowm5@gmail.com>
---
 vllm_omni/platforms/npu/platform.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py
index e2a576cfd..cce70c7ec 100644
--- a/vllm_omni/platforms/npu/platform.py
+++ b/vllm_omni/platforms/npu/platform.py
@@ -44,6 +44,12 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
     @classmethod
     def set_device(cls, device: torch.device) -> None:
         super().set_device(device)
+
+        # Register vllm_ascend custom ops (torch.ops._C_ascend.*).
+        from vllm_ascend.utils import enable_custom_op
+
+        enable_custom_op()
+
         # Ascend quantized weights are converted from ND to FRACTAL_NZ
         # after loading. Enable internal format so the NZ storage layout
         # is preserved for fused NPU kernels.