Metadata-Version: 2.5
Name: fa4-jax
Version: 0.1.0
Summary: FlashAttention-4 (CuTeDSL) for JAX: fwd/bwd/jvp with flex-attention mask mods, via cutejax
Project-URL: Repository, https://github.com/nshepperd/fa4-jax
Author-email: Emily Shepperd <nshepperd@gmail.com>
License-Expression: MIT AND BSD-3-Clause AND Apache-2.0
License-File: LICENSE
License-File: LICENSE.flash-attention
License-File: LICENSE.quack
Requires-Python: >=3.12
Requires-Dist: cutejax>=0.1.5
Requires-Dist: jax>=0.8.0
Requires-Dist: jaxtyping
Requires-Dist: nvidia-cutlass-dsl>=4.8
Provides-Extra: cu12
Requires-Dist: jax[cuda12]; extra == 'cu12'
Provides-Extra: cu13
Requires-Dist: jax[cuda13]; extra == 'cu13'
Requires-Dist: nvidia-cutlass-dsl[cu13]>=4.8; extra == 'cu13'
Description-Content-Type: text/markdown

# fa4-jax

FlashAttention-4 (Tri Dao's CuTe DSL kernels) as a JAX library, for the
SM80/SM120 family (Ampere → consumer Blackwell), SM90 (Hopper: H100/H200,
via the wgmma + TMA kernels) and SM100 (datacenter Blackwell: B200/GB200,
via the tcgen05 kernels), built on [cutejax](https://github.com/nshepperd/cutejax).

```sh
pip install "fa4-jax[cu13]"   # jax[cuda13] + the CuTe DSL's CUDA 13 runtime libraries
```

What it provides over `jax.nn.dot_product_attention`:

- **Native two-sided sliding windows** — `window_size=(left, right)` takes the
  kernel's local path and *skips* out-of-band KV blocks (cuDNN's fused kernel
  only takes a left window).
- **Flex-attention mask mods** — a `MaskMod` is a callable pytree: its class
  and `static()` fields are baked into the kernel as code via cutejax's
  static-pytree-leaf codegen; its array fields (e.g. document ids) stay
  runtime data, so new values reuse the compiled kernel.
- **Both differentiation modes** — reverse mode runs FA4's backward kernels;
  forward mode (`jax.jvp` / `jax.linearize`) runs a fused tangent kernel
  (`fa4_jax/cute/flash_jvp.py`, ours — upstream has no JVP). Built on custom
  primitives with a linear tangent + transpose rule, not `jax.custom_vjp`,
  which would forbid forward mode.
- **Varlen (packed) attention** — `flash_attn_varlen` takes a batch of
  different-length sequences concatenated along the token axis with no
  padding, delimited by int32 `cu_seqlens` offset arrays (the flash-attention
  varlen convention). One compile serves every packing: token counts *and*
  batch size stay symbolic.
- **A pure-JAX reference backend** with identical semantics (including the
  differentiation wiring) — the ground truth the kernels are tested against,
  and `flash_attn(..., backend="reference")` works anywhere.

```python
import jax.numpy as jnp
from fa4_jax import flash_attn, flash_attn_varlen, document_mask

out = flash_attn(q, k, v, causal=True)                       # (b, s, h, d) bf16/fp16
out = flash_attn(q, k, v, window_size=(128, 128))            # two-sided local
out = flash_attn(q, k, v, mask_mod=document_mask(doc_ids))   # flex: block-diag by (b, s) int32 id
out, lse = flash_attn(q, k, v, return_lse=True)

# Packed varlen: q (total_q, h, d), sequence b is q[cu_q[b]:cu_q[b+1]].
cu = jnp.asarray([0, 17, 300, 428], jnp.int32)               # 3 sequences
out = flash_attn_varlen(q, k, v, cu, cu, causal=True)        # (total_q, h, d)
```

Custom masks are frozen dataclasses subclassing `MaskMod`; array fields are
pytree leaves (runtime), anything else is declared `static()`:

```python
from dataclasses import dataclass
from fa4_jax import MaskMod, static

@dataclass(frozen=True)
class PrefixLM(MaskMod):
    prefix_len: jax.Array          # (b,) runtime data
    window: int = static()         # compile-time constant
    def __call__(self, b, h, q, kv):
        return (kv < self.prefix_len[b]) | ((kv <= q) & (q - kv <= self.window))
```

`causal` and `window_size` are bottom-right aligned when `s_q != s_kv` (the
flash-attention convention: the last query lines up with the last key, so
decode / chunked prefill just works). Mask mods see raw indices, so
`mask_mod=causal_mask()` is top-left and only equals `causal=True` when the
lengths match.

Varlen semantics: attention never crosses sequences; `causal`/`window_size`
apply per sequence, bottom-right aligned within it (`cu_seqlens_q` and
`cu_seqlens_k` may pack differently — e.g. one query per sequence against
long caches). Mask mods see the sequence index as `b` and sequence-local
`q_idx`/`kv_idx`. The LSE comes out `(heads, total_q)`. `cu_seqlens` contents
are runtime data and unchecked: they must be nondecreasing, start at 0, and
end at the token count (uncovered tokens get unspecified outputs/gradients).
Zero-length sequences are fine. Optional `seqused_q`/`seqused_k` (int32
`(batch,)`) limit each slot to its first `seqused[b]` tokens — the KV-cache
case, where slots are allocated larger than currently filled; unused K is
invisible, unused Q rows get unspecified outputs/gradients, and the
band alignment uses the used lengths. Reverse mode runs the FA4 varlen
backward and forward mode the fused varlen tangent kernel; `vmap` over a
varlen call is rejected — the packing *is* the batch.

Restrictions (kernel-side, all checked with clear errors): fp16/bf16, head_dim
a multiple of 8 and ≤ 256 — on the 99 KB parts (sm_86/89, sm_120) via smaller
tiles above 128 (64×64 forward, 64×32 backward); on sm_90 the tiles come from
upstream's per-head-dim table and a config is rejected only when its shared
memory doesn't fit. Two exceptions stay at ≤ 128: sm_100, where larger heads
would need 2-CTA MMAs, i.e. cluster launches, which this wrapper doesn't do,
and forward-mode (JVP), whose six tiles don't fit 99 KB beyond that — no cotangent
through `lse` in reverse mode. GQA works in all three passes (the backward
accumulates dK/dV in fp32 across the q heads of each kv head). On sm_100 the
backward runs the tcgen05 kernel for equal head dims and falls back to the
SM80-MMA kernel for mixed ones (which miscompile in upstream b26); on sm_90 it
runs the wgmma kernel whenever upstream's config for the head dim fits (GQA
there needs equal head dims) and otherwise the same SM80-MMA fallback. Mask
mods work in all three passes — the SM80/SM120 backward's mask_mod support is
a local addition to the vendored kernel (the SM90 and SM100 ones have it
upstream).
Combining
causal/window_size with a mask_mod keeps the native band's KV-block skipping
and runs the mod as a second masking pass over the visited tiles (also a
local mask.py change; upstream's branch structure silently drops the mod in
that combination). The mod itself can't skip blocks — that would need flex
block sparsity, which is not wired up.

## Layout

- `src/fa4_jax/attention.py` — public API, custom primitives, jvp/transpose rules.
- `src/fa4_jax/kernels.py` — the cute backend: cutejax plumbing around the kernels.
- `src/fa4_jax/masks.py` — dual-backend `MaskMod` (one predicate, kernel codegen + jnp).
- `src/fa4_jax/reference.py` — executable spec; ground truth for the tests.
- `src/fa4_jax/cute/` — vendored `flash_attn.cute` (see its `VENDOR.md`) plus
  our `flash_jvp.py` tangent kernel.

## Tests

```sh
uv run pytest              # kernels vs the jnp reference (fwd/grad/jvp/masks) + API checks; needs a CUDA device
```

Kernel-facing tests run the device's own kernels: SM80-MMA on sm_8x and
sm_12x, wgmma on sm_90, tcgen05 on sm_100 (the latter two backwards fall
back to SM80-MMA for configs their native kernels can't take). Covering every
generation means running the suite on each kind of hardware.

`conftest.py` sets `XLA_PYTHON_CLIENT_ALLOCATOR=cuda_async` and
`XLA_PYTHON_CLIENT_MEM_FRACTION=0.2` so the suite shares the GPU.

## Releasing

Tags drive the version (hatch-vcs). Before tagging, check the generations
no local machine has (Hopper, datacenter Blackwell). First for free:
`ci/cross_compile.py` builds every kernel for sm_90a/sm_100a on any GPU box
without launching anything, which catches tracing and compile errors (e.g. a
new cutlass-dsl changing an API). Then the real suite on rented RunPod pods,
with the API key injected by keyward:

```sh
uv run python ci/cross_compile.py 90 && uv run python ci/cross_compile.py 100
kw run -e RUNPOD_API_KEY=RUNPOD_API_KEY -- python3 ci/runpod_ci.py
git tag v0.2.0 && git push origin v0.2.0
```

The tag push runs `.github/workflows/release.yml`: build the wheel and sdist,
publish to PyPI (trusted publishing), then create a GitHub release with both
attached.

## License

fa4-jax's own code is MIT (`LICENSE`). The vendored kernels keep their
upstream licenses: flash-attention's CuTe DSL kernels in `src/fa4_jax/cute/`
are BSD-3-Clause (`LICENSE.flash-attention`), and quack's helpers in
`src/fa4_jax/cute/quack/` are Apache-2.0 (`LICENSE.quack`).
