forked from ccf-ai-infra/GPUCodeForces
68 lines
1.9 KiB
Python
68 lines
1.9 KiB
Python
import time
|
|
|
|
import torch
|
|
|
|
from geglu_cuda import ModelNew
|
|
from geglu_torch import Model, get_init_inputs, get_inputs
|
|
|
|
|
|
def _to_cuda(values):
|
|
return [x.cuda() if isinstance(x, torch.Tensor) else x for x in values]
|
|
|
|
|
|
def run_benchmark():
|
|
if not torch.cuda.is_available():
|
|
print("CUDA is not available.")
|
|
return False, 0.0
|
|
|
|
init_inputs = _to_cuda(get_init_inputs())
|
|
inputs = _to_cuda(get_inputs())
|
|
|
|
torch_model = Model(*init_inputs).cuda().eval()
|
|
cuda_model = ModelNew(*init_inputs).cuda().eval()
|
|
|
|
with torch.no_grad():
|
|
output_torch = torch_model(*inputs)
|
|
output_cuda = cuda_model(*inputs)
|
|
|
|
max_diff = (output_torch - output_cuda).abs().max().item()
|
|
mean_diff = (output_torch - output_cuda).abs().mean().item()
|
|
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-4, atol=1e-4)
|
|
|
|
print("-------------------- precision check --------------------")
|
|
print(f"max diff: {max_diff:.8f}")
|
|
print(f"mean diff: {mean_diff:.8f}")
|
|
print(f"allclose: {precision_flag}")
|
|
|
|
for _ in range(20):
|
|
torch_model(*inputs)
|
|
cuda_model(*inputs)
|
|
|
|
num_iterations = 200
|
|
|
|
torch.cuda.synchronize()
|
|
start = time.time()
|
|
for _ in range(num_iterations):
|
|
torch_model(*inputs)
|
|
torch.cuda.synchronize()
|
|
torch_time = (time.time() - start) / num_iterations
|
|
|
|
torch.cuda.synchronize()
|
|
start = time.time()
|
|
for _ in range(num_iterations):
|
|
cuda_model(*inputs)
|
|
torch.cuda.synchronize()
|
|
cuda_time = (time.time() - start) / num_iterations
|
|
|
|
speedup = torch_time / cuda_time if cuda_time > 0 else 0.0
|
|
print("-------------------- performance check --------------------")
|
|
print(f"PyTorch GEGLU average time: {torch_time:.6f} s")
|
|
print(f"Custom CUDA GEGLU average time: {cuda_time:.6f} s")
|
|
print(f"Speedup: {speedup:.2f}x")
|
|
|
|
return precision_flag, speedup
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_benchmark()
|