Flash Attention Integration AI Prompt

Integrate Flash Attention into this transformer model to reduce memory and improve speed.

1. Installation and compatibility check:
   - Install flash-attn: pip install flash-attn --no-build-isolation
   - Verify: requires GPU with compute capability ≥ 8.0 (A100, H100, 3090, 4090)
   - Check PyTorch version compatibility

2. Drop-in replacement:
   - Replace standard scaled dot-product attention with flash_attn_func or flash_attn_varlen_func
   - For HuggingFace models: set attn_implementation='flash_attention_2' in from_pretrained

3. Expected improvements:
   - Memory: O(N) instead of O(N²) in sequence length — enables much longer sequences
   - Speed: 2–4× faster than standard attention on A100
   - No approximation: exact same output as standard attention (not approximate)

4. Sequence length scaling:
   - Benchmark max sequence length with standard attention vs Flash Attention on the same GPU memory budget
   - Demonstrate quadratic vs linear memory scaling

5. Causal vs bidirectional:
   - For decoder models: set causal=True in flash_attn_func
   - For encoder models: causal=False

6. Variable-length sequences:
   - Use flash_attn_varlen_func with cu_seqlens to handle variable-length batches without padding waste
   - Compute cumulative sequence lengths from attention masks

7. Fallback:
   - Check if Flash Attention is available at runtime; fall back to scaled_dot_product_attention if not

Return: Flash Attention integration code, before/after memory and speed benchmark, and fallback implementation.