[NVIDIA][CuTe,Fwd,sm120] Implement Pack-GQA on SM120 (+ graceful Spli… #98
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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" |