flash-attention/.github/workflows/ci.yml at 5835c733e7e9c07606b045255768e8a7e9e851bd · Dao-AILab/flash-attention · 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
name: CI

on:
  push:
    branches: [main, ci-fix]

permissions:
  contents: read

env:
  CI_WORK_DIR: ${{ vars.CI_WORK_DIR || format('/scratch/user/{0}', github.actor) }}
  # The mla_absorbed cases exercise the MLA backward kernels (flash_bwd_mla_sm100.py +
  # dq_dqv + dk); without them CI runs no MLA test at all. We cover the distinct backward
  # paths: sparse (kv_sparsity=True) non-causal + causal, dense (kv_sparsity=False), and
  # shared_kv=True. (deterministic/hdim/mha_type are single-valued for this test.)
  FA4_TEST_FILTER: >-
    1024-1024-128-True-0-0.0-False-False-False-mha-dtype0
    or 1024-1024-128-False-0-0.0-False-False-False-mha-dtype0
    or test_flash_attn_ex2_emu_decode_prefill_consistency
    or test_flash_attn_mla_absorbed and 256-256-False-True-64-False-0-False-False-mqa-dtype0
    or test_flash_attn_mla_absorbed and 256-256-False-True-64-True-0-False-False-mqa-dtype0
    or test_flash_attn_mla_absorbed and 256-256-False-False-64-False-0-False-False-mqa-dtype0
    or test_flash_attn_mla_absorbed and 256-256-True-True-64-False-0-False-False-mqa-dtype0

jobs:
  lint:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - name: Install ruff
        run: pip install ruff
      - name: Ruff check
        run: ruff check flash_attn/cute/ --extend-exclude "flash_attn/cute/flash_bwd.py,flash_attn/cute/flash_fwd.py,flash_attn/cute/flash_fwd_sm100.py,flash_attn/cute/interface.py"
      - name: Ruff format
        run: ruff format --check flash_attn/cute/ --exclude "flash_attn/cute/flash_bwd.py,flash_attn/cute/flash_fwd.py,flash_attn/cute/flash_fwd_sm100.py,flash_attn/cute/interface.py"

  fa4-correctness-and-benchmark:
    strategy:
      fail-fast: false
      matrix:
        gpu: [b200]
    runs-on: [self-hosted, '${{ matrix.gpu }}']
    name: fa4-correctness-and-benchmark (${{ matrix.gpu }})
    timeout-minutes: 60
    steps:
      - uses: actions/checkout@v4
      - name: Log in to Docker Hub
        uses: docker/login-action@v3
        with:
          username: ${{ secrets.DOCKERHUB_USERNAME }}
          password: ${{ secrets.DOCKERHUB_TOKEN }}
      - uses: ./.github/actions/gpu-test
        with:
          test-filter: ${{ env.FA4_TEST_FILTER }}
          fa4_image_cu129: "togethercomputer/training-performance:flash-attn-cu12.9-26.03.25@sha256:304a5c3d2b3a75b151cd2a964cd26d444e0d8b5686d63943df13378c9705f943"
          fa4_image_cu130: "togethercomputer/training-performance:flash-attn-cu13.0-26.06.27@sha256:2c31843c7137cbe909611eb70a9b6940581a2e3018d2c8b629f84dfd04b9bfd7"