[Bugfix] Fix w8a8 benchmarks for int8 case (#5643)

This commit is contained in:
Tyler Michael Smith 2024-06-18 20:33:25 -04:00 committed by GitHub
parent b23ce92032
commit 6820724e51
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -120,9 +120,8 @@ def bench_int8(dtype: torch.dtype, m: int, k: int, n: int, label: str,
# cutlass impl
timers.append(
bench_fn(a, b, scale_a.to(device="cpu"), scale_b.to(device="cpu"),
torch.bfloat16, label, sub_label, cutlass_impl,
"cutlass_i8_i8_bf16_scaled_mm"))
bench_fn(a, b, scale_a, scale_b, torch.bfloat16, label, sub_label,
cutlass_impl, "cutlass_i8_i8_bf16_scaled_mm"))
return timers