Fix tritonbench integration issue

Author: yf225

benchmarks/run.py

@@ -629,19 +629,9 @@ def _inner() -> Callable[..., Any] | object:
             ["--input-id", str(start_idx), "--num-inputs", str(shard_size)]
         )
 
-    # Re-parse args with the new input range
-    tb_args, unknown_args = tb_parser.parse_known_args(tritonbench_args)
-
-    # Use the public API to load and run the operator
-    from tritonbench.operators import load_opbench_by_name
-
-    op_class = load_opbench_by_name(operator_name)
-    benchmark = op_class(tb_args=tb_args, extra_args=unknown_args)
-    benchmark.run()
+    from tritonbench.run import run as tritonbench_run
 
-    # Print results if available
-    if hasattr(benchmark, "output"):
-        print(benchmark.output)
+    tritonbench_run(tritonbench_args)
 
     # Force garbage collection multiple times to ensure memory is freed
     for _ in range(3):