{"owner":"Dao-AILab","repo":"flash-attention","hasSkills":true,"totalSkillsCount":2,"totalTokensCount":3510,"categories":["claude-rule","plugin-manifest"],"hasMcp":false,"mcpConfig":null,"found":["CLAUDE.md","tools/ci/README.md"],"skills":{"CLAUDE.md":"# CLAUDE.md\n\nThis file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.\n\n## Project Overview\n\nFlashAttention-4 (FA4) — fast, memory-efficient exact attention kernels written in Python using CuTeDSL (NVIDIA CUTLASS DSL). Kernels are compiled to PTX/CUBIN at runtime. Targets Hopper (SM90) and Blackwell (SM100/SM110) GPUs. Package name: `flash-attn-4`.\n\nThe repository also contains older generations (FA2 in top-level `csrc/`, FA3 in `hopper/`) but active development is on FA4 in `flash_attn/cute/`.\n\n## Agent Scratch Space\n\nUse `agent_space/` for project-local scratch work such as lab notes, profiling outputs, temporary repro scripts, and experiment artifacts. Treat it as disposable workspace rather than product code.\n\n## Build & Install\n\n```bash\npip install flash-attn-4\n# or dev install:\npip install -e \"flash_attn/cute[dev]\"\n```\n\nDependencies: `nvidia-cutlass-dsl>=4.6.2`, `torch`, `einops`, `apache-tvm-ffi`, `quack-kernels>=0.5.3`.\n\n## Running Tests\n\n```bash\npytest tests/cute/test_flash_attn.py\npytest tests/cute/test_flash_attn.py -k \"test_flash_attn_output\" -x  # single test\npytest tests/cute/test_flash_attn_mla.py  # MLA (qv): dense / sparse / paged, 1CTA (FLASH_ATTENTION_MLA_1CTA=1) and 2CTA\npytest tests/cute/test_flash_attn_varlen.py\npytest tests/cute/test_mask_mod.py\npytest tests/cute/test_score_mod.py\npytest tests/cute/test_block_sparsity.py\n```\n\n### Fast two-pass testing\n\nCompilation dominates test time. The fast workflow separates compilation (parallel, no GPU work) from execution (uses cached binaries). The compilation pass still opens a CUDA context per xdist worker (a few hundred MB each), so set `CUDA_VISIBLE_DEVICES` for both passes:\n\n```bash\n# Pass 1: compile all kernels in parallel using FakeTensorMode (no GPU memory allocation)\nFLASH_ATTENTION_FAKE_TENSOR=1 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -n 64 -x tests/cute/test_flash_attn.py\n\n# Pass 2: run tests using cached compiled kernels\nFLASH_ATTENTION_FAKE_TENSOR=0 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -x tests/cute/test_flash_attn.py\n```\n\n- `FLASH_ATTENTION_FAKE_TENSOR=1` — uses PyTorch FakeTensorMode to compile kernels without allocating GPU memory or running them.\n- `FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1` — enables persistent disk cache at `/tmp/${USER}/flash_attention_cute_dsl_cache/`.\n- `-n 256` — pytest-xdist parallel workers (only useful in the compilation pass).\n- `tools/two_pass_tests.sh -k EXPR -g 0,1 [-f FILES]` runs both passes, both limited to the `-g` GPUs. The execution pass runs with\n  `FLASH_ATTENTION_TEST_EXPECT_CACHED=1`, which fails any test that still compiles a kernel.\n  `FLASH_ATTENTION_TEST_COUNT_COMPILES=1` only reports the compiles.\n- Tests must support the compilation pass: decorate with\n  `@maybe_fake_tensor_mode(USE_FAKE_TENSOR)`, issue every kernel variant (every env /\n  monkeypatch knob) before the first data-dependent check, then `if is_fake_mode(): return`.\n- `FLASH_ATTENTION_TEST_RECORD_KEYS=path` records every compile-cache lookup, for diffing the\n  kernel set of a selection across changes.\n- **Most test files do not support the compilation pass yet.** Only `test_flash_attn_mla.py`\n  is fully converted; `test_flash_attn.py` mostly is, and `test_flash_attn_fast.py` and\n  `test_flash_attn_combine.py` partly. The mask_mod, score_mod, block-sparsity, varlen, hd256\n  scheduler and race-condition files are not converted at all. Their tests run with real tensors in pass 1 and\n  can compile again in pass 2, so `FLASH_ATTENTION_TEST_EXPECT_CACHED=1` (and\n  `tools/two_pass_tests.sh`) reports them as failures even when the results are correct. Run\n  those files in a single pass, or check that a failure is \"kernel(s) compiled during the\n  execution pass\" before treating it as a regression.\n\nTests are parametrized over dtype (fp16/bf16), head dimension (64, 96, 128), sequence length, causal/non-causal, and MHA/GQA/MQA.\n\nIf you get OOM errors running tests or benchmarks, use `nvidia-smi` to find a free GPU and select it with `CUDA_VISIBLE_DEVICES=<id>`.\n\n## Linting\n\nPre-commit uses ruff on `flash_attn/cute/` files. Large kernel files (`flash_bwd.py`, `flash_fwd.py`, `flash_fwd_sm100.py`, `interface.py`) are excluded from auto-formatting.\n\n```bash\nruff check flash_attn/cute/ --fix\nruff format flash_attn/cute/\n```\n\n## Code Architecture\n\n### Public API (`flash_attn/cute/interface.py`)\n\nTwo entry points exported from `flash_attn/cute/__init__.py`:\n- `flash_attn_func(q, k, v, ...)` — standard attention\n- `flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)` — variable-length\n\nKey parameters: `causal`, `window_size_left/right`, `softmax_scale`, `softcap`, `score_mod`, `mask_mod`, `block_sparse_tensors`, `num_splits`, `pack_gqa`, `m_block_size`, `n_block_size`, `num_threads`.\n\nTensor layout: `(batch, seqlen, num_heads, head_dim)`, last dim contiguous, 16-byte aligned.\n\n### Forward Kernels\n\n- `flash_fwd.py` — `FlashAttentionForwardSm90`: Hopper forward. No SplitKV or paged KV.\n- `flash_fwd_sm100.py` — `FlashAttentionForwardSm100`: Blackwell forward. Full features including SplitKV, paged KV cache, persistent kernels, 2CTA instructions.\n- `flash_fwd_combine.py` — `FlashAttentionForwardCombine`: merges SplitKV partial results.\n\n### Backward Kernels\n\n- `flash_bwd.py` — `FlashAttentionBackwardSm80`: Ampere backward (base).\n- `flash_bwd_sm90.py` — `FlashAttentionBackwardSm90`: Hopper backward.\n- `flash_bwd_sm100.py` — `FlashAttentionBackwardSm100`: Blackwell backward with 2CTA and block sparse support.\n- `flash_bwd_preprocess.py` / `flash_bwd_postprocess.py` — auxiliary backward kernels.\n\n### Core Abstractions\n\n- `softmax.py` — Online softmax with row_max/row_sum tracking, score modifier support.\n- `mask.py` — `AttentionMask`: causal, local/sliding window, block sparse, mask_mod application.\n- `block_info.py` — `BlockInfo`: tile dimensions, n/m block range computation for causal/local masking.\n- `seqlen_info.py` — `SeqlenInfoQK`: sequence length and offset tracking for varlen.\n- `pipeline.py` — `PipelineStateSimple`: circular buffer index/phase management for pipelined loads.\n- `tile_scheduler.py` — Tile scheduling strategies (single tile, varlen-aware, persistent).\n- `copy_utils.py` — Type-converting copies, shared-to-register loads, TMA copy atoms.\n- `named_barrier.py` — Named barrier enums for warp synchronization.\n\n### Architecture-Specific Helpers\n\n- `hopper_helpers.py` — SM90 warp-group GEMM, shared memory layout creation, fence/commit/wait.\n- `blackwell_helpers.py` — SM100 UMMA-based GEMM, PTX-optimized paths, 2CTA support.\n- `mma_sm100_desc.py` — Hardware MMA descriptor enums (formats, saturation, scaling).\n\n### Other Components\n\n- `pack_gqa.py` — Packs multiple Q heads per KV head for efficient GQA.\n- `paged_kv.py` — `PagedKVManager`: paged KV cache with TMA support.\n- `fast_math.py` — exp2 polynomial coefficients, softcap score_mod creation.\n- `utils.py` — Hash functions for compile cache keys, warp reductions, predicates.\n- `cache_utils.py` — JIT compilation cache management.\n- `cute_dsl_utils.py` — Patched `cute.compile` that optionally dumps SASS.\n\n### Compilation & Caching\n\nKernels are JIT-compiled. Cache key includes dtype, head_dim, causal, mask/score_mod hashes, architecture, block sizes. Caching levels: in-memory LRU + optional disk cache via `get_jit_cache()`.\n\nEnv vars: `CUTE_CUBIN_PATH` (dump CUBIN/SASS), `CUTE_DSL_KEEP_PTX=1` (inspect PTX), `CUTE_DSL_PTXAS_PATH` (custom ptxas).\n\n## Key Patterns\n\n- Compile-time constants use `cutlass.Constexpr[type]` for kernel specialization.\n- Score/mask modifiers are user-defined `@cute.jit` callables injected into the kernel at compile time.\n- Forward execution: load Q tile → loop over K/V blocks (pipelined) → online softmax accumulation → store O and LSE.\n- 2CTA instructions (SM100, hdim=128): both CTAs in a cluster coordinate via shared mbarriers; tx_count must be multiplied by `cta_group_size`.\n\n## Debugging GPU Kernels\n\n**Before proposing a root cause for any hang, deadlock, illegal-address trap, Xid fault, sanitizer report, or numerical mismatch that is not visible in the CuteDSL source, read `AI/DEBUG_METHODOLOGY.md` and follow its protocol** (falsifiable-prediction discipline, evidence tiers, fix-validation hygiene, hypothesis ledger in `agent_space/`).\n\nTactical docs in `AI/`:\n- `DEBUG_2CTA.md` — kernel hang/deadlock debugging (printf bisection, pipeline barrier analysis, 2CTA pitfalls).\n- `RACECHECK_TMA_HAZARD.md` — `compute-sanitizer` false positives with `cp.async.bulk` (repro scripts: `racecheck_repro_1d_*.py`).\n- `CLC_TRACE_DEBUG.md` — visualization of CLC scheduling (`parse_clc_log.py`).\n- `SASS_MMA_ANALYSIS.md` — dumping SASS and analyzing HGMMA instruction mix.\n- `SM90_BLOCK_SIZE_TUNING.md` — choosing tile sizes/MMA configs on Hopper (`sm90_config_search.py`).\n- `SM90_R2P_MASKING_SASS.md` — SASS-level analysis of R2P predicate masking in SM90 forward.\n- `VARLEN_PREPROCESS_TILE_BUG.md` — post-mortem: varlen preprocess tile-size mismatch and padded-offset layout.\n\nKey tools:\n- `cute.printf` with thread guards (`tidx % 32 == 0`, `elect_one()`) for targeted output\n- `compute-sanitizer --tool=racecheck` (beware false positives with raw TMA)\n- `CUTE_DSL_KEEP_PTX=1` and `CUTE_DSL_LINEINFO=1` for PTX inspection and sanitizer source mapping\n","tools/ci/README.md":"# FA4 CI\n\nCI runs on a self-hosted GPU runner inside an Apptainer (SIF) container. The container is either\npulled from Docker Hub (default) or, when the runner has one, a SIF built locally from\n`tools/ci/docker/Dockerfile` (registry-free mode, see below). Triggered on every push to `main`.\n\n## Two-pass test strategy\n\n- **Pass 1** — compile kernels in parallel via `FakeTensorMode` (no GPU memory needed)\n- **Pass 2** — run tests using cached compiled kernels on real GPU\n\nSee `run_fa4_ci.py` for the shared logic used by both CI and `test_ci_local.sh`.\n\nBefore Pass 1 the driver provisions a disk-backed Apptainer overlay (replaces the SIF-baked cutlass-dsl /\nquack / FA4 with the versions pinned in `flash_attn/cute/pyproject.toml` plus the checked-out FA4) and\n**verifies it from a fresh session**: one distribution per package, `flash_attn.cute` from the checkout,\nand the CuTeDSL runtime owned by the installed cutlass-dsl. On failure it re-provisions once, then fails\nwith `OVERLAY VERIFY FAILED`. The verified runtime is pinned via `CUTE_DSL_LIBS`. This guards against\nan apptainer session hand-off race where the SIF-baked packages reappear next to the installed ones.\n\n## Required GitHub secrets / variables\n\n| Name | Kind | Value |\n|------|------|-------|\n| `DOCKERHUB_USERNAME` | Secret | Docker Hub username |\n| `DOCKERHUB_TOKEN` | Secret | Docker Hub access token |\n| `CI_WORK_DIR` | Variable | Large-disk path on runner, e.g. `/scratch/user/johnson` |\n| `FA4_LOCAL_SIF` | Variable | Optional. Absolute path of a runner-local SIF; enables registry-free mode |\n\n`CI_WORK_DIR` is used for SIF caching and Apptainer temp files. Falls back to `/scratch/user/<github-actor>` if unset.\n\n## Registry-free mode (runner-local SIF)\n\nWhen `FA4_LOCAL_SIF` is set, the GPU job skips the Docker Hub login and pull and runs in that SIF, so\nCI needs no registry credentials. The Docker Hub secrets are then unused.\n\n1. On the runner: `tools/ci/docker/build_local_sif.sh` builds `tools/ci/docker/Dockerfile` and writes a SIF\n   to `CI_WORK_DIR/local-sif/` (keep it there: registry mode prunes `*.sif` directly under `CI_WORK_DIR`).\n2. Set the repo variable `FA4_LOCAL_SIF` to the printed path. A missing file fails the job explicitly.\n\nThe image only provides OS, Python and torch; `run_fa4_ci.py` installs cutlass-dsl, quack and FA4 at job\ntime, so DSL pin bumps never need a rebuild. Rebuild when the Dockerfile changes.\n\n## Updating the container image (registry mode)\n\n1. Build and push a new image via `tools/ci/docker/build.sh` + `tag_and_push.sh`.\n2. Update `FA4_IMAGE` in `.github/workflows/ci.yml` with the new tag and `sha256` digest.\n3. The old SIF is automatically deleted from the runner on the next CI run.\n\n## Expanding test coverage\n\nEdit `FA4_TEST_FILTER` in `.github/workflows/ci.yml`. To run the full suite, set it to an empty string and increase `compile-workers` in the `gpu-test` action call.\n\nAlternatively, edit `run_fa4_ci.py` to change `DEFAULT_TEST_TARGET` or worker defaults — changes there apply to both CI and local runs.\n\n## FA2 import isolation\n\nTests run inside the Apptainer container. The repo's `flash_attn/__init__.py` imports the FA2 C extension (`flash_attn_2_cuda`) which is absent in the container. `run_fa4_ci.py` works around this by:\n\n1. Installing FA4 from the current repo into the container at runtime (`uv pip install -e flash_attn/cute`).\n2. Running pytest from `/tmp` with absolute test paths — this keeps the repo root out of `sys.path[0]` so the installed FA4 package is found instead of the FA2 `__init__.py`.\n\n`flash_attn/__init__.py` is intentionally not modified; isolation is handled entirely in CI.\n\n## Adding a new runner / GPU type\n\n1. Register a self-hosted runner on the machine with the desired label (e.g. `h100`).\n2. Add the label to the `gpu` matrix in `.github/workflows/ci.yml`.\n3. Set `CI_WORK_DIR` for the new machine if its scratch path differs.\n"},"files":{"CLAUDE.md":"# CLAUDE.md\n\nThis file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.\n\n## Project Overview\n\nFlashAttention-4 (FA4) — fast, memory-efficient exact attention kernels written in Python using CuTeDSL (NVIDIA CUTLASS DSL). Kernels are compiled to PTX/CUBIN at runtime. Targets Hopper (SM90) and Blackwell (SM100/SM110) GPUs. Package name: `flash-attn-4`.\n\nThe repository also contains older generations (FA2 in top-level `csrc/`, FA3 in `hopper/`) but active development is on FA4 in `flash_attn/cute/`.\n\n## Agent Scratch Space\n\nUse `agent_space/` for project-local scratch work such as lab notes, profiling outputs, temporary repro scripts, and experiment artifacts. Treat it as disposable workspace rather than product code.\n\n## Build & Install\n\n```bash\npip install flash-attn-4\n# or dev install:\npip install -e \"flash_attn/cute[dev]\"\n```\n\nDependencies: `nvidia-cutlass-dsl>=4.6.2`, `torch`, `einops`, `apache-tvm-ffi`, `quack-kernels>=0.5.3`.\n\n## Running Tests\n\n```bash\npytest tests/cute/test_flash_attn.py\npytest tests/cute/test_flash_attn.py -k \"test_flash_attn_output\" -x  # single test\npytest tests/cute/test_flash_attn_mla.py  # MLA (qv): dense / sparse / paged, 1CTA (FLASH_ATTENTION_MLA_1CTA=1) and 2CTA\npytest tests/cute/test_flash_attn_varlen.py\npytest tests/cute/test_mask_mod.py\npytest tests/cute/test_score_mod.py\npytest tests/cute/test_block_sparsity.py\n```\n\n### Fast two-pass testing\n\nCompilation dominates test time. The fast workflow separates compilation (parallel, no GPU work) from execution (uses cached binaries). The compilation pass still opens a CUDA context per xdist worker (a few hundred MB each), so set `CUDA_VISIBLE_DEVICES` for both passes:\n\n```bash\n# Pass 1: compile all kernels in parallel using FakeTensorMode (no GPU memory allocation)\nFLASH_ATTENTION_FAKE_TENSOR=1 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -n 64 -x tests/cute/test_flash_attn.py\n\n# Pass 2: run tests using cached compiled kernels\nFLASH_ATTENTION_FAKE_TENSOR=0 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -x tests/cute/test_flash_attn.py\n```\n\n- `FLASH_ATTENTION_FAKE_TENSOR=1` — uses PyTorch FakeTensorMode to compile kernels without allocating GPU memory or running them.\n- `FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1` — enables persistent disk cache at `/tmp/${USER}/flash_attention_cute_dsl_cache/`.\n- `-n 256` — pytest-xdist parallel workers (only useful in the compilation pass).\n- `tools/two_pass_tests.sh -k EXPR -g 0,1 [-f FILES]` runs both passes, both limited to the `-g` GPUs. The execution pass runs with\n  `FLASH_ATTENTION_TEST_EXPECT_CACHED=1`, which fails any test that still compiles a kernel.\n  `FLASH_ATTENTION_TEST_COUNT_COMPILES=1` only reports the compiles.\n- Tests must support the compilation pass: decorate with\n  `@maybe_fake_tensor_mode(USE_FAKE_TENSOR)`, issue every kernel variant (every env /\n  monkeypatch knob) before the first data-dependent check, then `if is_fake_mode(): return`.\n- `FLASH_ATTENTION_TEST_RECORD_KEYS=path` records every compile-cache lookup, for diffing the\n  kernel set of a selection across changes.\n- **Most test files do not support the compilation pass yet.** Only `test_flash_attn_mla.py`\n  is fully converted; `test_flash_attn.py` mostly is, and `test_flash_attn_fast.py` and\n  `test_flash_attn_combine.py` partly. The mask_mod, score_mod, block-sparsity, varlen, hd256\n  scheduler and race-condition files are not converted at all. Their tests run with real tensors in pass 1 and\n  can compile again in pass 2, so `FLASH_ATTENTION_TEST_EXPECT_CACHED=1` (and\n  `tools/two_pass_tests.sh`) reports them as failures even when the results are correct. Run\n  those files in a single pass, or check that a failure is \"kernel(s) compiled during the\n  execution pass\" before treating it as a regression.\n\nTests are parametrized over dtype (fp16/bf16), head dimension (64, 96, 128), sequence length, causal/non-causal, and MHA/GQA/MQA.\n\nIf you get OOM errors running tests or benchmarks, use `nvidia-smi` to find a free GPU and select it with `CUDA_VISIBLE_DEVICES=<id>`.\n\n## Linting\n\nPre-commit uses ruff on `flash_attn/cute/` files. Large kernel files (`flash_bwd.py`, `flash_fwd.py`, `flash_fwd_sm100.py`, `interface.py`) are excluded from auto-formatting.\n\n```bash\nruff check flash_attn/cute/ --fix\nruff format flash_attn/cute/\n```\n\n## Code Architecture\n\n### Public API (`flash_attn/cute/interface.py`)\n\nTwo entry points exported from `flash_attn/cute/__init__.py`:\n- `flash_attn_func(q, k, v, ...)` — standard attention\n- `flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)` — variable-length\n\nKey parameters: `causal`, `window_size_left/right`, `softmax_scale`, `softcap`, `score_mod`, `mask_mod`, `block_sparse_tensors`, `num_splits`, `pack_gqa`, `m_block_size`, `n_block_size`, `num_threads`.\n\nTensor layout: `(batch, seqlen, num_heads, head_dim)`, last dim contiguous, 16-byte aligned.\n\n### Forward Kernels\n\n- `flash_fwd.py` — `FlashAttentionForwardSm90`: Hopper forward. No SplitKV or paged KV.\n- `flash_fwd_sm100.py` — `FlashAttentionForwardSm100`: Blackwell forward. Full features including SplitKV, paged KV cache, persistent kernels, 2CTA instructions.\n- `flash_fwd_combine.py` — `FlashAttentionForwardCombine`: merges SplitKV partial results.\n\n### Backward Kernels\n\n- `flash_bwd.py` — `FlashAttentionBackwardSm80`: Ampere backward (base).\n- `flash_bwd_sm90.py` — `FlashAttentionBackwardSm90`: Hopper backward.\n- `flash_bwd_sm100.py` — `FlashAttentionBackwardSm100`: Blackwell backward with 2CTA and block sparse support.\n- `flash_bwd_preprocess.py` / `flash_bwd_postprocess.py` — auxiliary backward kernels.\n\n### Core Abstractions\n\n- `softmax.py` — Online softmax with row_max/row_sum tracking, score modifier support.\n- `mask.py` — `AttentionMask`: causal, local/sliding window, block sparse, mask_mod application.\n- `block_info.py` — `BlockInfo`: tile dimensions, n/m block range computation for causal/local masking.\n- `seqlen_info.py` — `SeqlenInfoQK`: sequence length and offset tracking for varlen.\n- `pipeline.py` — `PipelineStateSimple`: circular buffer index/phase management for pipelined loads.\n- `tile_scheduler.py` — Tile scheduling strategies (single tile, varlen-aware, persistent).\n- `copy_utils.py` — Type-converting copies, shared-to-register loads, TMA copy atoms.\n- `named_barrier.py` — Named barrier enums for warp synchronization.\n\n### Architecture-Specific Helpers\n\n- `hopper_helpers.py` — SM90 warp-group GEMM, shared memory layout creation, fence/commit/wait.\n- `blackwell_helpers.py` — SM100 UMMA-based GEMM, PTX-optimized paths, 2CTA support.\n- `mma_sm100_desc.py` — Hardware MMA descriptor enums (formats, saturation, scaling).\n\n### Other Components\n\n- `pack_gqa.py` — Packs multiple Q heads per KV head for efficient GQA.\n- `paged_kv.py` — `PagedKVManager`: paged KV cache with TMA support.\n- `fast_math.py` — exp2 polynomial coefficients, softcap score_mod creation.\n- `utils.py` — Hash functions for compile cache keys, warp reductions, predicates.\n- `cache_utils.py` — JIT compilation cache management.\n- `cute_dsl_utils.py` — Patched `cute.compile` that optionally dumps SASS.\n\n### Compilation & Caching\n\nKernels are JIT-compiled. Cache key includes dtype, head_dim, causal, mask/score_mod hashes, architecture, block sizes. Caching levels: in-memory LRU + optional disk cache via `get_jit_cache()`.\n\nEnv vars: `CUTE_CUBIN_PATH` (dump CUBIN/SASS), `CUTE_DSL_KEEP_PTX=1` (inspect PTX), `CUTE_DSL_PTXAS_PATH` (custom ptxas).\n\n## Key Patterns\n\n- Compile-time constants use `cutlass.Constexpr[type]` for kernel specialization.\n- Score/mask modifiers are user-defined `@cute.jit` callables injected into the kernel at compile time.\n- Forward execution: load Q tile → loop over K/V blocks (pipelined) → online softmax accumulation → store O and LSE.\n- 2CTA instructions (SM100, hdim=128): both CTAs in a cluster coordinate via shared mbarriers; tx_count must be multiplied by `cta_group_size`.\n\n## Debugging GPU Kernels\n\n**Before proposing a root cause for any hang, deadlock, illegal-address trap, Xid fault, sanitizer report, or numerical mismatch that is not visible in the CuteDSL source, read `AI/DEBUG_METHODOLOGY.md` and follow its protocol** (falsifiable-prediction discipline, evidence tiers, fix-validation hygiene, hypothesis ledger in `agent_space/`).\n\nTactical docs in `AI/`:\n- `DEBUG_2CTA.md` — kernel hang/deadlock debugging (printf bisection, pipeline barrier analysis, 2CTA pitfalls).\n- `RACECHECK_TMA_HAZARD.md` — `compute-sanitizer` false positives with `cp.async.bulk` (repro scripts: `racecheck_repro_1d_*.py`).\n- `CLC_TRACE_DEBUG.md` — visualization of CLC scheduling (`parse_clc_log.py`).\n- `SASS_MMA_ANALYSIS.md` — dumping SASS and analyzing HGMMA instruction mix.\n- `SM90_BLOCK_SIZE_TUNING.md` — choosing tile sizes/MMA configs on Hopper (`sm90_config_search.py`).\n- `SM90_R2P_MASKING_SASS.md` — SASS-level analysis of R2P predicate masking in SM90 forward.\n- `VARLEN_PREPROCESS_TILE_BUG.md` — post-mortem: varlen preprocess tile-size mismatch and padded-offset layout.\n\nKey tools:\n- `cute.printf` with thread guards (`tidx % 32 == 0`, `elect_one()`) for targeted output\n- `compute-sanitizer --tool=racecheck` (beware false positives with raw TMA)\n- `CUTE_DSL_KEEP_PTX=1` and `CUTE_DSL_LINEINFO=1` for PTX inspection and sanitizer source mapping\n","tools/ci/README.md":"# FA4 CI\n\nCI runs on a self-hosted GPU runner inside an Apptainer (SIF) container. The container is either\npulled from Docker Hub (default) or, when the runner has one, a SIF built locally from\n`tools/ci/docker/Dockerfile` (registry-free mode, see below). Triggered on every push to `main`.\n\n## Two-pass test strategy\n\n- **Pass 1** — compile kernels in parallel via `FakeTensorMode` (no GPU memory needed)\n- **Pass 2** — run tests using cached compiled kernels on real GPU\n\nSee `run_fa4_ci.py` for the shared logic used by both CI and `test_ci_local.sh`.\n\nBefore Pass 1 the driver provisions a disk-backed Apptainer overlay (replaces the SIF-baked cutlass-dsl /\nquack / FA4 with the versions pinned in `flash_attn/cute/pyproject.toml` plus the checked-out FA4) and\n**verifies it from a fresh session**: one distribution per package, `flash_attn.cute` from the checkout,\nand the CuTeDSL runtime owned by the installed cutlass-dsl. On failure it re-provisions once, then fails\nwith `OVERLAY VERIFY FAILED`. The verified runtime is pinned via `CUTE_DSL_LIBS`. This guards against\nan apptainer session hand-off race where the SIF-baked packages reappear next to the installed ones.\n\n## Required GitHub secrets / variables\n\n| Name | Kind | Value |\n|------|------|-------|\n| `DOCKERHUB_USERNAME` | Secret | Docker Hub username |\n| `DOCKERHUB_TOKEN` | Secret | Docker Hub access token |\n| `CI_WORK_DIR` | Variable | Large-disk path on runner, e.g. `/scratch/user/johnson` |\n| `FA4_LOCAL_SIF` | Variable | Optional. Absolute path of a runner-local SIF; enables registry-free mode |\n\n`CI_WORK_DIR` is used for SIF caching and Apptainer temp files. Falls back to `/scratch/user/<github-actor>` if unset.\n\n## Registry-free mode (runner-local SIF)\n\nWhen `FA4_LOCAL_SIF` is set, the GPU job skips the Docker Hub login and pull and runs in that SIF, so\nCI needs no registry credentials. The Docker Hub secrets are then unused.\n\n1. On the runner: `tools/ci/docker/build_local_sif.sh` builds `tools/ci/docker/Dockerfile` and writes a SIF\n   to `CI_WORK_DIR/local-sif/` (keep it there: registry mode prunes `*.sif` directly under `CI_WORK_DIR`).\n2. Set the repo variable `FA4_LOCAL_SIF` to the printed path. A missing file fails the job explicitly.\n\nThe image only provides OS, Python and torch; `run_fa4_ci.py` installs cutlass-dsl, quack and FA4 at job\ntime, so DSL pin bumps never need a rebuild. Rebuild when the Dockerfile changes.\n\n## Updating the container image (registry mode)\n\n1. Build and push a new image via `tools/ci/docker/build.sh` + `tag_and_push.sh`.\n2. Update `FA4_IMAGE` in `.github/workflows/ci.yml` with the new tag and `sha256` digest.\n3. The old SIF is automatically deleted from the runner on the next CI run.\n\n## Expanding test coverage\n\nEdit `FA4_TEST_FILTER` in `.github/workflows/ci.yml`. To run the full suite, set it to an empty string and increase `compile-workers` in the `gpu-test` action call.\n\nAlternatively, edit `run_fa4_ci.py` to change `DEFAULT_TEST_TARGET` or worker defaults — changes there apply to both CI and local runs.\n\n## FA2 import isolation\n\nTests run inside the Apptainer container. The repo's `flash_attn/__init__.py` imports the FA2 C extension (`flash_attn_2_cuda`) which is absent in the container. `run_fa4_ci.py` works around this by:\n\n1. Installing FA4 from the current repo into the container at runtime (`uv pip install -e flash_attn/cute`).\n2. Running pytest from `/tmp` with absolute test paths — this keeps the repo root out of `sys.path[0]` so the installed FA4 package is found instead of the FA2 `__init__.py`.\n\n`flash_attn/__init__.py` is intentionally not modified; isolation is handled entirely in CI.\n\n## Adding a new runner / GPU type\n\n1. Register a self-hosted runner on the machine with the desired label (e.g. `h100`).\n2. Add the label to the `gpu` matrix in `.github/workflows/ci.yml`.\n3. Set `CI_WORK_DIR` for the new machine if its scratch path differs.\n"},"items":[{"name":"CLAUDE.md","path":"CLAUDE.md","rawUrl":"https://raw.githubusercontent.com/Dao-AILab/flash-attention/HEAD/CLAUDE.md","title":"Claude Agent Guidelines & System Prompt","category":"claude-rule","format":"markdown","content":"# CLAUDE.md\n\nThis file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.\n\n## Project Overview\n\nFlashAttention-4 (FA4) — fast, memory-efficient exact attention kernels written in Python using CuTeDSL (NVIDIA CUTLASS DSL). Kernels are compiled to PTX/CUBIN at runtime. Targets Hopper (SM90) and Blackwell (SM100/SM110) GPUs. Package name: `flash-attn-4`.\n\nThe repository also contains older generations (FA2 in top-level `csrc/`, FA3 in `hopper/`) but active development is on FA4 in `flash_attn/cute/`.\n\n## Agent Scratch Space\n\nUse `agent_space/` for project-local scratch work such as lab notes, profiling outputs, temporary repro scripts, and experiment artifacts. Treat it as disposable workspace rather than product code.\n\n## Build & Install\n\n```bash\npip install flash-attn-4\n# or dev install:\npip install -e \"flash_attn/cute[dev]\"\n```\n\nDependencies: `nvidia-cutlass-dsl>=4.6.2`, `torch`, `einops`, `apache-tvm-ffi`, `quack-kernels>=0.5.3`.\n\n## Running Tests\n\n```bash\npytest tests/cute/test_flash_attn.py\npytest tests/cute/test_flash_attn.py -k \"test_flash_attn_output\" -x  # single test\npytest tests/cute/test_flash_attn_mla.py  # MLA (qv): dense / sparse / paged, 1CTA (FLASH_ATTENTION_MLA_1CTA=1) and 2CTA\npytest tests/cute/test_flash_attn_varlen.py\npytest tests/cute/test_mask_mod.py\npytest tests/cute/test_score_mod.py\npytest tests/cute/test_block_sparsity.py\n```\n\n### Fast two-pass testing\n\nCompilation dominates test time. The fast workflow separates compilation (parallel, no GPU work) from execution (uses cached binaries). The compilation pass still opens a CUDA context per xdist worker (a few hundred MB each), so set `CUDA_VISIBLE_DEVICES` for both passes:\n\n```bash\n# Pass 1: compile all kernels in parallel using FakeTensorMode (no GPU memory allocation)\nFLASH_ATTENTION_FAKE_TENSOR=1 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -n 64 -x tests/cute/test_flash_attn.py\n\n# Pass 2: run tests using cached compiled kernels\nFLASH_ATTENTION_FAKE_TENSOR=0 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -x tests/cute/test_flash_attn.py\n```\n\n- `FLASH_ATTENTION_FAKE_TENSOR=1` — uses PyTorch FakeTensorMode to compile kernels without allocating GPU memory or running them.\n- `FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1` — enables persistent disk cache at `/tmp/${USER}/flash_attention_cute_dsl_cache/`.\n- `-n 256` — pytest-xdist parallel workers (only useful in the compilation pass).\n- `tools/two_pass_tests.sh -k EXPR -g 0,1 [-f FILES]` runs both passes, both limited to the `-g` GPUs. The execution pass runs with\n  `FLASH_ATTENTION_TEST_EXPECT_CACHED=1`, which fails any test that still compiles a kernel.\n  `FLASH_ATTENTION_TEST_COUNT_COMPILES=1` only reports the compiles.\n- Tests must support the compilation pass: decorate with\n  `@maybe_fake_tensor_mode(USE_FAKE_TENSOR)`, issue every kernel variant (every env /\n  monkeypatch knob) before the first data-dependent check, then `if is_fake_mode(): return`.\n- `FLASH_ATTENTION_TEST_RECORD_KEYS=path` records every compile-cache lookup, for diffing the\n  kernel set of a selection across changes.\n- **Most test files do not support the compilation pass yet.** Only `test_flash_attn_mla.py`\n  is fully converted; `test_flash_attn.py` mostly is, and `test_flash_attn_fast.py` and\n  `test_flash_attn_combine.py` partly. The mask_mod, score_mod, block-sparsity, varlen, hd256\n  scheduler and race-condition files are not converted at all. Their tests run with real tensors in pass 1 and\n  can compile again in pass 2, so `FLASH_ATTENTION_TEST_EXPECT_CACHED=1` (and\n  `tools/two_pass_tests.sh`) reports them as failures even when the results are correct. Run\n  those files in a single pass, or check that a failure is \"kernel(s) compiled during the\n  execution pass\" before treating it as a regression.\n\nTests are parametrized over dtype (fp16/bf16), head dimension (64, 96, 128), sequence length, causal/non-causal, and MHA/GQA/MQA.\n\nIf you get OOM errors running tests or benchmarks, use `nvidia-smi` to find a free GPU and select it with `CUDA_VISIBLE_DEVICES=<id>`.\n\n## Linting\n\nPre-commit uses ruff on `flash_attn/cute/` files. Large kernel files (`flash_bwd.py`, `flash_fwd.py`, `flash_fwd_sm100.py`, `interface.py`) are excluded from auto-formatting.\n\n```bash\nruff check flash_attn/cute/ --fix\nruff format flash_attn/cute/\n```\n\n## Code Architecture\n\n### Public API (`flash_attn/cute/interface.py`)\n\nTwo entry points exported from `flash_attn/cute/__init__.py`:\n- `flash_attn_func(q, k, v, ...)` — standard attention\n- `flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)` — variable-length\n\nKey parameters: `causal`, `window_size_left/right`, `softmax_scale`, `softcap`, `score_mod`, `mask_mod`, `block_sparse_tensors`, `num_splits`, `pack_gqa`, `m_block_size`, `n_block_size`, `num_threads`.\n\nTensor layout: `(batch, seqlen, num_heads, head_dim)`, last dim contiguous, 16-byte aligned.\n\n### Forward Kernels\n\n- `flash_fwd.py` — `FlashAttentionForwardSm90`: Hopper forward. No SplitKV or paged KV.\n- `flash_fwd_sm100.py` — `FlashAttentionForwardSm100`: Blackwell forward. Full features including SplitKV, paged KV cache, persistent kernels, 2CTA instructions.\n- `flash_fwd_combine.py` — `FlashAttentionForwardCombine`: merges SplitKV partial results.\n\n### Backward Kernels\n\n- `flash_bwd.py` — `FlashAttentionBackwardSm80`: Ampere backward (base).\n- `flash_bwd_sm90.py` — `FlashAttentionBackwardSm90`: Hopper backward.\n- `flash_bwd_sm100.py` — `FlashAttentionBackwardSm100`: Blackwell backward with 2CTA and block sparse support.\n- `flash_bwd_preprocess.py` / `flash_bwd_postprocess.py` — auxiliary backward kernels.\n\n### Core Abstractions\n\n- `softmax.py` — Online softmax with row_max/row_sum tracking, score modifier support.\n- `mask.py` — `AttentionMask`: causal, local/sliding window, block sparse, mask_mod application.\n- `block_info.py` — `BlockInfo`: tile dimensions, n/m block range computation for causal/local masking.\n- `seqlen_info.py` — `SeqlenInfoQK`: sequence length and offset tracking for varlen.\n- `pipeline.py` — `PipelineStateSimple`: circular buffer index/phase management for pipelined loads.\n- `tile_scheduler.py` — Tile scheduling strategies (single tile, varlen-aware, persistent).\n- `copy_utils.py` — Type-converting copies, shared-to-register loads, TMA copy atoms.\n- `named_barrier.py` — Named barrier enums for warp synchronization.\n\n### Architecture-Specific Helpers\n\n- `hopper_helpers.py` — SM90 warp-group GEMM, shared memory layout creation, fence/commit/wait.\n- `blackwell_helpers.py` — SM100 UMMA-based GEMM, PTX-optimized paths, 2CTA support.\n- `mma_sm100_desc.py` — Hardware MMA descriptor enums (formats, saturation, scaling).\n\n### Other Components\n\n- `pack_gqa.py` — Packs multiple Q heads per KV head for efficient GQA.\n- `paged_kv.py` — `PagedKVManager`: paged KV cache with TMA support.\n- `fast_math.py` — exp2 polynomial coefficients, softcap score_mod creation.\n- `utils.py` — Hash functions for compile cache keys, warp reductions, predicates.\n- `cache_utils.py` — JIT compilation cache management.\n- `cute_dsl_utils.py` — Patched `cute.compile` that optionally dumps SASS.\n\n### Compilation & Caching\n\nKernels are JIT-compiled. Cache key includes dtype, head_dim, causal, mask/score_mod hashes, architecture, block sizes. Caching levels: in-memory LRU + optional disk cache via `get_jit_cache()`.\n\nEnv vars: `CUTE_CUBIN_PATH` (dump CUBIN/SASS), `CUTE_DSL_KEEP_PTX=1` (inspect PTX), `CUTE_DSL_PTXAS_PATH` (custom ptxas).\n\n## Key Patterns\n\n- Compile-time constants use `cutlass.Constexpr[type]` for kernel specialization.\n- Score/mask modifiers are user-defined `@cute.jit` callables injected into the kernel at compile time.\n- Forward execution: load Q tile → loop over K/V blocks (pipelined) → online softmax accumulation → store O and LSE.\n- 2CTA instructions (SM100, hdim=128): both CTAs in a cluster coordinate via shared mbarriers; tx_count must be multiplied by `cta_group_size`.\n\n## Debugging GPU Kernels\n\n**Before proposing a root cause for any hang, deadlock, illegal-address trap, Xid fault, sanitizer report, or numerical mismatch that is not visible in the CuteDSL source, read `AI/DEBUG_METHODOLOGY.md` and follow its protocol** (falsifiable-prediction discipline, evidence tiers, fix-validation hygiene, hypothesis ledger in `agent_space/`).\n\nTactical docs in `AI/`:\n- `DEBUG_2CTA.md` — kernel hang/deadlock debugging (printf bisection, pipeline barrier analysis, 2CTA pitfalls).\n- `RACECHECK_TMA_HAZARD.md` — `compute-sanitizer` false positives with `cp.async.bulk` (repro scripts: `racecheck_repro_1d_*.py`).\n- `CLC_TRACE_DEBUG.md` — visualization of CLC scheduling (`parse_clc_log.py`).\n- `SASS_MMA_ANALYSIS.md` — dumping SASS and analyzing HGMMA instruction mix.\n- `SM90_BLOCK_SIZE_TUNING.md` — choosing tile sizes/MMA configs on Hopper (`sm90_config_search.py`).\n- `SM90_R2P_MASKING_SASS.md` — SASS-level analysis of R2P predicate masking in SM90 forward.\n- `VARLEN_PREPROCESS_TILE_BUG.md` — post-mortem: varlen preprocess tile-size mismatch and padded-offset layout.\n\nKey tools:\n- `cute.printf` with thread guards (`tidx % 32 == 0`, `elect_one()`) for targeted output\n- `compute-sanitizer --tool=racecheck` (beware false positives with raw TMA)\n- `CUTE_DSL_KEEP_PTX=1` and `CUTE_DSL_LINEINFO=1` for PTX inspection and sanitizer source mapping\n","isInternal":false,"tokens":2490,"sizeBytes":9410},{"name":"README.md","path":"tools/ci/README.md","rawUrl":"https://raw.githubusercontent.com/Dao-AILab/flash-attention/HEAD/tools/ci/README.md","title":"ci Documentation","category":"plugin-manifest","format":"markdown","content":"# FA4 CI\n\nCI runs on a self-hosted GPU runner inside an Apptainer (SIF) container. The container is either\npulled from Docker Hub (default) or, when the runner has one, a SIF built locally from\n`tools/ci/docker/Dockerfile` (registry-free mode, see below). Triggered on every push to `main`.\n\n## Two-pass test strategy\n\n- **Pass 1** — compile kernels in parallel via `FakeTensorMode` (no GPU memory needed)\n- **Pass 2** — run tests using cached compiled kernels on real GPU\n\nSee `run_fa4_ci.py` for the shared logic used by both CI and `test_ci_local.sh`.\n\nBefore Pass 1 the driver provisions a disk-backed Apptainer overlay (replaces the SIF-baked cutlass-dsl /\nquack / FA4 with the versions pinned in `flash_attn/cute/pyproject.toml` plus the checked-out FA4) and\n**verifies it from a fresh session**: one distribution per package, `flash_attn.cute` from the checkout,\nand the CuTeDSL runtime owned by the installed cutlass-dsl. On failure it re-provisions once, then fails\nwith `OVERLAY VERIFY FAILED`. The verified runtime is pinned via `CUTE_DSL_LIBS`. This guards against\nan apptainer session hand-off race where the SIF-baked packages reappear next to the installed ones.\n\n## Required GitHub secrets / variables\n\n| Name | Kind | Value |\n|------|------|-------|\n| `DOCKERHUB_USERNAME` | Secret | Docker Hub username |\n| `DOCKERHUB_TOKEN` | Secret | Docker Hub access token |\n| `CI_WORK_DIR` | Variable | Large-disk path on runner, e.g. `/scratch/user/johnson` |\n| `FA4_LOCAL_SIF` | Variable | Optional. Absolute path of a runner-local SIF; enables registry-free mode |\n\n`CI_WORK_DIR` is used for SIF caching and Apptainer temp files. Falls back to `/scratch/user/<github-actor>` if unset.\n\n## Registry-free mode (runner-local SIF)\n\nWhen `FA4_LOCAL_SIF` is set, the GPU job skips the Docker Hub login and pull and runs in that SIF, so\nCI needs no registry credentials. The Docker Hub secrets are then unused.\n\n1. On the runner: `tools/ci/docker/build_local_sif.sh` builds `tools/ci/docker/Dockerfile` and writes a SIF\n   to `CI_WORK_DIR/local-sif/` (keep it there: registry mode prunes `*.sif` directly under `CI_WORK_DIR`).\n2. Set the repo variable `FA4_LOCAL_SIF` to the printed path. A missing file fails the job explicitly.\n\nThe image only provides OS, Python and torch; `run_fa4_ci.py` installs cutlass-dsl, quack and FA4 at job\ntime, so DSL pin bumps never need a rebuild. Rebuild when the Dockerfile changes.\n\n## Updating the container image (registry mode)\n\n1. Build and push a new image via `tools/ci/docker/build.sh` + `tag_and_push.sh`.\n2. Update `FA4_IMAGE` in `.github/workflows/ci.yml` with the new tag and `sha256` digest.\n3. The old SIF is automatically deleted from the runner on the next CI run.\n\n## Expanding test coverage\n\nEdit `FA4_TEST_FILTER` in `.github/workflows/ci.yml`. To run the full suite, set it to an empty string and increase `compile-workers` in the `gpu-test` action call.\n\nAlternatively, edit `run_fa4_ci.py` to change `DEFAULT_TEST_TARGET` or worker defaults — changes there apply to both CI and local runs.\n\n## FA2 import isolation\n\nTests run inside the Apptainer container. The repo's `flash_attn/__init__.py` imports the FA2 C extension (`flash_attn_2_cuda`) which is absent in the container. `run_fa4_ci.py` works around this by:\n\n1. Installing FA4 from the current repo into the container at runtime (`uv pip install -e flash_attn/cute`).\n2. Running pytest from `/tmp` with absolute test paths — this keeps the repo root out of `sys.path[0]` so the installed FA4 package is found instead of the FA2 `__init__.py`.\n\n`flash_attn/__init__.py` is intentionally not modified; isolation is handled entirely in CI.\n\n## Adding a new runner / GPU type\n\n1. Register a self-hosted runner on the machine with the desired label (e.g. `h100`).\n2. Add the label to the `gpu` matrix in `.github/workflows/ci.yml`.\n3. Set `CI_WORK_DIR` for the new machine if its scratch path differs.\n","isInternal":false,"tokens":1020,"sizeBytes":3922}],"systemPromptSnippet":"<yakaai_skills repo=\"Dao-AILab/flash-attention\">\n<!-- File: CLAUDE.md (Tokens: ~2490 | Category: claude-rule) -->\n# CLAUDE.md\n\nThis file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.\n\n## Project Overview\n\nFlashAttention-4 (FA4) — fast, memory-efficient exact attention kernels written in Python using CuTeDSL (NVIDIA CUTLASS DSL). Kernels are compiled to PTX/CUBIN at runtime. Targets Hopper (SM90) and Blackwell (SM100/SM110) GPUs. Package name: `flash-attn-4`.\n\nThe repository also contains older generations (FA2 in top-level `csrc/`, FA3 in `hopper/`) but active development is on FA4 in `flash_attn/cute/`.\n\n## Agent Scratch Space\n\nUse `agent_space/` for project-local scratch work such as lab notes, profiling outputs, temporary repro scripts, and experiment artifacts. Treat it as disposable workspace rather than product code.\n\n## Build & Install\n\n```bash\npip install flash-attn-4\n# or dev install:\npip install -e \"flash_attn/cute[dev]\"\n```\n\nDependencies: `nvidia-cutlass-dsl>=4.6.2`, `torch`, `einops`, `apache-tvm-ffi`, `quack-kernels>=0.5.3`.\n\n## Running Tests\n\n```bash\npytest tests/cute/test_flash_attn.py\npytest tests/cute/test_flash_attn.py -k \"test_flash_attn_output\" -x  # single test\npytest tests/cute/test_flash_attn_mla.py  # MLA (qv): dense / sparse / paged, 1CTA (FLASH_ATTENTION_MLA_1CTA=1) and 2CTA\npytest tests/cute/test_flash_attn_varlen.py\npytest tests/cute/test_mask_mod.py\npytest tests/cute/test_score_mod.py\npytest tests/cute/test_block_sparsity.py\n```\n\n### Fast two-pass testing\n\nCompilation dominates test time. The fast workflow separates compilation (parallel, no GPU work) from execution (uses cached binaries). The compilation pass still opens a CUDA context per xdist worker (a few hundred MB each), so set `CUDA_VISIBLE_DEVICES` for both passes:\n\n```bash\n# Pass 1: compile all kernels in parallel using FakeTensorMode (no GPU memory allocation)\nFLASH_ATTENTION_FAKE_TENSOR=1 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -n 64 -x tests/cute/test_flash_attn.py\n\n# Pass 2: run tests using cached compiled kernels\nFLASH_ATTENTION_FAKE_TENSOR=0 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -x tests/cute/test_flash_attn.py\n```\n\n- `FLASH_ATTENTION_FAKE_TENSOR=1` — uses PyTorch FakeTensorMode to compile kernels without allocating GPU memory or running them.\n- `FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1` — enables persistent disk cache at `/tmp/${USER}/flash_attention_cute_dsl_cache/`.\n- `-n 256` — pytest-xdist parallel workers (only useful in the compilation pass).\n- `tools/two_pass_tests.sh -k EXPR -g 0,1 [-f FILES]` runs both passes, both limited to the `-g` GPUs. The execution pass runs with\n  `FLASH_ATTENTION_TEST_EXPECT_CACHED=1`, which fails any test that still compiles a kernel.\n  `FLASH_ATTENTION_TEST_COUNT_COMPILES=1` only reports the compiles.\n- Tests must support the compilation pass: decorate with\n  `@maybe_fake_tensor_mode(USE_FAKE_TENSOR)`, issue every kernel variant (every env /\n  monkeypatch knob) before the first data-dependent check, then `if is_fake_mode(): return`.\n- `FLASH_ATTENTION_TEST_RECORD_KEYS=path` records every compile-cache lookup, for diffing the\n  kernel set of a selection across changes.\n- **Most test files do not support the compilation pass yet.** Only `test_flash_attn_mla.py`\n  is fully converted; `test_flash_attn.py` mostly is, and `test_flash_attn_fast.py` and\n  `test_flash_attn_combine.py` partly. The mask_mod, score_mod, block-sparsity, varlen, hd256\n  scheduler and race-condition files are not converted at all. Their tests run with real tensors in pass 1 and\n  can compile again in pass 2, so `FLASH_ATTENTION_TEST_EXPECT_CACHED=1` (and\n  `tools/two_pass_tests.sh`) reports them as failures even when the results are correct. Run\n  those files in a single pass, or check that a failure is \"kernel(s) compiled during the\n  execution pass\" before treating it as a regression.\n\nTests are parametrized over dtype (fp16/bf16), head dimension (64, 96, 128), sequence length, causal/non-causal, and MHA/GQA/MQA.\n\nIf you get OOM errors running tests or benchmarks, use `nvidia-smi` to find a free GPU and select it with `CUDA_VISIBLE_DEVICES=<id>`.\n\n## Linting\n\nPre-commit uses ruff on `flash_attn/cute/` files. Large kernel files (`flash_bwd.py`, `flash_fwd.py`, `flash_fwd_sm100.py`, `interface.py`) are excluded from auto-formatting.\n\n```bash\nruff check flash_attn/cute/ --fix\nruff format flash_attn/cute/\n```\n\n## Code Architecture\n\n### Public API (`flash_attn/cute/interface.py`)\n\nTwo entry points exported from `flash_attn/cute/__init__.py`:\n- `flash_attn_func(q, k, v, ...)` — standard attention\n- `flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)` — variable-length\n\nKey parameters: `causal`, `window_size_left/right`, `softmax_scale`, `softcap`, `score_mod`, `mask_mod`, `block_sparse_tensors`, `num_splits`, `pack_gqa`, `m_block_size`, `n_block_size`, `num_threads`.\n\nTensor layout: `(batch, seqlen, num_heads, head_dim)`, last dim contiguous, 16-byte aligned.\n\n### Forward Kernels\n\n- `flash_fwd.py` — `FlashAttentionForwardSm90`: Hopper forward. No SplitKV or paged KV.\n- `flash_fwd_sm100.py` — `FlashAttentionForwardSm100`: Blackwell forward. Full features including SplitKV, paged KV cache, persistent kernels, 2CTA instructions.\n- `flash_fwd_combine.py` — `FlashAttentionForwardCombine`: merges SplitKV partial results.\n\n### Backward Kernels\n\n- `flash_bwd.py` — `FlashAttentionBackwardSm80`: Ampere backward (base).\n- `flash_bwd_sm90.py` — `FlashAttentionBackwardSm90`: Hopper backward.\n- `flash_bwd_sm100.py` — `FlashAttentionBackwardSm100`: Blackwell backward with 2CTA and block sparse support.\n- `flash_bwd_preprocess.py` / `flash_bwd_postprocess.py` — auxiliary backward kernels.\n\n### Core Abstractions\n\n- `softmax.py` — Online softmax with row_max/row_sum tracking, score modifier support.\n- `mask.py` — `AttentionMask`: causal, local/sliding window, block sparse, mask_mod application.\n- `block_info.py` — `BlockInfo`: tile dimensions, n/m block range computation for causal/local masking.\n- `seqlen_info.py` — `SeqlenInfoQK`: sequence length and offset tracking for varlen.\n- `pipeline.py` — `PipelineStateSimple`: circular buffer index/phase management for pipelined loads.\n- `tile_scheduler.py` — Tile scheduling strategies (single tile, varlen-aware, persistent).\n- `copy_utils.py` — Type-converting copies, shared-to-register loads, TMA copy atoms.\n- `named_barrier.py` — Named barrier enums for warp synchronization.\n\n### Architecture-Specific Helpers\n\n- `hopper_helpers.py` — SM90 warp-group GEMM, shared memory layout creation, fence/commit/wait.\n- `blackwell_helpers.py` — SM100 UMMA-based GEMM, PTX-optimized paths, 2CTA support.\n- `mma_sm100_desc.py` — Hardware MMA descriptor enums (formats, saturation, scaling).\n\n### Other Components\n\n- `pack_gqa.py` — Packs multiple Q heads per KV head for efficient GQA.\n- `paged_kv.py` — `PagedKVManager`: paged KV cache with TMA support.\n- `fast_math.py` — exp2 polynomial coefficients, softcap score_mod creation.\n- `utils.py` — Hash functions for compile cache keys, warp reductions, predicates.\n- `cache_utils.py` — JIT compilation cache management.\n- `cute_dsl_utils.py` — Patched `cute.compile` that optionally dumps SASS.\n\n### Compilation & Caching\n\nKernels are JIT-compiled. Cache key includes dtype, head_dim, causal, mask/score_mod hashes, architecture, block sizes. Caching levels: in-memory LRU + optional disk cache via `get_jit_cache()`.\n\nEnv vars: `CUTE_CUBIN_PATH` (dump CUBIN/SASS), `CUTE_DSL_KEEP_PTX=1` (inspect PTX), `CUTE_DSL_PTXAS_PATH` (custom ptxas).\n\n## Key Patterns\n\n- Compile-time constants use `cutlass.Constexpr[type]` for kernel specialization.\n- Score/mask modifiers are user-defined `@cute.jit` callables injected into the kernel at compile time.\n- Forward execution: load Q tile → loop over K/V blocks (pipelined) → online softmax accumulation → store O and LSE.\n- 2CTA instructions (SM100, hdim=128): both CTAs in a cluster coordinate via shared mbarriers; tx_count must be multiplied by `cta_group_size`.\n\n## Debugging GPU Kernels\n\n**Before proposing a root cause for any hang, deadlock, illegal-address trap, Xid fault, sanitizer report, or numerical mismatch that is not visible in the CuteDSL source, read `AI/DEBUG_METHODOLOGY.md` and follow its protocol** (falsifiable-prediction discipline, evidence tiers, fix-validation hygiene, hypothesis ledger in `agent_space/`).\n\nTactical docs in `AI/`:\n- `DEBUG_2CTA.md` — kernel hang/deadlock debugging (printf bisection, pipeline barrier analysis, 2CTA pitfalls).\n- `RACECHECK_TMA_HAZARD.md` — `compute-sanitizer` false positives with `cp.async.bulk` (repro scripts: `racecheck_repro_1d_*.py`).\n- `CLC_TRACE_DEBUG.md` — visualization of CLC scheduling (`parse_clc_log.py`).\n- `SASS_MMA_ANALYSIS.md` — dumping SASS and analyzing HGMMA instruction mix.\n- `SM90_BLOCK_SIZE_TUNING.md` — choosing tile sizes/MMA configs on Hopper (`sm90_config_search.py`).\n- `SM90_R2P_MASKING_SASS.md` — SASS-level analysis of R2P predicate masking in SM90 forward.\n- `VARLEN_PREPROCESS_TILE_BUG.md` — post-mortem: varlen preprocess tile-size mismatch and padded-offset layout.\n\nKey tools:\n- `cute.printf` with thread guards (`tidx % 32 == 0`, `elect_one()`) for targeted output\n- `compute-sanitizer --tool=racecheck` (beware false positives with raw TMA)\n- `CUTE_DSL_KEEP_PTX=1` and `CUTE_DSL_LINEINFO=1` for PTX inspection and sanitizer source mapping\n\n<!-- File: tools/ci/README.md (Tokens: ~1020 | Category: plugin-manifest) -->\n# FA4 CI\n\nCI runs on a self-hosted GPU runner inside an Apptainer (SIF) container. The container is either\npulled from Docker Hub (default) or, when the runner has one, a SIF built locally from\n`tools/ci/docker/Dockerfile` (registry-free mode, see below). Triggered on every push to `main`.\n\n## Two-pass test strategy\n\n- **Pass 1** — compile kernels in parallel via `FakeTensorMode` (no GPU memory needed)\n- **Pass 2** — run tests using cached compiled kernels on real GPU\n\nSee `run_fa4_ci.py` for the shared logic used by both CI and `test_ci_local.sh`.\n\nBefore Pass 1 the driver provisions a disk-backed Apptainer overlay (replaces the SIF-baked cutlass-dsl /\nquack / FA4 with the versions pinned in `flash_attn/cute/pyproject.toml` plus the checked-out FA4) and\n**verifies it from a fresh session**: one distribution per package, `flash_attn.cute` from the checkout,\nand the CuTeDSL runtime owned by the installed cutlass-dsl. On failure it re-provisions once, then fails\nwith `OVERLAY VERIFY FAILED`. The verified runtime is pinned via `CUTE_DSL_LIBS`. This guards against\nan apptainer session hand-off race where the SIF-baked packages reappear next to the installed ones.\n\n## Required GitHub secrets / variables\n\n| Name | Kind | Value |\n|------|------|-------|\n| `DOCKERHUB_USERNAME` | Secret | Docker Hub username |\n| `DOCKERHUB_TOKEN` | Secret | Docker Hub access token |\n| `CI_WORK_DIR` | Variable | Large-disk path on runner, e.g. `/scratch/user/johnson` |\n| `FA4_LOCAL_SIF` | Variable | Optional. Absolute path of a runner-local SIF; enables registry-free mode |\n\n`CI_WORK_DIR` is used for SIF caching and Apptainer temp files. Falls back to `/scratch/user/<github-actor>` if unset.\n\n## Registry-free mode (runner-local SIF)\n\nWhen `FA4_LOCAL_SIF` is set, the GPU job skips the Docker Hub login and pull and runs in that SIF, so\nCI needs no registry credentials. The Docker Hub secrets are then unused.\n\n1. On the runner: `tools/ci/docker/build_local_sif.sh` builds `tools/ci/docker/Dockerfile` and writes a SIF\n   to `CI_WORK_DIR/local-sif/` (keep it there: registry mode prunes `*.sif` directly under `CI_WORK_DIR`).\n2. Set the repo variable `FA4_LOCAL_SIF` to the printed path. A missing file fails the job explicitly.\n\nThe image only provides OS, Python and torch; `run_fa4_ci.py` installs cutlass-dsl, quack and FA4 at job\ntime, so DSL pin bumps never need a rebuild. Rebuild when the Dockerfile changes.\n\n## Updating the container image (registry mode)\n\n1. Build and push a new image via `tools/ci/docker/build.sh` + `tag_and_push.sh`.\n2. Update `FA4_IMAGE` in `.github/workflows/ci.yml` with the new tag and `sha256` digest.\n3. The old SIF is automatically deleted from the runner on the next CI run.\n\n## Expanding test coverage\n\nEdit `FA4_TEST_FILTER` in `.github/workflows/ci.yml`. To run the full suite, set it to an empty string and increase `compile-workers` in the `gpu-test` action call.\n\nAlternatively, edit `run_fa4_ci.py` to change `DEFAULT_TEST_TARGET` or worker defaults — changes there apply to both CI and local runs.\n\n## FA2 import isolation\n\nTests run inside the Apptainer container. The repo's `flash_attn/__init__.py` imports the FA2 C extension (`flash_attn_2_cuda`) which is absent in the container. `run_fa4_ci.py` works around this by:\n\n1. Installing FA4 from the current repo into the container at runtime (`uv pip install -e flash_attn/cute`).\n2. Running pytest from `/tmp` with absolute test paths — this keeps the repo root out of `sys.path[0]` so the installed FA4 package is found instead of the FA2 `__init__.py`.\n\n`flash_attn/__init__.py` is intentionally not modified; isolation is handled entirely in CI.\n\n## Adding a new runner / GPU type\n\n1. Register a self-hosted runner on the machine with the desired label (e.g. `h100`).\n2. Add the label to the `gpu` matrix in `.github/workflows/ci.yml`.\n3. Set `CI_WORK_DIR` for the new machine if its scratch path differs.\n</yakaai_skills>"}