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

6820724e · Tyler Michael Smith · GitHub · b23ce920 · 6820724e
Unverified Commit 6820724e authored Jun 18, 2024 by Tyler Michael Smith Committed by GitHub Jun 19, 2024
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 3 deletions

benchmarks/cutlass_benchmarks/w8a8_benchmarks.py benchmarks/cutlass_benchmarks/w8a8_benchmarks.py +2 -3

No files found.
--- a/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py
+++ b/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py
@@ -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