Add flag for nvshmem/rocshmem on release name.

End of hash, 'a' means rocshmem, 'b' means nvshmem.
Signed-off-by: lijian <lijian6@sugon.com>

build.sh

@@ -139,6 +139,11 @@ hipcc -Wno-unused-result -Wsign-compare -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack
 # build whl
 echo "Using Python: $(which python3)"
 python3 --version
-python setup.py bdist_wheel
+if [ "$USE_NVSHMEM" == "ON" ]; then
+    python setup.py bdist_wheel --shmem=nv
+fi
+if [ "$USE_ROCSHMEM" == "ON" ]; then
+    python setup.py bdist_wheel --shmem=rocm
+fi
 echo "✅ Build complete:"
 ls -lh dist/

setup.py

-import os
+import os, sys
 import subprocess
 import setuptools
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
@@ -6,6 +6,23 @@ from datetime import datetime
 
 date_tag = datetime.now().strftime("%Y%m%d")
 
+# -----------------------
+# 解析自定义参数
+# -----------------------
+shmem = None
+other = []
+for arg in sys.argv:
+    if arg.startswith("--shmem="):
+        shmem = arg.split("=", 1)[1]
+        if shmem == "rocm":
+            shmem = "a"
+        elif shmem == "nv":
+            shmem = "b"
+    else:
+        other.append(arg)
+
+sys.argv = other
+
 if __name__ == '__main__':
     try:
         cmd = ['git', 'rev-parse', '--short', 'HEAD']
@@ -15,7 +32,7 @@ if __name__ == '__main__':
 
     setuptools.setup(
         name='deep_ep',
-        version='1.0.0' + revision + '.' + date_tag,
+        version='1.0.0' + revision + shmem + '.' + date_tag,
         packages=setuptools.find_packages(include=['deep_ep']),
         include_package_data=True,
         package_data={"deep_ep": ["deep_ep_cpp.cpython-310-x86_64-linux-gnu.so"]},