{"owner":"Dao-AILab","repo":"flash-attention","hasSkills":true,"hasMcp":false,"mcpConfig":null,"found":["CLAUDE.md"],"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.5.2`, `torch`, `einops`, `apache-tvm-ffi`, `quack-kernels>=0.5.0`.\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_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 needed) from execution (uses cached binaries):\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\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"}}