kernel-builder/.github/workflows/build_kernel.yaml at main · sywangyi/kernel-builder · 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
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 examples/relu && nix build .\#redistributable.torch29-cxx11-cu126-x86_64-linux )
      - name: Copy relu kernel
        run: cp -rL examples/relu/result relu-kernel

      - name: Build relu kernel (CPU)
        run: ( cd examples/relu && nix build .\#redistributable.torch29-cxx11-cpu-x86_64-linux )
      - name: Copy relu kernel (CPU)
        run: cp -rL examples/relu/result relu-kernel-cpu

      - name: Build cutlass GEMM kernel
        run: ( cd examples/cutlass-gemm && nix build .\#redistributable.torch29-cxx11-cu126-x86_64-linux )
      - name: Copy cutlass GEMM kernel
        run: cp -rL examples/cutlass-gemm/result cutlass-gemm-kernel

      - name: Build relu-backprop-compile kernel
        run: ( cd examples/relu-backprop-compile && nix build .\#redistributable.torch29-cxx11-cu126-x86_64-linux )
      - name: Copy relu-backprop-compile kernel
        run: cp -rL 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 examples/relu-specific-torch && nix build . )

      - name: Build relu kernel (compiler flags)
        run: ( cd 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 examples/relu && nix build .#devShells.x86_64-linux.test )

      - name: Build silu-and-mul kernel
        run: ( cd examples/silu-and-mul && nix build .\#redistributable.torch-cuda )
      - name: Copy silu-and-mul kernel
        run: cp -rL 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