Move --skip-decode to model cli helper

Signed-off-by: aviator19941 <avinash.sharma@amd.com>
nod-ai · Nov 16, 2024 · 2244b9e · 2244b9e
1 parent 9b3f92a
commit 2244b9e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 11 deletions.
diff --git a/sharktank/sharktank/examples/export_paged_llm_v1.py b/sharktank/sharktank/examples/export_paged_llm_v1.py
@@ -54,17 +54,6 @@ def main():
         help="Enables strictness during export",
         action="store_true",
     )
-    parser.add_argument(
-        "--attention-kernel",
-        type=str,
-        default="decomposed",
-        choices=["decomposed", "torch"],
-    )
-    parser.add_argument(
-        "--skip-decode",
-        help="Enables prefill only, skips decode",
-        action="store_true",
-    )
 
     cli.add_quantization_options(parser)
     cli.add_model_options(parser)

diff --git a/sharktank/sharktank/utils/cli.py b/sharktank/sharktank/utils/cli.py
@@ -69,6 +69,11 @@ def add_model_options(parser: argparse.ArgumentParser):
         default="decomposed",
         choices=["decomposed", "torch"],
     )
+    parser.add_argument(
+        "--skip-decode",
+        help="Enables prefill only, skips decode",
+        action="store_true",
+    )
 
 
 def add_quantization_options(parser: argparse.ArgumentParser):