Skip to content

Commit 642836c

Browse files
authored
[Benchmark] Fix arg parsing issue in tritonbench integration (#417)
1 parent b403f2e commit 642836c

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

benchmarks/run.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -477,10 +477,15 @@ def _inner() -> Callable[..., Any] | object:
477477

478478
def main() -> None:
479479
# Parse command line arguments
480-
parser = argparse.ArgumentParser(description="Run Helion kernels with tritonbench")
480+
parser = argparse.ArgumentParser(
481+
description="Run Helion kernels with tritonbench",
482+
allow_abbrev=False, # Disable prefix matching to prevent --k from matching --kernel
483+
)
481484
parser.add_argument(
482485
"--kernel",
486+
"--op",
483487
type=str,
488+
dest="kernel",
484489
help="Name(s) of the Helion kernel module(s) to run. Can be a single kernel or comma-separated list (e.g., vector_add or vector_add,rms_norm). If not specified, runs all kernels.",
485490
)
486491
parser.add_argument(

0 commit comments

Comments
 (0)