FlashAttention

https://github.com/dao-ailab/flash-attention
Admin
FlashAttention is a fast and memory-efficient exact attention implementation with IO-awareness for...
Tokens:58,895
Snippets:357
Trust Score:7.8
License:BSD-3-Clause
Update:1 month ago
Show doc for...
Tokens:
Raw
### Install flash-attn-4

Source: https://github.com/dao-ailab/flash-attention/blob/main/CLAUDE.md

Install the flash-attn-4 package from PyPI, or perform a development install from the flash_attn/cute directory with the dev extras.

```bash
pip install flash-attn-4
# or dev install:
pip install -e "flash_attn/cute[dev]"
```

--------------------------------

### Install FlashAttention-2 from source

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Install FlashAttention-2 from source. The first command uses pip with --no-build-isolation; the second compiles directly from the source tree.

```sh
pip install flash-attn --no-build-isolation
```

```sh
python setup.py install
```

--------------------------------

### Install FlashAttention-3 beta with setup.py

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Install the FlashAttention-3 beta package from the hopper subdirectory. Requires an H100/H800 GPU and CUDA >= 12.3; CUDA 12.8 is recommended for best performance.

```sh
cd hopper
python setup.py install
```

--------------------------------

### Install FlashAttention

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Installs the FlashAttention package for fast and memory-efficient exact attention. This removes the need for activation checkpointing.

```sh
pip install flash-attn
```

--------------------------------

### Install FlashAttention with Triton backend

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Install FlashAttention with the Triton backend for AMD ROCm. Requires PyTorch for ROCm installed first. The FLASH_ATTENTION_TRITON_AMD_ENABLE environment variable enables the Triton backend.

```sh
cd flash-attention
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pip install --no-build-isolation .
```

--------------------------------

### Install FlashAttention-3 with uv in pyproject.toml

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Install FlashAttention-3 using uv by adding the dependency and source configuration to pyproject.toml. The source points to the hopper subdirectory of the flash-attention repository.

```toml
[project]
dependencies = [
    "flash-attn-3"
]

[tool.uv]
no-build-isolation = true

[tool.uv.sources]
flash-attn-3 = { git = "https://github.com/Dao-AILab/flash-attention", subdirectory = "hopper" }
```

--------------------------------

### Install with system Triton

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Install FlashAttention with the Triton backend while keeping the system Triton. Set AITER_USE_SYSTEM_TRITON=1 to avoid replacing the existing Triton on a ROCm PyTorch image.

```sh
cd flash-attention
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" AITER_USE_SYSTEM_TRITON=1 pip install --no-build-isolation .
```

--------------------------------

### Install FlashAttention-4 with pip

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Install FlashAttention-4. The second command adds the cu13 extra for best performance on CUDA 13.

```sh
pip install flash-attn-4
```

```sh
pip install "flash-attn-4[cu13]"
```

--------------------------------

### Install optimized cross-entropy loss

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Installs the optimized cross-entropy loss kernel. Supports bfloat16 and in-place backward to save memory.

```sh
cd ../csrc/xentropy && pip install .
```

--------------------------------

### Import and call flash_attn_interface.flash_attn_func

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Import and call the FlashAttention-3 interface after installation. The flash_attn_func() call is shown without arguments.

```python
from flash_attn_3 import flash_attn_interface
flash_attn_interface.flash_attn_func()
```

--------------------------------

### Install fused dense library

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Installs the fused matmul+bias and matmul+bias+gelu kernels from the fused_dense_lib directory. For best performance with bfloat16, use CUDA >= 11.8.

```sh
cd ../csrc/fused_dense_lib && pip install .
```

--------------------------------

### Install fused dropout+residual+LayerNorm

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Installs the fused dropout+residual+LayerNorm kernel. Supports dimensions divisible by 8, up to 6144, and works for both pre-norm and post-norm architectures.

```sh
cd ../csrc/layer_norm && pip install .
```

--------------------------------

### Install fused rotary embedding

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Installs the fused rotary embedding kernel.

```sh
cd ../csrc/rotary && pip install .
```

--------------------------------

### RS vs SS HGMMA examples for dQ = dS @ K

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/SASS_MMA_ANALYSIS.md

Compares register-source (RS) and shared-memory-source (SS) modes for the same GEMM.

```text
HGMMA.64x192x16.F32.BF16 R24, R232, gdesc[UR4], ...    # RS: A=R232 (dS in regs), B=gdesc (K in smem)
```

```text
HGMMA.64x192x16.F32.BF16 R24, gdesc[UR16], gdesc[UR4], ...   # SS: both from smem
```

--------------------------------

### Using Flash Attention 2 and 3 with the kernels library

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Use the kernels library to load Flash Attention 2 and 3 kernels. Requires installing the kernels package (pip install kernels). The get_kernel function fetches the specified kernel module, and the flash_attn_func is then accessed from the module.

```python
# pip install kernels

from kernels import get_kernel

# FA2
fa_module = get_kernel("kernels-community/flash-attn2", version=1)
flash_attn_func = fa_module.flash_attn_func

# FA3
fa3_module = get_kernel("kernels-community/flash-attn3", version=1)
flash_attn_func = fa3_module.flash_attn_func
```

--------------------------------

### Build against system aiter

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Build FlashAttention against an already installed aiter package instead of the bundled third_party/aiter submodule.

```sh
cd flash-attention
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" FLASH_ATTENTION_USE_SYSTEM_AITER="TRUE" pip install --no-build-isolation .
```

--------------------------------

### Example HGMMA instruction with operand annotations

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/SASS_MMA_ANALYSIS.md

Shows the destination register and source operand A in an HGMMA instruction.

```text
HGMMA.64x96x16.F32.BF16 R72, gdesc[UR16], RZ, ...
                         ^^^  ^^^^^^^^^^^
                         |    operand A (see RS vs SS below)
                         destination register = accumulator
```

--------------------------------

### Import flash_attn_qkvpacked_func and flash_attn_func

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Import the two main FlashAttention functions for scaled dot product attention.

```python
from flash_attn import flash_attn_qkvpacked_func, flash_attn_func
```

--------------------------------

### Resume training with resume=True

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Resumes a training run at exactly the same batch. Set a name to the run and then set resume=True when resuming.

```sh
python run.py experiment=pile/gpt3s-flash trainer.devices=8 name=pile-gpt3s-flash resume=True
```

--------------------------------

### Run sm90_config_search.py for SM90 block size tuning

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/SM90_BLOCK_SIZE_TUNING.md

Commands to enumerate feasible SM90 configurations for FlashAttention on Hopper. Use --headdim to specify head dimension(s), --mode to select forward/backward, and --tile-m/--tile-n to restrict tile choices.

```bash
python flash_attn/cute/sm90_config_search.py --headdim 128
```

```bash
python flash_attn/cute/sm90_config_search.py --mode fwd --headdim 192-128
```

```bash
python flash_attn/cute/sm90_config_search.py --mode bwd --headdim 192 --tile-m 64,80 --tile-n 64,96
```

--------------------------------

### Prepare Openwebtext dataset with pytest

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Sets PYTHONPATH and runs the Openwebtext dataset preparation test. Takes about 1 hour on a 64-core CPU and produces a 17GB processed dataset.

```sh
export PYTHONPATH=$PWD:$PYTHONPATH
pytest -q -s tests/datamodules/test_language_modeling_hf.py -k "openwebtext"
```

--------------------------------

### Train GPT2 on Openwebtext with 8 GPUs

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Runs GPT2 training experiments on Openwebtext with 8 GPUs. Default parameters are set for 8 x A100 80GB. Add trainer.precision=bf16 to train with bf16 instead of fp16.

```sh
python run.py experiment=owt/gpt2s-flash trainer.devices=8  # 125M
python run.py experiment=owt/gpt2m-flash trainer.devices=8  # 355M
python run.py experiment=owt/gpt2l-flash trainer.devices=8  # 760M
python run.py experiment=owt/gpt2xl-flash trainer.devices=8  # 1.6B
```

--------------------------------

### Run FlashAttention-4 tests

Source: https://github.com/dao-ailab/flash-attention/blob/main/CLAUDE.md

Run the test suite for FlashAttention-4. Use the -k flag to run a single test and -x to stop on the first failure.

```bash
pytest tests/cute/test_flash_attn.py
pytest tests/cute/test_flash_attn.py -k "test_flash_attn_output" -x  # single test
pytest tests/cute/test_flash_attn_varlen.py
pytest tests/cute/test_mask_mod.py
pytest tests/cute/test_score_mod.py
pytest tests/cute/test_block_sparsity.py
```

--------------------------------

### PTX comparison: hazard vs clean cp.async.bulk instructions

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/RACECHECK_TMA_HAZARD.md

Shows the PTX instructions for the hazard (bulk) and clean (tensor) reproducers. The hazard uses `cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes` while the clean version uses `cp.async.bulk.tensor.1d.shared::cta.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint`. The difference is the instruction type that racecheck instruments.

```ptx
# racecheck_repro_1d_bulk.py  (HAZARD)
cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes
    [%r42], [%rd12], %r43, [%r6+-16];

# racecheck_repro_1d_tensor.py  (CLEAN)
cp.async.bulk.tensor.1d.shared::cta.global.tile.mbarrier::complete_tx::bytes.L2::cache_hint
    [%r43], [%rd1, {%r71}], [%r6+-16], %rd8;
```

--------------------------------

### Build and run Docker image

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Build and run the Docker image for FlashAttention with Triton backend. The run command includes necessary device and security flags for ROCm.

```sh
docker build -t flash-attn-triton .
docker run -it --network=host --user root --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ipc=host --shm-size 16G --device=/dev/kfd --device=/dev/dri flash-attn-triton
```

--------------------------------

### Lint and format flash_attn/cute/ with ruff

Source: https://github.com/dao-ailab/flash-attention/blob/main/CLAUDE.md

Run ruff linting and formatting on the flash_attn/cute/ directory. Large kernel files (flash_bwd.py, flash_fwd.py, flash_fwd_sm100.py, interface.py) are excluded from auto-formatting.

```bash
ruff check flash_attn/cute/ --fix
ruff format flash_attn/cute/
```

--------------------------------

### Run Triton backend tests

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Run the Triton backend tests. Note that the full suite takes hours to complete.

```sh
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pytest tests/test_flash_attn_triton_amd.py
```

--------------------------------

### flash_attn_qkvpacked_func signature and docstring

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Function signature and docstring for flash_attn_qkvpacked_func. Use when Q, K, V are already stacked into one tensor; it is faster than flash_attn_func because the backward pass avoids explicit concatenation of gradients. Set dropout_p to 0.0 during evaluation. Supports sliding window local attention when window_size is not (-1, -1).

```python
flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False,
                          window_size=(-1, -1), alibi_slopes=None, deterministic=False):
"""dropout_p should be set to 0.0 during evaluation
If Q, K, V are already stacked into 1 tensor, this function will be faster than
calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
of the gradients of Q, K, V.
If window_size != (-1, -1), implements sliding window local attention. Query at position i
will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
Arguments:
    qkv: (batch_size, seqlen, 3, nheads, headdim)
    dropout_p: float. Dropout probability.
    softmax_scale: float. The scaling of QK^T before applying softmax.
        Default to 1 / sqrt(headdim).
    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
    window_size: (left, right). If not (-1, -1), implements sliding window local attention.
    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to
        the attention score of query i and key j.
    deterministic: bool. Whether to use the deterministic implementation of the backward pass,
        which is slightly slower and uses more memory. The forward pass is always deterministic.
Return:
    out: (batch_size, seqlen, nheads, headdim).
"""
```

--------------------------------

### Fix: pass tile to SeqlenInfo.create

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/VARLEN_PREPROCESS_TILE_BUG.md

The fixed call passes tile=self.tile_m to SeqlenInfo.create in flash_bwd_preprocess.py line 216, which matches the backward kernel's tile size.

```python
# Before:
seqlen = SeqlenInfo.create(batch_idx, mO.shape[1], mCuSeqlensQ, mSeqUsedQ)
# After:
seqlen = SeqlenInfo.create(batch_idx, mO.shape[1], mCuSeqlensQ, mSeqUsedQ, tile=self.tile_m)
```

--------------------------------

### Parse CLC trace (HTML view)

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/CLC_TRACE_DEBUG.md

Runs the parser with --html to generate an HTML view of the trace, saved to agent_space/clc_trace.html.

```bash
python AI/parse_clc_log.py agent_space/clc_trace.log --html -o agent_space/clc_trace.html
```

--------------------------------

### Dockerfile for Triton backend

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Dockerfile to build FlashAttention with the Triton backend on a ROCm PyTorch image.

```dockerfile
FROM rocm/pytorch:latest

WORKDIR /workspace

# build flash attention with triton backend
RUN git clone https://github.com/Dao-AILab/flash-attention &&\ 
    cd flash-attention &&\
    FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pip install --no-build-isolation .

# set working dir
WORKDIR /workspace/flash-attention

# set env variable to use triton backend
ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"
```

--------------------------------

### Dump SASS and extract HGMMA instructions

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/SASS_MMA_ANALYSIS.md

Compile with cubin output, find the cubin, disassemble, and grep for HGMMA instructions.

```bash
# Compile with cubin output
CUTE_DSL_KEEP_CUBIN=1 python -c "..."

# Find the cubin (saved in cwd with long name)
ls *.cubin

# Disassemble and extract HGMMA instructions
nvdisasm kernel.sm_90a.cubin | grep "HGMMA\."
```

--------------------------------

### Parse CLC trace (text summary)

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/CLC_TRACE_DEBUG.md

Runs the lightweight parser in AI/parse_clc_log.py to produce a compact per-SM / per-CTA summary of the trace.

```bash
python AI/parse_clc_log.py agent_space/clc_trace.log
```

--------------------------------

### Run racecheck reproducer commands for AI/

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/RACECHECK_TMA_HAZARD.md

Three shell commands: first validates correctness, second runs racecheck with CUTE_DSL_LINEINFO=1 showing one error, third runs racecheck on tensor variant showing zero hazards.

```bash
python AI/racecheck_repro_1d_bulk.py                                              # correctness
CUTE_DSL_LINEINFO=1 compute-sanitizer --tool=racecheck python AI/racecheck_repro_1d_bulk.py   # 1 error
compute-sanitizer --tool=racecheck python AI/racecheck_repro_1d_tensor.py         # 0 hazards
```

--------------------------------

### flash_attn_with_kvcache signature and docstring

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Function signature and docstring for flash_attn_with_kvcache. Use for incremental decoding with a pre-allocated KV cache. When k and v are provided, k_cache and v_cache are updated inplace with the new values, and attention is computed with the updated cache in one kernel. Ensure the cache is large enough to hold the new values; use cache_seqlens to track current sequence lengths.

```python
def flash_attn_with_kvcache(
    q,
    k_cache,
    v_cache,
    k=None,
    v=None,
    rotary_cos=None,
    rotary_sin=None,
    cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,
    cache_batch_idx: Optional[torch.Tensor] = None,
    block_table: Optional[torch.Tensor] = None,
    softmax_scale=None,
    causal=False,
    window_size=(-1, -1),  # -1 means infinite context window
    rotary_interleaved=True,
    alibi_slopes=None,
):
    """
    If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
    k and v. This is useful for incremental decoding: you can pass in the cached keys/values from
    the previous step, and update them with the new keys/values from the current step, and do
    attention with the updated cache, all in 1 kernel.

    If you pass in k / v, you must make sure that the cache is large enough to hold the new values.
    For example, the KV cache could be pre-allocated with the max sequence length, and you can use
    cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.
    """
```

--------------------------------

### Capture CLC trace with FA_LOG_LEVEL=3 and FA_CLC=1

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/CLC_TRACE_DEBUG.md

Minimal repro pattern to capture a CLC trace. Requires FA_LOG_LEVEL=3 for device-side prints and FA_CLC=1 to request CLC; the kernel may still fall back if the shape/features disable it. Redirects stdout/stderr to agent_space/clc_trace.log.

```bash
FA_LOG_LEVEL=3 FA_CLC=1 CUDA_VISIBLE_DEVICES=0 python - <<'PY' \
  > agent_space/clc_trace.log 2>&1
import torch
from flash_attn.cute.interface import flash_attn_func

torch.manual_seed(0)
q = torch.randn(1, 512, 16, 128, device='cuda', dtype=torch.bfloat16)
k = torch.randn(1, 512, 1, 128, device='cuda', dtype=torch.bfloat16)
v = torch.randn(1, 512, 1, 128, device='cuda', dtype=torch.bfloat16)
flash_attn_func(q, k, v, causal=True)
torch.cuda.synchronize()
PY
```

--------------------------------

### Prepare The Pile dataset with pytest

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Sets PYTHONPATH and runs The Pile dataset preparation test. Takes about 20 hours on a 64-core CPU and produces a 699GB processed dataset.

```sh
export PYTHONPATH=$PWD:$PYTHONPATH
pytest -q -s tests/datamodules/test_language_modeling_hf.py -k "pile"
```

--------------------------------

### Run FlashAttention-3 tests with pytest

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Run the FlashAttention-3 test suite. Set PYTHONPATH to the current directory before invoking pytest.

```sh
export PYTHONPATH=$PWD
pytest -q -s test_flash_attn.py
```

--------------------------------

### flash_attn_qkvpacked_func signature

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Function signature for FlashAttention-2 when inputs have the same sequence lengths in the same batch. It is simpler and faster to use this function in that case.

```python
flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False)
```

--------------------------------

### SASS HGMMA breakdown for BWD SM90 case study

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/SASS_MMA_ANALYSIS.md

Lists all 47 HGMMA instructions grouped by GEMM, with destination registers, sources, and modes.

```text
#1-12   64x112x16  dst=R24              src=gdesc[UR..]           SS  ×12  →  S = Q @ K.T
#13-20  64x112x16  dst=R24              src=gdesc[UR..]           SS  ×8   →  dP = dO @ V.T
#21-28  64x112x16  dst=R176/R120 alt    src=gdesc[UR..].tnsp      SS  ×4ea →  dV = P.T @ dO
#29-35  64x192x16  dst=R24              src=R232..R248            RS  ×7   →  dQ = dS @ K
#36-47  64x112x16  dst=R136/R80/R24 cyc src=gdesc[UR8].tnsp       SS  ×4ea →  dK = dS.T @ Q
```

--------------------------------

### Run FlashAttention-2 CK tests with pytest

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Runs the FlashAttention-2 test suite for the CK (Composable Kernel) backend. This is a separate test command from the main FlashAttention tests.

```sh
pytest tests/test_flash_attn_ck.py
```

--------------------------------

### Run FlashAttention tests with pytest

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Runs the FlashAttention test suite, comparing output and gradients against a reference implementation up to a numerical tolerance. The maximum numerical error is checked to be at most twice that of a PyTorch baseline across different head dimensions, input dtypes, sequence lengths, and causal/non-causal settings.

```sh
pytest -q -s tests/test_flash_attn.py
```

--------------------------------

### Train GPT3 on The Pile with 8 GPUs

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Runs GPT3 training experiments on The Pile with 8 GPUs. Default parameters are set for 8 x A100 80GB. To train with rotary embedding, run the experiments pile/gpt3{s,m,l,xl}-flash-rotary.

```sh
python run.py experiment=pile/gpt3s-flash trainer.devices=8  # 125M
python run.py experiment=pile/gpt3m-flash trainer.devices=8  # 355M
python run.py experiment=pile/gpt3l-flash trainer.devices=8  # 760M
python run.py experiment=pile/gpt3xl-flash trainer.devices=8  # 1.3B
python run.py experiment=pile/gpt3-2.7B-flash-hdim128 trainer.devices=8  # 2.7B
```

--------------------------------

### Fast two-pass testing with compilation cache

Source: https://github.com/dao-ailab/flash-attention/blob/main/CLAUDE.md

Two-pass testing workflow to avoid compilation overhead. Pass 1 compiles all kernels in parallel with FakeTensorMode (no GPU memory allocation) and caches them; Pass 2 runs the tests using the cached binaries. Set FLASH_ATTENTION_FAKE_TENSOR=1 for compilation and 0 for execution, and enable the persistent disk cache with FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1.

```bash
# Pass 1: compile all kernels in parallel using FakeTensorMode (no GPU memory allocation)
FLASH_ATTENTION_FAKE_TENSOR=1 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -n 64 -x tests/cute/test_flash_attn.py

# Pass 2: run tests using cached compiled kernels
FLASH_ATTENTION_FAKE_TENSOR=0 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -x tests/cute/test_flash_attn.py
```

--------------------------------

### Bug ledger template (agent_space/ledger_<bug-slug>.md)

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/DEBUG_METHODOLOGY.md

Template for a per-bug debugging ledger stored at agent_space/ledger_<bug-slug>.md. Records toolchain, deterministic repro command, evidence table, and a theory entry with prediction, falsification, cost, status, and rescue count. Use one file per bug and append as results land.

```markdown
# fp16 mismatch, local attention hdim64 — ledger
Toolchain: nvidia-cutlass-dsl 4.5.2, ptxas 13.0, driver 580.xx
Repro: CUDA_VISIBLE_DEVICES=3 pytest tests/cute/test_flash_attn.py -k "..." (deterministic, fails every run)

## Evidence
| Observed (artifact, file:line) | Inferred (causal claim) |
|---|---|
| diff.log:12 — first divergence at (b=0, h=2, q=191, d=17); q=191 is the last row of its m-block | tile-edge mask handling? |

## Theories
### T1: local-window mask off by one on the diagonal n-block
Predicts:      the set of divergent q-rows moves when n_block_size goes 128 → 64
Falsified by:  divergent-row set unchanged across n_block_size
Cost to check: 10 min (one recompile, diff the mismatch indices)
Status:        CONFIRMED — set shifted exactly with the tile edge (diff_n64.log:3)
Rescues:       0
```

--------------------------------

### Use flash_attn_func from flash_attn.cute

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Use FlashAttention-4 in Python. The flash_attn_func takes query, key, and value tensors and returns the attention output. The causal=True flag enables causal masking.

```python
from flash_attn.cute import flash_attn_func

out = flash_attn_func(q, k, v, causal=True)
```

--------------------------------

### flash_attn_qkvpacked_func

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Computes scaled dot product attention when Q, K, V are packed into a single tensor. This is faster than using separate Q, K, V tensors in the backward pass.

```APIDOC
## flash_attn_qkvpacked_func

### Description
Computes scaled dot product attention (softmax(Q @ K^T * softmax_scale) @ V) when Q, K, V are packed into a single tensor. This is faster than calling flash_attn_func on separate Q, K, V tensors because the backward pass avoids explicit concatenation of gradients.

### Method
Python function

### Endpoint
flash_attn.flash_attn_qkvpacked_func

### Parameters
#### Path Parameters
- None

#### Query Parameters
- None

#### Request Body
- **qkv** (tensor) - Required - (batch_size, seqlen, 3, nheads, headdim) - Packed Q, K, V tensor.
- **dropout_p** (float) - Optional - Dropout probability. Should be set to 0.0 during evaluation. Default: 0.0.
- **softmax_scale** (float) - Optional - Scaling of QK^T before softmax. Default: 1 / sqrt(headdim).
- **causal** (bool) - Optional - Whether to apply causal attention mask. Default: False.
- **window_size** (tuple) - Optional - (left, right) sliding window local attention. If (-1, -1), no window. Default: (-1, -1).
- **alibi_slopes** (tensor) - Optional - (nheads,) or (batch_size, nheads) fp32. Adds bias (-alibi_slope * |i - j|) to attention scores.
- **deterministic** (bool) - Optional - Use deterministic backward pass (slower, more memory). Forward is always deterministic. Default: False.

### Request Example
```python
from flash_attn import flash_attn_qkvpacked_func
out = flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False, window_size=(-1, -1), alibi_slopes=None, deterministic=False)
```

### Response
#### Success Response (Return)
- **out** (tensor) - (batch_size, seqlen, nheads, headdim) - Attention output.

#### Response Example
```python
# out shape: (batch_size, seqlen, nheads, headdim)
```
```

--------------------------------

### dK GEMM SASS showing 3 parts of 64 rows each

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/SASS_MMA_ANALYSIS.md

Illustrates how a 192-row output is split into three parts with their own accumulators.

```text
HGMMA.64x96x16 dst=R120   # part 0, K-step 0
HGMMA.64x96x16 dst=R72    # part 1, K-step 0
HGMMA.64x96x16 dst=R24    # part 2, K-step 0
HGMMA.64x96x16 dst=R120   # part 0, K-step 1
HGMMA.64x96x16 dst=R72    # part 1, K-step 1
HGMMA.64x96x16 dst=R24    # part 2, K-step 1
...  (4 K-steps total)
```

--------------------------------

### Host log line confirming CLC selection

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/CLC_TRACE_DEBUG.md

Host log line that explicitly confirms CLC was selected. Check for scheduling_mode=CLC to verify the shape actually used the CLC path.

```text
[FA] TileScheduler=SingleTileLPTScheduler, scheduling_mode=CLC, USE_2CTA=False
```

--------------------------------

### Use specific aiter commit

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Use a specific aiter commit for testing or development. Replace <commit-sha> with the desired commit hash.

```sh
cd flash-attention
cd third_party/aiter && git fetch origin && git checkout <commit-sha> && cd ../..
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pip install --no-build-isolation .
```

--------------------------------

### Set custom Triton config for attn_fwd

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Set a single Triton config for attn_fwd, overriding the hardcoded defaults. Use when not autotuning.

```sh
FLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON='{"BLOCK_M":128,"BLOCK_N":64,"waves_per_eu":1,"PRE_LOAD_V":false,"num_stages":1,"num_warps":8}'
```

--------------------------------

### Cite FlashAttention and FlashAttention-2

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

BibTeX entries for citing FlashAttention and FlashAttention-2. Use these when referencing the codebase or the papers in academic work.

```bibtex
@inproceedings{dao2022flashattention,
  title={Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},
  author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2022}
}
@inproceedings{dao2023flashattention2,
  title={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning},
  author={Dao, Tri},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2024}
}
```

--------------------------------

### Padded offset calculation with tile-dependent alignment

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/VARLEN_PREPROCESS_TILE_BUG.md

padded_offset_q uses tile-aligned gaps between sequences; the tile size must match the kernel's allocation. With tile_m=64 vs 128, the computed offset for the same batch differs, causing the preprocess to write to invalid locations.

```python
padded_offset_q = ((offset_q + batch_idx * tile_m) // tile_m) * tile_m
```

--------------------------------

### flash_attn_func signature and docstring

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Function signature and docstring for flash_attn_func. Use when Q, K, V are separate tensors. Supports multi-query and grouped-query attention (MQA/GQA) by passing KV with fewer heads than Q; the number of Q heads must be divisible by the number of KV heads. Set dropout_p to 0.0 during evaluation.

```python
flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False,
                window_size=(-1, -1), alibi_slopes=None, deterministic=False):
"""dropout_p should be set to 0.0 during evaluation
Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
If window_size != (-1, -1), implements sliding window local attention. Query at position i
will only attend to keys between
[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.

Arguments:
    q: (batch_size, seqlen, nheads, headdim)
    k: (batch_size, seqlen, nheads_k, headdim)
    v: (batch_size, seqlen, nheads_k, headdim)
    dropout_p: float. Dropout probability.
    softmax_scale: float. The scaling of QK^T before applying softmax.
        Default to 1 / sqrt(headdim).
    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
    window_size: (left, right). If not (-1, -1), implements sliding window local attention.
    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
        is added to the attention score of query i and key j.
    deterministic: bool. Whether to use the deterministic implementation of the backward pass,
        which is slightly slower and uses more memory. The forward pass is always deterministic.
Return:
    out: (batch_size, seqlen, nheads, headdim).
"""
```

--------------------------------

### CLC scheduler trace line format

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/CLC_TRACE_DEBUG.md

Trace line emitted by SM100 forward kernels at FA_LOG_LEVEL=3 for each scheduler-warp query. Fields: sm is the SM id, cta is blockIdx.x, m_blk/h/b/s are logical work coordinates, and valid indicates whether the returned work tile is valid (1) or the scheduler is exhausted (0).

```text
[CLC] query sm=<smid> cta=<blockIdx.x> (m_blk=<m>,h=<h>,b=<b>,s=<s>) valid=<0|1>
```

--------------------------------

### flash_attn_with_kvcache

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Computes attention with an in-place updatable KV cache, useful for incremental decoding. Supports rotary embeddings, cache sequence lengths, and block tables.

```APIDOC
## flash_attn_with_kvcache

### Description
Computes attention with an in-place updatable KV cache. If k and v are provided, k_cache and v_cache are updated in-place with the new values, enabling incremental decoding in a single kernel call.

### Method
Python function

### Endpoint
flash_attn.flash_attn_with_kvcache

### Parameters
#### Path Parameters
- None

#### Query Parameters
- None

#### Request Body
- **q** (tensor) - Required - Query tensor.
- **k_cache** (tensor) - Required - Key cache tensor. Updated in-place if k is provided.
- **v_cache** (tensor) - Required - Value cache tensor. Updated in-place if v is provided.
- **k** (tensor) - Optional - New key values to update the cache. If provided, k_cache is updated in-place.
- **v** (tensor) - Optional - New value values to update the cache. If provided, v_cache is updated in-place.
- **rotary_cos** (tensor) - Optional - Cosine values for rotary positional embeddings.
- **rotary_sin** (tensor) - Optional - Sine values for rotary positional embeddings.
- **cache_seqlens** (int or tensor) - Optional - Current sequence lengths of each sequence in the batch. Used to track cache positions.
- **cache_batch_idx** (tensor) - Optional - Batch indices for cache updates.
- **block_table** (tensor) - Optional - Block table for paged attention.
- **softmax_scale** (float) - Optional - Scaling of QK^T before softmax. Default: 1 / sqrt(headdim).
- **causal** (bool) - Optional - Whether to apply causal attention mask. Default: False.
- **window_size** (tuple) - Optional - (left, right) sliding window local attention. -1 means infinite context window. Default: (-1, -1).
- **rotary_interleaved** (bool) - Optional - Whether rotary embeddings are interleaved. Default: True.
- **alibi_slopes** (tensor) - Optional - ALiBi slopes for attention bias.

### Request Example
```python
from flash_attn import flash_attn_with_kvcache
out = flash_attn_with_kvcache(q, k_cache, v_cache, k=None, v=None, rotary_cos=None, rotary_sin=None, cache_seqlens=None, cache_batch_idx=None, block_table=None, softmax_scale=None, causal=False, window_size=(-1, -1), rotary_interleaved=True, alibi_slopes=None)
```

### Response
#### Success Response (Return)
- **out** (tensor) - Attention output.

#### Response Example
```python
# out shape: (batch_size, seqlen, nheads, headdim)
```
```

--------------------------------

### R2P SASS pattern for masking with leftover bit handling

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/SM90_R2P_MASKING_SASS.md

Shows the SASS instruction pattern where R2P converts 7 bits of a byte into predicates P0-P6, followed by FSEL instructions. Leftover bits (bit 7, 15, 23, 31) are handled with LOP3.LUT or ISETP instructions. The pattern repeats for each byte of the 32-bit mask.

```sass
SHF.R.U32.HI R9, RZ, R9, R16    ; shift to create bitmask
R2P PR, R9, 0x7f                  ; byte 0 → predicates P0-P6
FSEL R15, R36, -INF, P6           ; apply P6: keep or mask to -inf
R2P PR, R9.B1, 0x7f              ; byte 1 → predicates P0-P6
FSEL R52, R52, -INF, P6           ; apply P6
R2P PR, R9.B2, 0x7f              ; byte 2
...
R2P PR, R9.B3, 0x7f              ; byte 3
```

```sass
R2P PR, R12,     0x7f           ; bits 0-6   → P0-P6  (7 elements)
  14× FSEL using P0-P6           ; apply to 7 cols × 2 rows
LOP3.LUT P0, RZ, R12, 0x80, ... ; test bit 7  (1 element)
  2× FSEL using P0

R2P PR, R12.B1,  0x7f           ; bits 8-14  → P0-P6  (7 elements)
  14× FSEL using P0-P6
LOP3.LUT P1, RZ, R12, 0x8000, ..; test bit 15 (1 element)
  2× FSEL using P1

R2P PR, R12.B2,  0x7f           ; bits 16-22 → P0-P6  (7 elements)
  14× FSEL using P0-P6
LOP3.LUT P0, RZ, R12, 0x800000,..; test bit 23 (1 element)
  2× FSEL using P0

R2P PR, R12.B3,  0x7f           ; bits 24-30 → P0-P6  (7 elements)
  14× FSEL using P0-P6
ISETP.GT P0, R12, -1            ; test bit 31 (sign bit) (1 element)
  2× FSEL using P0
```

--------------------------------

### Printf guards to avoid print storms in CUTLASS DSL

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/DEBUG_2CTA.md

Use these guards to limit printf output to one thread per warp, one thread per CTA, or a specific thread. The elect_one context manager is not a bool; use it as a context manager.

```python
# One thread per warp:
if cute.arch.thread_idx()[0] % 32 == 0:
    cute.printf("...")

# One thread per CTA (elect_one is a context manager, not a bool):
with cute.arch.elect_one():
    cute.printf("...")

# One specific thread:
if tidx == 0:
    cute.printf("...")
```

--------------------------------

### Limit compilation jobs with MAX_JOBS

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Limit parallel compilation jobs with MAX_JOBS to avoid exhausting RAM on machines with less than 96GB and many CPU cores.

```sh
MAX_JOBS=4 pip install flash-attn --no-build-isolation
```

--------------------------------

### Fix: change copy_stats to use TMA descriptor

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/RACECHECK_TMA_HAZARD.md

Replace the `copy_stats` definition in the load function to use `cpasync.make_tiled_tma_atom` with `CopyBulkTensorTileG2SOp` instead of `cute.make_copy_atom` with `cpasync.CopyBulkG2SOp`. This generates `cp.async.bulk.tensor.1d` instead of `cp.async.bulk`, which racecheck does not instrument.

```python
copy_atom_stats = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), Float32)
copy_stats = partial(cute.copy, copy_atom_stats)
```

```python
copy_atom_stats = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), Float32)
copy_stats = partial(cute.copy, copy_atom_stats)
```

--------------------------------

### Construct GPT3-1.3B model with rotary embedding

Source: https://github.com/dao-ailab/flash-attention/blob/main/training/README.md

Constructs a GPT3-1.3B model with rotary embedding using GPT2Config and GPTLMHeadModel. Sets seqlen=2048, hidden_dim=2048, nheads=16, n_layer=24, and rotary_emb_fraction=0.5. Enables flash attention, fused MLP, fused bias FC, and fused dropout+add+LN.

```python
from transformers.models.gpt2.configuration_gpt2 import GPT2Config
from flash_attn.models.gpt import GPTLMHeadModel

seqlen = 2048
hidden_dim = 2048
nheads = 16
n_layer = 24
rotary_emb_fraction = 0.5
config = GPT2Config(vocab_size=50257, n_positions=seqlen, n_embd=hidden_dim,
                    n_layer=n_layer, n_head=nheads, 
                    scale_attn_by_inverse_layer_idx=True, 
                    rotary_emb_fraction=rotary_emb_fraction,
                    use_flash_attn=True, fused_mlp=True,
                    fused_bias_fc=True, fused_dropout_add_ln=True, 
                    pad_vocab_size_multiple=8)
model = GPTLMHeadModel(config)
```

--------------------------------

### Theory prediction template

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/DEBUG_METHODOLOGY.md

A template for documenting a root-cause theory with its prediction, falsifier, cost to check, and status. Use it to ensure a theory is testable before investing in a patch.

```text
Theory:      <mechanism, one sentence>
Predicts:    <observation that must hold if true>
Falsified by: <observation that would kill it>
Cost to check: <minutes>
Status:      UNTESTED | CONFIRMED | FALSIFIED
```

--------------------------------

### flash_attn_func signature

Source: https://github.com/dao-ailab/flash-attention/blob/main/README.md

Function signature for FlashAttention-2 when inputs have the same sequence lengths in the same batch. It is simpler and faster to use this function in that case.

```python
flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)
```

--------------------------------

### Racecheck error output for hazard reproducer

Source: https://github.com/dao-ailab/flash-attention/blob/main/AI/RACECHECK_TMA_HAZARD.md

Output from Compute Sanitizer racecheck showing a race between a write at 0x430 and multiple reads at 0x770-0x800 in `racecheck_repro_1d_bulk.py`. The write is the `cp.async.bulk` instruction on line 46, and the reads are four `ld.shared.b32` on line 55.

```text
Error: Race reported between Write access at ...+0x430 in racecheck_repro_1d_bulk.py:46
    and Read access at ...+0x770 in racecheck_repro_1d_bulk.py:55 [248 hazards]
    and Read access at ...+0x7a0 in racecheck_repro_1d_bulk.py:55 [248 hazards]
    and Read access at ...+0x7d0 in racecheck_repro_1d_bulk.py:55 [248 hazards]
    and Read access at ...+0x800 in racecheck_repro_1d_bulk.py:55 [248 hazards]
```
Tokens:
Raw