Merge kernel-builder into the kernels repo

.github/workflows/build_documentation.yaml

@@ -12,6 +12,7 @@ jobs:
     uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@main
     with:
       commit_sha: ${{ github.sha }}
+      package_path: kernels/
       package: kernels
     secrets:
       hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}

.github/workflows/build_kernel.yaml

@@ -0,0 +1,106 @@
+name: "Build and test kernel"
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+    types: [opened, synchronize, reopened] # trigger on PRs
+  workflow_dispatch:
+
+jobs:
+  build:
+    name: Build kernels
+    runs-on:
+      group: aws-highmemory-32-plus-nix
+    steps:
+      - uses: actions/checkout@v6
+      - uses: DeterminateSystems/nix-installer-action@main
+        with:
+          extra-conf: |
+            max-jobs = 4
+            cores = 12
+            sandbox-fallback = false
+      - uses: cachix/cachix-action@v16
+        with:
+          name: huggingface
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+        env:
+          USER: runner
+      - name: Nix info
+        run: nix-shell -p nix-info --run "nix-info -m"
+      - name: Build relu kernel
+        run: ( cd builder/examples/relu && nix build .\#redistributable.torch29-cxx11-cu126-x86_64-linux )
+      - name: Copy relu kernel
+        run: cp -rL builder/examples/relu/result relu-kernel
+
+      - name: Build relu kernel (CPU)
+        run: ( cd builder/examples/relu && nix build .\#redistributable.torch29-cxx11-cpu-x86_64-linux )
+      - name: Copy relu kernel (CPU)
+        run: cp -rL builder/examples/relu/result relu-kernel-cpu
+
+      - name: Build cutlass GEMM kernel
+        run: ( cd builder/examples/cutlass-gemm && nix build .\#redistributable.torch29-cxx11-cu126-x86_64-linux )
+      - name: Copy cutlass GEMM kernel
+        run: cp -rL builder/examples/cutlass-gemm/result cutlass-gemm-kernel
+
+      - name: Build relu-backprop-compile kernel
+        run: ( cd builder/examples/relu-backprop-compile && nix build .\#redistributable.torch29-cxx11-cu126-x86_64-linux )
+      - name: Copy relu-backprop-compile kernel
+        run: cp -rL builder/examples/relu-backprop-compile/result relu-backprop-compile-kernel
+
+      # Just test that we build with the extra torchVersions argument.
+      - name: Build relu kernel (specific Torch version)
+        run: ( cd builder/examples/relu-specific-torch && nix build . )
+
+      - name: Build relu kernel (compiler flags)
+        run: ( cd builder/examples/relu-compiler-flags && nix build .\#redistributable.torch29-cxx11-cu126-x86_64-linux )
+
+      - name: Test that we can build a test shell (e.g. that gcc corresponds to CUDA-required)
+        run: ( cd builder/examples/relu && nix build .#devShells.x86_64-linux.test )
+
+      - name: Build silu-and-mul kernel
+        run: ( cd builder/examples/silu-and-mul && nix build .\#redistributable.torch-cuda )
+      - name: Copy silu-and-mul kernel
+        run: cp -rL builder/examples/silu-and-mul/result silu-and-mul-kernel
+
+      - name: Upload kernel artifacts
+        uses: actions/upload-artifact@v6
+        with:
+          name: built-kernels
+          path: |
+            activation-kernel
+            cutlass-gemm-kernel
+            relu-kernel
+            relu-kernel-cpu
+            relu-backprop-compile-kernel
+            silu-and-mul-kernel
+
+  test:
+    name: Test kernels
+    needs: build
+    runs-on:
+      group: aws-g6-12xlarge-plus
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Download kernel artifacts
+        uses: actions/download-artifact@v7
+        with:
+          name: built-kernels
+          path: .
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      - name: Build Docker image
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: tests/Dockerfile.test-kernel
+          platforms: linux/amd64
+          load: true
+          push: false
+          tags: kernel-builder:latest
+
+      - name: Run Tests
+        run: |
+          docker run --gpus all kernel-builder:latest

.github/workflows/build_kernel_macos.yaml

@@ -0,0 +1,31 @@
+name: "Build and test kernel (macOS)"
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+    types: [opened, synchronize, reopened] # trigger on PRs
+  workflow_dispatch:
+
+jobs:
+  build:
+    name: Build kernel
+    runs-on: macos-26
+    steps:
+      - name: "Select Xcode"
+        run: sudo xcrun xcode-select -s /Applications/Xcode_26.0.app
+      - name: "Install Metal Toolchain"
+        run: xcodebuild -downloadComponent metalToolchain
+      - uses: actions/checkout@v6
+      - uses: cachix/install-nix-action@v31
+      - uses: cachix/cachix-action@v15
+        with:
+          name: huggingface
+          #authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+      # For now we only test that there are no regressions in building macOS
+      # kernels. Also run tests once we have a macOS runner.
+      - name: Build relu kernel
+        run: ( cd builder/examples/relu && nix build .\#redistributable.torch29-metal-aarch64-darwin -L )
+
+      - name: Build relu metal cpp kernel
+        run: ( cd builder/examples/relu-metal-cpp && nix build .\#redistributable.torch29-metal-aarch64-darwin -L )

.github/workflows/build_kernel_rocm.yaml

@@ -0,0 +1,37 @@
+name: "Build and test kernel (ROCm)"
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+    types: [opened, synchronize, reopened] # trigger on PRs
+  workflow_dispatch:
+
+jobs:
+  build:
+    name: Build kernel
+    runs-on:
+      group: aws-highmemory-32-plus-nix
+    steps:
+      - uses: actions/checkout@v6
+      - uses: DeterminateSystems/nix-installer-action@main
+        with:
+          extra-conf: |
+            max-jobs = 4
+            cores = 12
+            sandbox-fallback = false
+      - uses: cachix/cachix-action@v14
+        with:
+          name: huggingface
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+        env:
+          USER: runner
+      - name: Nix info
+        run: nix-shell -p nix-info --run "nix-info -m"
+      # For now we only test that there are no regressions in building ROCm
+      # kernels. Also run tests once we have a ROCm runner.
+      - name: Build relu kernel
+        run: ( cd buidler/examples/relu && nix build .\#redistributable.torch29-cxx11-rocm63-x86_64-linux -L )
+
+      - name: Build relu kernel (compiler flags)
+        run: ( cd builder/examples/relu-compiler-flags && nix build .\#redistributable.torch29-cxx11-rocm63-x86_64-linux )

.github/workflows/build_kernel_windows.yaml

@@ -0,0 +1,73 @@
+name: "Build and test kernel - Windows"
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+    types: [opened, synchronize, reopened] # trigger on PRs
+  workflow_dispatch:
+
+jobs:
+  build:
+    strategy:
+      matrix:
+        os: [ windows-2022 ]
+        python: [ 3.12 ]
+        torch: [
+#          { version: '2.9.1', cuda: '12.6.3', wheel: '126' },
+          { version: '2.9.1', cuda: '12.8.1', wheel: '128' },
+#          { version: '2.9.1', cuda: '13.0.1', wheel: '130' }
+        ]
+
+    name: Build kernel
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - uses: actions/cache@v5
+        with:
+          key: cuda-toolkit-v${{ matrix.cuda }}-${{ matrix.os }}
+          path: |
+            C:\Program Files\NVIDIA GPU Computing Toolkit
+            ~/.cargo/registry
+            ~/.cargo/git
+
+      - uses: actions/checkout@v6
+
+      # CUDA environment setup
+      - uses: Jimver/cuda-toolkit@v0.2.29
+        id: setup-cuda-toolkit
+        with:
+          cuda: ${{ matrix.torch.cuda }}  # TODO(mfuntowicz): How can we test multiple CUDA versions than align with torch?
+      - name: "NVCC checks"
+        run: nvcc -V
+
+      # Rust build environment setup
+      - uses: actions-rs/toolchain@v1
+        with:
+          toolchain: stable
+          profile: minimal
+          override: true
+
+      - name: Build build2cmake
+        run: ( cd build2cmake && cargo build --release )
+
+      # Python environment setup
+      - uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python }}
+          cache: 'pip'
+
+      - name: Install PyTorch
+        run: pip install torch --index-url https://download.pytorch.org/whl/cu${{ matrix.torch.wheel }}
+
+      - name: Build cutlass GEMM kernel
+        run: ( builder\scripts\windows\builder.ps1 -SourceFolder builder/examples/cutlass-gemm -BuildConfig Release -Backend cuda -Build -Force )
+
+      - name: Build relu kernel
+        run: ( builder\scripts\windows\builder.ps1 -SourceFolder builder/examples/relu -BuildConfig Release -Backend cuda -Build -Force )
+
+      - name: Build relu-backprop-compile kernel
+        run: ( builder\scripts\windows\builder.ps1 -SourceFolder builder/examples/relu-backprop-compile -BuildConfig Release -Backend cuda -Build -Force  )
+
+      - name: Build silu-and-mul kernel
+        run: ( builder\scripts\windows\builder.ps1 -SourceFolder builder/examples/silu-and-mul -BuildConfig Release -Backend cuda -Build -Force)

.github/workflows/build_kernel_xpu.yaml

@@ -0,0 +1,37 @@
+name: "Build and test kernel (XPU)"
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+    types: [opened, synchronize, reopened] # trigger on PRs
+  workflow_dispatch:
+
+jobs:
+  build:
+    name: Build kernel
+    runs-on:
+      group: aws-highmemory-32-plus-nix
+    steps:
+      - uses: actions/checkout@v6
+      - uses: DeterminateSystems/nix-installer-action@main
+        with:
+          extra-conf: |
+            max-jobs = 4
+            cores = 12
+            sandbox-fallback = false
+      - uses: cachix/cachix-action@v16
+        with:
+          name: huggingface
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+        env:
+          USER: runner
+      - name: Nix info
+        run: nix-shell -p nix-info --run "nix-info -m"
+      # For now we only test that there are no regressions in building XPU
+      # kernels. Also run tests once we have a XPU runner.
+      - name: Build relu kernel
+        run: ( cd builder/examples/relu && nix build .\#redistributable.torch29-cxx11-xpu20252-x86_64-linux -L )
+
+      - name: Build relu kernel (compiler flags)
+        run: ( cd builder/examples/relu-compiler-flags && nix build .\#redistributable.torch29-cxx11-xpu20252-x86_64-linux )

.github/workflows/build_pr_documentation.yaml

@@ -1,6 +1,11 @@
 name: Build PR Documentation
 
-on: pull_request
+on:
+  pull_request:
+    paths:
+      - "docs/**"
+      - "kernels/**"
+      - ".github/workflows/build_pr_documentation.yml"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -12,4 +17,5 @@ jobs:
     with:
       commit_sha: ${{ github.event.pull_request.head.sha }}
       pr_number: ${{ github.event.number }}
+      package_path: kernels/kernels/
       package: kernels

.github/workflows/check_variants.yaml

@@ -0,0 +1,38 @@
+name: "Check build variants"
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+    types: [opened, synchronize, reopened] # trigger on PRs
+  workflow_dispatch:
+
+jobs:
+  build:
+    name: Check build variants
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: cachix/install-nix-action@v27
+        with:
+          nix_path: nixpkgs=channel:nixos-unstable
+      - name: Generate variants JSON
+        run: nix eval --raw .#lib.allBuildVariantsJSON | nix run nixpkgs#jq 'walk(if type == "array" then sort else . end)' > build-variants.json
+      - name: Check if variants JSON is up-to-date
+        run: |
+          if git diff --exit-code build-variants.json; then
+            echo "✅ variants.json is up-to-date"
+          else
+            echo "🛑 regenerate variants.json: nix eval --raw .#lib.allBuildVariantsJSON | nix run nixpkgs#jq 'walk(if type == "array" then sort else . end)' > build-variants.json"
+            exit 1
+          fi
+      - name: Generate variants Markdown
+        run: nix run nixpkgs#python3 builder/scripts/gen_variants_markdown.py
+      - name: Check if variants Markdown is up-to-date
+        run: |
+          if git diff --exit-code builder/docs/build-variants.md; then
+            echo "✅ docs/build-variants.md is up-to-date"
+          else
+            echo "🛑 regenerate docs/build-variants: nix run nixpkgs#python3 builder/scripts/gen_variants_markdown.py"
+            exit 1
+          fi