Add MXFP8 Support and Consolidated stable scaled mm API (#26) #1
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: Build Wheels | |
| on: | |
| push: | |
| branches: | |
| - main | |
| tags: | |
| - "v*" | |
| pull_request: | |
| workflow_dispatch: | |
| permissions: | |
| id-token: write | |
| contents: read | |
| jobs: | |
| build_wheels_linux: | |
| name: Build Linux wheels | |
| runs-on: ubuntu-22.04 | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| - name: Set up Python 3.12 | |
| uses: actions/setup-python@v5 | |
| with: | |
| python-version: "3.12" | |
| # Build CUDA wheel using cibuildwheel with manylinux container | |
| - name: Build CUDA wheel (manylinux) | |
| uses: pypa/cibuildwheel@v2.22.0 | |
| env: | |
| # Use NVIDIA's CUDA devel image based on UBI8 (manylinux_2_28 compatible) | |
| # CUDA 13.0 will be installed inside the container | |
| CIBW_MANYLINUX_X86_64_IMAGE: quay.io/pypa/manylinux_2_28_x86_64 | |
| # Build for Python 3.10, 3.11, and 3.12 | |
| # 3.10/3.11 will be version-specific, 3.12 will be abi3 (stable ABI) | |
| CIBW_BUILD: "cp310-manylinux_x86_64 cp311-manylinux_x86_64 cp312-manylinux_x86_64" | |
| CIBW_SKIP: "*-musllinux_*" | |
| # Install CUDA 13.0 toolkit inside the container | |
| CIBW_BEFORE_ALL_LINUX: | | |
| # Install dnf plugins for config-manager command | |
| dnf install -y dnf-plugins-core | |
| # Add NVIDIA CUDA repository for RHEL 8 | |
| dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo | |
| dnf clean all | |
| # Install ONLY the minimal CUDA packages needed for compilation (not the full toolkit) | |
| dnf install -y cuda-nvcc-13-0 cuda-cudart-devel-13-0 libcublas-devel-13-0 libnvjitlink-devel-13-0 | |
| # Install Python development headers (fallback for CMake FindPython) | |
| dnf install -y python3-devel | |
| # Verify CUDA installation | |
| /usr/local/cuda-13.0/bin/nvcc --version | |
| # Set CUDA environment and install build dependencies | |
| CIBW_ENVIRONMENT_LINUX: > | |
| CUDA_HOME=/usr/local/cuda-13.0 | |
| PATH=/usr/local/cuda-13.0/bin:$PATH | |
| LD_LIBRARY_PATH=/usr/local/cuda-13.0/lib64:$LD_LIBRARY_PATH | |
| LIBRARY_PATH=/usr/local/cuda-13.0/lib64/stubs:$LIBRARY_PATH | |
| # Install Python build dependencies | |
| CIBW_BEFORE_BUILD: "pip install nanobind cmake setuptools wheel" | |
| # Repair wheel for manylinux compliance, excluding CUDA runtime libs | |
| # Users must have CUDA installed on their system | |
| CIBW_REPAIR_WHEEL_COMMAND_LINUX: > | |
| auditwheel repair -w {dest_dir} {wheel} | |
| --exclude libcuda.so.1 | |
| --exclude libcudart.so.13 | |
| --exclude libnvrtc.so.13 | |
| --exclude libcublas.so.13 | |
| --exclude libcublasLt.so.13 | |
| --exclude libnvJitLink.so.13 | |
| with: | |
| output-dir: wheelhouse | |
| - name: Build CPU-only wheel (py3-none-any) | |
| shell: bash | |
| run: | | |
| python -m pip install --upgrade pip | |
| pip install build setuptools wheel | |
| # Build CPU-only variant (pure Python, no CUDA) | |
| python setup.py bdist_wheel --no-cuda | |
| # Move to wheelhouse directory | |
| mv dist/*py3-none-any.whl wheelhouse/ | |
| - name: List built wheels | |
| shell: bash | |
| run: | | |
| ls -lh wheelhouse/ | |
| echo "Built wheels:" | |
| ls -1 wheelhouse/*.whl | |
| - name: Upload Linux CUDA wheel | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: wheels-linux-cuda | |
| path: wheelhouse/*manylinux*.whl | |
| if-no-files-found: error | |
| - name: Upload CPU-only wheel | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: wheels-cpu-only | |
| path: wheelhouse/*py3-none-any.whl | |
| if-no-files-found: error | |
| build_wheels_linux_arm64: | |
| name: Build Linux ARM64 wheels | |
| runs-on: ubuntu-22.04-arm | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| - name: Set up Python 3.12 | |
| uses: actions/setup-python@v5 | |
| with: | |
| python-version: "3.12" | |
| # Build CUDA wheel using cibuildwheel with manylinux container for ARM64 | |
| - name: Build CUDA wheel (manylinux ARM64) | |
| uses: pypa/cibuildwheel@v2.22.0 | |
| env: | |
| # Use manylinux_2_28 for ARM64 | |
| CIBW_MANYLINUX_AARCH64_IMAGE: quay.io/pypa/manylinux_2_28_aarch64 | |
| # Build for Python 3.10, 3.11, and 3.12 on ARM64 | |
| CIBW_BUILD: "cp310-manylinux_aarch64 cp311-manylinux_aarch64 cp312-manylinux_aarch64" | |
| CIBW_SKIP: "*-musllinux_*" | |
| # Native ARM64 runner - no emulation needed | |
| CIBW_ARCHS_LINUX: aarch64 | |
| # Install CUDA 13.0 toolkit inside the container (using sbsa repo for ARM64) | |
| CIBW_BEFORE_ALL_LINUX: | | |
| # Install dnf plugins for config-manager command | |
| dnf install -y dnf-plugins-core | |
| # Add NVIDIA CUDA repository for RHEL 8 ARM64 (sbsa = server-based system architecture) | |
| dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo | |
| dnf clean all | |
| # Install ONLY the minimal CUDA packages needed for compilation (not the full toolkit) | |
| dnf install -y cuda-nvcc-13-0 cuda-cudart-devel-13-0 libcublas-devel-13-0 libnvjitlink-devel-13-0 | |
| # Install Python development headers (fallback for CMake FindPython) | |
| dnf install -y python3-devel | |
| # Verify CUDA installation | |
| /usr/local/cuda-13.0/bin/nvcc --version | |
| # Set CUDA environment and install build dependencies | |
| CIBW_ENVIRONMENT_LINUX: > | |
| CUDA_HOME=/usr/local/cuda-13.0 | |
| PATH=/usr/local/cuda-13.0/bin:$PATH | |
| LD_LIBRARY_PATH=/usr/local/cuda-13.0/lib64:$LD_LIBRARY_PATH | |
| LIBRARY_PATH=/usr/local/cuda-13.0/lib64/stubs:$LIBRARY_PATH | |
| # Install Python build dependencies | |
| CIBW_BEFORE_BUILD: "pip install nanobind cmake setuptools wheel" | |
| # Repair wheel for manylinux compliance, excluding CUDA runtime libs | |
| # Users must have CUDA installed on their system | |
| CIBW_REPAIR_WHEEL_COMMAND_LINUX: > | |
| auditwheel repair -w {dest_dir} {wheel} | |
| --exclude libcuda.so.1 | |
| --exclude libcudart.so.13 | |
| --exclude libnvrtc.so.13 | |
| --exclude libcublas.so.13 | |
| --exclude libcublasLt.so.13 | |
| --exclude libnvJitLink.so.13 | |
| with: | |
| output-dir: wheelhouse | |
| - name: List built wheels | |
| shell: bash | |
| run: | | |
| ls -lh wheelhouse/ | |
| echo "Built wheels:" | |
| ls -1 wheelhouse/*.whl | |
| - name: Upload Linux ARM64 CUDA wheel | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: wheels-linux-arm64-cuda | |
| path: wheelhouse/*manylinux*.whl | |
| if-no-files-found: error | |
| build_wheels_windows: | |
| name: Build Windows wheel (Python ${{ matrix.python-version }}) | |
| runs-on: windows-2022 | |
| strategy: | |
| matrix: | |
| python-version: ["3.10", "3.11", "3.12"] | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| - name: Install CUDA Toolkit 13.0.2 | |
| uses: Jimver/cuda-toolkit@v0.2.29 | |
| id: cuda-toolkit | |
| with: | |
| cuda: '13.0.2' | |
| method: 'network' | |
| use-github-cache: 'true' | |
| log-file-suffix: 'windows-2022.txt' | |
| - name: Set additional CUDA environment variables | |
| shell: pwsh | |
| run: | | |
| echo "CUDA_PATH_V13_0=$env:CUDA_PATH" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append | |
| - name: Verify CUDA installation | |
| shell: bash | |
| run: | | |
| echo "CUDA_HOME: $CUDA_HOME" | |
| echo "CUDA_PATH: $CUDA_PATH" | |
| nvcc --version | |
| - name: Setup MSVC | |
| uses: ilammy/msvc-dev-cmd@v1 | |
| - name: Setup CUDA MSBuild integration | |
| shell: pwsh | |
| run: | | |
| # Copy CUDA MSBuild integration files to Visual Studio directory | |
| $cudaPath = $env:CUDA_PATH | |
| $cudaMSBuildSource = Join-Path $cudaPath "extras\visual_studio_integration\MSBuildExtensions" | |
| if (-not (Test-Path $cudaMSBuildSource)) { | |
| Write-Error "CUDA MSBuild extensions not found at $cudaMSBuildSource" | |
| exit 1 | |
| } | |
| # Find Visual Studio installation | |
| $vsBasePaths = @( | |
| "C:\Program Files (x86)\Microsoft Visual Studio", | |
| "C:\Program Files\Microsoft Visual Studio" | |
| ) | |
| $vsVersions = @("2022", "2019") | |
| $vsEditions = @("BuildTools", "Enterprise", "Professional", "Community") | |
| $toolsets = @("v170", "v160") | |
| $destinations = @() | |
| foreach ($basePath in $vsBasePaths) { | |
| if (Test-Path $basePath) { | |
| foreach ($version in $vsVersions) { | |
| foreach ($edition in $vsEditions) { | |
| foreach ($toolset in $toolsets) { | |
| $vsPath = Join-Path $basePath "$version\$edition\MSBuild\Microsoft\VC\$toolset\BuildCustomizations" | |
| if (Test-Path $vsPath) { | |
| $destinations += $vsPath | |
| } | |
| } | |
| } | |
| } | |
| } | |
| } | |
| if ($destinations.Count -eq 0) { | |
| Write-Error "Could not find Visual Studio MSBuild directory" | |
| exit 1 | |
| } | |
| # CUDA 13.0 uses versioned filenames that need to be copied as both versioned and standard names | |
| $fileMapping = @{ | |
| "CUDA 13.0.props" = @("CUDA 13.0.props", "CUDA.props") | |
| "CUDA 13.0.targets" = @("CUDA 13.0.targets", "CUDA.targets") | |
| "CUDA 13.0.xml" = @("CUDA 13.0.xml", "CUDA.xml") | |
| "Nvda.Build.CudaTasks.v13.0.dll" = @("Nvda.Build.CudaTasks.v13.0.dll") | |
| } | |
| $copySuccess = $false | |
| foreach ($destDir in $destinations) { | |
| try { | |
| foreach ($srcFileName in $fileMapping.Keys) { | |
| $srcFile = Join-Path $cudaMSBuildSource $srcFileName | |
| if (Test-Path $srcFile) { | |
| foreach ($destFileName in $fileMapping[$srcFileName]) { | |
| $destFile = Join-Path $destDir $destFileName | |
| Copy-Item -Path $srcFile -Destination $destFile -Force | |
| } | |
| } | |
| } | |
| $copySuccess = $true | |
| } catch { | |
| Write-Warning "Failed to copy to $destDir : $_" | |
| } | |
| } | |
| if (-not $copySuccess) { | |
| Write-Error "Failed to copy CUDA MSBuild files to any destination" | |
| exit 1 | |
| } | |
| Write-Host "CUDA MSBuild integration configured successfully" | |
| - name: Set up Python ${{ matrix.python-version }} | |
| uses: actions/setup-python@v5 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - name: Create python3.lib for stable ABI | |
| if: matrix.python-version == '3.12' | |
| shell: pwsh | |
| run: | | |
| # Create python3.lib for stable ABI support | |
| $pythonDir = Split-Path -Parent (Get-Command python).Source | |
| $libsDir = Join-Path $pythonDir "libs" | |
| $pythonLib = Get-ChildItem -Path $libsDir -Filter "python3*.lib" | Where-Object { $_.Name -match "python3\d+\.lib" } | Select-Object -First 1 | |
| if ($pythonLib) { | |
| $python3Lib = Join-Path $libsDir "python3.lib" | |
| if (-not (Test-Path $python3Lib)) { | |
| Copy-Item -Path $pythonLib.FullName -Destination $python3Lib | |
| Write-Host "Created python3.lib for stable ABI" | |
| } | |
| } else { | |
| Write-Error "Could not find python3XX.lib" | |
| exit 1 | |
| } | |
| - name: Install CMake | |
| uses: lukka/get-cmake@latest | |
| with: | |
| cmakeVersion: "~3.27.0" | |
| - name: Install build dependencies | |
| shell: bash | |
| run: | | |
| python -m pip install --upgrade pip | |
| pip install build setuptools wheel nanobind | |
| - name: Build CUDA wheel | |
| shell: bash | |
| run: | | |
| # Build with default platform-specific CUDA architectures from setup.py | |
| # Python 3.12 will produce abi3 wheel, 3.10/3.11 will be version-specific | |
| python setup.py bdist_wheel | |
| - name: List built wheels | |
| shell: bash | |
| run: | | |
| ls -lh dist/ | |
| echo "Built wheels:" | |
| ls -1 dist/*.whl | |
| - name: Upload Windows wheel | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: wheels-windows-cuda-${{ matrix.python-version }} | |
| path: dist/*.whl | |
| if-no-files-found: error | |
| test_linux: | |
| name: Test on Linux (Python ${{ matrix.python-version }}) | |
| needs: build_wheels_linux | |
| runs-on: ubuntu-22.04 | |
| strategy: | |
| matrix: | |
| python-version: ["3.10", "3.11", "3.12"] | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| - name: Set up Python ${{ matrix.python-version }} | |
| uses: actions/setup-python@v5 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - name: Download CPU-only wheel | |
| uses: actions/download-artifact@v4 | |
| with: | |
| name: wheels-cpu-only | |
| path: dist/ | |
| - name: Install wheel and test dependencies | |
| shell: bash | |
| run: | | |
| python -m pip install --upgrade pip | |
| # Install CPU-only PyTorch first (smaller download, no CUDA needed for CI) | |
| pip install torch --index-url https://download.pytorch.org/whl/cpu | |
| # Install the CPU-only wheel (py3-none-any works everywhere) | |
| pip install dist/*py3-none-any.whl | |
| pip install pytest ruff | |
| - name: Ruff check | |
| shell: bash | |
| run: | | |
| ruff check . | |
| - name: Run tests | |
| shell: bash | |
| run: | | |
| # Run all tests - conftest handles backend availability | |
| # Tests requiring CUDA will be skipped on CPU-only runners | |
| python -m pytest tests/ -v --tb=short | |
| test_windows: | |
| name: Test on Windows (Python ${{ matrix.python-version }}) | |
| needs: build_wheels_windows | |
| runs-on: windows-2022 | |
| strategy: | |
| matrix: | |
| python-version: ["3.10", "3.11", "3.12"] | |
| steps: | |
| - name: Checkout code | |
| uses: actions/checkout@v4 | |
| - name: Set up Python ${{ matrix.python-version }} | |
| uses: actions/setup-python@v5 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - name: Download Windows wheel | |
| uses: actions/download-artifact@v4 | |
| with: | |
| name: wheels-windows-cuda-${{ matrix.python-version }} | |
| path: dist/ | |
| - name: Install wheel and test dependencies | |
| shell: bash | |
| run: | | |
| python -m pip install --upgrade pip | |
| # Install CPU-only PyTorch first (smaller download, no CUDA needed for CI) | |
| pip install torch --index-url https://download.pytorch.org/whl/cpu | |
| # Windows: install the CUDA wheel (will work in CPU-only mode) | |
| pip install dist/*.whl | |
| pip install pytest ruff | |
| - name: Ruff check | |
| shell: bash | |
| run: | | |
| ruff check . | |
| - name: Run tests | |
| shell: bash | |
| run: | | |
| # Run all tests - conftest handles backend availability | |
| # Tests requiring CUDA will be skipped on CPU-only runners | |
| python -m pytest tests/ -v --tb=short | |
| publish: | |
| name: Publish to PyPI | |
| needs: [build_wheels_linux, build_wheels_linux_arm64, build_wheels_windows, test_linux, test_windows] | |
| runs-on: ubuntu-latest | |
| if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v') | |
| environment: pypi | |
| steps: | |
| - name: Download Linux CUDA wheel | |
| uses: actions/download-artifact@v4 | |
| with: | |
| name: wheels-linux-cuda | |
| path: dist/ | |
| - name: Download Linux ARM64 CUDA wheel | |
| uses: actions/download-artifact@v4 | |
| with: | |
| name: wheels-linux-arm64-cuda | |
| path: dist/ | |
| - name: Download Windows CUDA wheel (3.10) | |
| uses: actions/download-artifact@v4 | |
| with: | |
| name: wheels-windows-cuda-3.10 | |
| path: dist/ | |
| - name: Download Windows CUDA wheel (3.11) | |
| uses: actions/download-artifact@v4 | |
| with: | |
| name: wheels-windows-cuda-3.11 | |
| path: dist/ | |
| - name: Download Windows CUDA wheel (3.12) | |
| uses: actions/download-artifact@v4 | |
| with: | |
| name: wheels-windows-cuda-3.12 | |
| path: dist/ | |
| - name: Download CPU-only wheel | |
| uses: actions/download-artifact@v4 | |
| with: | |
| name: wheels-cpu-only | |
| path: dist/ | |
| - name: List wheels to publish | |
| run: | | |
| echo "Wheels to publish:" | |
| ls -la dist/ | |
| - name: Publish to PyPI | |
| uses: pypa/gh-action-pypi-publish@release/v1 | |
| with: | |
| password: ${{ secrets.PYPI_TOKEN }} | |
| skip-existing: true | |
| verbose: true |