Attention Kernel Comparison: FlashAttention vs FlashInfer vs xformers vs Triton

Your choice of attention kernel can leave 40% of your GPU’s potential on the table. I learned this when we profiled our 70B serving stack and found the decode phase running at 58% of memory bandwidth — not because the workload was fundamentally limited, but because we’d defaulted to FlashAttention for everything, including paged KV cache decode where it’s simply not designed to excel. Switching to FlashInfer for decode alone recovered 18% throughput. The attention kernel is not a detail; it’s the difference between underutilizing your hardware and saturating it. Four major implementations compete for LLM inference workloads, and each one makes fundamentally different tradeoffs that matter when you’re serving production traffic.

The Contenders

FlashAttention-2/3

The gold standard for prefill. Tiled attention that keeps intermediates in SRAM, reducing HBM traffic from $O(N^2)$ to $O(N)$. FA-3 adds Hopper-specific optimizations (WGMMA, TMA, warp specialization) for ~1.3x over FA-2 on H100.

Limitation: Designed for contiguous Q, K, V tensors. Does not natively support paged KV cache (non-contiguous blocks). Serving systems use FA for prefill (contiguous KV) and a separate paged kernel for decode.

FlashInfer

Built for serving. Provides both contiguous and paged attention kernels:

import flashinfer

# Prefill: contiguous KV
output = flashinfer.single_prefill_with_kv_cache(
    q, k, v, causal=True
)

# Decode: paged KV cache
output = flashinfer.batch_decode_with_padded_kv_cache(
    q, kv_cache, page_table, seq_lens
)

Advantage over FlashAttention: native paged KV cache support means vLLM/SGLang can use one library for both prefill and decode instead of switching between FlashAttention and a custom paged kernel.

xformers Memory-Efficient Attention

Meta’s library. Provides memory_efficient_attention() that automatically selects the best backend:

from xformers.ops import memory_efficient_attention

output = memory_efficient_attention(query, key, value, attn_bias=causal_mask)

Advantage: cleanest PyTorch API, good for research. Disadvantage: slower than FlashAttention for production workloads, limited paged KV support.

Custom Triton Kernels

For novel attention patterns (sliding window, block sparse, linear attention), write a custom Triton kernel:

import triton
import triton.language as tl

@triton.jit
def sliding_window_attention_kernel(
    Q, K, V, Out,
    stride_qb, stride_qh, stride_qs, stride_qd,
    window_size: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    # Custom attention with sliding window
    # Each query attends to only [max(0, pos-W), pos] keys
    pid = tl.program_id(0)
    # ... kernel implementation

Advantage: maximum flexibility, Python-like syntax. Disadvantage: 20-40% slower than hand-tuned CUDA for standard patterns, harder to optimize for specific GPU architectures.

Throughput Benchmarks

When Each Wins