[build] fix aphrodite-kernels wheel installation for pypi compat (#1606)

* [build] fix aphrodite-kernels wheel installation for pypi compat

Signed-off-by: AlpinDale <alpindale@gmail.com>

* address comment

Signed-off-by: AlpinDale <alpindale@gmail.com>

* some cleanup

Signed-off-by: AlpinDale <alpindale@gmail.com>

---------

Signed-off-by: AlpinDale <alpindale@gmail.com>

Author: AlpinDale

aphrodite/platforms/cuda.py

@@ -3,12 +3,33 @@
 """
 
 import os
+import sys
 from collections.abc import Callable
 from functools import cache, wraps
 from typing import TYPE_CHECKING, TypeVar
 
+from aphrodite.utils.system_utils import get_kernels_install_command
+
 # import custom ops, trigger op registration
-import aphrodite_kernels._C  # noqa
+try:
+    import aphrodite_kernels._C  # noqa
+except ImportError:
+    install_cmd = get_kernels_install_command()
+    error_msg = f"Failed to import aphrodite_kernels. Please install it using:\n  {install_cmd}"
+
+    # Override excepthook to suppress traceback for this specific error
+    original_excepthook = sys.excepthook
+
+    def clean_excepthook(exc_type, exc_value, exc_traceback):
+        # Check if this is our kernels import error by checking the message
+        if exc_type is ImportError and exc_value and "Failed to import aphrodite_kernels" in str(exc_value):
+            print(str(exc_value), file=sys.stderr)
+            sys.exit(1)
+        else:
+            original_excepthook(exc_type, exc_value, exc_traceback)
+
+    sys.excepthook = clean_excepthook
+    raise ImportError(error_msg) from None
 import torch
 from typing_extensions import ParamSpec
 

aphrodite/utils/system_utils.py

@@ -7,9 +7,11 @@
 import sys
 from collections.abc import Callable, Iterator
 from pathlib import Path
+from shutil import which
 from typing import TextIO
 
 import psutil
+import regex as re
 
 import aphrodite.envs as envs
 from aphrodite.logger import init_logger
@@ -297,3 +299,85 @@ def set_ulimit(target_soft_limit: int = 65535):
                 current_soft,
                 e,
             )
+
+
+def get_kernels_install_command() -> str:
+    """Get the install command for aphrodite-kernels based on the current environment."""
+
+    base_url = "https://downloads.pygmalion.chat/whl"
+    install_page = "https://aphrodite.pygmalion.chat/installation/installation/"
+
+    # Detect if uv is available
+    pip_cmd = "uv pip" if which("uv") else "pip"
+
+    try:
+        import torch
+
+        # Check for CPU build - no pre-built wheels available
+        if torch.version.cuda is None and torch.version.hip is None:
+            return (
+                "Pre-built wheels for CPU are not available.\n"
+                f"Please build aphrodite-kernels from source. See: {install_page}"
+            )
+
+        # Check for ROCm - no pre-built wheels available
+        if torch.version.hip is not None:
+            return (
+                "Pre-built wheels for ROCm are not available.\n"
+                f"Please build aphrodite-kernels from source. See: {install_page}"
+            )
+
+        # Check for CUDA
+        if torch.version.cuda:
+            cuda_major, cuda_minor = torch.version.cuda.split(".")
+            cuda_version_str = f"{cuda_major}{cuda_minor}"
+
+            # Support CUDA 12.8 and 12.9
+            if cuda_version_str == "129":
+                extra_index_url = f"{base_url}/cu129"
+                return f"{pip_cmd} install --extra-index-url {extra_index_url} aphrodite-kernels==0.0.1"
+            elif cuda_version_str == "128":
+                # Default (12.8) uses base URL without /cu128 suffix
+                extra_index_url = base_url
+                return f"{pip_cmd} install --extra-index-url {extra_index_url} aphrodite-kernels==0.0.1"
+
+            # Try to detect from nvcc if available
+            try:
+                import subprocess
+
+                from torch.utils.cpp_extension import CUDA_HOME
+
+                if CUDA_HOME:
+                    nvcc_output = subprocess.check_output(
+                        [f"{CUDA_HOME}/bin/nvcc", "-V"], universal_newlines=True, stderr=subprocess.DEVNULL
+                    )
+                    version_match = re.search(r"release (\d+\.\d+)", nvcc_output)
+                    if version_match:
+                        nvcc_version = version_match.group(1)
+                        nvcc_major, nvcc_minor = nvcc_version.split(".")
+                        nvcc_version_str = f"{nvcc_major}{nvcc_minor}"
+                        if nvcc_version_str == "129":
+                            extra_index_url = f"{base_url}/cu129"
+                            return f"{pip_cmd} install --extra-index-url {extra_index_url} aphrodite-kernels==0.0.1"
+                        elif nvcc_version_str == "128":
+                            extra_index_url = base_url
+                            return f"{pip_cmd} install --extra-index-url {extra_index_url} aphrodite-kernels==0.0.1"
+            except (subprocess.CalledProcessError, FileNotFoundError, ImportError) as e:
+                logger.warning(
+                    "nvcc-based CUDA version detection failed. This is expected if nvcc is not in your PATH. Error: %s",
+                    e,
+                )
+                pass
+
+            # Unsupported CUDA version - direct to build from source
+            return (
+                f"Your CUDA version ({torch.version.cuda}) is not supported by pre-built wheels.\n"
+                f"Only CUDA 12.8 and 12.9 are supported. "
+                f"Please build aphrodite-kernels from source. See: {install_page}"
+            )
+    except ImportError:
+        pass
+
+    # Fallback - default to 12.8
+    extra_index_url = base_url
+    return f"{pip_cmd} install --extra-index-url {extra_index_url} aphrodite-kernels==0.0.1"

requirements/cu128.txt

@@ -1 +1,2 @@
-aphrodite-kernels @ https://downloads.pygmalion.chat/whl/aphrodite_kernels/aphrodite_kernels-0.0.1-cp38-abi3-manylinux1_x86_64.whl
+--extra-index-url https://downloads.pygmalion.chat/whl/
+aphrodite-kernels == 0.0.1

requirements/cu129.txt

@@ -1,2 +1,2 @@
-aphrodite-kernels @ https://downloads.pygmalion.chat/whl/cu129/aphrodite_kernels/aphrodite_kernels-0.0.1-cp38-abi3-manylinux1_x86_64.whl
-
+--extra-index-url https://downloads.pygmalion.chat/whl/cu129/
+aphrodite-kernels == 0.0.1

setup.py

@@ -512,54 +512,6 @@ def get_aphrodite_version() -> str:
     return version
 
 
-def get_kernels_extra_index_url() -> str | None:
-    """Get the extra index URL for aphrodite-kernels wheels based on target device."""
-    base_url = "https://downloads.pygmalion.chat/whl"
-
-    if _is_cpu():
-        return f"{base_url}/cpu"
-    elif _is_hip():
-        return f"{base_url}/rocm"
-    elif _is_cuda():
-        # Determine CUDA version for wheel URL
-        cuda_version = None
-        if torch.version.cuda:
-            cuda_major, cuda_minor = torch.version.cuda.split(".")
-            cuda_version_str = f"{cuda_major}{cuda_minor}"
-            # Support CUDA 12.6, 12.8, and 12.9
-            if cuda_version_str in ["126", "128", "129"]:
-                cuda_version = cuda_version_str
-            else:
-                # Fallback to detected version or default
-                try:
-                    nvcc_version = get_nvcc_cuda_version()
-                    if nvcc_version >= Version("12.9"):
-                        cuda_version = "129"
-                    elif nvcc_version >= Version("12.8"):
-                        cuda_version = "128"
-                    elif nvcc_version >= Version("12.6"):
-                        cuda_version = "126"
-                except Exception:
-                    pass
-
-        # Use MAIN_CUDA_VERSION as fallback
-        if cuda_version is None:
-            main_cuda = envs.APHRODITE_MAIN_CUDA_VERSION
-            if main_cuda == "12.9":
-                cuda_version = "129"
-            elif main_cuda == "12.8":
-                cuda_version = "128"
-            elif main_cuda == "12.6":
-                cuda_version = "126"
-            else:
-                # Default to 12.8
-                cuda_version = "128"
-
-        return f"{base_url}/cu{cuda_version}"
-
-    return None
-
-
 def get_requirements() -> list[str]:
     """Get Python package dependencies from requirements.txt."""
     requirements_dir = ROOT_DIR / "requirements"
@@ -589,48 +541,6 @@ def _read_requirements(filename: str) -> list[str]:
             modified_requirements.append(req)
         requirements = modified_requirements
 
-        # Add CUDA-specific kernels requirements file based on detected CUDA version
-        cuda_version = None
-        if torch.version.cuda:
-            cuda_major, cuda_minor = torch.version.cuda.split(".")
-            cuda_version_str = f"{cuda_major}{cuda_minor}"
-            if cuda_version_str == "128":
-                cuda_version = "128"
-            elif cuda_version_str == "129":
-                cuda_version = "129"
-            else:
-                # Try to detect from nvcc
-                try:
-                    nvcc_version = get_nvcc_cuda_version()
-                    if nvcc_version >= Version("12.9"):
-                        cuda_version = "129"
-                    elif nvcc_version >= Version("12.8"):
-                        cuda_version = "128"
-                except Exception:
-                    pass
-
-        # Use MAIN_CUDA_VERSION as fallback
-        if cuda_version is None:
-            main_cuda = envs.APHRODITE_MAIN_CUDA_VERSION
-            if main_cuda == "12.9":
-                cuda_version = "129"
-            elif main_cuda == "12.8":
-                cuda_version = "128"
-
-        # Add kernels requirements file if CUDA version is supported
-        if cuda_version == "128":
-            requirements += _read_requirements("cu128.txt")
-        elif cuda_version == "129":
-            requirements += _read_requirements("cu129.txt")
-        else:
-            # If cuda_version is None or unsupported, don't add kernels
-            # User will need to build from source
-            logger.warning(
-                "No pre-built aphrodite-kernels wheels available for CUDA version %s. "
-                "Only CUDA 12.8 and 12.9 are supported. "
-                "You will need to build aphrodite-kernels from source.",
-                torch.version.cuda if torch.version.cuda else "unknown",
-            )
     elif _is_hip():
         requirements = _read_requirements("rocm.txt")
     elif _is_hpu():
@@ -644,27 +554,15 @@ def _read_requirements(filename: str) -> list[str]:
     else:
         raise ValueError("Unsupported platform, please use CUDA, ROCm, or CPU.")
 
+    # Filter out aphrodite-kernels from install_requires
+    # Users will get a helpful error message when they try to import it
+    requirements = [req for req in requirements if "aphrodite-kernels" not in req.lower()]
+
     return requirements
 
 
 ext_modules = []
 
-# Determine extra index URL for aphrodite-kernels wheels
-kernels_extra_index_url = get_kernels_extra_index_url()
-if kernels_extra_index_url:
-    current_extra_index = os.environ.get("PIP_EXTRA_INDEX_URL", "")
-    if kernels_extra_index_url not in current_extra_index:
-        if current_extra_index:
-            os.environ["PIP_EXTRA_INDEX_URL"] = f"{current_extra_index} {kernels_extra_index_url}"
-        else:
-            os.environ["PIP_EXTRA_INDEX_URL"] = kernels_extra_index_url
-        logger.info(
-            "Added extra index URL for aphrodite-kernels: %s\n"
-            "To install with this index URL, use: pip install --extra-index-url %s .",
-            kernels_extra_index_url,
-            kernels_extra_index_url,
-        )
-
 package_data = {
     "aphrodite": [
         "endpoints/kobold/klite.embd",