"""Kernel-level benchmark: triton_fa2 vs torch SDPA at the real Mistral shapes. Run (pin ONE free device; this is a light job and must NOT take the heavy lock): CUDA_VISIBLE_DEVICES= uv run python scripts/bench_flash_attention.py Shapes are Mistral-7B-v0.3's: b=4, 32 query heads : 8 KV heads, head_dim 128, causal, bf16, seq {512, 1024, 2048}. The forward column is compared against the baseline pinned in the plan artifact `tmp/sdpa_baseline.json`, measured on this same box: 0.189 / 0.624 / 2.282 ms. Purpose is to bound a decision, not to win: research (FA2 paper 2307.08691, flash-attention#1556) says Triton loses to CUTLASS at head_dim 128 on Ampere, ~1.3-1.5x forward and ~2x backward, and head_dim 128 is exactly Mistral's. At seq 1024 attention is a few percent of layer FLOPs, so even a 2x backward loss is a small single-digit share of step time. This script turns that into a measured number so the SDPA-default decision is evidence rather than vibes. """ from __future__ import annotations import json import os import sys sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from align.utils.guard import apply_resource_guards # Light job: no heavy lock, no 7B model. Still honors the path policy so no # cache lands off /apps. apply_resource_guards(cpu_threads=8, heavy=False, label="bench-flash-attention") import torch from torch.nn.attention import SDPBackend, sdpa_kernel from align.kernels.flash_attention import flash_attention REPO = "/apps/development/rlhf/align" OUT = f"{REPO}/docs/bench_flash_attention.json" #: Measured on this box during planning; see the module docstring. PINNED_SDPA_FWD_MS = {512: 0.189, 1024: 0.624, 2048: 2.282} B, HQ, HKV, D = 4, 32, 8, 128 SEQS = (512, 1024, 2048) def _bench(fn, iters=30, warmup=10): for _ in range(warmup): fn() torch.cuda.synchronize() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() for _ in range(iters): fn() end.record() torch.cuda.synchronize() return start.elapsed_time(end) / iters def bench_layer_share(seq: int = 1024, batch: int = 2) -> dict: """Put attention back in context: one full Mistral decoder layer, fwd+bwd. Benchmark 1 gives the ratio ON the attention kernel; this converts it into the number the SDPA-vs-triton decision actually turns on -- share of a training step. A layer is built from Mistral-7B-v0.3's real dimensions (hidden 4096, intermediate 14336, 32:8 heads, head_dim 128) so the MLP and projection GEMMs are correctly sized against the attention core. This is NOT the section-11/12 end-to-end number the section also asks for; that one needs the real 7B weights and a GPU window, and is deferred to it. A single layer times 32 is a tight bound on the attention share because the per-layer work is identical across layers. """ import torch.nn as nn H, I = 4096, 14336 dev_t, dt = "cuda", torch.bfloat16 lin = lambda a, b: nn.Linear(a, b, bias=False, device=dev_t, dtype=dt) wq, wk, wv, wo = lin(H, HQ * D), lin(H, HKV * D), lin(H, HKV * D), lin(HQ * D, H) wg, wu, wd = lin(H, I), lin(H, I), lin(I, H) x = torch.randn(batch * seq, H, device=dev_t, dtype=dt, requires_grad=True) def layer(attn_fn): def run(): x.grad = None q = wq(x).view(batch, seq, HQ, D).transpose(1, 2) k = wk(x).view(batch, seq, HKV, D).transpose(1, 2) v = wv(x).view(batch, seq, HKV, D).transpose(1, 2) a = attn_fn(q, k, v) h = x + wo(a.transpose(1, 2).reshape(batch * seq, H)) out = h + wd(torch.nn.functional.silu(wg(h)) * wu(h)) out.sum().backward() return run def sdpa_attn(q, k, v): with sdpa_kernel(SDPBackend.FLASH_ATTENTION): return torch.nn.functional.scaled_dot_product_attention( q, k, v, is_causal=True, enable_gqa=True ) t_sdpa = _bench(layer(sdpa_attn), iters=20, warmup=8) t_tri = _bench(layer(lambda q, k, v: flash_attention(q, k, v, causal=True)), iters=20, warmup=8) out = { "seq": seq, "batch": batch, "hidden": H, "intermediate": I, "n_layers": 32, "layer_sdpa_ms": t_sdpa, "layer_triton_ms": t_tri, "layer_ratio": t_tri / t_sdpa, "step_penalty_pct": 100.0 * (t_tri - t_sdpa) / t_sdpa, "est_32layer_sdpa_s": t_sdpa * 32 / 1000.0, "est_32layer_triton_s": t_tri * 32 / 1000.0, } print(f"\nlayer-level (Mistral decoder layer, b={batch} seq={seq}, fwd+bwd):") print(f" SDPA {t_sdpa:7.2f} ms/layer -> {out['est_32layer_sdpa_s']:.2f} s / 32 layers") print(f" triton {t_tri:7.2f} ms/layer -> {out['est_32layer_triton_s']:.2f} s / 32 layers") print(f" cost of choosing triton_fa2: {out['step_penalty_pct']:+.1f}% of the 32-layer body") return out def main() -> int: if not torch.cuda.is_available(): print("no CUDA device visible", file=sys.stderr) return 1 if torch.cuda.device_count() != 1: print( f"{torch.cuda.device_count()} devices visible; pin exactly one " "with CUDA_VISIBLE_DEVICES (shared box)", file=sys.stderr, ) return 1 dev = torch.cuda.get_device_name(0) cc = ".".join(str(x) for x in torch.cuda.get_device_capability(0)) print(f"device: {dev} (sm{cc.replace('.', '')}) torch {torch.__version__}") print(f"shapes: b={B}, {HQ}q:{HKV}kv heads, head_dim={D}, causal, bf16\n") results = {} for seq in SEQS: q = torch.randn(B, HQ, seq, D, device="cuda", dtype=torch.bfloat16) k = torch.randn(B, HKV, seq, D, device="cuda", dtype=torch.bfloat16) v = torch.randn(B, HKV, seq, D, device="cuda", dtype=torch.bfloat16) qg, kg, vg = (t.clone().requires_grad_() for t in (q, k, v)) dout = torch.randn_like(qg) def sdpa_fwd(): with sdpa_kernel(SDPBackend.FLASH_ATTENTION): return torch.nn.functional.scaled_dot_product_attention( q, k, v, is_causal=True, enable_gqa=True ) def tri_fwd(): return flash_attention(q, k, v, causal=True) def sdpa_fwd_bwd(): for t in (qg, kg, vg): t.grad = None with sdpa_kernel(SDPBackend.FLASH_ATTENTION): o = torch.nn.functional.scaled_dot_product_attention( qg, kg, vg, is_causal=True, enable_gqa=True ) o.backward(dout) def tri_fwd_bwd(): for t in (qg, kg, vg): t.grad = None flash_attention(qg, kg, vg, causal=True).backward(dout) row = { "sdpa_fwd_ms": _bench(sdpa_fwd), "triton_fwd_ms": _bench(tri_fwd), "sdpa_fwd_bwd_ms": _bench(sdpa_fwd_bwd), "triton_fwd_bwd_ms": _bench(tri_fwd_bwd), "pinned_sdpa_fwd_ms": PINNED_SDPA_FWD_MS[seq], } row["fwd_ratio_triton_over_sdpa"] = row["triton_fwd_ms"] / row["sdpa_fwd_ms"] row["fwd_bwd_ratio_triton_over_sdpa"] = row["triton_fwd_bwd_ms"] / row["sdpa_fwd_bwd_ms"] row["sdpa_fwd_drift_vs_pinned"] = row["sdpa_fwd_ms"] / row["pinned_sdpa_fwd_ms"] results[str(seq)] = row # Rebind to None rather than `del`: the bench closures above hold these # names as free variables, and `del` empties the closure cells (ruff # F821; a later call would NameError). Rebinding frees the tensors just # the same, so empty_cache() can actually reclaim. q = k = v = qg = kg = vg = dout = None torch.cuda.empty_cache() hdr = ( f"{'seq':>6} | {'SDPA fwd':>9} {'pinned':>8} | {'triton fwd':>11} " f"{'ratio':>7} | {'SDPA f+b':>9} {'triton f+b':>11} {'ratio':>7}" ) print(hdr) print("-" * len(hdr)) for seq in SEQS: r = results[str(seq)] print( f"{seq:>6} | {r['sdpa_fwd_ms']:>9.3f} {r['pinned_sdpa_fwd_ms']:>8.3f} | " f"{r['triton_fwd_ms']:>11.3f} {r['fwd_ratio_triton_over_sdpa']:>6.2f}x | " f"{r['sdpa_fwd_bwd_ms']:>9.3f} {r['triton_fwd_bwd_ms']:>11.3f} " f"{r['fwd_bwd_ratio_triton_over_sdpa']:>6.2f}x" ) layer = bench_layer_share() payload = { "device": dev, "compute_capability": cc, "torch": torch.__version__, "shape": { "batch": B, "q_heads": HQ, "kv_heads": HKV, "head_dim": D, "causal": True, "dtype": "bfloat16", }, "results": results, "layer_level": layer, } os.makedirs(os.path.dirname(OUT), exist_ok=True) with open(OUT, "w") as f: json.dump(payload, f, indent=2) f.write("\n") print(f"\nwrote {OUT}") return 0 if __name__ == "__main__": raise SystemExit(main())