| import torch |
|
|
| from kernels.benchmark import Benchmark |
|
|
|
|
| _original_allclose = torch.allclose |
|
|
|
|
| def _fp8_dtype() -> torch.dtype: |
| if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"): |
| return torch.float8_e4m3fnuz |
| return torch.float8_e4m3fn |
|
|
|
|
| def _fp8_max() -> float: |
| return 240.0 if torch.version.hip is not None else 448.0 |
|
|
|
|
| def _is_fp8_dtype(dtype: torch.dtype) -> bool: |
| return dtype == torch.float8_e4m3fn or ( |
| hasattr(torch, "float8_e4m3fnuz") and dtype == torch.float8_e4m3fnuz |
| ) |
|
|
|
|
| def _flashrt_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False): |
| if _is_fp8_dtype(input.dtype) or _is_fp8_dtype(other.dtype): |
| return _original_allclose( |
| input.float(), |
| other.float(), |
| rtol=0, |
| atol=0, |
| equal_nan=equal_nan, |
| ) |
| return _original_allclose(input, other, rtol=rtol, atol=atol, equal_nan=equal_nan) |
|
|
|
|
| torch.allclose = _flashrt_allclose |
|
|
|
|
| QUANT_SHAPES = [ |
| ("decode_m1", 1, 4096), |
| ("decode_m2", 2, 4096), |
| ("decode_m4", 4, 4096), |
| ("decode_m8", 8, 4096), |
| ("small_m16", 16, 4096), |
| ("small_m32", 32, 4096), |
| ("prefill_m64", 64, 4096), |
| ("prefill_m128", 128, 4096), |
| ("prefill_m256", 256, 4096), |
| ("wide_n8192_m16", 16, 8192), |
| ("wide_n8192_m128", 128, 8192), |
| ("vla_n12288_m16", 16, 12288), |
| ("vla_n12288_m64", 64, 12288), |
| ("vla_n16384_m16", 16, 16384), |
| ("vla_n16384_m64", 64, 16384), |
| ] |
|
|
|
|
| class BiasGeluFp8QuantizeBenchmark(Benchmark): |
| seed = 1 |
|
|
| def _setup_shape(self, m: int, n: int) -> None: |
| self.input = torch.randn((m, n), device=self.device, dtype=torch.bfloat16) |
| self.bias = torch.randn((n,), device=self.device, dtype=torch.bfloat16) |
| self.scale = torch.tensor([0.25], device=self.device, dtype=torch.float32) |
| self.out = torch.empty_like(self.input, dtype=_fp8_dtype()) |
|
|
| def _reference(self) -> torch.Tensor: |
| y = self.input.float() + self.bias.float() |
| y = torch.nn.functional.gelu(y, approximate="tanh") |
| limit = _fp8_max() |
| y = torch.clamp(y / self.scale.float(), -limit, limit) |
| return y.to(_fp8_dtype()) |
|
|
| def setup_decode_m1(self) -> None: |
| self._setup_shape(1, 4096) |
|
|
| def benchmark_decode_m1(self) -> None: |
| self.kernel.bias_gelu_quantize_fp8_static_bf16( |
| self.input, self.bias, self.scale, out=self.out |
| ) |
|
|
| def verify_decode_m1(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_decode_m8(self) -> None: |
| self._setup_shape(8, 4096) |
|
|
| def benchmark_decode_m8(self) -> None: |
| self.kernel.bias_gelu_quantize_fp8_static_bf16( |
| self.input, self.bias, self.scale, out=self.out |
| ) |
|
|
| def verify_decode_m8(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_small_m16(self) -> None: |
| self._setup_shape(16, 4096) |
|
|
| def benchmark_small_m16(self) -> None: |
| self.kernel.bias_gelu_quantize_fp8_static_bf16( |
| self.input, self.bias, self.scale, out=self.out |
| ) |
|
|
| def verify_small_m16(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_prefill_m64(self) -> None: |
| self._setup_shape(64, 4096) |
|
|
| def benchmark_prefill_m64(self) -> None: |
| self.kernel.bias_gelu_quantize_fp8_static_bf16( |
| self.input, self.bias, self.scale, out=self.out |
| ) |
|
|
| def verify_prefill_m64(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_prefill_m128(self) -> None: |
| self._setup_shape(128, 4096) |
|
|
| def benchmark_prefill_m128(self) -> None: |
| self.kernel.bias_gelu_quantize_fp8_static_bf16( |
| self.input, self.bias, self.scale, out=self.out |
| ) |
|
|
| def verify_prefill_m128(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_wide_n8192_m16(self) -> None: |
| self._setup_shape(16, 8192) |
|
|
| def benchmark_wide_n8192_m16(self) -> None: |
| self.kernel.bias_gelu_quantize_fp8_static_bf16( |
| self.input, self.bias, self.scale, out=self.out |
| ) |
|
|
| def verify_wide_n8192_m16(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_wide_n8192_m128(self) -> None: |
| self._setup_shape(128, 8192) |
|
|
| def benchmark_wide_n8192_m128(self) -> None: |
| self.kernel.bias_gelu_quantize_fp8_static_bf16( |
| self.input, self.bias, self.scale, out=self.out |
| ) |
|
|
| def verify_wide_n8192_m128(self) -> torch.Tensor: |
| return self._reference() |
|
|
|
|
| def _register_bias_gelu_shapes() -> None: |
| for label, m, n in QUANT_SHAPES: |
|
|
| def setup(self, m=m, n=n) -> None: |
| self._setup_shape(m, n) |
|
|
| def benchmark(self) -> None: |
| self.kernel.bias_gelu_quantize_fp8_static_bf16( |
| self.input, self.bias, self.scale, out=self.out |
| ) |
|
|
| def verify(self) -> torch.Tensor: |
| return self._reference() |
|
|
| setattr(BiasGeluFp8QuantizeBenchmark, f"setup_{label}", setup) |
| setattr(BiasGeluFp8QuantizeBenchmark, f"benchmark_{label}", benchmark) |
| setattr(BiasGeluFp8QuantizeBenchmark, f"verify_{label}", verify) |
|
|
|
|
| class GeluFp8QuantizeBenchmark(Benchmark): |
| seed = 3 |
|
|
| def _setup_shape(self, m: int, n: int) -> None: |
| self.input = torch.randn((m, n), device=self.device, dtype=torch.bfloat16) |
| self.scale = torch.tensor([0.25], device=self.device, dtype=torch.float32) |
| self.out = torch.empty_like(self.input, dtype=_fp8_dtype()) |
|
|
| def _reference(self) -> torch.Tensor: |
| y = torch.nn.functional.gelu(self.input.float(), approximate="tanh") |
| limit = _fp8_max() |
| y = torch.clamp(y / self.scale.float(), -limit, limit) |
| return y.to(_fp8_dtype()) |
|
|
| def setup_decode_m1(self) -> None: |
| self._setup_shape(1, 4096) |
|
|
| def benchmark_decode_m1(self) -> None: |
| self.kernel.gelu_quantize_fp8_static_bf16( |
| self.input, self.scale, out=self.out |
| ) |
|
|
| def verify_decode_m1(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_decode_m8(self) -> None: |
| self._setup_shape(8, 4096) |
|
|
| def benchmark_decode_m8(self) -> None: |
| self.kernel.gelu_quantize_fp8_static_bf16( |
| self.input, self.scale, out=self.out |
| ) |
|
|
| def verify_decode_m8(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_small_m16(self) -> None: |
| self._setup_shape(16, 4096) |
|
|
| def benchmark_small_m16(self) -> None: |
| self.kernel.gelu_quantize_fp8_static_bf16( |
| self.input, self.scale, out=self.out |
| ) |
|
|
| def verify_small_m16(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_prefill_m64(self) -> None: |
| self._setup_shape(64, 4096) |
|
|
| def benchmark_prefill_m64(self) -> None: |
| self.kernel.gelu_quantize_fp8_static_bf16( |
| self.input, self.scale, out=self.out |
| ) |
|
|
| def verify_prefill_m64(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_prefill_m128(self) -> None: |
| self._setup_shape(128, 4096) |
|
|
| def benchmark_prefill_m128(self) -> None: |
| self.kernel.gelu_quantize_fp8_static_bf16( |
| self.input, self.scale, out=self.out |
| ) |
|
|
| def verify_prefill_m128(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_wide_n8192_m16(self) -> None: |
| self._setup_shape(16, 8192) |
|
|
| def benchmark_wide_n8192_m16(self) -> None: |
| self.kernel.gelu_quantize_fp8_static_bf16( |
| self.input, self.scale, out=self.out |
| ) |
|
|
| def verify_wide_n8192_m16(self) -> torch.Tensor: |
| return self._reference() |
|
|
| def setup_wide_n8192_m128(self) -> None: |
| self._setup_shape(128, 8192) |
|
|
| def benchmark_wide_n8192_m128(self) -> None: |
| self.kernel.gelu_quantize_fp8_static_bf16( |
| self.input, self.scale, out=self.out |
| ) |
|
|
| def verify_wide_n8192_m128(self) -> torch.Tensor: |
| return self._reference() |
|
|
|
|
| def _register_gelu_shapes() -> None: |
| for label, m, n in QUANT_SHAPES: |
|
|
| def setup(self, m=m, n=n) -> None: |
| self._setup_shape(m, n) |
|
|
| def benchmark(self) -> None: |
| self.kernel.gelu_quantize_fp8_static_bf16( |
| self.input, self.scale, out=self.out |
| ) |
|
|
| def verify(self) -> torch.Tensor: |
| return self._reference() |
|
|
| setattr(GeluFp8QuantizeBenchmark, f"setup_{label}", setup) |
| setattr(GeluFp8QuantizeBenchmark, f"benchmark_{label}", benchmark) |
| setattr(GeluFp8QuantizeBenchmark, f"verify_{label}", verify) |
|
|
|
|
| _register_bias_gelu_shapes() |
| _register_gelu_shapes() |
|
|
