Triton FlashAttention-2 forward#
Track 07 - GPU · Notebook 04 · Runtime: ≈25 min on Ampere+ GPU
Prerequisites:
07_gpu/03(Triton tiled matmul).Hardware: Ampere or newer (compute capability ≥ 8.0). T4 is not supported by the Triton FlashAttention kernel because it relies on features introduced in Ampere’s tensor cores. No CUDA → every kernel check is skipped.
Papers:
Dao et al. 2022, FlashAttention.
Dao 2023, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.
What#
A Triton kernel for the FlashAttention-2 forward pass. The two characteristic properties:
Tiling over the Q rows instead of the K rows. FA1 tiled K,V to keep SRAM resident; FA2’s key change is to also move the Q iteration outside so each program handles one
BLOCK_M × dQ tile. Result: the parallel dimension matches the number of Q tiles, which is the same shape as the output - much higher occupancy.Online softmax with running max and denominator. Attention is
softmax(QK^T) V. FlashAttention never materialises the fullT × Tscore matrix; it streams blocks of K/V, updates a running(m, l, O)state per row, and rescales on each new block so the finalOis the correct softmax-weighted sum.
Target: match F.scaled_dot_product_attention to 1e-3 max abs error
at B=1, H=8, T=1024, D=64, and hit ≥3× naive attention speed at
T=2048.
from llm_systems_cookbook.nb import bootstrap
import math
import time
import torch
import torch.nn.functional as F
info = hardware_check()
print(info)
IS_CUDA = torch.cuda.is_available()
IS_AMPERE_PLUS = IS_CUDA and torch.cuda.get_device_capability(0) >= (8, 0)
if IS_CUDA and not IS_AMPERE_PLUS:
print("WARNING: CUDA detected but compute capability < 8.0; kernel checks will be skipped.")
s = bootstrap("07_gpu_04_triton_flashattention")
Smoke check: attention shape algebra#
Even without a GPU we can verify the shape invariants of the attention operation we’re about to accelerate. This ensures the kernel-free portion of the notebook does something useful, and that the FLOP count we target below is correct.
def attention_flops(batch: int, heads: int, seq: int, head_dim: int) -> int:
# Q @ K^T (softmax) score @ V each is B*H*T*T*D FLOPs with the 2x for mul+add.
return 4 * batch * heads * seq * seq * head_dim
def attention_bytes(batch: int, heads: int, seq: int, head_dim: int, dtype_bytes: int = 2) -> int:
# Q read + K read + V read + O write.
return dtype_bytes * batch * heads * head_dim * (seq + 2 * seq + seq)
ai = attention_flops(1, 8, 1024, 64) / attention_bytes(1, 8, 1024, 64)
print(f"attention AI at B=1 H=8 T=1024 D=64 = {ai:.1f} FLOPs/byte")
s.check(
"attention_ai_above_128",
lambda: ai > 128, # T=1024 attention at modest head count crosses A100 ridge.
msg=f"AI = {ai:.1f}",
)
The kernel#
One program per (batch, head, Q-tile) triple. Inside the program we
loop over K-tiles, maintaining three running state tensors:
m_i: row-wise max so far.l_i: row-wise softmax denominator so far.acc: runningsoftmax(...) Voutput, FP32.
For each new K-tile we compute its scores, take the new row max, rescale
the running acc and l_i by exp(m_old - m_new), and add the new
block’s contribution. That’s online softmax in six lines.
The tl.dot calls use FP32 accumulation; the final acc / l_i is cast
back to the input dtype for the store.
triton = None
tl = None
flash_attn_triton = None
if IS_AMPERE_PLUS:
try:
import triton
import triton.language as tl
@triton.jit
def _fa2_fwd(
Q, K, V, O,
sm_scale,
stride_qb, stride_qh, stride_qm, stride_qk,
stride_kb, stride_kh, stride_kn, stride_kk,
stride_vb, stride_vh, stride_vn, stride_vk,
stride_ob, stride_oh, stride_om, stride_ok,
Z, H, N_CTX,
BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr,
CAUSAL: tl.constexpr,
):
start_m = tl.program_id(0)
off_hz = tl.program_id(1)
off_z = off_hz // H
off_h = off_hz % H
q_offset = off_z * stride_qb + off_h * stride_qh
k_offset = off_z * stride_kb + off_h * stride_kh
v_offset = off_z * stride_vb + off_h * stride_vh
o_offset = off_z * stride_ob + off_h * stride_oh
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, HEAD_DIM)
q_ptrs = Q + q_offset + offs_m[:, None] * stride_qm + offs_k[None, :] * stride_qk
q = tl.load(q_ptrs, mask=offs_m[:, None] < N_CTX, other=0.0)
m_i = tl.zeros((BLOCK_M,), dtype=tl.float32) - float("inf")
l_i = tl.zeros((BLOCK_M,), dtype=tl.float32)
acc = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32)
end_n = (start_m + 1) * BLOCK_M if CAUSAL else N_CTX
for start_n in range(0, end_n, BLOCK_N):
cur_n = start_n + offs_n
k_ptrs = K + k_offset + cur_n[:, None] * stride_kn + offs_k[None, :] * stride_kk
v_ptrs = V + v_offset + cur_n[:, None] * stride_vn + offs_k[None, :] * stride_vk
k_blk = tl.load(k_ptrs, mask=cur_n[:, None] < N_CTX, other=0.0)
v_blk = tl.load(v_ptrs, mask=cur_n[:, None] < N_CTX, other=0.0)
qk = tl.dot(q, tl.trans(k_blk)) * sm_scale
if CAUSAL:
qk = tl.where(offs_m[:, None] >= cur_n[None, :], qk, float("-inf"))
qk = tl.where(cur_n[None, :] < N_CTX, qk, float("-inf"))
m_new = tl.maximum(m_i, tl.max(qk, axis=1))
alpha = tl.exp(m_i - m_new)
p = tl.exp(qk - m_new[:, None])
l_i = l_i * alpha + tl.sum(p, axis=1)
acc = acc * alpha[:, None] + tl.dot(p.to(v_blk.dtype), v_blk)
m_i = m_new
acc = acc / l_i[:, None]
o_ptrs = O + o_offset + offs_m[:, None] * stride_om + offs_k[None, :] * stride_ok
tl.store(o_ptrs, acc.to(O.dtype.element_ty), mask=offs_m[:, None] < N_CTX)
def flash_attn_triton(
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = True
) -> torch.Tensor:
B, H, T, D = q.shape
o = torch.empty_like(q)
BLOCK_M = 64
BLOCK_N = 64
grid = (triton.cdiv(T, BLOCK_M), B * H)
_fa2_fwd[grid](
q, k, v, o,
1.0 / math.sqrt(D),
q.stride(0), q.stride(1), q.stride(2), q.stride(3),
k.stride(0), k.stride(1), k.stride(2), k.stride(3),
v.stride(0), v.stride(1), v.stride(2), v.stride(3),
o.stride(0), o.stride(1), o.stride(2), o.stride(3),
B, H, T,
BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=D,
CAUSAL=causal,
)
return o
print("Triton FA2 kernel compiled")
except Exception as e: # noqa: BLE001
print(f"Triton unavailable: {type(e).__name__}: {e}")
flash_attn_triton = None
else:
print("CUDA compute capability < 8.0; kernel will not be compiled.")
Correctness#
Compare against F.scaled_dot_product_attention, which dispatches to
cuDNN’s own fused attention. Max abs error threshold: 1e-3 at
B=1, H=8, T=1024, D=64 in FP16.
if IS_AMPERE_PLUS and flash_attn_triton is not None:
torch.manual_seed(0)
B, H, T, D = 1, 8, 1024, 64
q = torch.randn((B, H, T, D), device="cuda", dtype=torch.float16)
k = torch.randn((B, H, T, D), device="cuda", dtype=torch.float16)
v = torch.randn((B, H, T, D), device="cuda", dtype=torch.float16)
with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]):
o_ref = F.scaled_dot_product_attention(q, k, v, is_causal=True)
o_triton = flash_attn_triton(q, k, v, causal=True)
max_err = (o_ref - o_triton).abs().max().item()
print(f"max abs error vs SDPA = {max_err:.3e}")
s.check(
"triton_fa2_matches_sdpa_within_1e_minus_3",
lambda: max_err < 1e-3,
msg=f"max abs err = {max_err:.3e}",
)
else:
s.skip("triton_fa2_matches_sdpa_within_1e_minus_3", "no Ampere+ / no Triton")
Performance#
Target: 3× naive attention at T=2048. Naive attention here is the
MATH backend of scaled_dot_product_attention, which materialises
the full T × T score matrix.
def bench_attn(fn, q, k, v, n_iter: int = 20) -> float:
for _ in range(3):
fn(q, k, v)
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(n_iter):
fn(q, k, v)
torch.cuda.synchronize()
return n_iter / (time.perf_counter() - t0)
if IS_AMPERE_PLUS and flash_attn_triton is not None:
torch.manual_seed(0)
B, H, T, D = 1, 8, 2048, 64
q = torch.randn((B, H, T, D), device="cuda", dtype=torch.float16)
k = torch.randn((B, H, T, D), device="cuda", dtype=torch.float16)
v = torch.randn((B, H, T, D), device="cuda", dtype=torch.float16)
def naive(q, k, v):
with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]):
return F.scaled_dot_product_attention(q, k, v, is_causal=True)
def ours(q, k, v):
return flash_attn_triton(q, k, v, causal=True)
its_naive = bench_attn(naive, q, k, v)
its_ours = bench_attn(ours, q, k, v)
print(f"naive SDPA (MATH) = {its_naive:6.1f} it/s")
print(f"triton FA2 = {its_ours:6.1f} it/s")
s.benchmark(
"triton_fa2_at_least_3x_naive",
measured=its_ours,
baseline=its_naive,
must_beat=3.0,
unit="it/s",
higher_is_better=True,
)
else:
s.skip("triton_fa2_at_least_3x_naive", "no Ampere+ / no Triton")
The O(N^2) vs O(N) memory story#
Naive attention materialises a T x T score matrix, so peak memory
grows quadratically with sequence length. FA2 never does - it streams
K/V blocks and keeps only the running (m, l, O) state per Q row,
so memory stays linear. Plot both peak memory and latency against
T to see the two curves diverge.
import matplotlib.pyplot as plt
if IS_AMPERE_PLUS and flash_attn_triton is not None:
Ts = [256, 512, 1024, 2048]
B, H, D = 1, 8, 64
naive_mem, flash_mem = [], []
naive_ms, flash_ms = [], []
for T in Ts:
q = torch.randn((B, H, T, D), device="cuda", dtype=torch.float16)
k = torch.randn((B, H, T, D), device="cuda", dtype=torch.float16)
v = torch.randn((B, H, T, D), device="cuda", dtype=torch.float16)
def _naive(q, k, v):
with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]):
return F.scaled_dot_product_attention(q, k, v, is_causal=True)
for label, fn, mem, lat in [("naive", _naive, naive_mem, naive_ms),
("flash", lambda q, k, v: flash_attn_triton(q, k, v, causal=True), flash_mem, flash_ms)]:
torch.cuda.reset_peak_memory_stats()
torch.cuda.synchronize()
its = bench_attn(fn, q, k, v, n_iter=10)
lat.append(1000.0 / its)
mem.append(torch.cuda.max_memory_allocated() / 1024**2)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.3))
ax1.plot(Ts, naive_mem, "o-", label="naive (MATH)", color="tab:gray")
ax1.plot(Ts, flash_mem, "o-", label="flash (triton)", color="tab:blue")
ax1.set_xlabel("seq len T"); ax1.set_ylabel("peak mem (MiB)")
ax1.set_title("memory: O(T^2) vs O(T)")
ax1.set_xscale("log"); ax1.set_yscale("log"); ax1.legend(); ax1.grid(True, which="both", alpha=0.3)
ax2.plot(Ts, naive_ms, "o-", label="naive", color="tab:gray")
ax2.plot(Ts, flash_ms, "o-", label="flash", color="tab:blue")
ax2.set_xlabel("seq len T"); ax2.set_ylabel("latency (ms/iter)")
ax2.set_title("latency vs seq length")
ax2.set_xscale("log"); ax2.set_yscale("log"); ax2.legend(); ax2.grid(True, which="both", alpha=0.3)
fig.tight_layout(); plt.show()
else:
print("skipped - no Ampere+ / no Triton.")
Exercises#
Add a
BLOCK_Mautotune sweep over{32, 64, 128}and measure how occupancy interacts with register pressure.Implement the backward pass (Dao 2023 §3.2). The trick: recompute
P = softmax(S)inside the backward so you never store aT × Tprobability matrix.Extend to variable-length sequences via a
cu_seqlensindex tensorstandard in real serving stacks to pack a batch of different-length requests into one kernel launch.
References#
flash-attnsource:flash_attn/flash_attn_triton.pyfor the production version of this kernel.Dao 2023 §2 for the m/l/O recursion; §3 for the FA2 partitioning change we implemented here.
s.summary()
s.save()