Fix CI after do_bench refactor in pytorch inductor (pytorch#242)

Summary:
We are relying on some private APIs from inductor and a recent refactor: pytorch/pytorch#125736 broken the do_bench
API we rely on for autoquant, maybe we should use our own do_bench or rely on triton's directly?

Test Plan:
regression tests
python test/integration/test_integration.py -k test_autoquant_one_input_29_cuda

Reviewers:

Subscribers:

Tasks:

Tags:

torchao/quantization/autoquant.py

@@ -15,6 +15,8 @@
 except:
     from torch._inductor.runtime.runtime_utils import do_bench
 
+from .utils import TORCH_VERSION_AFTER_2_4
+
 aten = torch.ops.aten
 
 AUTOQUANT_CACHE = {}
@@ -197,7 +199,11 @@ def do_autoquant_bench(op, *args, **kwargs):
         graph = torch.cuda.CUDAGraph()
         with torch.cuda.graph(graph, stream=stream):
             op(*args, **kwargs)
-        res = do_bench(lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median")
+        if TORCH_VERSION_AFTER_2_4:
+            from torch._inductor.runtime.runtime_utils import do_bench_gpu
+            res = do_bench_gpu(lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median")
+        else:
+            res = do_bench(lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median")
     return res
 
 def _is_interpolate_mode(mode):