forked from ccf-ai-infra/TileOPs-Metax
247 lines
7.9 KiB
Python
247 lines
7.9 KiB
Python
"""Skipped benchmarks for unsupported fp8 elementwise ops (e4m3fn, e5m2).
|
|
|
|
Keeps unsupported fp8 benchmark cases visible without turning the nightly
|
|
benchmark suite red. Current elementwise dtype contracts reject these fp8
|
|
inputs; remove the skip marks when the corresponding ops add fp8 support.
|
|
"""
|
|
|
|
from math import prod
|
|
from typing import Optional
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
|
|
from tileops.ops.elementwise import AddFwdOp, ExpFwdOp, ReluFwdOp, SiluAndMulFwdOp
|
|
from workloads.workload_base import FixtureBase
|
|
|
|
# Shapes modeled on real LLM workloads: (batch, seq_len, hidden_dim).
|
|
# Small: (1, 2048, 4096) - single-batch inference, LLaMA-7B hidden.
|
|
# Medium: (8, 2048, 4096) - multi-batch inference.
|
|
# Large: (4, 4096, 8192) - training, LLaMA-70B hidden.
|
|
# A non-pow2 hidden (LLaMA-7B intermediate=11008) is added in the
|
|
# unary/binary sweep to exercise tail handling.
|
|
_SHAPES = (
|
|
(1, 2048, 4096),
|
|
(8, 2048, 4096),
|
|
(4, 4096, 8192),
|
|
(1, 2048, 11008),
|
|
)
|
|
_FP8_DTYPES = [torch.float8_e4m3fn, torch.float8_e5m2]
|
|
_UNSUPPORTED_FP8_SKIP = pytest.mark.skip(
|
|
reason=(
|
|
"TileOPs elementwise ops currently reject fp8 dtypes; "
|
|
"benchmark is kept as an explicit unsupported case"
|
|
)
|
|
)
|
|
|
|
|
|
def _shape_id(shape: tuple[int, ...]) -> str:
|
|
return "x".join(str(s) for s in shape)
|
|
|
|
|
|
# Helpers
|
|
|
|
|
|
class Fp8UnaryBenchCase:
|
|
def __init__(self, shape: tuple[int, ...], dtype: torch.dtype):
|
|
self.shape = shape
|
|
self.n_total = prod(shape)
|
|
self.dtype = dtype
|
|
|
|
def gen_inputs(self) -> tuple[torch.Tensor]:
|
|
x = (torch.randn(*self.shape, dtype=torch.float16, device="cuda") * 2.0)
|
|
return (x.to(self.dtype),)
|
|
|
|
|
|
class Fp8UnaryBenchmark(BenchmarkBase[Fp8UnaryBenchCase]):
|
|
def calculate_flops(self) -> Optional[float]:
|
|
return self.workload.n_total
|
|
|
|
def calculate_memory(self) -> Optional[float]:
|
|
# fp8 in (1B) + fp8 out (1B) per element
|
|
return self.workload.n_total * 2
|
|
|
|
|
|
class Fp8BinaryBenchCase:
|
|
def __init__(self, shape: tuple[int, ...], dtype: torch.dtype):
|
|
self.shape = shape
|
|
self.n_total = prod(shape)
|
|
self.dtype = dtype
|
|
|
|
def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor]:
|
|
a = (torch.randn(*self.shape, dtype=torch.float16, device="cuda") * 0.5).to(self.dtype)
|
|
b = (torch.randn(*self.shape, dtype=torch.float16, device="cuda") * 0.5).to(self.dtype)
|
|
return a, b
|
|
|
|
|
|
class Fp8BinaryBenchmark(BenchmarkBase[Fp8BinaryBenchCase]):
|
|
def calculate_flops(self) -> Optional[float]:
|
|
return self.workload.n_total
|
|
|
|
def calculate_memory(self) -> Optional[float]:
|
|
# fp8 in a (1B) + fp8 in b (1B) + fp8 out (1B)
|
|
return self.workload.n_total * 3
|
|
|
|
|
|
class Fp8FusedGatedBenchCase:
|
|
def __init__(self, shape: tuple[int, int], dtype: torch.dtype):
|
|
# ``shape`` is the *output* shape (M, N). The input has 2*N
|
|
# along the trailing axis for the gate/value split.
|
|
self.shape = shape
|
|
self.M, self.N = shape
|
|
self.dtype = dtype
|
|
|
|
def gen_inputs(self) -> tuple[torch.Tensor]:
|
|
x = (torch.randn(self.M, 2 * self.N, dtype=torch.float16, device="cuda") * 0.5)
|
|
return (x.to(self.dtype),)
|
|
|
|
|
|
class Fp8FusedGatedBenchmark(BenchmarkBase[Fp8FusedGatedBenchCase]):
|
|
def calculate_flops(self) -> Optional[float]:
|
|
# FIXME(staged-rollout): hardcoded silu FLOPs in Fp8FusedGatedBenchmark
|
|
#
|
|
# Broken invariant: calculate_flops assumes silu (5 FLOPs/elem), wrong for other activations
|
|
# Why: only silu is benchmarked currently, other activations not yet added
|
|
# Cleanup: implement per-activation FLOPs lookup when benchmarking gelu/other activations
|
|
return self.workload.M * self.workload.N * 5
|
|
|
|
def calculate_memory(self) -> Optional[float]:
|
|
# Read x (M*2N*1B) + write y (M*N*1B)
|
|
return (self.workload.M * 2 * self.workload.N + self.workload.M * self.workload.N)
|
|
|
|
|
|
# Unary fp8 benchmarks: relu, exp
|
|
|
|
_unary_params = []
|
|
for _op_name, _op_cls, _bl_fn in [
|
|
("relu_fp8", ReluFwdOp, torch.relu),
|
|
("exp_fp8", ExpFwdOp, torch.exp),
|
|
]:
|
|
for _shape in _SHAPES:
|
|
for _dt in _FP8_DTYPES:
|
|
_unary_params.append(pytest.param(
|
|
_op_name, _shape, _dt, _op_cls, _bl_fn,
|
|
marks=_UNSUPPORTED_FP8_SKIP,
|
|
id=f"{_op_name}-{_shape_id(_shape)}-{_dt}",
|
|
))
|
|
|
|
|
|
class Fp8UnaryBenchFixture(FixtureBase):
|
|
PARAMS = [("op_name, shape, dtype, op_cls, baseline_fn", _unary_params)]
|
|
|
|
|
|
@Fp8UnaryBenchFixture
|
|
def test_fp8_unary_bench(op_name, shape, dtype, op_cls, baseline_fn):
|
|
test = Fp8UnaryBenchCase(shape=shape, dtype=dtype)
|
|
bm = Fp8UnaryBenchmark(test)
|
|
inputs = test.gen_inputs()
|
|
|
|
n_total = prod(shape)
|
|
op = op_cls(N_total=n_total, dtype=dtype)
|
|
result = bm.profile(op, *inputs)
|
|
BenchmarkReport.record(
|
|
op_name, {"shape": shape, "dtype": dtype}, result, tag="tileops",
|
|
)
|
|
|
|
# Baseline: PyTorch fp16 compute then cast back to fp8
|
|
def baseline(*args):
|
|
return baseline_fn(args[0].to(torch.float16)).to(dtype)
|
|
|
|
result_bl = bm.profile(baseline, *inputs)
|
|
BenchmarkReport.record(
|
|
op_name, {"shape": shape, "dtype": dtype}, result_bl, tag="torch",
|
|
)
|
|
|
|
|
|
# Binary fp8 benchmark: add
|
|
|
|
_binary_params = []
|
|
for _shape in _SHAPES:
|
|
for _dt in _FP8_DTYPES:
|
|
_binary_params.append(pytest.param(
|
|
"add_fp8", _shape, _dt,
|
|
marks=_UNSUPPORTED_FP8_SKIP,
|
|
id=f"add_fp8-{_shape_id(_shape)}-{_dt}",
|
|
))
|
|
|
|
|
|
class Fp8BinaryBenchFixture(FixtureBase):
|
|
PARAMS = [("op_name, shape, dtype", _binary_params)]
|
|
|
|
|
|
@Fp8BinaryBenchFixture
|
|
def test_fp8_binary_bench(op_name, shape, dtype):
|
|
test = Fp8BinaryBenchCase(shape=shape, dtype=dtype)
|
|
bm = Fp8BinaryBenchmark(test)
|
|
inputs = test.gen_inputs()
|
|
|
|
op = AddFwdOp(a_shape=shape, b_shape=shape, dtype=dtype)
|
|
result = bm.profile(op, *inputs)
|
|
BenchmarkReport.record(
|
|
op_name, {"shape": shape, "dtype": dtype}, result, tag="tileops",
|
|
)
|
|
|
|
def baseline(a, b):
|
|
return (a.to(torch.float16) + b.to(torch.float16)).to(dtype)
|
|
|
|
result_bl = bm.profile(baseline, *inputs)
|
|
BenchmarkReport.record(
|
|
op_name, {"shape": shape, "dtype": dtype}, result_bl, tag="torch",
|
|
)
|
|
|
|
|
|
# Fused gated fp8 benchmark: silu_and_mul
|
|
|
|
# Fused gated output shapes: (batch * seq_len, intermediate_dim).
|
|
# LLaMA-7B: hidden=4096, intermediate=11008 (non-pow2)
|
|
# LLaMA-13B: hidden=5120, intermediate=13824 (non-pow2)
|
|
# LLaMA-70B: hidden=8192, intermediate=28672
|
|
_GATED_SHAPES = [
|
|
(1 * 2048, 11008), # LLaMA-7B single-batch inference
|
|
(8 * 2048, 11008), # LLaMA-7B multi-batch inference
|
|
(4 * 4096, 28672), # LLaMA-70B training
|
|
]
|
|
_gated_params = []
|
|
for _shape in _GATED_SHAPES:
|
|
for _dt in _FP8_DTYPES:
|
|
_gated_params.append(pytest.param(
|
|
"silu_and_mul_fp8", _shape, _dt,
|
|
marks=_UNSUPPORTED_FP8_SKIP,
|
|
id=f"silu_and_mul_fp8-{_shape_id(_shape)}-{_dt}",
|
|
))
|
|
|
|
|
|
class Fp8FusedGatedBenchFixture(FixtureBase):
|
|
PARAMS = [("op_name, shape, dtype", _gated_params)]
|
|
|
|
|
|
@Fp8FusedGatedBenchFixture
|
|
def test_fp8_fused_gated_bench(op_name, shape, dtype):
|
|
test = Fp8FusedGatedBenchCase(shape=shape, dtype=dtype)
|
|
bm = Fp8FusedGatedBenchmark(test)
|
|
inputs = test.gen_inputs()
|
|
|
|
M, N = shape
|
|
op = SiluAndMulFwdOp(M=M, N=N, dtype=dtype)
|
|
result = bm.profile(op, *inputs)
|
|
BenchmarkReport.record(
|
|
op_name, {"shape": shape, "dtype": dtype}, result, tag="tileops",
|
|
)
|
|
|
|
def baseline(x):
|
|
x_fp16 = x.to(torch.float16)
|
|
gate = x_fp16[:, :N]
|
|
value = x_fp16[:, N:]
|
|
return (F.silu(gate) * value).to(dtype)
|
|
|
|
result_bl = bm.profile(baseline, *inputs)
|
|
BenchmarkReport.record(
|
|
op_name, {"shape": shape, "dtype": dtype}, result_bl, tag="torch-ref",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-vvs"])
|