flash-attention_skill

Overview

This skill optimizes transformer attention using Flash Attention to deliver 2–4x speedups and 10–20x memory reductions for long-context models. It provides practical integration paths: PyTorch native SDPA (PyTorch ≥2.2), the flash-attn library for advanced features, and FP8 kernels for H100. The content focuses on checklists, code snippets, and benchmarking steps to verify performance and correctness.

How this skill works

Flash Attention replaces the standard softmax-based attention with an IO-aware tiled algorithm and recomputation that reduces intermediate memory and improves GPU utilization. The skill shows how to enable the PyTorch backend hook or call flash-attn kernels directly, and how to enable features such as causal attention, sliding-window/local attention, multi-query attention, and H100 FP8 acceleration. It also explains verification steps: profiling, numeric comparisons, and common fixes for CUDA or dtype issues.

When to use it

• Training or inference with sequence lengths >512 tokens where attention dominates compute.
• When GPUs run out of memory (OOM) due to attention states or KV caches.
• When you need 2–4x throughput improvements for long-context models (>2K tokens yields larger gains).
• To leverage H100 FP8 kernels for maximal FP8 throughput and lower memory on H100-class hardware.
• If you need sliding-window/local attention or multi-query attention for latency/memory tradeoffs.

Best practices

• Confirm environment: PyTorch ≥2.2 for native SDPA or install flash-attn with --no-build-isolation for CUDA toolkits.
• Benchmark with realistic batch/sequence sizes and use torch.utils.benchmark or wall-clock timing after warmup.
• Validate numeric parity: compare Flash Attention outputs to standard attention on float16/bfloat16 and check max difference (<1e-3 typical).
• Force backend selection during testing (torch.backends.cuda.sdp_kernel) to isolate speed/accuracy impacts.
• Use sliding-window or MQA when model architecture benefits from reduced KV state size; enable causal True for autoregressive models.

Example use cases

• Enable Flash Attention in an existing PyTorch transformer to remove attention OOMs and double throughput for 2K+ token training batches.
• Switch to flash-attn library to use sliding-window attention for long-document retrieval or local-context models.
• Deploy autoregressive LLM inference on H100 and enable FP8 kernels to maximize throughput and cut memory footprint.
• Benchmark and compare standard vs Flash Attention on a production GPU to quantify latency and VRAM savings before rollout.
• Integrate Flash Attention into Hugging Face model layers for Llama/GPT-style models to speed long-context fine-tuning.