| import torch | |
| import quiptools_cuda | |
| def benchmark(): | |
| torch.manual_seed(42) | |
| M = 1 | |
| N = 12288 | |
| K = 4096 | |
| x = torch.randn((M, K), dtype=torch.float32, device="cuda") | |
| Qidxs = torch.randint(1<<15, (N, K//8), dtype=torch.int16, device="cuda") | |
| codebook = torch.randint(0x7FFFFFFFFFFFFFFF, (256,), dtype=torch.int64, device="cuda") | |
| # start_event = torch.cuda.Event(enable_timing=True) | |
| # end_event = torch.cuda.Event(enable_timing=True) | |
| # start_event.record() | |
| x = quiptools_cuda.decode_matmul_e8p(x, Qidxs - 0x8000, codebook) | |
| # end_event.record() | |
| # torch.cuda.synchronize() | |
| # elapsed_time_ms = start_event.elapsed_time(end_event) | |
| # print(f"Elapsed: {elapsed_time_ms:.4f}ms") | |
| if __name__ == "__main__": | |
| benchmark() | |