ci/cd: add github workflows to publish flashinfer-cubin wheel to pypi

.github/workflows/release-pypi-flashinfer-cubin-wheel.yml

@@ -0,0 +1,138 @@
+name: Publish flashinfer-cubin wheel to PyPI
+
+on:
+  workflow_dispatch:
+    inputs:
+      tag:
+        description: 'Tag (e.g., v1.2.3) to build and publish'
+        required: true
+        type: string
+
+jobs:
+  build-and-upload:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Validate tag format
+        run: |
+          if [[ ! "${{ inputs.tag }}" =~ ^v[0-9]+\.[0-9]+\.[0-9]+([a-z0-9]+)?$ ]]; then
+            echo "Error: Tag '${{ inputs.tag }}' does not match the expected format (e.g., v1.2.3 or v1.2.3.post1 or v1.2.3rc1)"
+            exit 1
+          fi
+          echo "✓ Tag format is valid: ${{ inputs.tag }}"
+
+      - name: Check out tag
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.tag }}
+          submodules: true
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+
+      - name: Verify tag matches version.txt (CRITICAL)
+        run: |
+          # Extract version from tag (remove 'v' prefix)
+          TAG_VERSION="${{ inputs.tag }}"
+          TAG_VERSION="${TAG_VERSION#v}"
+
+          # Check version.txt FIRST - this is the source of truth
+          if [ ! -f "version.txt" ]; then
+            echo "Error: version.txt file not found!"
+            exit 1
+          fi
+
+          VERSION_TXT=$(cat version.txt | tr -d '[:space:]')
+
+          if [ "$TAG_VERSION" != "$VERSION_TXT" ]; then
+            echo "❌ CRITICAL ERROR: version.txt does not match tag!"
+            echo "  Tag version: $TAG_VERSION"
+            echo "  version.txt: $VERSION_TXT"
+            echo ""
+            echo "Please update version.txt to match the release version before creating a release."
+            echo "The tag should be 'v$VERSION_TXT' (e.g., if version.txt contains '1.2.3', tag should be 'v1.2.3')"
+            exit 1
+          fi
+
+          echo "✓ version.txt matches tag version: $VERSION_TXT"
+
+      - name: Install build dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install build twine wheel
+          # Install dependencies required for building flashinfer-cubin
+          pip install setuptools>=61.0 requests filelock torch tqdm
+
+      - name: Build flashinfer-cubin wheel
+        run: |
+          echo "Building flashinfer-cubin wheel..."
+          cd flashinfer-cubin
+
+          # Clean any previous builds
+          rm -rf dist build *.egg-info
+
+          # Build the wheel using the build module for better isolation
+          python -m build --wheel
+
+          echo "✓ Build completed"
+          ls -lh dist/
+
+          # Move the wheel to the root dist directory for consistency
+          mkdir -p ../dist
+          cp dist/*.whl ../dist/
+          cd ..
+
+      - name: Check wheel contents
+        run: |
+          echo "Verifying wheel contents..."
+          # Create a temp directory to extract and check wheel
+          mkdir -p temp_wheel_check
+          cd temp_wheel_check
+
+          # Extract wheel (wheels are just zip files)
+          unzip -l ../dist/flashinfer_cubin*.whl | head -30
+          echo "..."
+
+          # Check if cubins are included
+          unzip -l ../dist/flashinfer_cubin*.whl | grep -c "\.cubin" || true
+
+          cd ..
+          rm -rf temp_wheel_check
+          echo "✓ Wheel archive created successfully"
+
+      - name: Store build artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+          retention-days: 7
+
+      - name: Test installation from wheel
+        run: |
+          echo "Testing installation from built wheel..."
+          python -m venv test-env
+          source test-env/bin/activate
+
+          # Install the wheel
+          pip install dist/flashinfer_cubin*.whl
+
+          # Test import and check for cubins
+          python -c "import flashinfer_cubin"
+          deactivate
+          rm -rf test-env
+
+      - name: Check package with twine
+        run: |
+          echo "Running twine check..."
+          twine check dist/flashinfer_cubin*.whl
+          echo "✓ Package validation passed"
+
+      - name: Upload to PyPI
+        env:
+          TWINE_USERNAME: __token__
+          TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
+        run: |
+          echo "Uploading flashinfer-cubin wheel to PyPI..."
+          twine upload --verbose --non-interactive dist/flashinfer_cubin*.whl
+          echo "✓ Successfully uploaded to PyPI"

flashinfer-cubin/.gitignore

@@ -1,2 +1,3 @@
 build/
 flashinfer_cubin/cubins/
+flashinfer_cubin/_build_meta.py

flashinfer-cubin/build_backend.py

@@ -0,0 +1,102 @@
+"""
+Custom build backend that downloads cubins before building the package.
+"""
+
+import os
+import sys
+from pathlib import Path
+from setuptools import build_meta as _orig
+from setuptools.build_meta import *
+
+# Add parent directory to path to import artifacts module
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+# add flashinfer._build_meta if not there, it should exist in Path(__file__).parent.parent / "flashinfer" / "_build_meta.py"
+build_meta_file = Path(__file__).parent.parent / "flashinfer" / "_build_meta.py"
+if not build_meta_file.exists():
+    version_file = Path(__file__).parent.parent / "version.txt"
+    if version_file.exists():
+        with open(version_file, "r") as f:
+            version = f.read().strip()
+    with open(build_meta_file, "w") as f:
+        f.write('"""Build metadata for flashinfer package."""\n')
+        f.write(f'__version__ = "{version}"\n')
+
+
+def _download_cubins():
+    """Download cubins to the source directory before building."""
+    from flashinfer.artifacts import download_artifacts
+
+    # Create cubins directory in the source tree
+    cubin_dir = Path(__file__).parent / "flashinfer_cubin" / "cubins"
+    cubin_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_dir)
+
+    try:
+        print(f"Downloading cubins to {cubin_dir}...")
+        download_artifacts()
+        print(f"Successfully downloaded cubins to {cubin_dir}")
+
+        # Count the downloaded files
+        cubin_files = list(cubin_dir.rglob("*.cubin"))
+        print(f"Downloaded {len(cubin_files)} cubin files")
+
+    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)
+
+
+def _create_build_metadata():
+    """Create build metadata file with version information."""
+    version_file = Path(__file__).parent.parent / "version.txt"
+    if version_file.exists():
+        with open(version_file, "r") as f:
+            version = f.read().strip()
+    else:
+        version = "0.0.0+unknown"
+
+    # Create build metadata in the source tree
+    package_dir = Path(__file__).parent / "flashinfer_cubin"
+    build_meta_file = package_dir / "_build_meta.py"
+
+    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}")
+    return version
+
+
+def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
+    """Build a wheel, downloading cubins first."""
+    _download_cubins()
+    _create_build_metadata()
+    return _orig.build_wheel(wheel_directory, config_settings, metadata_directory)
+
+
+def build_sdist(sdist_directory, config_settings=None):
+    """Build a source distribution, downloading cubins first."""
+    _download_cubins()
+    _create_build_metadata()
+    return _orig.build_sdist(sdist_directory, config_settings)
+
+
+def build_editable(wheel_directory, config_settings=None, metadata_directory=None):
+    """Build an editable install, downloading cubins first."""
+    _download_cubins()
+    _create_build_metadata()
+    return _orig.build_editable(wheel_directory, config_settings, metadata_directory)
+
+
+# Pass through all other hooks
+get_requires_for_build_wheel = _orig.get_requires_for_build_wheel
+get_requires_for_build_sdist = _orig.get_requires_for_build_sdist
+get_requires_for_build_editable = _orig.get_requires_for_build_editable
+prepare_metadata_for_build_wheel = _orig.prepare_metadata_for_build_wheel
+prepare_metadata_for_build_editable = _orig.prepare_metadata_for_build_editable

flashinfer-cubin/pyproject.toml

@@ -1,6 +1,7 @@
 [build-system]
-requires = ["setuptools>=61.0", "wheel", "requests", "filelock", "torch", "tqdm"] # NOTE(Zihao): we should remove torch once https://github.com/flashinfer-ai/flashinfer/pull/1641 merged
-build-backend = "setuptools.build_meta"
+requires = ["setuptools>=61.0", "wheel", "requests", "filelock", "torch", "tqdm", "numpy"] # NOTE(Zihao): we should remove torch once https://github.com/flashinfer-ai/flashinfer/pull/1641 merged
+build-backend = "build_backend"
+backend-path = ["."]
 
 [project]
 name = "flashinfer-cubin"
@@ -26,6 +27,7 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 dependencies = [
@@ -47,8 +49,4 @@ include-package-data = true
 version = {attr = "flashinfer_cubin.__version__"}
 
 [tool.setuptools.package-data]
-flashinfer_cubin = ["cubins/**/*"]
-
-[tool.setuptools.cmdclass]
-build_py = "setup.DownloadAndBuildPy"
-sdist = "setup.CustomSdist"
+flashinfer_cubin = ["cubins/**/*", "_build_meta.py"]