import torch

from kernels.benchmark import Benchmark


class NunchakuLiteKernelsBenchmark(Benchmark):
    def setup(self):
        self.m = 16
        self.k = 1536
        self.n = 3072
        self.input = torch.randn(self.m, self.k, device=self.device, dtype=torch.float16)
        self.kernel_weights = torch.zeros(self.n // 4, (self.k // 64) * 32, device=self.device, dtype=torch.int32)
        self.scales = torch.ones(self.k // 64, self.n, device=self.device, dtype=torch.float16) * 0.01
        self.zeros = torch.zeros_like(self.scales)

    def benchmark_base(self):
        self.kernel.awq_gemm_w4a16_g64_int32(self.input, self.kernel_weights, self.scales, self.zeros)

    def verify_base(self) -> torch.Tensor:
        return torch.zeros(self.m, self.n, device=self.device, dtype=torch.float16)