build: Add CachedWheel

Signed-off-by: oliver könig <[email protected]>

Author: ko3n1g

.github/workflows/publish.yml

@@ -82,7 +82,8 @@ jobs:
           pip install torch --index-url https://download.pytorch.org/whl/cpu
       - name: Build core package
         env:
-          GROUPED_GEMM_SKIP_CUDA_BUILD: "TRUE"
+          DEEP_GEMM_NO_LOCAL_VERSION: "TRUE"
+          DEEP_GEMM_SKIP_CUDA_BUILD: "TRUE"
         run: |
           python setup.py sdist --dist-dir=dist
       - name: Deploy

setup.py

@@ -1,3 +1,4 @@
+from calendar import c
 import os
 import setuptools
 import shutil
@@ -9,11 +10,84 @@
 from torch.utils.cpp_extension import CUDAExtension, CUDA_HOME
 from pathlib import Path
 import subprocess
+import sys
+import platform
+from packaging.version import parse
+from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
+import urllib
 
 SKIP_CUDA_BUILD = os.getenv("DEEP_GEMM_SKIP_CUDA_BUILD", "FALSE") == "TRUE"
+NO_LOCAL_VERSION = os.getenv("DEEP_GEMM_NO_LOCAL_VERSION", "FALSE") == "TRUE"
+FORCE_BUILD = os.getenv("DEEP_GEMM_FORCE_BUILD", "FALSE") == "TRUE"
 
-
+BASE_WHEEL_URL = (
+    "https://github.com/DeepSeek-AI/DeepGEMM/releases/download/{tag_name}/{wheel_name}"
+)
+PACKAGE_NAME = "deep_gemm"
 current_dir = os.path.dirname(os.path.realpath(__file__))
+
+
+def get_package_version():
+    with open(Path(current_dir) / "deep_gemm" / "__init__.py", "r") as f:
+        version_match = re.search(r"^__version__\s*=\s*(.*)$", f.read(), re.MULTILINE)
+    public_version = ast.literal_eval(version_match.group(1))
+    revision = ""
+
+    if not NO_LOCAL_VERSION:
+        try:
+            cmd = ["git", "rev-parse", "--short", "HEAD"]
+            revision = "+" + subprocess.check_output(cmd).decode("ascii").rstrip()
+        except:
+            revision = ""
+
+    return f"{public_version}{revision}"
+
+
+def get_platform():
+    """
+    Returns the platform name as used in wheel filenames.
+    """
+    if sys.platform.startswith("linux"):
+        return f"linux_{platform.uname().machine}"
+    elif sys.platform == "darwin":
+        mac_version = ".".join(platform.mac_ver()[0].split(".")[:2])
+        return f"macosx_{mac_version}_x86_64"
+    elif sys.platform == "win32":
+        return "win_amd64"
+    else:
+        raise ValueError("Unsupported platform: {}".format(sys.platform))
+
+
+def get_wheel_url():
+    torch_version_raw = parse(torch.__version__)
+    python_version = f"cp{sys.version_info.major}{sys.version_info.minor}"
+    platform_name = get_platform()
+    grouped_gemm_version = get_package_version()
+    torch_version = f"{torch_version_raw.major}.{torch_version_raw.minor}"
+    cxx11_abi = str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()
+
+    # Determine the version numbers that will be used to determine the correct wheel
+    # We're using the CUDA version used to build torch, not the one currently installed
+    # _, cuda_version_raw = get_cuda_bare_metal_version(CUDA_HOME)
+    torch_cuda_version = parse(torch.version.cuda)
+    # For CUDA 11, we only compile for CUDA 11.8, and for CUDA 12 we only compile for CUDA 12.3
+    # to save CI time. Minor versions should be compatible.
+    torch_cuda_version = (
+        parse("11.8") if torch_cuda_version.major == 11 else parse("12.3")
+    )
+    # cuda_version = f"{cuda_version_raw.major}{cuda_version_raw.minor}"
+    cuda_version = f"{torch_cuda_version.major}"
+
+    # Determine wheel URL based on CUDA version, torch version, python version and OS
+    wheel_filename = f"{PACKAGE_NAME}-{grouped_gemm_version}+cu{cuda_version}torch{torch_version}cxx11abi{cxx11_abi}-{python_version}-{python_version}-{platform_name}.whl"
+
+    wheel_url = BASE_WHEEL_URL.format(
+        tag_name=f"v{grouped_gemm_version}", wheel_name=wheel_filename
+    )
+
+    return wheel_url, wheel_filename
+
+
 cxx_flags = [
     "-std=c++17",
     "-O3",
@@ -101,23 +175,39 @@ def prepare_includes(self):
     ext_modules = []
 
 
-NO_LOCAL_VERSION = os.getenv("DEEP_GEMM_NO_LOCAL_VERSION", "FALSE") == "TRUE"
+class CachedWheelsCommand(_bdist_wheel):
+    """
+    The CachedWheelsCommand plugs into the default bdist wheel, which is ran by pip when it cannot
+    find an existing wheel (which is currently the case for all grouped gemm installs). We use
+    the environment parameters to detect whether there is already a pre-built version of a compatible
+    wheel available and short-circuits the standard full build pipeline.
+    """
 
+    def run(self):
+        if FORCE_BUILD:
+            return super().run()
 
-def get_package_version():
-    with open(Path(current_dir) / "deep_gemm" / "__init__.py", "r") as f:
-        version_match = re.search(r"^__version__\s*=\s*(.*)$", f.read(), re.MULTILINE)
-    public_version = ast.literal_eval(version_match.group(1))
-    revision = ""
-
-    if not NO_LOCAL_VERSION:
+        wheel_url, wheel_filename = get_wheel_url()
+        print("Guessing wheel URL: ", wheel_url)
         try:
-            cmd = ["git", "rev-parse", "--short", "HEAD"]
-            revision = "+" + subprocess.check_output(cmd).decode("ascii").rstrip()
-        except:
-            revision = ""
+            urllib.request.urlretrieve(wheel_url, wheel_filename)
 
-    return f"{public_version}{revision}"
+            # Make the archive
+            # Lifted from the root wheel processing command
+            # https://github.com/pypa/wheel/blob/cf71108ff9f6ffc36978069acb28824b44ae028e/src/wheel/bdist_wheel.py#LL381C9-L381C85
+            if not os.path.exists(self.dist_dir):
+                os.makedirs(self.dist_dir)
+
+            impl_tag, abi_tag, plat_tag = self.get_tag()
+            archive_basename = f"{self.wheel_dist_name}-{impl_tag}-{abi_tag}-{plat_tag}"
+
+            wheel_path = os.path.join(self.dist_dir, archive_basename + ".whl")
+            print("Raw wheel path", wheel_path)
+            os.rename(wheel_filename, wheel_path)
+        except (urllib.error.HTTPError, urllib.error.URLError):
+            print("Precompiled wheel not found. Building from source...")
+            # If the wheel could not be downloaded, build from source
+            super().run()
 
 
 if __name__ == "__main__":
@@ -135,7 +225,5 @@ def get_package_version():
         },
         ext_modules=ext_modules,
         zip_safe=False,
-        cmdclass={
-            "build_py": CustomBuildPy,
-        },
+        cmdclass={"build_py": CustomBuildPy, "bdist_wheel": CachedWheelsCommand},
     )