File size: 2,656 Bytes
52ab2d1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | 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()
class FP8SwiGLUMlpBenchmark(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.gate_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.gate_up_w = _quantize_fp8(
torch.randn((2 * H, K), device=self.device, dtype=torch.bfloat16),
self.gate_up_scale,
)
self.down_w = _quantize_fp8(
torch.randn((N, H), device=self.device, dtype=torch.bfloat16),
self.down_scale,
)
self.gate_up_bf16 = torch.empty((M, 2 * 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:
gate_up = (
_dequant_fp8(self.x, self.x_scale)
@ _dequant_fp8(self.gate_up_w, self.gate_up_scale).T
).to(torch.bfloat16)
gate, up = gate_up.float().chunk(2, dim=1)
hidden_fp8 = _quantize_fp8(torch.nn.functional.silu(gate) * up, self.hidden_scale)
return (
_dequant_fp8(hidden_fp8, self.hidden_scale)
@ _dequant_fp8(self.down_w, self.down_scale).T
).to(torch.bfloat16)
def setup_smoke_mlp(self) -> None:
self._setup_shape(10, 1024, 4096, 1024)
def benchmark_smoke_mlp(self) -> None:
self.kernel.fp8_swiglu_mlp_bf16(
self.x,
self.gate_up_w,
self.down_w,
self.x_scale,
self.gate_up_scale,
self.hidden_scale,
self.down_scale,
gate_up_bf16=self.gate_up_bf16,
hidden_fp8=self.hidden_fp8,
out=self.out,
)
def verify_smoke_mlp(self) -> torch.Tensor:
return self._reference()
|