Blame for code/scripts/bench_flash_attention.py

"""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=<free gpu> 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())