scripts/attention_kernel_bench.py

# /// script
# requires-python = ">=3.10"
# dependencies = ["torch"]
# ///
"""
Empirical validation of the roofline prediction in roofline_analysis.py
(Claim 3): measure real decode-attention kernel latency vs KV cache length L
and query block size D (MTP degree proxy) on a GPU, and check that:

  - at fixed D, latency is flat in L once L is large enough that KV-cache
    bytes dominate DRAM traffic in the memory-bound regime, and grows with L
    only through the (small) compute term -- i.e. bandwidth-bound behavior
  - at fixed L, increasing D increases latency sub-linearly relative to the
    FLOP growth (D is "free" until compute-bound), then above the predicted
    D* the added latency starts tracking FLOPs (compute-bound) rather than
    being nearly free -- i.e. we look for the knee in latency-vs-D matching
    the analytical D* from roofline_analysis.py

Uses torch.nn.functional.scaled_dot_product_attention (SDPA) as a real,
tensor-core-backed attention kernel (no custom CUDA needed). Two configs:
  - "mha": standard multi-head attention, KV cache sized n_h x d_h per token
  - "mla": KV cache compressed to (kv_lora_rank + rope_dim) per token, shared
    across heads, with a linear up-projection applied before SDPA to emulate
    the "absorbed" MLA compute (FLOPs of the up-projection are included in
    the timed region, matching the analytical FLOP accounting).

Run standalone: python attention_kernel_bench.py --out outputs/kernel_bench.json
"""
import argparse
import json
import time

import torch
import torch.nn.functional as F

N_HEADS = 128
HEAD_DIM = 128
KV_LORA_RANK = 512
ROPE_DIM = 64

L_VALUES = [512, 2048, 8192, 32768]
D_VALUES = [1, 2, 4, 8, 16, 32, 64]
N_WARMUP = 5
N_ITERS = 20


def bench_mha(L, D, device, dtype):
    q = torch.randn(1, N_HEADS, D, HEAD_DIM, device=device, dtype=dtype)
    k = torch.randn(1, N_HEADS, L, HEAD_DIM, device=device, dtype=dtype)
    v = torch.randn(1, N_HEADS, L, HEAD_DIM, device=device, dtype=dtype)

    def step():
        return F.scaled_dot_product_attention(q, k, v, is_causal=False)

    return time_it(step, device)


def bench_mla(L, D, device, dtype):
    latent_dim = KV_LORA_RANK + ROPE_DIM
    kv_latent = torch.randn(1, L, latent_dim, device=device, dtype=dtype)
    # absorbed query: W^UK already folded in, so scoring happens directly in
    # latent space against kv_latent -- no K up-projection needed.
    q = torch.randn(1, N_HEADS, D, latent_dim, device=device, dtype=dtype)
    # K latent is shared/broadcast across all heads (read once, reused n_h times).
    k_shared = kv_latent.unsqueeze(1).expand(1, N_HEADS, L, latent_dim)
    # V still needs a per-head up-projection back to head_dim before the output.
    w_uv = torch.randn(N_HEADS, latent_dim, HEAD_DIM, device=device, dtype=dtype)

    def step():
        v = torch.einsum("bld,hde->bhle", kv_latent, w_uv)
        return F.scaled_dot_product_attention(q, k_shared, v, is_causal=False)

    return time_it(step, device)


def time_it(fn, device):
    for _ in range(N_WARMUP):
        fn()
    if device == "cuda":
        torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(N_ITERS):
        fn()
    if device == "cuda":
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - start
    return (elapsed / N_ITERS) * 1e3  # ms/iter


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="outputs/kernel_bench.json")
    ap.add_argument("--quick", action="store_true", help="smoke-test: tiny sizes, few points")
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32
    print(f"device={device} dtype={dtype}")

    l_values, d_values = L_VALUES, D_VALUES
    if args.quick:
        l_values, d_values = [64, 256], [1, 4, 16]

    rows = []
    for L in l_values:
        for D in d_values:
            try:
                t_mha = bench_mha(L, D, device, dtype)
                t_mla = bench_mla(L, D, device, dtype)
            except RuntimeError as e:
                print(f"SKIP L={L} D={D}: {e}")
                continue
            row = {"L": L, "D": D, "mha_ms": t_mha, "mla_ms": t_mla}
            rows.append(row)
            print(f"L={L:<7} D={D:<4} mha={t_mha:8.4f}ms  mla={t_mla:8.4f}ms")

    import os
    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    with open(args.out, "w") as f:
        json.dump({"device": device, "dtype": str(dtype), "rows": rows}, f, indent=2)
    print(f"\nSaved {args.out}")


if __name__ == "__main__":
    main()