[Kernel] Add inventory guards and clean benchmark layout (#32788)
This commit is contained in:
Xiaoyu Zhang
@@ -0,0 +1,487 @@
|
||||
# Benchmark FP8 attention for FA4 (CuTe-DSL) on SM100.
|
||||
#
|
||||
# Run (recommended):
|
||||
# python benchmark/kernels/attention/bench_flash_attention_fp8.py
|
||||
#
|
||||
# Notes:
|
||||
# - This is intended to be used while bringing up FP8 support for SM100.
|
||||
# - FP8 correctness depends on descales + max-offset scaling being implemented in the SM100 kernel.
|
||||
# This script optionally checks output vs a BF16 PyTorch baseline on dequantized FP8 inputs.
|
||||
#
|
||||
# Adapted from: `hopper/benchmark_flash_attention_fp8.py`
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import inspect
|
||||
import math
|
||||
import time
|
||||
from typing import Iterable
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from fa4_benchmark_utils import benchmark_forward
|
||||
|
||||
from sglang.kernels.ops.attention.flash_attn.cute.interface import (
|
||||
_flash_attn_fwd as flash_attn_cute_fwd,
|
||||
)
|
||||
|
||||
try:
|
||||
import cudnn
|
||||
except ImportError:
|
||||
cudnn = None
|
||||
|
||||
|
||||
def _torch_float8_dtype(name: str) -> torch.dtype:
|
||||
if name in ("fp8", "fp8_e4m3", "fp8_e4m3fn"):
|
||||
return torch.float8_e4m3fn
|
||||
if name in ("fp8_e5m2", "fp8_e5m2fn"):
|
||||
return torch.float8_e5m2
|
||||
raise ValueError(f"Unsupported fp8 dtype name: {name}")
|
||||
|
||||
|
||||
def _parse_int_list(csv: str) -> list[int]:
|
||||
out: list[int] = []
|
||||
for part in csv.split(","):
|
||||
part = part.strip()
|
||||
if not part:
|
||||
continue
|
||||
out.append(int(part))
|
||||
return out
|
||||
|
||||
|
||||
def attention_pytorch(qkv: torch.Tensor, causal: bool) -> torch.Tensor:
|
||||
"""
|
||||
qkv: (batch, seqlen, 3, nheads, headdim)
|
||||
out: (batch, seqlen, nheads, headdim)
|
||||
"""
|
||||
batch_size, seqlen, _, nheads, d = qkv.shape
|
||||
q, k, v = qkv.unbind(dim=2)
|
||||
q = rearrange(q, "b t h d -> (b h) t d")
|
||||
k = rearrange(k, "b s h d -> (b h) d s")
|
||||
softmax_scale = 1.0 / math.sqrt(d)
|
||||
scores = torch.empty(
|
||||
batch_size * nheads, seqlen, seqlen, dtype=qkv.dtype, device=qkv.device
|
||||
)
|
||||
scores = rearrange(
|
||||
torch.baddbmm(scores, q, k, beta=0, alpha=softmax_scale),
|
||||
"(b h) t s -> b h t s",
|
||||
h=nheads,
|
||||
)
|
||||
if causal:
|
||||
causal_mask = torch.triu(
|
||||
torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1
|
||||
)
|
||||
scores = scores + causal_mask.to(dtype=scores.dtype)
|
||||
attention = torch.softmax(scores, dim=-1)
|
||||
output = torch.einsum("bhts,bshd->bthd", attention, v)
|
||||
return output.to(dtype=qkv.dtype)
|
||||
|
||||
|
||||
def flops(batch: int, seqlen: int, headdim: int, nheads: int, causal: bool) -> int:
|
||||
# Matches the hopper benchmark’s convention.
|
||||
return 4 * batch * seqlen**2 * nheads * headdim // (2 if causal else 1)
|
||||
|
||||
|
||||
def efficiency(flop: int, seconds: float) -> float:
|
||||
return (flop / seconds / 1e12) if not math.isnan(seconds) else 0.0
|
||||
|
||||
|
||||
def time_fwd(fn, *args, repeats: int, **kwargs) -> float:
|
||||
time.sleep(1) # reduce residual throttling effects between benchmarks
|
||||
_, m = benchmark_forward(fn, *args, repeats=repeats, verbose=False, **kwargs)
|
||||
return float(m.mean)
|
||||
|
||||
|
||||
def convert_to_cudnn_type(torch_type):
|
||||
if torch_type == torch.float16:
|
||||
return cudnn.data_type.HALF
|
||||
if torch_type == torch.bfloat16:
|
||||
return cudnn.data_type.BFLOAT16
|
||||
if torch_type == torch.float32:
|
||||
return cudnn.data_type.FLOAT
|
||||
if torch_type == torch.int32:
|
||||
return cudnn.data_type.INT32
|
||||
if torch_type == torch.int64:
|
||||
return cudnn.data_type.INT64
|
||||
if torch_type == torch.float8_e4m3fn:
|
||||
return cudnn.data_type.FP8_E4M3
|
||||
if torch_type == torch.float8_e5m2:
|
||||
return cudnn.data_type.FP8_E5M2
|
||||
raise ValueError("Unsupported tensor data type.")
|
||||
|
||||
|
||||
def cudnn_sdpa_fp8_setup(qkv: torch.Tensor, seqlen_q: int, seqlen_k: int, causal: bool):
|
||||
"""Minimal cudnn.fp8 sdpa runner (optional)."""
|
||||
assert cudnn is not None, "cudnn python bindings not available"
|
||||
b, _, _, nheads, headdim = qkv.shape
|
||||
o_gpu = torch.zeros(
|
||||
b, seqlen_q, nheads, headdim, dtype=qkv.dtype, device=qkv.device
|
||||
)
|
||||
o_gpu_transposed = torch.as_strided(
|
||||
o_gpu,
|
||||
[b, nheads, seqlen_q, headdim],
|
||||
[nheads * seqlen_q * headdim, headdim, nheads * headdim, 1],
|
||||
)
|
||||
amax_s_gpu = torch.empty(1, 1, 1, 1, dtype=torch.float32, device=qkv.device)
|
||||
amax_o_gpu = torch.empty(1, 1, 1, 1, dtype=torch.float32, device=qkv.device)
|
||||
|
||||
graph = cudnn.pygraph(
|
||||
io_data_type=convert_to_cudnn_type(qkv.dtype),
|
||||
intermediate_data_type=cudnn.data_type.FLOAT,
|
||||
compute_data_type=cudnn.data_type.FLOAT,
|
||||
)
|
||||
|
||||
new_q = torch.as_strided(
|
||||
qkv,
|
||||
[b, nheads, seqlen_q, headdim],
|
||||
[seqlen_q * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
|
||||
storage_offset=0,
|
||||
)
|
||||
q = graph.tensor(
|
||||
name="Q",
|
||||
dim=list(new_q.shape),
|
||||
stride=list(new_q.stride()),
|
||||
data_type=convert_to_cudnn_type(qkv.dtype),
|
||||
)
|
||||
|
||||
new_k = torch.as_strided(
|
||||
qkv,
|
||||
[b, nheads, seqlen_k, headdim],
|
||||
[seqlen_k * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
|
||||
storage_offset=nheads * headdim,
|
||||
)
|
||||
k = graph.tensor(
|
||||
name="K",
|
||||
dim=list(new_k.shape),
|
||||
stride=list(new_k.stride()),
|
||||
data_type=convert_to_cudnn_type(qkv.dtype),
|
||||
)
|
||||
|
||||
new_v = torch.as_strided(
|
||||
qkv,
|
||||
[b, nheads, seqlen_k, headdim],
|
||||
[seqlen_k * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
|
||||
storage_offset=nheads * headdim * 2,
|
||||
)
|
||||
v = graph.tensor(
|
||||
name="V",
|
||||
dim=list(new_v.shape),
|
||||
stride=list(new_v.stride()),
|
||||
data_type=convert_to_cudnn_type(qkv.dtype),
|
||||
)
|
||||
|
||||
def _scale_tensor():
|
||||
return graph.tensor(
|
||||
dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT
|
||||
)
|
||||
|
||||
default_scale_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float32, device="cuda")
|
||||
descale_q = _scale_tensor()
|
||||
descale_k = _scale_tensor()
|
||||
descale_v = _scale_tensor()
|
||||
descale_s = _scale_tensor()
|
||||
scale_s = _scale_tensor()
|
||||
scale_o = _scale_tensor()
|
||||
|
||||
o, _, amax_s, amax_o = graph.sdpa_fp8(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
descale_q=descale_q,
|
||||
descale_k=descale_k,
|
||||
descale_v=descale_v,
|
||||
descale_s=descale_s,
|
||||
scale_s=scale_s,
|
||||
scale_o=scale_o,
|
||||
is_inference=True,
|
||||
attn_scale=1.0 / math.sqrt(headdim),
|
||||
use_causal_mask=causal,
|
||||
name="sdpa",
|
||||
)
|
||||
o.set_output(True).set_dim(o_gpu_transposed.shape).set_stride(
|
||||
o_gpu_transposed.stride()
|
||||
)
|
||||
amax_s.set_output(False).set_dim(amax_s_gpu.shape).set_stride(amax_s_gpu.stride())
|
||||
amax_o.set_output(False).set_dim(amax_o_gpu.shape).set_stride(amax_o_gpu.stride())
|
||||
|
||||
graph.validate()
|
||||
graph.build_operation_graph()
|
||||
graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
|
||||
graph.check_support()
|
||||
graph.build_plans()
|
||||
|
||||
variant_pack = {
|
||||
q: new_q,
|
||||
k: new_k,
|
||||
v: new_v,
|
||||
descale_q: default_scale_gpu,
|
||||
descale_k: default_scale_gpu,
|
||||
descale_v: default_scale_gpu,
|
||||
descale_s: default_scale_gpu,
|
||||
scale_s: default_scale_gpu,
|
||||
scale_o: default_scale_gpu,
|
||||
o: o_gpu_transposed,
|
||||
amax_s: amax_s_gpu,
|
||||
amax_o: amax_o_gpu,
|
||||
}
|
||||
workspace = torch.empty(
|
||||
graph.get_workspace_size(), device="cuda", dtype=torch.uint8
|
||||
)
|
||||
|
||||
def run():
|
||||
graph.execute(variant_pack, workspace)
|
||||
return o_gpu
|
||||
|
||||
return run
|
||||
|
||||
|
||||
def _maybe_pass_descales(callable_, **kwargs):
|
||||
sig = inspect.signature(callable_)
|
||||
return {k: v for k, v in kwargs.items() if k in sig.parameters}
|
||||
|
||||
|
||||
def main(argv: Iterable[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--repeats", type=int, default=30)
|
||||
parser.add_argument("--dim", type=int, default=2048)
|
||||
parser.add_argument("--headdims", default="64,128")
|
||||
parser.add_argument("--dtype", default="fp8_e4m3fn")
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument(
|
||||
"--check",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Enable correctness checks vs BF16 PyTorch baseline.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--check-quantization-only",
|
||||
action="store_true",
|
||||
help="Check FP8 kernel vs dequantized-FP8 baseline (quantization error only).",
|
||||
)
|
||||
parser.add_argument("--atol-bf16", type=float, default=0.10)
|
||||
parser.add_argument("--rtol-bf16", type=float, default=0.10)
|
||||
parser.add_argument("--atol-fp8", type=float, default=0.50)
|
||||
parser.add_argument("--rtol-fp8", type=float, default=0.50)
|
||||
parser.add_argument("--run-cudnn", action="store_true")
|
||||
args = parser.parse_args(list(argv) if argv is not None else None)
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA is required")
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
if major != 10:
|
||||
raise RuntimeError(
|
||||
f"This benchmark is for SM100 (compute capability 10.x). Got {major}.{minor}."
|
||||
)
|
||||
|
||||
torch.manual_seed(args.seed)
|
||||
device = "cuda"
|
||||
fp8_dtype = _torch_float8_dtype(args.dtype)
|
||||
headdim_vals = _parse_int_list(args.headdims)
|
||||
bs_seqlen_vals = [
|
||||
(32, 512),
|
||||
(16, 1024),
|
||||
(8, 2048),
|
||||
(4, 4096),
|
||||
(2, 8192),
|
||||
(1, 16384),
|
||||
]
|
||||
|
||||
methods = ["Pytorch", "FA4-CuTe-BF16", "FA4-CuTe-FP8"] + (
|
||||
["cuDNN-FP8"] if args.run_cudnn and cudnn is not None else []
|
||||
)
|
||||
|
||||
fp8_failures = []
|
||||
|
||||
for headdim in headdim_vals:
|
||||
for causal in (False, True):
|
||||
for batch, seqlen in bs_seqlen_vals:
|
||||
torch.cuda.empty_cache()
|
||||
nheads = args.dim // headdim
|
||||
if args.dim % headdim != 0:
|
||||
raise ValueError(
|
||||
f"--dim must be divisible by headdim ({args.dim=} {headdim=})"
|
||||
)
|
||||
|
||||
q_bf16 = torch.randn(
|
||||
batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
|
||||
)
|
||||
k_bf16 = torch.randn(
|
||||
batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
|
||||
)
|
||||
v_bf16 = torch.randn(
|
||||
batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
|
||||
)
|
||||
qkv_bf16 = torch.stack([q_bf16, k_bf16, v_bf16], dim=2)
|
||||
|
||||
times = {}
|
||||
speeds = {}
|
||||
|
||||
out_ref_bf16 = None
|
||||
try:
|
||||
out_ref_bf16 = attention_pytorch(
|
||||
qkv_bf16, causal=causal
|
||||
) # warmup / reference
|
||||
t = time_fwd(
|
||||
attention_pytorch, qkv_bf16, causal=causal, repeats=args.repeats
|
||||
)
|
||||
times["Pytorch"] = t
|
||||
except RuntimeError as e:
|
||||
if "out of memory" in str(e).lower():
|
||||
times["Pytorch"] = float("nan")
|
||||
out_ref_bf16 = None
|
||||
else:
|
||||
raise
|
||||
|
||||
# FA4 / CuTe BF16 baseline
|
||||
try:
|
||||
softmax_scale = headdim**-0.5
|
||||
out_fa4_bf16, _ = flash_attn_cute_fwd(
|
||||
q_bf16,
|
||||
k_bf16,
|
||||
v_bf16,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=causal,
|
||||
) # warmup / compile
|
||||
t = time_fwd(
|
||||
flash_attn_cute_fwd,
|
||||
q_bf16,
|
||||
k_bf16,
|
||||
v_bf16,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=causal,
|
||||
repeats=args.repeats,
|
||||
)
|
||||
times["FA4-CuTe-BF16"] = t
|
||||
if args.check and out_ref_bf16 is not None:
|
||||
torch.testing.assert_close(
|
||||
out_fa4_bf16,
|
||||
out_ref_bf16,
|
||||
atol=args.atol_bf16,
|
||||
rtol=args.rtol_bf16,
|
||||
)
|
||||
except Exception as e:
|
||||
# Treat as fatal: BF16 kernel should be usable for basic sanity checking.
|
||||
raise RuntimeError("FA4-CuTe BF16 baseline failed") from e
|
||||
|
||||
# FA4 / CuTe FP8
|
||||
q_fp8 = q_bf16.to(fp8_dtype)
|
||||
k_fp8 = k_bf16.to(fp8_dtype)
|
||||
v_fp8 = v_bf16.to(fp8_dtype)
|
||||
|
||||
# Placeholder descales (FA3-style: per-(batch, kv_head)).
|
||||
q_descale = torch.ones(
|
||||
batch, nheads, device=device, dtype=torch.float32
|
||||
)
|
||||
k_descale = torch.ones(
|
||||
batch, nheads, device=device, dtype=torch.float32
|
||||
)
|
||||
v_descale = torch.ones(
|
||||
batch, nheads, device=device, dtype=torch.float32
|
||||
)
|
||||
|
||||
# Optional: FP8 reference baseline (dequantized FP8 -> PyTorch) for quantization-error-only checks
|
||||
out_ref_fp8 = None
|
||||
if args.check and args.check_quantization_only:
|
||||
try:
|
||||
# Dequantize FP8 inputs back to BF16 (applying descales)
|
||||
q_ref_fp8 = (
|
||||
q_fp8.to(torch.bfloat16) * q_descale[:, None, :, None]
|
||||
).to(torch.bfloat16)
|
||||
k_ref_fp8 = (
|
||||
k_fp8.to(torch.bfloat16) * k_descale[:, None, :, None]
|
||||
).to(torch.bfloat16)
|
||||
v_ref_fp8 = (
|
||||
v_fp8.to(torch.bfloat16) * v_descale[:, None, :, None]
|
||||
).to(torch.bfloat16)
|
||||
qkv_ref_fp8 = torch.stack(
|
||||
[q_ref_fp8, k_ref_fp8, v_ref_fp8], dim=2
|
||||
)
|
||||
out_ref_fp8 = attention_pytorch(qkv_ref_fp8, causal=causal)
|
||||
except RuntimeError as e:
|
||||
if "out of memory" in str(e).lower():
|
||||
out_ref_fp8 = None
|
||||
else:
|
||||
raise
|
||||
|
||||
fa4_kwargs = dict(softmax_scale=softmax_scale, causal=causal)
|
||||
fa4_kwargs.update(
|
||||
_maybe_pass_descales(
|
||||
flash_attn_cute_fwd,
|
||||
q_descale=q_descale,
|
||||
k_descale=k_descale,
|
||||
v_descale=v_descale,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
# Warmup/compile (will raise until FP8 is implemented)
|
||||
out_fa4_fp8, _ = flash_attn_cute_fwd(
|
||||
q_fp8, k_fp8, v_fp8, **fa4_kwargs
|
||||
)
|
||||
t = time_fwd(
|
||||
flash_attn_cute_fwd,
|
||||
q_fp8,
|
||||
k_fp8,
|
||||
v_fp8,
|
||||
repeats=args.repeats,
|
||||
**fa4_kwargs,
|
||||
)
|
||||
times["FA4-CuTe-FP8"] = t
|
||||
if args.check:
|
||||
# Choose baseline: quantization-only (dequantized FP8) or full (BF16)
|
||||
if args.check_quantization_only:
|
||||
ref_baseline = out_ref_fp8
|
||||
else:
|
||||
ref_baseline = out_ref_bf16
|
||||
|
||||
if ref_baseline is not None:
|
||||
torch.testing.assert_close(
|
||||
out_fa4_fp8,
|
||||
ref_baseline,
|
||||
atol=args.atol_fp8,
|
||||
rtol=args.rtol_fp8,
|
||||
)
|
||||
except Exception as e:
|
||||
fp8_failures.append((causal, headdim, batch, seqlen, repr(e)))
|
||||
times["FA4-CuTe-FP8"] = float("nan")
|
||||
|
||||
if args.run_cudnn and cudnn is not None:
|
||||
qkv_fp8 = qkv_bf16.to(fp8_dtype)
|
||||
runner = cudnn_sdpa_fp8_setup(
|
||||
qkv_fp8, seqlen, seqlen, causal=causal
|
||||
)
|
||||
_ = runner() # warmup
|
||||
t = time_fwd(lambda: runner(), repeats=args.repeats)
|
||||
times["cuDNN-FP8"] = t
|
||||
|
||||
print(
|
||||
f"### causal={causal}, headdim={headdim}, batch={batch}, seqlen={seqlen} ###"
|
||||
)
|
||||
for method in methods:
|
||||
t = times.get(method, float("nan"))
|
||||
speeds[method] = efficiency(
|
||||
flops(batch, seqlen, headdim, nheads, causal), t
|
||||
)
|
||||
if math.isnan(t):
|
||||
print(f"{method} fwd: (skipped)")
|
||||
else:
|
||||
print(
|
||||
f"{method} fwd: {speeds[method]:.2f} TFLOPs/s, {t * 1e3:.3f} ms"
|
||||
)
|
||||
if math.isnan(times.get("FA4-CuTe-FP8", float("nan"))):
|
||||
print("FA4-CuTe-FP8 status: FAILED")
|
||||
|
||||
if fp8_failures:
|
||||
print(f"\nFP8 failures: {len(fp8_failures)} (showing first 5)")
|
||||
for causal, headdim, batch, seqlen, err in fp8_failures[:5]:
|
||||
print(
|
||||
f"- causal={causal} headdim={headdim} batch={batch} seqlen={seqlen}: {err}"
|
||||
)
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user