CI: sandboxed builds and enable Cachix pushes

.github/workflows/build_kernel.yaml

@@ -9,20 +9,25 @@ on:
 
 jobs:
   build:
-    name: Build kernel
+    name: Build kernels
     runs-on:
-      group: aws-g6-12xlarge-plus
+      group: aws-highmemory-32-plus-nix
     steps:
       - uses: actions/checkout@v4
-      - uses: cachix/install-nix-action@v27
+      - uses: DeterminateSystems/nix-installer-action@main
         with:
-          nix_path: nixpkgs=channel:nixos-unstable
-      - uses: cachix/cachix-action@v14
+          extra-conf: |
+            max-jobs = 4
+            cores = 12
+            sandbox-fallback = false
+      - uses: cachix/cachix-action@v16
         with:
           name: huggingface
-          #authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
         env:
-          USER: github_runner
+          USER: runner
+      - name: Nix info
+        run: nix-shell -p nix-info --run "nix-info -m"
       - name: Build activation kernel
         run: ( cd examples/activation && nix build .\#redistributable.torch29-cxx11-cu126-x86_64-linux )
       - name: Copy activation kernel
@@ -58,6 +63,31 @@ jobs:
       - name: Copy silu-and-mul-universal kernel
         run: cp -rL examples/silu-and-mul-universal/result silu-and-mul-universal-kernel
 
+      - name: Upload kernel artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: built-kernels
+          path: |
+            activation-kernel
+            cutlass-gemm-kernel
+            relu-kernel
+            relu-backprop-compile-kernel
+            silu-and-mul-universal-kernel
+
+  test:
+    name: Test kernels
+    needs: build
+    runs-on:
+      group: aws-g6-12xlarge-plus
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Download kernel artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: built-kernels
+          path: .
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
       - name: Build Docker image

.github/workflows/build_kernel_rocm.yaml

@@ -11,18 +11,23 @@ jobs:
   build:
     name: Build kernel
     runs-on:
-      group: aws-g6-12xlarge-plus
+      group: aws-highmemory-32-plus-nix
     steps:
       - uses: actions/checkout@v4
-      - uses: cachix/install-nix-action@v27
+      - uses: DeterminateSystems/nix-installer-action@main
         with:
-          nix_path: nixpkgs=channel:nixos-unstable
+          extra-conf: |
+            max-jobs = 4
+            cores = 12
+            sandbox-fallback = false
       - uses: cachix/cachix-action@v14
         with:
           name: huggingface
-          #authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
         env:
-          USER: github_runner
+          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

.github/workflows/build_kernel_xpu.yaml

@@ -11,18 +11,23 @@ jobs:
   build:
     name: Build kernel
     runs-on:
-      group: aws-g6-12xlarge-plus
+      group: aws-highmemory-32-plus-nix
     steps:
       - uses: actions/checkout@v4
-      - uses: cachix/install-nix-action@v27
+      - uses: DeterminateSystems/nix-installer-action@main
         with:
-          nix_path: nixpkgs=channel:nixos-unstable
-      - uses: cachix/cachix-action@v14
+          extra-conf: |
+            max-jobs = 4
+            cores = 12
+            sandbox-fallback = false
+      - uses: cachix/cachix-action@v16
         with:
           name: huggingface
-          #authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
         env:
-          USER: github_runner
+          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

.github/workflows/test_extra_commands.yaml

@@ -10,18 +10,19 @@ on:
 jobs:
   build:
     name: Build kernel
-    runs-on:
-      group: aws-g6-12xlarge-plus
     steps:
       - uses: actions/checkout@v4
-      - uses: cachix/install-nix-action@v27
+      - uses: DeterminateSystems/nix-installer-action@main
         with:
-          nix_path: nixpkgs=channel:nixos-unstable
+          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: github_runner
+          USER: runner
       - name: Test nix run .#kernels
         run: ( cd examples/relu ; nix run .#kernels -- lock ../../tests/run-kernels )