Update Wall-X to 1.1.0 (#104)

setup.py

@@ -1,63 +1,78 @@
-import os
-import torch
 from pathlib import Path
-from setuptools import setup, find_packages
+
+from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
-cwd = Path(os.path.dirname(os.path.abspath(__file__)))
+ROOT = Path(__file__).resolve().parent
+OPS_DIR = Path("wall_x") / "model" / "core" / "ops"
+CSRC_DIR = OPS_DIR / "csrc"
 
-nvcc_flags = [
-    "-std=c++17",  # NOTE: CUTLASS requires c++17
-    "-DENABLE_BF16",  # Enable BF16 for cuda_version >= 11
+PUBLIC_SCRIPTS = [
+    "scripts/compute_norm_stats.py",
+    "scripts/draw_openloop_plot.py",
+    "scripts/fake_inference.py",
+    "scripts/infer_libero.py",
+    "scripts/merge_sharded_weights.py",
+    "scripts/merge_tokenizer.py",
+    "scripts/run_libero.sh",
+    "scripts/run_serving.sh",
 ]
 
-env_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST", None)
 
-if env_arch_list:
-    # Let PyTorch builder to choose device to target for.
-    device_capability = ""
-else:
-    device_capability = torch.cuda.get_device_capability()
-    device_capability = f"{device_capability[0]}{device_capability[1]}"
+def read_readme() -> str:
+    readme = ROOT / "README.md"
+    return readme.read_text(encoding="utf-8") if readme.exists() else ""
 
-if device_capability:
-    nvcc_flags.extend(
-        [
-            f"--generate-code=arch=compute_{device_capability},code=sm_{device_capability}",
-            f"-DGROUPED_GEMM_DEVICE_CAPABILITY={device_capability}",
-        ]
+
+def build_ext_modules():
+    binding = CSRC_DIR / "binding.cu"
+    if not (ROOT / binding).exists():
+        raise RuntimeError(
+            "CUDA operator sources are missing. The OSS export must include "
+            "wall_x/model/core/ops/csrc/binding.cu."
+        )
+
+    cuda_sources = [binding] + sorted(
+        path for path in (ROOT / CSRC_DIR).rglob("*.cu") if path.name != "binding.cu"
     )
+    cuda_sources = [
+        path if path == binding else path.relative_to(ROOT) for path in cuda_sources
+    ]
+
+    return [
+        CUDAExtension(
+            name="wall_x.model.core.ops._cuda_ext_bin",
+            sources=[str(path) for path in cuda_sources],
+            include_dirs=[str(CSRC_DIR), str(CSRC_DIR / "common")],
+            extra_compile_args={
+                "cxx": ["-O3", "-std=c++17"],
+                "nvcc": ["-O3", "--use_fast_math", "-std=c++17"],
+            },
+        )
+    ]
 
-ext_modules = [
-    CUDAExtension(
-        "wallx_csrc",
-        [
-            "csrc/ops.cu",
-            "csrc/dual_asym_grouped_gemm.cu",
-            "csrc/permute.cu",
-            "csrc/rope.cu",
-            "csrc/rope_index.cu",
-            "csrc/rot_pos.cu",
-            "csrc/window_index.cu",
-        ],
-        include_dirs=[f"{cwd}/3rdparty/cutlass/include/", f"{cwd}/csrc"],
-        extra_compile_args={
-            "cxx": ["-fopenmp", "-fPIC", "-Wno-strict-aliasing"],
-            "nvcc": nvcc_flags,
-        },
-    )
-]
 
 setup(
     name="wall_x",
-    version="1.0.1",
-    author="X2Robot Team",
+    version="1.1.0",
+    description="Training and inference code for WALL open-source embodied models.",
+    long_description=read_readme(),
+    long_description_content_type="text/markdown",
+    author="X-Square Robot",
+    url="https://github.com/X-Square-Robot/wall-x",
+    python_requires=">=3.10",
+    packages=find_packages(exclude=("tests", "tests.*")),
+    scripts=PUBLIC_SCRIPTS,
+    ext_modules=build_ext_modules(),
+    cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},
     classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: Apache Software License",
         "Programming Language :: Python :: 3",
-        "License :: OSI Approved :: BSD License",
-        "Operating System :: Unix",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
     ],
-    packages=find_packages(),
-    ext_modules=ext_modules,
-    cmdclass={"build_ext": BuildExtension},
 )