| import torch |
|
|
| from kernels.benchmark import Benchmark |
|
|
|
|
| def _quantize_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: |
| return torch.clamp(x.float() / scale.float(), -448.0, 448.0).to(torch.float8_e4m3fn) |
|
|
|
|
| def _dequant_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: |
| return x.float() * scale.float() |
|
|
|
|
| def _compiler_disable(fn): |
| compiler = getattr(torch, "compiler", None) |
| if compiler is not None and hasattr(compiler, "disable"): |
| return compiler.disable(fn) |
| return torch._dynamo.disable(fn) |
|
|
|
|
| def _gelu_quantize_fp8_boundary( |
| hidden: torch.Tensor, bias: torch.Tensor, scale: torch.Tensor |
| ) -> torch.Tensor: |
| hidden = torch.nn.functional.gelu( |
| hidden.float() + bias.float(), approximate="tanh" |
| ) |
| return _quantize_fp8(hidden, scale) |
|
|
|
|
| def _bf16_bias_add_boundary(out: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: |
| return (out.float() + bias.float()).to(torch.bfloat16) |
|
|
|
|
| _stable_gelu_quantize_fp8 = _compiler_disable(_gelu_quantize_fp8_boundary) |
| _stable_bf16_bias_add = _compiler_disable(_bf16_bias_add_boundary) |
|
|
|
|
| class FP8GeluMlpBenchmark(Benchmark): |
| seed = 17 |
|
|
| def _setup_shape(self, M: int, K: int, H: int, N: int) -> None: |
| self.M, self.K, self.H, self.N = M, K, H, N |
| self.x_scale = torch.tensor([0.05], device=self.device, dtype=torch.float32) |
| self.up_scale = torch.tensor([0.04], device=self.device, dtype=torch.float32) |
| self.hidden_scale = torch.tensor([0.25], device=self.device, dtype=torch.float32) |
| self.down_scale = torch.tensor([0.04], device=self.device, dtype=torch.float32) |
| self.x = _quantize_fp8( |
| torch.randn((M, K), device=self.device, dtype=torch.bfloat16), |
| self.x_scale, |
| ) |
| self.up_w = _quantize_fp8( |
| torch.randn((H, K), device=self.device, dtype=torch.bfloat16), |
| self.up_scale, |
| ) |
| self.down_w = _quantize_fp8( |
| torch.randn((N, H), device=self.device, dtype=torch.bfloat16), |
| self.down_scale, |
| ) |
| self.up_b = torch.randn((H,), device=self.device, dtype=torch.bfloat16) |
| self.down_b = torch.randn((N,), device=self.device, dtype=torch.bfloat16) |
| self.hidden = torch.empty((M, H), device=self.device, dtype=torch.bfloat16) |
| self.hidden_fp8 = torch.empty((M, H), device=self.device, dtype=torch.float8_e4m3fn) |
| self.out = torch.empty((M, N), device=self.device, dtype=torch.bfloat16) |
|
|
| def _reference(self) -> torch.Tensor: |
| hidden = ( |
| _dequant_fp8(self.x, self.x_scale) |
| @ _dequant_fp8(self.up_w, self.up_scale).T |
| ).to(torch.bfloat16) |
| hidden_fp8 = _stable_gelu_quantize_fp8( |
| hidden, self.up_b, self.hidden_scale |
| ) |
| out = ( |
| _dequant_fp8(hidden_fp8, self.hidden_scale) |
| @ _dequant_fp8(self.down_w, self.down_scale).T |
| ).to(torch.bfloat16) |
| return _stable_bf16_bias_add(out, self.down_b) |
|
|
| def setup_smoke_mlp(self) -> None: |
| self._setup_shape(16, 128, 256, 128) |
|
|
| def benchmark_smoke_mlp(self) -> None: |
| self.kernel.fp8_gelu_mlp_bf16( |
| self.x, |
| self.up_w, |
| self.up_b, |
| self.down_w, |
| self.down_b, |
| self.x_scale, |
| self.up_scale, |
| self.hidden_scale, |
| self.down_scale, |
| hidden_bf16=self.hidden, |
| hidden_fp8=self.hidden_fp8, |
| out=self.out, |
| ) |
|
|
| def verify_smoke_mlp(self) -> torch.Tensor: |
| return self._reference() |
|
|
