Spaces:
Running
Running
File size: 2,372 Bytes
e313085 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """Roofline floor for a task on a GPU: t_min = max(FLOPs / peak compute, bytes / peak bandwidth).
FLOPs come from torch.utils.flop_counter on the reference (matmul/conv/attention only, so
elementwise work counts as 0) and bytes are the compulsory traffic: inputs + parameters read once,
outputs written once. Both are lower bounds, so t_min is a floor no correct kernel can beat,
except tasks that legitimately skip input data (slicing), which KernelBench rarely has.
"""
from typing import Dict, Optional
# Dense datasheet peaks: (FP32 SIMT, TF32 tensor, BF16/FP16 tensor) in TFLOP/s, HBM in TB/s.
PEAKS = {
"H100 PCIe": (51.2, 378.0, 756.0, 2.0),
"H100": (66.9, 494.7, 989.4, 3.35), # SXM, "NVIDIA H100 80GB HBM3"
"H200": (66.9, 494.7, 989.4, 4.8),
"A100-SXM4-40GB": (19.5, 156.0, 312.0, 1.555),
"A100": (19.5, 156.0, 312.0, 2.039),
"L40S": (91.6, 183.0, 362.0, 0.864),
"L4": (30.3, 60.0, 121.0, 0.3),
}
def peaks(gpu: Optional[str]):
for key in sorted(PEAKS, key=len, reverse=True): # "L40S" before "L4", "H100 PCIe" before "H100"
if gpu and key in gpu:
return PEAKS[key]
return None
def roofline(flops, nbytes, dtype, gpu, t_ref_ms, t_new_ms=None) -> Optional[Dict]:
p = peaks(gpu)
if p is None or not nbytes or flops is None or not t_ref_ms:
return None
_, tf32, half, bw = p
# fp32 tasks: TF32 tensor cores are the roof, since tl.dot uses TF32 by default and passes atol=rtol=1e-2
compute = {"torch.float32": tf32, "torch.float16": half, "torch.bfloat16": half}.get(dtype)
if compute is None:
return None
t_compute = flops / (compute * 1e12) * 1e3
t_memory = nbytes / (bw * 1e12) * 1e3
t_min = max(t_compute, t_memory)
r = {
"arithmetic_intensity": flops / nbytes,
"bound": "compute" if t_compute > t_memory else "memory",
"t_min_ms": t_min,
"ref_roofline_frac": t_min / t_ref_ms,
}
if t_new_ms:
r["roofline_frac"] = t_min / t_new_ms
return r
def describe(r: Dict, gpu: str) -> str:
s = (f"Roofline on {gpu}: {r['bound']}-bound (arithmetic intensity {r['arithmetic_intensity']:.3g} FLOP/byte), "
f"floor {r['t_min_ms']:.4f} ms; PyTorch reaches {r['ref_roofline_frac']:.0%} of it")
if "roofline_frac" in r:
s += f", yours {r['roofline_frac']:.0%}"
return s + "."
|