Make padding-aware scheduling disableable (#771)

HabanaAI · Feb 3, 2025 · 85eb147 · 85eb147
2 parents e29b5f5 + 4deb8cf
commit 85eb147
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
@@ -465,7 +465,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parser.add_argument(
             '--use-padding-aware-scheduling',
             default=EngineArgs.use_padding_aware_scheduling,
-            action='store_true',
+            action=StoreBoolean,
+            nargs="?",
+            const="True",
             help=('Use padding-aware scheduling. If True, the scheduler '
                   'will consider padded tokens in prefill. '
                   'By default this is set to False on non-HPU devices. '))