CI: extend FA4 test matrix with causal/non-causal correctness and fwd+bwd benchmark seqlen 1K-32K (#2428)

Author: Johnsonms

test_ci_local.sh .github/scripts/test_ci_local.shtest_ci_local.sh renamed to .github/scripts/test_ci_local.sh

@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../../" && pwd)
 
 python3 "$SCRIPT_DIR/tools/ci/run_fa4_ci.py" \
   --repo-root "$SCRIPT_DIR" \

.github/workflows/ci.yml

@@ -9,7 +9,7 @@ permissions:
 
 env:
   CI_WORK_DIR: ${{ vars.CI_WORK_DIR || format('/scratch/user/{0}', github.actor) }}
-  FA4_TEST_FILTER: "1-1-128-True-0-0.0-False-False-False-mha-dtype0"
+  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"
 
 jobs:
   lint:
@@ -23,13 +23,13 @@ jobs:
       - 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"
 
-  test:
+  fa4-correctness-and-benchmark:
     strategy:
       fail-fast: false
       matrix:
         gpu: [b200]
     runs-on: [self-hosted, '${{ matrix.gpu }}']
-    name: test (${{ matrix.gpu }})
+    name: fa4-correctness-and-benchmark (${{ matrix.gpu }})
     timeout-minutes: 60
     steps:
       - uses: actions/checkout@v4

tools/ci/run_fa4_ci.py

@@ -87,11 +87,12 @@ def build_step_plan(
     ]
     if not skip_benchmark:
         steps.append(Step(
-            name="Benchmark (FA4 fwd, hdim=128, seqlen=8192)",
+            name="Benchmark (FA4 fwd, hdim=128, causal=both, seqlen=1K-32K)",
             command=[
                 "python3", "benchmarks/benchmark_attn.py",
-                "--backend", "fa4", "--fwd",
-                "--headdim", "128", "--seqlen", "8192",
+                "--backend", "fa4", "--fwd", "--bwd",
+                "--headdim", "128",
+                "--seqlen", "1024,2048,4096,8192,16384,32768",
                 "--causal", "both", "--warmup", "1", "--rep", "3",
             ],
             extra_env={"CUDA_VISIBLE_DEVICES": benchmark_visible_devices},