code/scripts/bench_flash_attention.py · viai957/align at main

align / code /scripts /bench_flash_attention.py

Add files using upload-large-folder tool

bcc8d35 verified 16 days ago

9.02 kB

"""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())

	"""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())