| import torch |
| import torch.nn.functional as F |
|
|
| from kernels.benchmark import Benchmark |
|
|
|
|
| class CausalConv1dBenchmark(Benchmark): |
| seed: int = 42 |
|
|
| def setup(self): |
| batch_size, dim, seqlen, width = 2, 64, 128, 4 |
| self.x = torch.randn( |
| batch_size, dim, seqlen, device=self.device, dtype=torch.float16 |
| ) |
| self.weight = torch.randn(dim, width, device=self.device, dtype=torch.float32) |
| self.bias = torch.randn(dim, device=self.device, dtype=torch.float32) |
| self.out = torch.empty( |
| batch_size, dim, seqlen, device=self.device, dtype=torch.float16 |
| ) |
| self.dim = dim |
| self.width = width |
| self.seqlen = seqlen |
|
|
| def benchmark_base(self): |
| self.out = self.kernel.causal_conv1d_fn(self.x, self.weight, self.bias) |
|
|
| def verify_base(self) -> torch.Tensor: |
| x_fp32 = self.x.to(self.weight.dtype) |
| out = F.conv1d( |
| x_fp32, |
| self.weight.unsqueeze(1), |
| self.bias, |
| padding=self.width - 1, |
| groups=self.dim, |
| ) |
| return out[..., : self.seqlen].to(self.x.dtype) |
|
|
| def setup_large(self): |
| batch_size, dim, seqlen, width = 8, 256, 512, 4 |
| self.x = torch.randn( |
| batch_size, dim, seqlen, device=self.device, dtype=torch.float16 |
| ) |
| self.weight = torch.randn(dim, width, device=self.device, dtype=torch.float32) |
| self.bias = torch.randn(dim, device=self.device, dtype=torch.float32) |
| self.out = torch.empty( |
| batch_size, dim, seqlen, device=self.device, dtype=torch.float16 |
| ) |
| self.dim = dim |
| self.width = width |
| self.seqlen = seqlen |
|
|
| def benchmark_large(self): |
| self.out = self.kernel.causal_conv1d_fn(self.x, self.weight, self.bias) |
|
|
| def verify_large(self) -> torch.Tensor: |
| x_fp32 = self.x.to(self.weight.dtype) |
| out = F.conv1d( |
| x_fp32, |
| self.weight.unsqueeze(1), |
| self.bias, |
| padding=self.width - 1, |
| groups=self.dim, |
| ) |
| return out[..., : self.seqlen].to(self.x.dtype) |
|
|
| def setup_xlarge(self): |
| batch_size, dim, seqlen, width = 16, 512, 1024, 4 |
| self.x = torch.randn( |
| batch_size, dim, seqlen, device=self.device, dtype=torch.float16 |
| ) |
| self.weight = torch.randn(dim, width, device=self.device, dtype=torch.float32) |
| self.bias = torch.randn(dim, device=self.device, dtype=torch.float32) |
| self.out = torch.empty( |
| batch_size, dim, seqlen, device=self.device, dtype=torch.float16 |
| ) |
| self.dim = dim |
| self.width = width |
| self.seqlen = seqlen |
|
|
| def benchmark_xlarge(self): |
| self.out = self.kernel.causal_conv1d_fn(self.x, self.weight, self.bias) |
|
|
| def verify_xlarge(self) -> torch.Tensor: |
| x_fp32 = self.x.to(self.weight.dtype) |
| out = F.conv1d( |
| x_fp32, |
| self.weight.unsqueeze(1), |
| self.bias, |
| padding=self.width - 1, |
| groups=self.dim, |
| ) |
| return out[..., : self.seqlen].to(self.x.dtype) |
|
|
