Skip to content

Commit 3c1b042

Browse files
committed
[Benchmark] Fix arg parsing issue in tritonbench integration
1 parent b7822e4 commit 3c1b042

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
@@ -478,10 +478,15 @@ def _inner() -> Callable[..., Any] | object:
478478

479479
def main() -> None:
480480
# Parse command line arguments
481-
parser = argparse.ArgumentParser(description="Run Helion kernels with tritonbench")
481+
parser = argparse.ArgumentParser(
482+
description="Run Helion kernels with tritonbench",
483+
allow_abbrev=False, # Disable prefix matching to prevent --k from matching --kernel
484+
)
482485
parser.add_argument(
483486
"--kernel",
487+
"--op",
484488
type=str,
489+
dest="kernel",
485490
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.",
486491
)
487492
parser.add_argument(

0 commit comments

Comments
 (0)