How I Beat PyTorch's Flash Attention: A CUDA Optimization Journey

In 2022, I started learning CUDA from scratch. Fast forward to today, and I've built a Flash Attention implementation that beats PyTorch's cuDNN backend at sequence lengths of 4096 and beyond. This is the story of that journey - from barely understanding GPU architecture to achieving 22.6 TFLOPS on consumer hardware.

The Challenge: Why Flash Attention?

Flash Attention is the algorithm powering modern LLMs like GPT-4, Claude, and Llama. Traditional attention has O(N²) memory complexity - for a 4096 token sequence, that's 67 million attention scores stored in memory. Flash Attention reduces this to O(N) by computing attention in tiles, never materializing the full attention matrix.

But here's the thing: implementing it efficiently requires mastering CUDA at the deepest level - tensor cores, shared memory, warp-level primitives, and PTX assembly. PyTorch's implementation is written by NVIDIA engineers. Could I actually compete?

The Results: Beating PyTorch

On my RTX 3080 Ti Laptop GPU (batch=1, heads=32, head_dim=128, causal=true):

At longer sequences - exactly where Flash Attention matters most for LLMs - my implementation wins. Here's how I got there.

Phase 1: The Tensor Core Learning Curve (1.75 TFLOPS)

My first working implementation achieved 1.75 TFLOPS - just 8% of PyTorch. But it was correct, and that's what mattered initially.

The hardest part was understanding tensor cores. These specialized units perform 16x8x16 matrix multiplications in a single instruction. But the data layout is... non-intuitive.

// The mma.sync.aligned.m16n8k16 instruction
// Performs: D = A * B + C where A is 16x16 (FP16), B is 16x8 (FP16), C/D are 16x8 (FP32)

// Fragment layout for A matrix (16x16, stored across 32 threads)
// Each thread holds 8 FP16 values in 4 registers
uint32_t A_frag[4];  // Holds 8 half values

// Fragment layout for B matrix (16x8, stored across 32 threads)
// Each thread holds 4 FP16 values in 2 registers
uint32_t B_frag[2];  // Holds 4 half values

// Accumulator (16x8 FP32, stored across 32 threads)
// Each thread holds 4 FP32 values
float C_frag[4];

// The actual PTX instruction
asm volatile(
    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
    "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n"
    : "=f"(C_frag[0]), "=f"(C_frag[1]), "=f"(C_frag[2]), "=f"(C_frag[3])
    : "r"(A_frag[0]), "r"(A_frag[1]), "r"(A_frag[2]), "r"(A_frag[3]),
      "r"(B_frag[0]), "r"(B_frag[1]),
      "f"(C_frag[0]), "f"(C_frag[1]), "f"(C_frag[2]), "f"(C_frag[3])
);

Getting the fragment mapping right took weeks. The NVIDIA documentation is sparse, and a single wrong index means garbage output. I wrote dozens of test cases to verify each element mapping.

Phase 2: The Tile Size Breakthrough (17.5 TFLOPS)

The 10x performance jump came from a counterintuitive discovery: smaller tiles are faster.

My initial implementation used 64x64 tiles for both Q and K/V. The math said bigger tiles = more compute per memory load = better efficiency. The reality was different.

// Before: 64x64 tiles = 1.75 TFLOPS
constexpr int BLOCK_M = 64;  // Q tile rows
constexpr int BLOCK_N = 64;  // K/V tile rows
// Shared memory: 64 * 128 * 2 * 2 = 32KB per tile
// Register pressure: High
// SM occupancy: Low (only 1-2 blocks per SM)

// After: 64x32 tiles = 17.5 TFLOPS
constexpr int BLOCK_M = 64;  // Q tile rows
constexpr int BLOCK_N = 32;  // K/V tile rows (halved!)
// Shared memory: 32 * 128 * 2 * 2 = 16KB per tile
// Register pressure: Lower
// SM occupancy: High (4+ blocks per SM)

The key insight: GPU occupancy matters more than per-block efficiency. With smaller K/V tiles, I could run more thread blocks concurrently, keeping all 58 SMs on my 3080 Ti busy.

Phase 3: Micro-Optimizations (18.4 TFLOPS)

Several micro-optimizations added incremental gains:

1. Shared Memory Bank Conflict Elimination

// Problem: 32 banks, 128 elements = 4 elements/bank = conflicts!
__shared__ half K_smem[BLOCK_N][HEAD_DIM];  // Stride = 128

// Solution: Add padding to offset banks
constexpr int K_STRIDE = HEAD_DIM + 8;  // Stride = 136
__shared__ half K_smem[BLOCK_N][K_STRIDE];
// Now threads access different banks

2. Warp Shuffle Reductions

// Before: Shared memory reduction (slow)
__shared__ float smem_max[BLOCK_M];
smem_max[threadIdx.x] = local_max;
__syncthreads();
// ... reduction loop with multiple syncs

// After: Warp shuffle reduction (fast, no shared memory)
float m_ij = local_max;
m_ij = fmaxf(m_ij, __shfl_xor_sync(0xffffffff, m_ij, 1));
m_ij = fmaxf(m_ij, __shfl_xor_sync(0xffffffff, m_ij, 2));
m_ij = fmaxf(m_ij, __shfl_xor_sync(0xffffffff, m_ij, 4));
m_ij = fmaxf(m_ij, __shfl_xor_sync(0xffffffff, m_ij, 8));
m_ij = fmaxf(m_ij, __shfl_xor_sync(0xffffffff, m_ij, 16));
// 5 instructions, zero shared memory, zero syncs

Warp shuffles are magical. They let threads exchange data within a warp (32 threads) in a single cycle, without touching shared memory.

Phase 4: The exp2f Optimization (20.9 TFLOPS)

This one optimization gave me a 13% speedup. Softmax requires computing millions of exponentials, and CUDA's `expf()` is slower than `exp2f()`.

// Before: Standard exponential
float p = expf(s - max_val);

// After: exp(x) = 2^(x * log2(e))
constexpr float LOG2E = 1.4426950408889634f;
float p = exp2f((s - max_val) * LOG2E);

// Why it's faster: The GPU's Special Function Unit (SFU) executes
// exp2f in fewer cycles than expf. The multiplication by LOG2E
// is essentially free compared to the SFU savings.

This is the kind of optimization you only learn from reading NVIDIA's internal documentation and GPU architecture papers. The compiler doesn't do this automatically.

Phase 5: Adaptive Tile Dispatch - Beating PyTorch (22.6 TFLOPS)

The final breakthrough came from realizing that different sequence lengths need different tile sizes.

My 64x32 kernel was great for long sequences, but underperformed at short ones:

The reason: at seq_len=512 with BLOCK_M=64, I launch only 8 blocks. My GPU has 58 SMs - most are idle! The smaller 32x32 kernel launches 16 blocks, better utilizing the hardware.

// The adaptive dispatcher
constexpr int SEQ_LEN_THRESHOLD = 768;

void flash_attention_forward(/* ... */) {
    if (seq_len < SEQ_LEN_THRESHOLD) {
        // Small tile: BLOCK_M=32, 64 threads
        // Higher occupancy, better for short sequences
        flash_attention_kernel_small<128, IS_CAUSAL>
            <<<grid_small, 64, smem_small>>>(Q, K, V, O, seq_len, scale);
    } else {
        // Large tile: BLOCK_M=64, 128 threads
        // Better tensor core utilization for long sequences
        flash_attention_kernel_large<128, IS_CAUSAL>
            <<<grid_large, 128, smem_large>>>(Q, K, V, O, seq_len, scale);
    }
}

This adaptive approach is why I beat PyTorch at long sequences while staying competitive at short ones. Production implementations need this kind of runtime dispatch.

What Didn't Work

Not every optimization succeeded. Here's what I tried and abandoned:

• •Double-Buffer Prefetch: Loading next K/V tile while computing current one. Result: 12 TFLOPS (33% slower). The kernel is compute-bound, not memory-bound.
• •Epilogue Fusion with half2 Stores: Vectorized 2-element writes. Result: 14-15 TFLOPS (20% slower). Register pressure from packing values hurt occupancy.
• •cp.async for Asynchronous Loads: Modern async memory copies. Result: Illegal memory access errors. Zero-fill paths for bounds checking were incompatible.

These failures taught me as much as the successes. GPU optimization is empirical - you have to measure everything.

The Complete Journey

Total improvement: 12.9x from initial to final.

Key Lessons Learned

• •Occupancy often beats efficiency: Running more blocks concurrently can outweigh per-block optimizations.
• •Profile everything: Intuition about what's slow is often wrong. Use Nsight Compute.
• •Read the PTX manual: The highest-performance code requires understanding assembly.
• •Test obsessively: One wrong index corrupts everything. Build comprehensive test suites.
• •Study the hardware: Understanding GPU architecture (SMs, warps, banks) is essential.

What's Next

The code is open source at GitHub. Future optimizations I'm exploring:

• •Multi-head parallelism across SMs
• •FP8 tensor cores on Ada Lovelace GPUs
• •Backward pass implementation
• •Sliding window attention for ultra-long sequences

This project proved that with enough persistence, you can compete with the best implementations in the world. The gap between 'good enough' and 'state-of-the-art' is smaller than most people think - it just requires going deep.

The best way to learn GPU programming is to pick a problem that matters, then refuse to stop until you've beaten the baseline. Flash Attention was that problem for me.