windflash 1.0.0

WindFlash — Autotuned Flash Attention v2 in pure Triton

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 written in OpenAI Triton. 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

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

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.

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 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.

Testing

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

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.