| |
| """Native/wrapper/eager/compile benchmark for padded FP8 producers.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| PACKAGE = Path(__file__).resolve().parents[1] |
| ROOT = PACKAGE.parent |
| sys.path.insert(0, str(PACKAGE / "tests")) |
| from test_padded_fp8_producers import load_source_ops |
|
|
| SHAPES = [ |
| ("decode", 1, 1, 1280, 16), |
| ("groot-dit", 1, 40, 1536, 64), |
| ("vla", 1, 51, 2048, 64), |
| ("vision", 1, 105, 1280, 128), |
| ("prefill", 2, 277, 2048, 320), |
| ] |
|
|
|
|
| def load_ops(backend: str, artifact: str | None): |
| if backend == "source": |
| return load_source_ops() |
| if artifact: |
| sys.path.insert(0, artifact) |
| try: |
| return importlib.import_module("padded_fp8_producers") |
| finally: |
| if artifact: |
| sys.path.remove(artifact) |
|
|
|
|
| def load_native(): |
| from torch.utils.cpp_extension import load |
|
|
| major, minor = torch.cuda.get_device_capability() |
| os.environ.setdefault( |
| "TORCH_CUDA_ARCH_LIST", "12.0a" if major >= 12 else f"{major}.{minor}" |
| ) |
| return load( |
| name="padded_fp8_producers_native_bench", |
| sources=[ |
| str(PACKAGE / "benchmarks" / "native_binding.cpp"), |
| str(PACKAGE / "csrc" / "padded_fp8_producers.cu"), |
| ], |
| extra_include_paths=[str(PACKAGE / "csrc")], |
| extra_cflags=["-O3"], |
| extra_cuda_cflags=[ |
| "-O3", |
| "--expt-relaxed-constexpr", |
| "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", |
| "-U__CUDA_NO_BFLOAT16_OPERATORS__", |
| "-U__CUDA_NO_BFLOAT162_OPERATORS__", |
| ], |
| verbose=False, |
| ) |
|
|
|
|
| def bench(fn, warmup=100, iterations=500) -> float: |
| 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(iterations): |
| fn() |
| end.record() |
| end.synchronize() |
| return start.elapsed_time(end) * 1000.0 / iterations |
|
|
|
|
| def eager_adaptive(x, weight, gamma, beta, scale, output): |
| batch, rows, _ = x.shape |
| norm = x.float() * torch.rsqrt( |
| x.float().square().mean(dim=-1, keepdim=True) + 1e-6 |
| ) |
| value = ( |
| (1.0 + gamma[:, None, :].float()) * norm * weight.float() |
| + beta[:, None, :].float() |
| ).bfloat16() |
| output[:, :rows].copy_( |
| torch.clamp(value.float() / scale, -448.0, 448.0).to( |
| torch.float8_e4m3fn |
| ) |
| ) |
| output[:, rows:].zero_() |
| return output |
|
|
|
|
| def eager_swiglu(gate, up, scale, output): |
| rows = gate.shape[0] |
| value = (F.silu(gate.float()) * up.float()).bfloat16() |
| output[:rows].copy_( |
| torch.clamp(value.float() / scale, -448.0, 448.0).to( |
| torch.float8_e4m3fn |
| ) |
| ) |
| output[rows:].zero_() |
| return output |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--backend", choices=("source", "installed"), default="source") |
| parser.add_argument("--artifact") |
| args = parser.parse_args() |
| ops = load_ops(args.backend, args.artifact) |
| native = load_native() |
| print("op,shape,native_us,wrapper_us,eager_us,compile_us,wrapper/native") |
|
|
| for label, batch, rows, dim, padded in SHAPES: |
| x = (torch.randn((batch, rows, dim), device="cuda") * 0.4).bfloat16() |
| weight = torch.ones((dim,), device="cuda", dtype=torch.bfloat16) |
| gamma = torch.zeros((batch, dim), device="cuda", dtype=torch.bfloat16) |
| beta = torch.zeros_like(gamma) |
| scale = torch.tensor([0.01], device="cuda", dtype=torch.float32) |
| native_out = torch.empty( |
| (batch, padded, dim), device="cuda", dtype=torch.float8_e4m3fn |
| ) |
| wrapper_out = torch.empty_like(native_out) |
| eager_out = torch.empty_like(native_out) |
| compile_out = torch.empty_like(native_out) |
|
|
| native_fn = lambda: native.adaptive( |
| x.data_ptr(), weight.data_ptr(), gamma.data_ptr(), beta.data_ptr(), |
| scale.data_ptr(), native_out.data_ptr(), batch, rows, padded, dim, |
| 1e-6 |
| ) |
| wrapper_fn = lambda: ops.adaptive_rms_norm_quant_fp8_padded_bf16( |
| x, weight, gamma, beta, scale, output=wrapper_out |
| ) |
| eager_fn = lambda: eager_adaptive( |
| x, weight, gamma, beta, scale, eager_out |
| ) |
| compiled_call = torch.compile(eager_adaptive, fullgraph=True) |
| compile_fn = lambda: compiled_call( |
| x, weight, gamma, beta, scale, compile_out |
| ) |
| native_fn() |
| wrapper_fn() |
| torch.testing.assert_close( |
| native_out.float(), wrapper_out.float(), rtol=0, atol=0 |
| ) |
| times = [bench(fn) for fn in (native_fn, wrapper_fn, eager_fn, compile_fn)] |
| print( |
| f"adaptive,{label}:B{batch}S{rows}D{dim}P{padded}," |
| f"{times[0]:.3f},{times[1]:.3f},{times[2]:.3f},{times[3]:.3f}," |
| f"{times[1] / times[0]:.3f}" |
| ) |
|
|
| gate = x.reshape(batch * rows, dim) |
| up = torch.randn_like(gate) |
| native_swiglu = torch.empty( |
| (batch * padded, dim), device="cuda", dtype=torch.float8_e4m3fn |
| ) |
| wrapper_swiglu = torch.empty_like(native_swiglu) |
| eager_swiglu_out = torch.empty_like(native_swiglu) |
| compile_swiglu_out = torch.empty_like(native_swiglu) |
| native_fn = lambda: native.swiglu( |
| gate.data_ptr(), up.data_ptr(), scale.data_ptr(), |
| native_swiglu.data_ptr(), batch * rows, batch * padded, dim |
| ) |
| wrapper_fn = lambda: ops.swiglu_quant_fp8_padded_bf16( |
| gate, up, scale, output=wrapper_swiglu |
| ) |
| eager_fn = lambda: eager_swiglu( |
| gate, up, scale, eager_swiglu_out |
| ) |
| compiled_call = torch.compile(eager_swiglu, fullgraph=True) |
| compile_fn = lambda: compiled_call( |
| gate, up, scale, compile_swiglu_out |
| ) |
| native_fn() |
| wrapper_fn() |
| torch.testing.assert_close( |
| native_swiglu.float(), wrapper_swiglu.float(), rtol=0, atol=0 |
| ) |
| times = [bench(fn) for fn in (native_fn, wrapper_fn, eager_fn, compile_fn)] |
| print( |
| f"swiglu,{label}:S{batch * rows}D{dim}P{batch * padded}," |
| f"{times[0]:.3f},{times[1]:.3f},{times[2]:.3f},{times[3]:.3f}," |
| f"{times[1] / times[0]:.3f}" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
