Metadata-Version: 2.4
Name: dew-flash-attn-cu13
Version: 0.1.0
Summary: FlashAttention-2 for JAX: Dao-AILab's CUDA kernels as XLA FFI targets
License-Expression: BSD-3-Clause
Project-URL: Source, https://github.com/AshishKumar4/dew/tree/main/kernels/flash_attn
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.11.2
Dynamic: license-file

# dew-flash-attn

FlashAttention-2 for JAX, built for [Dew](https://github.com/AshishKumar4/dew).
It packages Dao-AILab's CUDA kernels through flash_attn_jax's XLA FFI port, both
BSD-3, as a library with no Python ABI that registers two FFI targets with JAX.

There is one wheel for each CUDA major of your jax, for Linux x86_64 and any
Python from 3.11:

    pip install dew-flash-attn-cu12   # jax with jax-cuda12-plugin
    pip install dew-flash-attn-cu13   # jax with jax-cuda13-plugin

With it installed, Dew's `attention_impl="auto"` runs FlashAttention-2
on an A100 (sm80) for calls without a window, mask, bias or softcap and with
heads up to 256 wide. `attention_impl="flash"` asks for it on any sm8x GPU.

The kernels are compiled for sm80, which every compute capability 8.x runs.

Called directly:

```python
from dew_flash_attn import flash_mha

out = flash_mha(query, key, value, is_causal=True)  # [B, S, H, D], bf16 or fp16
```

Key heads may divide the query's (grouped-query attention). The backward pass
is the kernel's own.
