Optimizing Attention Flash

If you're working with transformer models and attention is your bottleneck, this implements Flash Attention to speed things up 2-4x while cutting memory usage by 10-20x. It uses IO-aware tiling and recomputation tricks under the hood. The skill shows you how to use PyTorch's native scaled_dot_product_attention (which auto-detects Flash Attention support) plus presumably the standalone flash-attn library for more control. Worth noting this has passed multiple security audits and comes from a repo with solid GitHub traction. The performance gains are real and measurable, especially on longer sequences where standard attention becomes prohibitively expensive.

skills.sh installs: 1K
GitHub stars: 14.4K
last updated: Aug 21, 2026

View on GitHub

Install

$ npx -y skills add orchestra-research/ai-research-skills --skill optimizing-attention-flash --agent claude-code

Installs into .claude/skills of the current project.

Files

Flash Attention - Fast Memory-Efficient Attention

Quick start

Flash Attention provides 2-4x speedup and 10-20x memory reduction for transformer attention through IO-aware tiling and recomputation.

PyTorch native (easiest, PyTorch 2.2+):

import torch
import torch.nn.functional as F

q = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)  # [batch, heads, seq, dim]
k = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)

# Automatically uses Flash Attention if available
out = F.scaled_dot_product_attention(q, k, v)

flash-attn library (more features):

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

from flash_attn import flash_attn_func

# q, k, v: [batch, seqlen, nheads, headdim]
out = flash_attn_func(q, k, v, dropout_p=0.0, causal=True)

Common workflows

Workflow 1: Enable in existing PyTorch model

Copy this checklist:

Flash Attention Integration:
- [ ] Step 1: Check PyTorch version (≥2.2)
- [ ] Step 2: Enable Flash Attention backend
- [ ] Step 3: Verify speedup with profiling
- [ ] Step 4: Test accuracy matches baseline

SKILL.md on GitHub