ci/cd: bring up flashinfer-cubin package

flashinfer-cubin/.gitignore

@@ -0,0 +1,2 @@
+build/
+flashinfer_cubin/cubins/

flashinfer-cubin/MANIFEST.in

@@ -0,0 +1,3 @@
+include README.md
+include LICENSE
+recursive-include flashinfer_cubin/cubins *

flashinfer-cubin/build_wheel.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+"""
+Copyright (c) 2025 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+
+def build_wheel():
+    """Build the flashinfer-cubin wheel."""
+
+    # Change to the flashinfer-cubin directory
+    script_dir = Path(__file__).parent
+    os.chdir(script_dir)
+
+    print("Building flashinfer-cubin wheel...")
+    print(f"Working directory: {script_dir}")
+
+    # Clean previous builds
+    dist_dir = script_dir / "dist"
+    build_dir = script_dir / "build"
+    egg_info_dir = script_dir / "flashinfer_cubin.egg-info"
+
+    for dir_to_clean in [dist_dir, build_dir, egg_info_dir]:
+        if dir_to_clean.exists():
+            print(f"Cleaning {dir_to_clean}")
+            import shutil
+
+            shutil.rmtree(dir_to_clean)
+
+    # Build wheel
+    try:
+        subprocess.run([sys.executable, "setup.py", "bdist_wheel"], check=True)
+
+        print("Wheel built successfully!")
+
+        # List built wheels
+        if dist_dir.exists():
+            wheels = list(dist_dir.glob("*.whl"))
+            if wheels:
+                print(f"Built wheel: {wheels[0]}")
+            else:
+                print("No wheel files found in dist/")
+
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to build wheel: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    build_wheel()

flashinfer-cubin/download_cubins.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+"""
+Copyright (c) 2025 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import os
+import sys
+import argparse
+from pathlib import Path
+
+# Add parent directory to path to import flashinfer modules
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from flashinfer.artifacts import download_artifacts
+from flashinfer.jit.cubin_loader import FLASHINFER_CUBINS_REPOSITORY
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Download FlashInfer cubins from artifactory"
+    )
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        type=str,
+        default="flashinfer_cubin/cubins",
+        help="Output directory for cubins (default: flashinfer_cubin/cubins)",
+    )
+    parser.add_argument(
+        "--threads",
+        "-t",
+        type=int,
+        default=4,
+        help="Number of download threads (default: 4)",
+    )
+    parser.add_argument(
+        "--repository",
+        "-r",
+        type=str,
+        default=None,
+        help="Override the cubins repository URL",
+    )
+
+    args = parser.parse_args()
+
+    # Set environment variables to control download behavior
+    if args.repository:
+        os.environ["FLASHINFER_CUBINS_REPOSITORY"] = args.repository
+
+    os.environ["FLASHINFER_CUBIN_DIR"] = str(Path(args.output_dir).absolute())
+    os.environ["FLASHINFER_CUBIN_DOWNLOAD_THREADS"] = str(args.threads)
+
+    print(f"Downloading cubins to {args.output_dir}")
+    print(
+        f"Repository: {os.environ.get('FLASHINFER_CUBINS_REPOSITORY', FLASHINFER_CUBINS_REPOSITORY)}"
+    )
+
+    # Use the existing download_artifacts function
+    try:
+        download_artifacts()
+        print("Download complete!")
+    except Exception as e:
+        print(f"Download failed: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

flashinfer-cubin/flashinfer_cubin/__init__.py

@@ -0,0 +1,67 @@
+"""
+Copyright (c) 2025 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import os
+from pathlib import Path
+
+# Get the path to the cubins directory within this package
+CUBIN_DIR = Path(__file__).parent / "cubins"
+
+
+def get_cubin_dir():
+    """Get the directory containing the cubins."""
+    return str(CUBIN_DIR)
+
+
+def list_cubins():
+    """List all available cubin files."""
+    if not CUBIN_DIR.exists():
+        return []
+
+    cubins = []
+    for root, _, files in os.walk(CUBIN_DIR):
+        for file in files:
+            if file.endswith(".cubin"):
+                rel_path = os.path.relpath(os.path.join(root, file), CUBIN_DIR)
+                cubins.append(rel_path)
+    return sorted(cubins)
+
+
+def get_cubin_path(relative_path):
+    """Get the absolute path to a specific cubin file."""
+    return str(CUBIN_DIR / relative_path)
+
+
+# Read version from build metadata or fallback to main flashinfer version.txt
+def _get_version():
+    # First try to read from build metadata (for wheel distributions)
+    try:
+        from . import _build_meta
+
+        return _build_meta.__version__
+    except ImportError:
+        pass
+
+    # Fallback to reading from the main flashinfer version.txt (for development)
+    version_file = Path(__file__).parent.parent.parent / "version.txt"
+    if version_file.exists():
+        with open(version_file, "r") as f:
+            return f.read().strip()
+    return "0.0.0"
+
+
+__version__ = _get_version()
+__all__ = ["get_cubin_dir", "list_cubins", "get_cubin_path", "CUBIN_DIR"]

flashinfer-cubin/setup.py

@@ -0,0 +1,158 @@
+"""
+Copyright (c) 2025 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import os
+import sys
+from pathlib import Path
+from setuptools import setup, find_packages
+from setuptools.command.build_py import build_py
+from setuptools.command.sdist import sdist
+
+# Add parent directory to path to import artifacts module
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from flashinfer.artifacts import download_artifacts
+
+
+def get_version():
+    version_file = Path(__file__).parent.parent / "version.txt"
+    if version_file.exists():
+        with open(version_file, "r") as f:
+            return f.read().strip()
+    return "0.0.0+unknown"
+
+
+class DownloadAndBuildPy(build_py):
+    """Custom build command that downloads cubins before building."""
+
+    def run(self):
+        print("Downloading cubins from artifactory...")
+
+        # Create a temporary directory for cubins within the package
+        cubin_package_dir = Path(self.build_lib) / "flashinfer_cubin" / "cubins"
+        cubin_package_dir.mkdir(parents=True, exist_ok=True)
+
+        # Set environment variable to download to our package directory
+        original_cubin_dir = os.environ.get("FLASHINFER_CUBIN_DIR")
+        os.environ["FLASHINFER_CUBIN_DIR"] = str(cubin_package_dir)
+
+        try:
+            # Download all cubins using the existing download_artifacts function
+            download_artifacts()
+            print(f"Cubins downloaded to {cubin_package_dir}")
+
+        finally:
+            # Restore original environment variable
+            if original_cubin_dir:
+                os.environ["FLASHINFER_CUBIN_DIR"] = original_cubin_dir
+            else:
+                os.environ.pop("FLASHINFER_CUBIN_DIR", None)
+
+        # Create build metadata file with version information
+        package_dir = Path(self.build_lib) / "flashinfer_cubin"
+        build_meta_file = package_dir / "_build_meta.py"
+        version = get_version()
+
+        with open(build_meta_file, "w") as f:
+            f.write('"""Build metadata for flashinfer-cubin package."""\n')
+            f.write(f'__version__ = "{version}"\n')
+
+        print(f"Created build metadata file with version {version}")
+
+        # Continue with normal build
+        super().run()
+
+
+class CustomSdist(sdist):
+    """Custom sdist command that includes downloaded cubins."""
+
+    def run(self):
+        # Download cubins first
+        print("Downloading cubins for source distribution...")
+
+        cubin_package_dir = (
+            Path(self.distribution.package_dir.get("", "."))
+            / "flashinfer_cubin"
+            / "cubins"
+        )
+        cubin_package_dir.mkdir(parents=True, exist_ok=True)
+
+        # Set environment variable to download to our package directory
+        original_cubin_dir = os.environ.get("FLASHINFER_CUBIN_DIR")
+        os.environ["FLASHINFER_CUBIN_DIR"] = str(cubin_package_dir)
+
+        try:
+            download_artifacts()
+        finally:
+            if original_cubin_dir:
+                os.environ["FLASHINFER_CUBIN_DIR"] = original_cubin_dir
+            else:
+                os.environ.pop("FLASHINFER_CUBIN_DIR", None)
+
+        # Create build metadata file with version information for sdist
+        package_dir = (
+            Path(self.distribution.package_dir.get("", ".")) / "flashinfer_cubin"
+        )
+        build_meta_file = package_dir / "_build_meta.py"
+        version = get_version()
+
+        with open(build_meta_file, "w") as f:
+            f.write('"""Build metadata for flashinfer-cubin package."""\n')
+            f.write(f'__version__ = "{version}"\n')
+
+        print(f"Created build metadata file with version {version}")
+
+        # Continue with normal sdist
+        super().run()
+
+
+setup(
+    name="flashinfer-cubin",
+    version=get_version(),
+    description="Pre-compiled CUDA kernels (cubins) for FlashInfer",
+    long_description="This package contains pre-compiled CUDA kernels (cubins) for FlashInfer. "
+    "It provides all necessary cubin files downloaded from the FlashInfer artifactory.",
+    long_description_content_type="text/plain",
+    author="FlashInfer team",
+    url="https://github.com/flashinfer-ai/flashinfer",
+    packages=find_packages(),
+    package_data={
+        "flashinfer_cubin": ["cubins/**/*"],
+    },
+    include_package_data=True,
+    python_requires=">=3.8",
+    install_requires=[
+        "requests",
+        "filelock",
+    ],
+    cmdclass={
+        "build_py": DownloadAndBuildPy,
+        "sdist": CustomSdist,
+    },
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: Apache Software License",
+        "Operating System :: OS Independent",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+)

flashinfer/artifacts.py

@@ -168,7 +168,7 @@ def download_artifacts():
         results = []
         with (
             patch_logger_for_tqdm(logger),
-            tqdm(total=len(futures), desc="Downloading cubins") as pbar,
+            tqdm.tqdm(total=len(futures), desc="Downloading cubins") as pbar,
         ):
             for ret in as_completed(futures):
                 result = ret.result()