Flash Attention vs Grouped Query Attention in PyTorch

You doubled your GPU memory, switched to an H100, and inference still OOMs at 32K context. The problem splits in two: the attention kernel's temporary SRAM usage during each forward pass, and the KV cache that accumulates one entry per layer per token per head. Flash attention vs grouped query attention addresses these two costs with independent mechanisms - a GPU kernel optimization and a model architecture change - yet most teams treat them as interchangeable.

Quick answer

FlashAttention-2 reduces peak SRAM usage during attention computation by tiling and fusing operations on-chip, cutting temporary memory and wall-clock time without changing model weights. Grouped query attention (GQA) shrinks the KV cache by sharing key-value heads across multiple query heads, reducing persistent memory that scales with sequence length. They target different bottlenecks and compose directly: running FlashAttention-2 on a GQA model like Llama 3 reduces both costs simultaneously.

Flash Attention vs grouped query attention target different bottlenecks

Standard multi-head attention (MHA) materializes the full Q×Kᵀ matrix in HBM, then reads it back for softmax and the V multiply. That round-trip dominates wall-clock time on bandwidth-bound GPUs. FlashAttention-2 eliminates the materialized matrix by tiling Q, K and V into SRAM blocks and fusing the softmax online, so the mathematical result stays identical but HBM traffic drops noticeably. No weights change. No retraining.

GQA operates at a different level. In MHA, every query head owns a private key head and value head. GQA groups multiple query heads behind a single shared KV pair - Llama 3 8B uses 8 KV heads for 32 query heads, a 4:1 ratio. The KV cache shrinks proportionally to the head-group ratio, and decode-time bandwidth drops with it.

Three memory pools are in play:

• Temporary attention memory - the intermediate matrices FlashAttention-2 eliminates.
• Model weights - unchanged by either technique.
• KV cache - persistent, grows with sequence length, and only GQA compresses it.

FlashAttention-2 attacks the first pool. GQA attacks the third. Assuming GQA replaces FlashAttention-2 leaves the quadratic temporary memory intact, which OOMs on long prefills even when the cache fits comfortably.

What grouped query attention is and how its head groups are arranged

Grouped query attention sits between multi-head attention and multi-query attention (MQA). GQA picks a point in between: several query heads share one KV pair, but more than one KV pair exists.

For a model with 32 query heads and 8 KV heads, each KV head serves a group of 4 query heads. Query heads 0-3 index into KV head 0, heads 4-7 into KV head 1, and so on. The grouping ratio must be an integer; Hugging Face Transformers enforces divisibility at config time.

Variant	KV heads	Grouping ratio	KV cache relative size
MHA	32	1:1	1×
GQA	8	4:1	0.25×
MQA	1	32:1	0.03×

Llama 3 8B and Mistral 7B both ship with 4:1 GQA. These ratios are baked into pretrained weights. You cannot flip a runtime flag on an MHA checkpoint and get equivalent-quality GQA - the KV projection matrices have different shapes, and the model needs uptraining to learn head sharing.

How GQA works during prefill, decoding, and KV-cache growth

During prefill, the model projects the full input through all 32 query heads but only 8 KV heads. Each query head computes attention scores against its assigned KV group. RoPE positional encodings apply before scoring. The output projection concatenates all 32 query-head outputs exactly as in MHA.

At decode time, each new token appends one key and one value vector per KV head per layer - not one per query head.

The KV-cache formula: 2 × layers × batch × seq_len × KV_heads × head_dim × bytes_per_element

For Llama 3 8B (32 layers, 8 KV heads, head_dim 128, bfloat16) at batch 1 and 8 192 tokens:

• GQA (8 KV heads): 2 × 32 × 1 × 8 192 × 8 × 128 × 2 ≈ 1.07 GB
• MHA (32 KV heads): 2 × 32 × 1 × 8 192 × 32 × 128 × 2 ≈ 4.29 GB

The ratio is 8/32 = 0.25 - a 75% reduction. Verify head counts in each model's config.json. Total VRAM also includes weights (~16 GB for 8B in bfloat16), CUDA allocator fragmentation, and framework overhead from PyTorch's attention layers. Cache compression extends the sequence length and batch size you can serve before hitting the memory wall.

Using FlashAttention-2 and grouped query attention together

Yes, they compose directly. FlashAttention-2's CUDA kernel accepts unequal query and KV head counts, repeating KV heads internally without expanding tensors in HBM. Each technique does its job independently inside the same forward pass.

During long-context prefill, FlashAttention-2 delivers its largest speedup because the temporary matrix scales quadratically. During decoding, GQA stays relevant because a 4× smaller cache means 4× less data read from HBM every step.

Compatibility checks before assuming both are active:

• GPU: sm_80+ (A100, H100). Turing cards fall back to a non-Flash kernel.
• dtype: bfloat16 or float16. float32 silently selects the math backend.
• Head dimension: ≤ 256 for flash-attn v2.5+; earlier versions cap at 128.
• Causal masking: pass is_causal=True explicitly.

Call torch.backends.cuda.flash_sdp_enabled() or check a profiler trace to verify selection. A config flag like attn_implementation="flash_attention_2" in Hugging Face requests the kernel but does not guarantee it - dtype mismatches cause silent fallback.

Implement GQA with PyTorch scaled dot product attention

PyTorch SDPA path for grouped query attention

PyTorch 2.5+ exposes enable_gqa=True in scaled_dot_product_attention, handling the head-count mismatch internally. Without that flag, you must repeat_interleave K and V, which allocates expanded tensors and defeats part of the memory savings.

import torch
import torch.nn.functional as F

B, S, H_Q, H_KV, D = 1, 2048, 32, 8, 128
Q = torch.randn(B, H_Q, S, D, dtype=torch.bfloat16, device="cuda")
K = torch.randn(B, H_KV, S, D, dtype=torch.bfloat16, device="cuda")
V = torch.randn(B, H_KV, S, D, dtype=torch.bfloat16, device="cuda")

out = F.scaled_dot_product_attention(
    Q, K, V,
    dropout_p=0.0,
    is_causal=True,
    enable_gqa=True,  # requires PyTorch >= 2.5
)

Set dropout_p=0.0 at inference - nonzero dropout disables the Flash kernel. All inputs must be contiguous bfloat16 or float16. To catch a silent fallback to the math backend:

torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(False)
torch.backends.cuda.enable_math_sdp(False)
# RuntimeError now fires if Flash cannot be selected

In Hugging Face Transformers, num_attention_heads and num_key_value_heads in config.json drive projection shapes. Setting attn_implementation="sdpa" in from_pretrained routes through SDPA. Verify correctness against an eager repeat_interleave reference with torch.allclose(out_sdpa, out_eager, atol=1e-2).

FlashAttention-2 path with unequal query and KV heads

The flash_attn package (v2.5+) accepts different head counts directly. The layout is (B, S, H, D), transposed from PyTorch's convention:

from flash_attn import flash_attn_func

Q_fa = Q.transpose(1, 2)  # (B, S, H_Q, D)
K_fa = K.transpose(1, 2)  # (B, S, H_KV, D)
V_fa = V.transpose(1, 2)

out_fa = flash_attn_func(Q_fa, K_fa, V_fa, causal=True)

The kernel broadcasts KV heads internally. I prefer the direct flash_attn dependency when I need deterministic kernel selection or variable-length batching via flash_attn_varlen_func. For standard serving behind vLLM or TGI, framework-managed integration is safer because those servers handle layout conversion and version pinning.

Benchmark standard attention, FlashAttention-2, and GQA on NVIDIA H100

A single latency number at one sequence length tells you almost nothing. Prefill is compute-bound; decoding is bandwidth-bound. Mixing them hides which technique earned the win.

Benchmark matrix and measurement controls

Fix dimensions to Llama 3 8B's shape: 32 layers, 32 query heads (MHA) or 8 KV heads (GQA), head_dim 128, bfloat16 on H100 SXM5 with CUDA 12.4 and PyTorch 2.5. Cross three kernel backends - eager matmul+softmax, PyTorch SDPA, FlashAttention-2 v2.5 - against MHA and GQA for six cells.

Run prefill at 512, 2 048, and 8 192 tokens. Run decoding with caches pre-filled to those lengths, generating 128 tokens. Use 10 warmup iterations, 50 timed trials, report p50, call torch.cuda.synchronize() before every timestamp, and reset torch.cuda.reset_peak_memory_stats() between configs. Record backend fallbacks as data points.

Latency, VRAM, and KV-cache comparison table

The KV-cache row uses exact arithmetic from the formula above. Latency and VRAM rows show qualitative direction - fill in measured values from your own hardware before making deployment decisions.

Metric	Eager MHA	Eager GQA	SDPA MHA	SDPA GQA	FA-2 MHA	FA-2 GQA
Prefill latency	baseline	≈ baseline	faster	faster	fastest	fastest
Decode latency	baseline	faster	faster	faster	faster	fastest
Peak VRAM	highest	high	moderate	moderate	low	lowest
KV-cache bytes/token	524 288	131 072	524 288	131 072	524 288	131 072
Backend support	all GPUs	all GPUs	sm_80+	sm_80+	sm_80+	sm_80+

Prefill speedups come almost entirely from the kernel - GQA barely changes prefill time because computation still fans out to 32 query heads. Decode speedups stack: FlashAttention-2 reduces kernel overhead, GQA reduces cache bandwidth. At 512 tokens, kernel launch overhead narrows differences to single-digit percentages.

When to choose PyTorch SDPA, FlashAttention-2, GQA, or both

Pick a GQA checkpoint when your workload hits the KV-cache wall - long contexts above 4K tokens, batch sizes above 8, or dozens of concurrent requests on a single GPU.

Pick PyTorch SDPA when portability matters. SDPA dispatches across Flash, memory-efficient, and math backends automatically, integrates with torch.compile, and adds zero dependencies. I reach for this on mixed GPU fleets.

Pick direct FlashAttention-2 when you control the hardware and long-context prefill dominates your serving profile. Below 2K tokens, kernel launch overhead eats most of the gain.

Combine GQA + FlashAttention-2 for production serving where both prefill compute and decode bandwidth constrain throughput. This is the default in vLLM and TGI when running Llama 3 or Mistral 7B on Ampere/Hopper GPUs.

Expect modest gains when sequences stay under 512 tokens, batches are 1, inputs arrive in float32, or the workload is dominated by MLP layers. Always verify kernel selection in a profiler trace - silent fallback erases the benefit entirely.

• Sequence length > 2K, batch > 4 → GQA checkpoint + FlashAttention-2
• Mixed GPU fleet → SDPA with enable_gqa=True
• Fixed Hopper hardware, long context → direct flash_attn
• Short prompts, low concurrency → weight quantization matters more
• MHA checkpoint → FlashAttention-2 still helps prefill; GQA requires retraining

FAQ

How does grouped query attention work

GQA assigns multiple query heads to share a single key-value head pair. During attention, each query head computes scores against only its assigned KV group rather than a private KV pair. This reduces distinct key and value projections stored in the KV cache, cutting cache memory proportionally to the grouping ratio - a 4:1 ratio (32 query heads, 8 KV heads) saves 75%.

What is grouped query attention (GQA)

Grouped query attention is a transformer attention variant between multi-head attention and multi-query attention. It uses fewer KV heads than query heads - Llama 3 8B and Mistral 7B both use 8 KV heads serving 32 query heads. The grouping ratio is fixed at training time via num_key_value_heads in the model config and cannot be changed at inference without retraining.