⚡ FlashAttention-2 MMA

Production-Grade Custom CUDA FlashAttention Implementation

---

📖 Overview

This repository contains a from-scratch implementation of FlashAttention-2, designed as a custom PyTorch C++ Extension. The architecture is engineered to support modern Large Language Model workloads, handling dynamic sequence lengths, causal masking, and advanced kernel dispatching strategies (Split-Q, Split-KV, Split-D).

Currently, the project features a Mathematically Verified Naive Baseline using standard CUDA cores, serving as a stable fallback to guarantee 100% correctness and deterministic output when validating against PyTorch's native scaled_dot_product_attention.

---

🏗️ Architecture & Dispatch Flow

The underlying C++ API intelligently routes the PyTorch Tensor inputs to specific CUDA kernels based on the optimal parallelization strategy for the given tensor dimensions:

graph TD
    A[PyTorch Python API] -->|torch.Tensor| B(C++ Binding Layer)
    B --> C{Dispatch Logic}
    
    C -->|seq_q == 1| D[Split-KV Kernel<br/>Decode Phase]
    C -->|head_dim > 256| E[Split-D Kernel<br/>Large Head Dim]
    C -->|Default| F[Split-Q Kernel<br/>Standard Prefill]
    
    D --> G[(Mathematical Backend<br/>CUDA Cores)]
    E --> G
    F --> G
    
    G --> H[Output Tensor]

Loading

Supported Dispatch Strategies

1. Split-Q (Prefill): Parallelizes across the Query sequence length. Optimal for standard forward passes.
2. Split-KV (Decode): Parallelizes the reduction across the Key/Value sequence length. Essential when $S_q = 1$.
3. Split-D (Large Head Dim): Partitions the head dimension ($D$) across thread blocks to handle massive projection spaces without shared memory bank conflicts.

---

📊 Performance Benchmarks

The project features two tiers of execution. The Naive Baseline is the functional scalar fallback, while the MMA Optimized utilizes NVIDIA's nvcuda::wmma API to compile down to raw hardware-level Tensor Core instructions.

Tensor Core MMA Speedup (D=128)

For standard LLM dimensions (head_dim = 128), the dispatch routes to the Tensor Core kernel, yielding ~14x Speedups over the baseline!

Configuration Context	Tensor Dimensions (B, H, Sq, Sk, D)	Naive Latency	MMA Latency	Speedup	TFLOPS (MMA)
LLaMA-7B Prefill	(1, 32, 4096, 4096, 128)	1802.30 ms	135.38 ms	13.3x 🚀	2.03
BERT-large	(8, 32, 2048, 2048, 128)	3827.81 ms	273.72 ms	14.0x 🚀	2.01

Complete Benchmark Matrix (NVIDIA T4)

Configuration Context	Tensor Dimensions (B, H, Sq, Sk, D)	TFLOPS	Latency (ms)
GPT-4 Style	(1, 16, 2048, 2048, 256)	0.14	476.95
Large Head	(1, 8, 1024, 1024, 512)	0.07	252.29
Massive Head	(1, 4, 512, 512, 1024)	0.04	113.67
Decode KV	(1, 8, 1, 8192, 128)	0.01	6.15

Note: The head_dim > 128 configurations automatically fallback to the mathematically verified scalar loop kernel (Naive), serving as a 100% accurate baseline for validation.

---

⚙️ Technical Details

Online Softmax (Safe Softmax)

To avoid out-of-memory (OOM) errors during the $N \times N$ attention matrix generation, the kernel implements a streaming (online) softmax computation algorithm: $$m_i = \max(m_i, s_{ij})$$ $$l_i = l_i \cdot e^{m_{old} - m_{new}} + e^{s_{ij} - m_{new}}$$

Causal Masking

The implementation fully supports autoregressive token generation by natively enforcing $k_{idx} \leq q_{idx}$ during the matrix multiplication loop, inherently bypassing redundant sequence padding.

---

🚀 Quickstart & Installation

1. Requirements

• PyTorch >= 2.0
• CUDA Toolkit >= 11.8
• Ninja Build System
• Google Colab / Linux Environment Recommended

2. Build the C++ Extension

Instead of dealing with isolated build environments, compile the CUDA libraries in-place:

python setup.py build_ext --inplace

3. Install Python Bindings

Link the package to your Python environment without triggering PEP-517 isolation:

pip install -e . --no-build-isolation

4. Verify Correctness

Run the PyTest suite to validate the custom extension against PyTorch's native C++ backend:

pytest tests/test_correctness.py -v

---

💻 Usage Example

import torch
from flash_attn_mma.functional import flash_attention

# 1. Initialize random tensors on CUDA
batch_size, num_heads, seq_len, head_dim = 2, 8, 1024, 128
q = torch.randn(batch_size, num_heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
k = torch.randn(batch_size, num_heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
v = torch.randn(batch_size, num_heads, seq_len, head_dim, dtype=torch.float16, device="cuda")

# 2. Run custom FlashAttention (causal masking enabled)
output = flash_attention(q, k, v, causal=True, strategy="auto")

print(f"Output shape: {output.shape}") 
# Output shape: torch.Size([2, 8, 1024, 128])