From 1569c8631d621753efa07e7d41e3e5e4b17616bc Mon Sep 17 00:00:00 2001
From: fanjibin <fanjibin@huawei.com>
Date: Tue, 2 Sep 2025 20:17:26 +0800
Subject: [PATCH] add ms_custom_ops submodule

---
 .gitmodules   |  4 ++++
 ms_custom_ops |  1 +
 setup.py      | 28 ++++++++++++++++++++++++++++
 3 files changed, 33 insertions(+)
 create mode 160000 ms_custom_ops

diff --git a/.gitmodules b/.gitmodules
index d5f8b94d..885ebe56 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -2,3 +2,7 @@
 	path = tests/mindformers
 	url = https://gitee.com/mindspore/mindformers.git
 	branch = master
+[submodule "ms_custom_ops"]
+	path = ms_custom_ops
+	url = https://gitee.com/mindspore/akg.git
+	branch = ms_custom_ops
diff --git a/ms_custom_ops b/ms_custom_ops
new file mode 160000
index 00000000..7fec7a3a
--- /dev/null
+++ b/ms_custom_ops
@@ -0,0 +1 @@
+Subproject commit 7fec7a3a66c9936c37be39f90431c1fc7349d4da
diff --git a/setup.py b/setup.py
index f7d37dcd..510cfdee 100644
--- a/setup.py
+++ b/setup.py
@@ -125,9 +125,35 @@ class CustomBuildExt(build_ext):
     def build_extension(self, ext):
         if ext.name == "vllm_mindspore._C_ops":
             self.build_c_ops(ext)
+        elif ext.name == "vllm_mindspore.ms_custom_ops":
+            self.build_ms_custom_ops(ext)
         else:
             raise ValueError(f"Unknown extension name: {ext.name}")
 
+    def build_ms_custom_ops(self, ext):
+        # "vllm_mindspore.ms_custom_ops" --> "ms_custom_ops"
+        ext_dir = os.path.dirname(
+            os.path.realpath(self.get_ext_fullpath(ext.name)))
+        build_cmd = build_cmd = (
+            "git submodule update --init ms_custom_ops && "
+            "cd ms_custom_ops && "
+            "python setup.py build && "
+            f"cp -r build/*/ms_custom_ops {ext_dir}/..")
+        try:
+            logger.info("Running build ms_custom_ops commands:\n%s", build_cmd)
+            result = subprocess.run(build_cmd,
+                                    cwd=self.ROOT_DIR,
+                                    text=True,
+                                    shell=True,
+                                    capture_output=False)
+            if result.returncode != 0:
+                raise RuntimeError(
+                    "Build ms_custom_ops failed with exit code {}".format(
+                        result.returncode))
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError("Failed to build {}: {}".format(
+                "ms_custom_ops", e)) from e
+
     def build_c_ops(self, ext):
         # "vllm_mindspore._C_ops" --> "_C_ops"
         ext_name = ext.name.split('.')[-1]
@@ -190,6 +216,8 @@ def _get_ext_modules():
     if os.path.exists(_get_ascend_home_path()):
         # sources are specified in CMakeLists.txt
         ext_modules.append(Extension("vllm_mindspore._C_ops", sources=[]))
+        ext_modules.append(
+            Extension("vllm_mindspore.ms_custom_ops", sources=[]))
     return ext_modules
 
 
-- 
Gitee