kernels/.github/workflows/build_kernel.yaml at 26824dbbe377dce6a41c83e075ac85acbed2d194 · huggingface/kernels · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
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 (${{ matrix.arch }})
    strategy:
      matrix:
        include:
          - arch: x86_64-linux
            runner: aws-highmemory-32-plus-nix
          - arch: aarch64-linux
            runner: aws-r8g-8xl-plus-nix
    runs-on:
      group: ${{ matrix.runner }}
    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-${{ matrix.arch }} )
      - name: Copy relu kernel
        run: cp -rL builder/examples/relu/result relu-kernel

      - name: Build relu-tvm-ffi kernel
        run: ( cd builder/examples/relu-tvm-ffi && nix build .\#redistributable.tvm-ffi01-cu126-${{ matrix.arch }} )
      - name: Copy relu-tvm-ffi kernel
        run: cp -rL builder/examples/relu-tvm-ffi/result relu-tvm-ffi-kernel

      - name: Build extra-data kernel
        run: ( cd builder/examples/extra-data && nix build .\#redistributable.torch29-cxx11-cu126-${{ matrix.arch }} )
      - name: Copy extra-data kernel
        run: cp -rL builder/examples/extra-data/result extra-data

      - name: Build relu kernel (CPU)
        run: ( cd builder/examples/relu && nix build .\#redistributable.torch29-cxx11-cpu-${{ matrix.arch }} )
      - 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-${{ matrix.arch }} )
      - 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-${{ matrix.arch }} )
      - 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-${{ matrix.arch }} )

      - 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.${{ matrix.arch }}.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-${{ matrix.arch }}
          path: |
            activation-kernel
            cutlass-gemm-kernel
            extra-data
            relu-kernel
            relu-tvm-ffi-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-x86_64-linux
          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: builder/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