From 70e476bf4191efac6f98cf19a61f9c84a6202336 Mon Sep 17 00:00:00 2001 From: Xu Zhao Date: Thu, 25 Apr 2024 12:45:09 -0400 Subject: [PATCH] Use softmax and limit number of batch --- userbenchmark/triton/ci.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/userbenchmark/triton/ci.py b/userbenchmark/triton/ci.py index fcabf298a6..20043c8b07 100644 --- a/userbenchmark/triton/ci.py +++ b/userbenchmark/triton/ci.py @@ -9,7 +9,7 @@ from torchbenchmark.util.triton_op import BenchmarkOperatorResult CI_TESTS = [ - ["--op", "launch_latency"], + ["--op", "softmax", "--num-inputs", "10"], ] def ci_result_to_userbenchmark_json(ci_metrics: List[BenchmarkOperatorResult]) -> Any: @@ -26,7 +26,7 @@ def run_ci(): metrics = _run(test_args, test_extra_args) ci_result.append(metrics) result = ci_result_to_userbenchmark_json(ci_result) - result_with_environ = get_output_json(result) + result_with_environ = get_output_json(BM_NAME, result) result_with_environ["environ"]["triton_version"] = triton.__version__ output_file = get_default_output_json_path(BM_NAME) json_str = json.dumps(result_with_environ, indent=4)