Metadata-Version: 2.4
Name: windflash
Version: 1.0.0
Summary: WindFlash — Autotuned Flash Attention v2 in pure Triton
Author: MapleCascade
License-Expression: MIT
License-File: LICENSE
Keywords: attention,flash-attention,gqa,transformer,triton
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: torch>=2.1
Requires-Dist: triton>=3.0
Provides-Extra: dev
Requires-Dist: pytest; extra == 'dev'
Description-Content-Type: text/markdown

# WindFlash

**Autotuned Flash Attention v2 in pure Triton — no C++ build step, no prebuilt wheels, just `pip install` and go.**

A from-scratch implementation of [Flash Attention v2](https://arxiv.org/abs/2307.08691)
written in [OpenAI Triton](https://triton-lang.org/). Originally started as
"Windows Flash Attention" — an attempt to get flash-attn working on Windows
where prebuilt wheels don't exist — and evolved into an independent kernel with
its own autotuning, GQA support, backward pass, and CUDA C++ fast path.

## Features

- **13 autotuned tile configs** — optimal tile sizes per (head_dim, seq_len, GPU), benchmarked once and cached to disk
- **Persistent autotuner cache** — first-run results saved to `~/.windflash/`, instant startup on subsequent runs
- **GQA / MQA** — native in-kernel head mapping, no KV expansion needed
- **Backward pass** — split dKV/dQ design (no atomics), supports standard `loss.backward()`
- **CUDA C++ short-seq kernel** — warp-shuffle kernel for N ≤ 128, compiled via NVRTC at first use
- **3-tier dispatch** — CUDA C++ → cached Triton → full autotuner, based on sequence length
- **bf16 / fp16** — native tensor-core paths
- **Causal masking** — fused in-kernel
- **Monkey-patch** — drop-in replacement for `torch.nn.functional.scaled_dot_product_attention`
- **SM 8.0+** — RTX 30xx, 40xx, A100, H100

## Install

```bash
pip install -e .
```

Or copy `src/windflash/__init__.py` into your project — it's a single file.

**Requirements:** Python ≥ 3.10, PyTorch ≥ 2.1, Triton ≥ 3.0

## Quick Start

```python
from windflash import flash_attention

# Standard MHA
out = flash_attention(q, k, v)                    # (B, H, N, D)

# GQA: 32 Q heads, 8 KV heads — no expansion needed
out = flash_attention(q, k_gqa, v_gqa)            # Q: (B, 32, N, D), KV: (B, 8, N, D)

# Causal masking
out = flash_attention(q, k, v, causal=True)

# Training
out = flash_attention(q, k, v)
loss = out.sum()
loss.backward()                                    # dQ, dK, dV via Triton

# Monkey-patch torch SDPA globally
from windflash import enable_windflash, disable_windflash
enable_windflash()
# All F.scaled_dot_product_attention() calls now route through WindFlash
```

## Benchmarks

RTX 4090, bf16, batch=1. Full results in [TECHNICAL_REPORT.md](TECHNICAL_REPORT.md).

### D=64, 14 heads (TFLOPS — higher is better)

| Seq Len | WindFlash | SDPA cuDNN | SDPA efficient | FlexAttention |
|---------|-----------|------------|----------------|---------------|
| 512     | 15.6      | 35.1       | 37.9           | 10.6          |
| 1024    | 56.7      | 110.9      | 75.4           | 40.5          |
| 2048    | **130.1** | 135.9      | 101.4          | 120.9         |

### D=128, 32 heads (TFLOPS)

| Seq Len | WindFlash | SDPA cuDNN | SDPA efficient | FlexAttention |
|---------|-----------|------------|----------------|---------------|
| 1024    | 148.4     | 153.8      | 99.0           | 135.3         |
| 2048    | 157.9     | 163.9      | 103.8          | 148.8         |
| 4096    | **159.7** | 166.4      | 103.5          | 150.7         |

**Headlines:**
- **96% of cuDNN** peak throughput at D=128, N=4096
- **50%+ faster** than SDPA efficient at D=128, N≥1024
- **O(N) memory** — 64× less than naive attention at N=2048

### RTX 3080 Ti

| Shape | WindFlash | cuDNN | efficient | Δ vs cuDNN |
|-------|-----------|-------|-----------|------------|
| 1×14×1024×64 | **0.077 ms** | 0.087 ms | 0.102 ms | **1.13× faster** |
| 1×14×2048×64 | 0.239 ms | 0.227 ms | 0.307 ms | 0.95× |
| 1×32×1024×128 | 0.323 ms | 0.264 ms | 0.409 ms | 0.82× |
| 1×32×4096×128 | 4.132 ms | 3.904 ms | 5.956 ms | 0.94× |

> WindFlash **beats cuDNN on RTX 3080 Ti** at 1×14×1024×64.
> Consistently **#2 overall**, ahead of xFormers/efficient on medium-to-long sequences.

### Real-World Validation

Tested end-to-end in a GQA speech synthesis model (16 Q heads, 8 KV heads, D=64, 28 layers):

| Scenario | Result |
|----------|--------|
| **Without** `torch.compile` | **+22.6%** faster than stock PyTorch SDPA |
| **With** `torch.compile` | **−8% to −13%** (degrades — see [When NOT to Use](#when-not-to-use)) |

## When to Use

- **Inference without `torch.compile`** — best single-optimization speedup (+22.6%) for GQA models
- **Long sequences (N ≥ 1024) at D=128** — near-cuDNN throughput, 50%+ over SDPA efficient
- **Triton-only environments** — no CUDA compiler, no cuDNN, no prebuilt binaries
- **GQA/MQA models** — native head mapping avoids KV expansion overhead
- **Research & learning** — ~970 lines of readable Triton with persistent autotuner, 3-tier dispatch, split backward — great for understanding Flash Attention internals
- **Can't install flash-attn** — dependency hell on Windows, missing nvcc, wheel version conflicts

## When NOT to Use

- **`torch.compile` is enabled** — graph tracing conflicts with Triton monkey-patching, causing recompilation and loss of graph fusions. Net −8% to −13% vs compile + native SDPA. **This is the #1 disqualifier.**
- **Short sequences (N < 512, D=64)** — SDPA efficient and cuDNN are 2–4× faster due to lower dispatch overhead
- **cuDNN is available and N ≥ 4096** — cuDNN peaks at 166 TFLOPS vs WindFlash's 160
- **FP32 / TF32** — WindFlash only supports bf16/fp16
- **Custom attention masks** — causal only; for padding, sliding window, or prefix masks, use FlexAttention

### Decision Tree

```
Can you use torch.compile?
  ├─ YES → Use PyTorch native SDPA (cuDNN auto-selected)
  │        WindFlash adds no value and may hurt.
  │
  └─ NO → Can you install flash-attn?
           ├─ YES → Use flash-attn for maximum throughput
           │
           └─ NO → WindFlash ✓
                    • Zero deps beyond Triton
                    • +22–50% over SDPA efficient
                    • Near-cuDNN at long sequences
```

> **Honest take:** I would not recommend WindFlash for most production use cases.
> PyTorch's built-in SDPA backends with `torch.compile` will serve you better.
> WindFlash is useful in *specific* situations — dependency constraints, research,
> environments where flash-attn won't install, or non-compiled inference pipelines.

## Architecture

```
                    ┌─────────────────────────────────┐
                    │         flash_attention()        │
                    └──────────────┬──────────────────┘
                                   │
              ┌────────────────────┼────────────────────┐
              │                    │                     │
        N ≤ 128 &            Cached config          Full Triton
        CUDA ext avail?      or N ≤ 256?            autotuner
              │                    │                     │
        CUDA C++ kernel      Direct Triton call     13-config sweep
        (~12 µs, warp        (zero overhead)        (cached to disk)
         shuffle)
```

Single-file kernel: `src/windflash/__init__.py` (~970 lines)
- **Forward:** online softmax with pre-scaled Q, tiled K/V blocks, fused GQA head mapping
- **Backward:** split dKV/dQ kernels (no atomics), adaptive block sizes per head_dim
- **Cache:** `~/.windflash/{gpu_hash}.json` — persists across restarts

For the full architecture deep-dive, competitive analysis, and version history,
see [TECHNICAL_REPORT.md](TECHNICAL_REPORT.md).

## Testing

```bash
python tests/test_windflash.py
```

75 tests across 15 categories: forward (MHA/GQA/MQA, causal, fp16), backward
(MHA/GQA/D=128), config cache, mode switching, monkey-patch, fallback paths,
CUDA C++ short-seq, custom sm_scale, and realistic model shapes. All pass.

## Benchmarking

```bash
python bench/bench_flash.py                 # all GPUs, all configs
python bench/bench_flash.py --gpu 0 --quick # single GPU, quick run
python bench/bench_flash.py --save          # save CSV
```

## License

MIT — see [LICENSE](LICENSE).
