@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 020f58abcd
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 02:15 UTC

[Core] Support multiple tasks per model (#20771)

Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>

Nicolò Lucchesi committed Jul 13, 2025 at 04:40 UTC 020f58abcdea65302225663130d08fd8f4dd755a
8 files changed +278 -147
tests/test_config.py
+46 -3
@@ -54,7 +54,7 @@ def test_get_field():
54 ("jason9693/Qwen2.5-1.5B-apeach", "pooling", "classify"),
55 ("cross-encoder/ms-marco-MiniLM-L-6-v2", "pooling", "classify"),
56 ("Qwen/Qwen2.5-Math-RM-72B", "pooling", "reward"),
57 - ("openai/whisper-small", "transcription", "transcription"),
57 + ("openai/whisper-small", "generate", "transcription"),
58 ],
59 )
60 def test_auto_task(model_id, expected_runner_type, expected_task):
@@ -69,7 +69,11 @@ def test_auto_task(model_id, expected_runner_type, expected_task):
69 )
70
71 assert config.runner_type == expected_runner_type
72 - assert config.task == expected_task
72 +
73 + if config.runner_type == "pooling":
74 + assert config.task == expected_task
75 + else:
76 + assert expected_task in config.supported_tasks
77
78
79 @pytest.mark.parametrize(
@@ -98,11 +102,50 @@ def test_score_task(model_id, expected_runner_type, expected_task):
102 assert config.task == expected_task
103
104
105 +@pytest.mark.parametrize(("model_id", "expected_runner_type", "expected_task"),
106 + [
107 + ("Qwen/Qwen2.5-1.5B-Instruct", "draft", "auto"),
108 + ])
109 +def test_draft_task(model_id, expected_runner_type, expected_task):
110 + config = ModelConfig(
111 + model_id,
112 + runner="draft",
113 + tokenizer=model_id,
114 + seed=0,
115 + dtype="float16",
116 + )
117 +
118 + assert config.runner_type == expected_runner_type
119 + assert config.task == expected_task
120 +
121 +
122 +@pytest.mark.parametrize(
123 + ("model_id", "expected_runner_type", "expected_task"),
124 + [
125 + ("openai/whisper-small", "generate", "transcription"),
126 + ],
127 +)
128 +def test_transcription_task(model_id, expected_runner_type, expected_task):
129 + config = ModelConfig(
130 + model_id,
131 + task="transcription",
132 + tokenizer=model_id,
133 + tokenizer_mode="auto",
134 + trust_remote_code=False,
135 + seed=0,
136 + dtype="float16",
137 + )
138 +
139 + assert config.runner_type == expected_runner_type
140 + assert config.task == expected_task
141 +
142 +
143 @pytest.mark.parametrize(("model_id", "bad_task"), [
144 ("Qwen/Qwen2.5-Math-RM-72B", "generate"),
145 + ("Qwen/Qwen3-0.6B", "transcription"),
146 ])
147 def test_incorrect_task(model_id, bad_task):
105 - with pytest.raises(ValueError, match=r"does not support the .* task"):
148 + with pytest.raises(ValueError, match=r"does not support task=.*"):
149 ModelConfig(
150 model_id,
151 task=bad_task,
vllm/config.py
+160 -96
@@ -91,24 +91,19 @@ logger = init_logger(__name__)
91 ConfigT = TypeVar("ConfigT", bound=ConfigType)
92
93 TaskOption = Literal["auto", "generate", "embedding", "embed", "classify",
94 - "score", "reward", "transcription"]
94 + "score", "reward", "transcription", "draft"]
95
96 -_ResolvedTask = Literal["generate", "embed", "classify", "reward", "draft",
97 - "transcription"]
96 +_ResolvedTask = Literal["generate", "transcription", "pooling", "embed",
97 + "classify", "reward", "draft"]
98
99 -RunnerType = Literal["generate", "pooling", "draft", "transcription"]
99 +RunnerOption = Literal["auto", "generate", "pooling", "draft"]
100
101 -_RUNNER_TASKS: dict[RunnerType, list[_ResolvedTask]] = {
102 - "generate": ["generate"],
103 - "pooling": ["embed", "classify", "reward"],
104 - "draft": ["draft"],
105 - "transcription": ["transcription"],
106 -}
101 +RunnerType = Literal["generate", "pooling", "draft"]
102
108 -_TASK_RUNNER: dict[_ResolvedTask, RunnerType] = {
109 - task: runner
110 - for runner, tasks in _RUNNER_TASKS.items()
111 - for task in tasks
103 +_RUNNER_TASKS: dict[RunnerType, list[_ResolvedTask]] = {
104 + "generate": ["generate", "transcription"],
105 + "pooling": ["pooling", "embed", "classify", "reward"],
106 + "draft": [],
107 }
108
109
@@ -234,11 +229,14 @@ class ModelConfig:
229 """Name or path of the Hugging Face model to use. It is also used as the
230 content for `model_name` tag in metrics output when `served_model_name` is
231 not specified."""
237 - task: Literal[TaskOption, Literal["draft"]] = "auto"
238 - """The task to use the model for. Each vLLM instance only supports one
239 - task, even if the same model can be used for multiple tasks. When the model
240 - only supports one task, "auto" can be used to select it; otherwise, you
241 - must specify explicitly which task to use."""
232 + runner: RunnerOption = "auto"
233 + """The type of model runner to use. Each vLLM instance only supports one
234 + model runner, even if the same model can be used for multiple types."""
235 + task: TaskOption = "auto"
236 + """The task to use the model for. If the model supports more than one
237 + model runner, this is used to select which model runner to run.
238 +
239 + Note that the model may support other tasks using the same model runner."""
240 tokenizer: SkipValidation[str] = None # type: ignore
241 """Name or path of the Hugging Face tokenizer to use. If unspecified, model
242 name or path will be used."""
@@ -553,10 +551,41 @@ class ModelConfig:
551 self.hf_image_processor_config = get_hf_image_processor_config(
552 self.model, hf_token=self.hf_token, revision=self.revision)
553
556 - supported_tasks, task = self._resolve_task(self.task)
557 - self.supported_tasks = supported_tasks
558 - self.task = task
559 - if self.task in ("draft", "generate"):
554 + # For pooling models, self.task is used to indicate the
555 + # user-selected task
556 + if self.task == "score":
557 + if self.registry.is_cross_encoder_model(self.architectures):
558 + self.task = "classify"
559 + else:
560 + self.task = "embed"
561 + elif self.task == "embedding":
562 + msg = ("The 'embedding' task has been renamed to 'embed', please "
563 + "use the new name. The old name will be removed in v1.0.")
564 + warnings.warn(msg, DeprecationWarning, stacklevel=2)
565 +
566 + self.task = "embed"
567 +
568 + all_supported_tasks = self._get_supported_tasks(self.task)
569 + logger.debug("Tasks supported by runner type: %s", all_supported_tasks)
570 + supported_runner_types = self._get_supported_runner_types(
571 + all_supported_tasks)
572 + runner_type = self._resolve_runner(self.runner, self.task,
573 + supported_runner_types,
574 + all_supported_tasks)
575 +
576 + logger.debug("Selected runner type: %s", runner_type)
577 + # For pooling models, self.task is used to indicate the
578 + # user-selected task
579 + if runner_type == "pooling" and self.task == "auto":
580 + selected_task = all_supported_tasks[runner_type][-1]
581 + assert selected_task != "pooling"
582 + self.task = selected_task
583 + self.supported_runner_types = supported_runner_types
584 + self.runner_type = runner_type
585 + self.supported_tasks = all_supported_tasks[runner_type]
586 +
587 + if self.runner_type in ("draft",
588 + "generate") and self.task != "transcription":
589 self.truncation_side = "left"
590 else:
591 self.truncation_side = "right"
@@ -780,11 +809,10 @@ class ModelConfig:
809 f"one of {get_args(TokenizerMode)}.")
810 self.tokenizer_mode = tokenizer_mode
811
783 - def _get_preferred_task(
812 + def _get_preferred_pooling_task(
813 self,
814 architectures: list[str],
786 - supported_tasks: set[_ResolvedTask],
787 - ) -> Optional[_ResolvedTask]:
815 + ) -> _ResolvedTask:
816 model_id = self.model
817 if get_pooling_config(model_id, self.revision):
818 return "embed"
@@ -795,92 +823,136 @@ class ModelConfig:
823
824 suffix_to_preferred_task: list[tuple[str, _ResolvedTask]] = [
825 # Other models follow this pattern
798 - ("ForCausalLM", "generate"),
799 - ("ForConditionalGeneration", "generate"),
826 ("ForSequenceClassification", "classify"),
801 - ("ChatModel", "generate"),
802 - ("LMHeadModel", "generate"),
827 ("EmbeddingModel", "embed"),
828 ("RewardModel", "reward"),
829 ]
830 _, arch = self.registry.inspect_model_cls(architectures)
831
832 for suffix, pref_task in suffix_to_preferred_task:
809 - if arch.endswith(suffix) and pref_task in supported_tasks:
833 + if arch.endswith(suffix):
834 return pref_task
835
812 - return None
836 + return "embed"
837
814 - def _resolve_task(
838 + def _get_supported_generation_tasks(
839 self,
816 - task_option: Literal[TaskOption, Literal["draft"]],
817 - ) -> tuple[set[_ResolvedTask], _ResolvedTask]:
818 - if task_option == "draft":
819 - return {"draft"}, "draft"
840 + task_option: TaskOption,
841 + ) -> list[_ResolvedTask]:
842 + registry = self.registry
843 + architectures = self.architectures
844 +
845 + if registry.is_transcription_only_model(architectures):
846 + return ["transcription"]
847 +
848 + supported_tasks = list[_ResolvedTask]()
849 + if registry.is_text_generation_model(architectures):
850 + supported_tasks.append("generate")
851 +
852 + if registry.is_transcription_model(architectures):
853 + supported_tasks.append("transcription")
854 +
855 + return supported_tasks
856
857 + def _get_supported_pooling_tasks(
858 + self,
859 + task_option: TaskOption,
860 + ) -> list[_ResolvedTask]:
861 registry = self.registry
862 architectures = self.architectures
863
824 - runner_support: dict[RunnerType, bool] = {
825 - # NOTE: Listed from highest to lowest priority,
826 - # in case the model supports multiple of them
827 - "transcription": registry.is_transcription_model(architectures),
828 - "generate": registry.is_text_generation_model(architectures),
829 - "pooling": registry.is_pooling_model(architectures),
864 + supported_tasks = list[_ResolvedTask]()
865 + if registry.is_pooling_model(architectures):
866 + supported_tasks.append("pooling")
867 +
868 + # For now, users must specify the task (other than "pooling")
869 + # to use for pooling models
870 + if task_option == "auto":
871 + preferred_task = self._get_preferred_pooling_task(
872 + architectures)
873 +
874 + supported_tasks.append(preferred_task)
875 + elif task_option in _RUNNER_TASKS["pooling"]:
876 + supported_tasks.append(cast(_ResolvedTask, task_option))
877 +
878 + return supported_tasks
879 +
880 + def _get_supported_tasks(
881 + self,
882 + task_option: TaskOption,
883 + ) -> dict[RunnerType, list[_ResolvedTask]]:
884 + return {
885 + "generate": self._get_supported_generation_tasks(task_option),
886 + "pooling": self._get_supported_pooling_tasks(task_option),
887 + "draft": ["draft"]
888 }
831 - supported_runner_types_lst: list[RunnerType] = [
832 - runner_type
833 - for runner_type, is_supported in runner_support.items()
834 - if is_supported
835 - ]
889
837 - supported_tasks_lst: list[_ResolvedTask] = [
838 - task for runner_type in supported_runner_types_lst
839 - for task in _RUNNER_TASKS[runner_type]
840 - ]
841 - supported_tasks = set(supported_tasks_lst)
890 + def _get_supported_runner_types(
891 + self,
892 + supported_tasks: dict[RunnerType, list[_ResolvedTask]],
893 + ) -> set[RunnerType]:
894 + return {
895 + runner
896 + for runner, runner_tasks in supported_tasks.items()
897 + if len(runner_tasks) > 0
898 + }
899
843 - if task_option == "auto":
844 - selected_task = next(iter(supported_tasks_lst))
900 + def _resolve_runner(
901 + self,
902 + runner_option: RunnerOption,
903 + task_option: TaskOption,
904 + supported_runner_types: set[RunnerType],
905 + supported_tasks: dict[RunnerType, list[_ResolvedTask]],
906 + ) -> RunnerType:
907 + if not supported_runner_types:
908 + raise ValueError("This model does not support any model runners!")
909 +
910 + if runner_option != "auto":
911 + if runner_option not in supported_runner_types:
912 + raise ValueError(
913 + f"This model does not support runner={runner_option!r}. "
914 + f"Available runners: {supported_runner_types}")
915
846 - if len(supported_tasks_lst) > 1:
847 - preferred_task = self._get_preferred_task(
848 - architectures, supported_tasks)
849 - if preferred_task is not None:
850 - selected_task = preferred_task
916 + return runner_option
917
852 - logger.info(
853 - "This model supports multiple tasks: %s. "
854 - "Defaulting to '%s'.", supported_tasks, selected_task)
855 - else:
856 - if task_option == "score":
857 - if not runner_support["pooling"]:
858 - msg = (f"This model does not support the '{task_option}' "
859 - f"task. Supported tasks: {supported_tasks}")
860 - raise ValueError(msg)
861 - if self.registry.is_cross_encoder_model(architectures):
862 - task_option = "classify"
863 - else:
864 - task_option = "embed"
918 + if task_option != "auto":
919 + for runner, runner_tasks in supported_tasks.items():
920 + if task_option in runner_tasks:
921 + return runner
922 else:
866 - # Aliases
867 - if task_option == "embedding":
868 - msg = ("The 'embedding' task has been renamed to "
869 - "'embed', please use the new name. The old name "
870 - "will be removed in v1.0.")
871 - warnings.warn(msg, DeprecationWarning, stacklevel=2)
923 + task_runner: RunnerType = next(
924 + runner for runner, tasks in _RUNNER_TASKS.items()
925 + if task_option in tasks)
926 + raise ValueError(
927 + f"This model does not support task={task_option!r}. "
928 + f"Available tasks for runner={task_runner!r}: "
929 + f"{supported_tasks[task_runner]}")
930
873 - task_option = "embed"
931 + suffix_to_preferred_runner: list[tuple[str, RunnerType]] = [
932 + ("ForCausalLM", "generate"),
933 + ("ForConditionalGeneration", "generate"),
934 + ("ChatModel", "generate"),
935 + ("LMHeadModel", "generate"),
936 + ("ForSequenceClassification", "pooling"),
937 + ("EmbeddingModel", "pooling"),
938 + ("RewardModel", "pooling"),
939 + ]
940 + _, arch = self.registry.inspect_model_cls(self.architectures)
941
875 - if task_option not in supported_tasks:
876 - msg = (
877 - f"This model does not support the '{task_option}' task. "
878 - f"Supported tasks: {supported_tasks}")
879 - raise ValueError(msg)
942 + for suffix, pref_runner in suffix_to_preferred_runner:
943 + if arch.endswith(suffix) and pref_runner in supported_runner_types:
944 + return pref_runner
945
881 - selected_task = task_option
946 + if "classify" in supported_tasks.get("pooling", []):
947 + # When multiple pooling tasks are present, default to
948 + # pooling (eg cross-encoder) for non-standard architectures.
949 + return "pooling"
950 + if "generate" in supported_runner_types:
951 + return "generate"
952 + if "pooling" in supported_runner_types:
953 + return "pooling"
954
883 - return supported_tasks, selected_task
955 + raise AssertionError("This line should not be reached")
956
957 def _parse_quant_hf_config(self):
958 quant_cfg = getattr(self.hf_config, "quantization_config", None)
@@ -1449,14 +1521,6 @@ class ModelConfig:
1521 def use_mla(self) -> bool:
1522 return self.is_deepseek_mla and not envs.VLLM_MLA_DISABLE
1523
1452 - @property
1453 - def supported_runner_types(self) -> set[RunnerType]:
1454 - return {_TASK_RUNNER[task] for task in self.supported_tasks}
1455 -
1456 - @property
1457 - def runner_type(self) -> RunnerType:
1458 - return _TASK_RUNNER[cast(_ResolvedTask, self.task)]
1459 -
1524 @property
1525 def is_v1_compatible(self) -> bool:
1526 architectures = getattr(self.hf_config, "architectures", [])
@@ -2694,7 +2758,7 @@ class SpeculativeConfig:
2758 if self.model is not None:
2759 self.draft_model_config = ModelConfig(
2760 model=self.model,
2697 - task="draft",
2761 + runner="draft",
2762 tokenizer=self.target_model_config.tokenizer,
2763 tokenizer_mode=self.target_model_config.tokenizer_mode,
2764 trust_remote_code=self.target_model_config.
vllm/entrypoints/llm.py
+31 -30
@@ -454,20 +454,19 @@ class LLM:
454 considered legacy and may be deprecated in the future. You should
455 instead pass them via the `inputs` parameter.
456 """
457 - runner_type = self.llm_engine.model_config.runner_type
458 - if runner_type not in ["generate", "transcription"]:
457 + model_config = self.llm_engine.model_config
458 + runner_type = model_config.runner_type
459 + if runner_type != "generate":
460 messages = [
460 - "LLM.generate() is only supported for (conditional) generation "
461 - "models (XForCausalLM, XForConditionalGeneration).",
461 + "LLM.generate() is only supported for generative models."
462 ]
463
464 - supported_runner_types = self.llm_engine.model_config \
465 - .supported_runner_types
466 - if "generate" in supported_runner_types:
464 + if "generate" in model_config.supported_runner_types:
465 messages.append(
466 "Your model supports the 'generate' runner, but is "
467 f"currently initialized for the '{runner_type}' runner. "
470 - "Please initialize vLLM using `--task generate`.")
468 + "Please initialize vLLM using `--task generate` or "
469 + "`--task transcription`.")
470
471 raise ValueError(" ".join(messages))
472
@@ -1091,13 +1090,12 @@ class LLM:
1090 considered legacy and may be deprecated in the future. You should
1091 instead pass them via the `inputs` parameter.
1092 """
1094 - runner_type = self.llm_engine.model_config.runner_type
1093 + model_config = self.llm_engine.model_config
1094 + runner_type = model_config.runner_type
1095 if runner_type != "pooling":
1096 messages = ["LLM.encode() is only supported for pooling models."]
1097
1098 - supported_runner_types = self.llm_engine.model_config \
1099 - .supported_runner_types
1100 - if "pooling" in supported_runner_types:
1098 + if "pooling" in model_config.supported_runner_types:
1099 messages.append(
1100 "Your model supports the 'pooling' runner, but is "
1101 f"currently initialized for the '{runner_type}' runner. "
@@ -1119,13 +1117,13 @@ class LLM:
1117 # Use default pooling params.
1118 pooling_params = PoolingParams()
1119 elif isinstance(pooling_params, PoolingParams):
1122 - pooling_params.verify(self.llm_engine.model_config)
1120 + pooling_params.verify(model_config)
1121 else:
1122 for pooling_param in pooling_params:
1125 - pooling_param.verify(self.llm_engine.model_config)
1123 + pooling_param.verify(model_config)
1124
1127 - tokenization_kwargs: dict[str, Any] = {}
1128 - _validate_truncation_size(self.llm_engine.model_config.max_model_len,
1125 + tokenization_kwargs = dict[str, Any]()
1126 + _validate_truncation_size(model_config.max_model_len,
1127 truncate_prompt_tokens, tokenization_kwargs)
1128
1129 self._validate_and_add_requests(
@@ -1178,9 +1176,10 @@ class LLM:
1176 A list of `EmbeddingRequestOutput` objects containing the
1177 embedding vectors in the same order as the input prompts.
1178 """
1181 - if self.llm_engine.model_config.task != "embed":
1182 - raise ValueError(
1183 - "Embedding API is only enabled for `--task embed`")
1179 + model_config = self.llm_engine.model_config
1180 + if "embed" not in model_config.supported_tasks:
1181 + raise ValueError("Embedding API is not supported by this model. "
1182 + "Please set `--task embed`.")
1183
1184 items = self.encode(prompts,
1185 truncate_prompt_tokens=truncate_prompt_tokens,
@@ -1223,9 +1222,11 @@ class LLM:
1222 A list of `ClassificationRequestOutput` objects containing the
1223 embedding vectors in the same order as the input prompts.
1224 """
1226 - if self.llm_engine.model_config.task != "classify":
1225 + model_config = self.llm_engine.model_config
1226 + if "classify" not in model_config.supported_tasks:
1227 raise ValueError(
1228 - "Classification API is only enabled for `--task classify`")
1228 + "Classification API is not supported by this model. "
1229 + "Please set `--task classify`.")
1230
1231 items = self.encode(prompts,
1232 use_tqdm=use_tqdm,
@@ -1392,13 +1393,12 @@ class LLM:
1393 A list of `ScoringRequestOutput` objects containing the
1394 generated scores in the same order as the input prompts.
1395 """
1395 - runner_type = self.llm_engine.model_config.runner_type
1396 + model_config = self.llm_engine.model_config
1397 + runner_type = model_config.runner_type
1398 if runner_type != "pooling":
1399 messages = ["LLM.score() is only supported for pooling models."]
1400
1399 - supported_runner_types = self.llm_engine.model_config \
1400 - .supported_runner_types
1401 - if "pooling" in supported_runner_types:
1401 + if "pooling" in model_config.supported_runner_types:
1402 messages.append(
1403 "Your model supports the 'pooling' runner, but is "
1404 f"currently initialized for the '{runner_type}' runner. "
@@ -1407,12 +1407,13 @@ class LLM:
1407
1408 raise ValueError(" ".join(messages))
1409
1410 - if self.llm_engine.model_config.task not in ("embed", "classify"):
1411 - raise ValueError("Score API is only enabled for "
1412 - "`--task embed or --task classify`.")
1410 + if all(t not in model_config.supported_tasks
1411 + for t in ("embed", "classify")):
1412 + raise ValueError("Score API is not supported by this model. "
1413 + "Please set `--task embed` or `--task classify`.")
1414
1414 - if (self.llm_engine.model_config.task == "classify"
1415 - and self.llm_engine.model_config.hf_config.num_labels != 1):
1415 + if (model_config.task == "classify"
1416 + and getattr(model_config.hf_config, "num_labels", 0) != 1):
1417 raise ValueError("Score API is only enabled for num_labels == 1.")
1418
1419 # the tokenizer for models such as
vllm/entrypoints/openai/api_server.py
+14 -12
@@ -1520,7 +1520,7 @@ async def init_app_state(
1520 reasoning_parser=args.reasoning_parser,
1521 enable_prompt_tokens_details=args.enable_prompt_tokens_details,
1522 enable_force_include_usage=args.enable_force_include_usage,
1523 - ) if model_config.runner_type == "generate" else None
1523 + ) if "generate" in model_config.supported_tasks else None
1524 state.openai_serving_chat = OpenAIServingChat(
1525 engine_client,
1526 model_config,
@@ -1537,7 +1537,7 @@ async def init_app_state(
1537 reasoning_parser=args.reasoning_parser,
1538 enable_prompt_tokens_details=args.enable_prompt_tokens_details,
1539 enable_force_include_usage=args.enable_force_include_usage,
1540 - ) if model_config.runner_type == "generate" else None
1540 + ) if "generate" in model_config.supported_tasks else None
1541 state.openai_serving_completion = OpenAIServingCompletion(
1542 engine_client,
1543 model_config,
@@ -1545,7 +1545,7 @@ async def init_app_state(
1545 request_logger=request_logger,
1546 return_tokens_as_token_ids=args.return_tokens_as_token_ids,
1547 enable_force_include_usage=args.enable_force_include_usage,
1548 - ) if model_config.runner_type == "generate" else None
1548 + ) if "generate" in model_config.supported_tasks else None
1549 state.openai_serving_pooling = OpenAIServingPooling(
1550 engine_client,
1551 model_config,
@@ -1553,7 +1553,7 @@ async def init_app_state(
1553 request_logger=request_logger,
1554 chat_template=resolved_chat_template,
1555 chat_template_content_format=args.chat_template_content_format,
1556 - ) if model_config.runner_type == "pooling" else None
1556 + ) if "pooling" in model_config.supported_tasks else None
1557 state.openai_serving_embedding = OpenAIServingEmbedding(
1558 engine_client,
1559 model_config,
@@ -1561,22 +1561,24 @@ async def init_app_state(
1561 request_logger=request_logger,
1562 chat_template=resolved_chat_template,
1563 chat_template_content_format=args.chat_template_content_format,
1564 - ) if model_config.task == "embed" else None
1564 + ) if "embed" in model_config.supported_tasks else None
1565 state.openai_serving_classification = ServingClassification(
1566 engine_client,
1567 model_config,
1568 state.openai_serving_models,
1569 request_logger=request_logger,
1570 - ) if model_config.task == "classify" else None
1570 + ) if "classify" in model_config.supported_tasks else None
1571
1572 - enable_serving_reranking = (model_config.task == "classify" and getattr(
1573 - model_config.hf_config, "num_labels", 0) == 1)
1572 + enable_serving_reranking = ("classify" in model_config.supported_tasks
1573 + and getattr(model_config.hf_config,
1574 + "num_labels", 0) == 1)
1575 state.openai_serving_scores = ServingScores(
1576 engine_client,
1577 model_config,
1578 state.openai_serving_models,
1578 - request_logger=request_logger) if (
1579 - model_config.task == "embed" or enable_serving_reranking) else None
1579 + request_logger=request_logger,
1580 + ) if ("embed" in model_config.supported_tasks
1581 + or enable_serving_reranking) else None
1582
1583 state.openai_serving_tokenization = OpenAIServingTokenization(
1584 engine_client,
@@ -1591,13 +1593,13 @@ async def init_app_state(
1593 model_config,
1594 state.openai_serving_models,
1595 request_logger=request_logger,
1594 - ) if model_config.runner_type == "transcription" else None
1596 + ) if "transcription" in model_config.supported_tasks else None
1597 state.openai_serving_translation = OpenAIServingTranslation(
1598 engine_client,
1599 model_config,
1600 state.openai_serving_models,
1601 request_logger=request_logger,
1600 - ) if model_config.runner_type == "transcription" else None
1602 + ) if "transcription" in model_config.supported_tasks else None
1603 state.task = model_config.task
1604
1605 state.enable_server_load_tracking = args.enable_server_load_tracking
vllm/entrypoints/openai/run_batch.py
+8 -6
@@ -348,7 +348,7 @@ async def main(args):
348 chat_template=None,
349 chat_template_content_format="auto",
350 enable_prompt_tokens_details=args.enable_prompt_tokens_details,
351 - ) if model_config.runner_type == "generate" else None
351 + ) if "generate" in model_config.supported_tasks else None
352 openai_serving_embedding = OpenAIServingEmbedding(
353 engine,
354 model_config,
@@ -356,17 +356,19 @@ async def main(args):
356 request_logger=request_logger,
357 chat_template=None,
358 chat_template_content_format="auto",
359 - ) if model_config.task == "embed" else None
359 + ) if "embed" in model_config.supported_tasks else None
360
361 - enable_serving_reranking = (model_config.task == "classify" and getattr(
362 - model_config.hf_config, "num_labels", 0) == 1)
361 + enable_serving_reranking = ("classify" in model_config.supported_tasks
362 + and getattr(model_config.hf_config,
363 + "num_labels", 0) == 1)
364
364 - openai_serving_scores = (ServingScores(
365 + openai_serving_scores = ServingScores(
366 engine,
367 model_config,
368 openai_serving_models,
369 request_logger=request_logger,
369 - ) if (model_config.task == "embed" or enable_serving_reranking) else None)
370 + ) if ("embed" in model_config.supported_tasks
371 + or enable_serving_reranking) else None
372
373 tracker = BatchProgressTracker()
374 logger.info("Reading batch from %s...", args.input_file)
vllm/model_executor/models/interfaces.py
+6
@@ -694,6 +694,12 @@ class SupportsTranscription(Protocol):
694
695 supports_transcription: ClassVar[Literal[True]] = True
696
697 + supports_transcription_only: ClassVar[bool] = False
698 + """
699 + Transcription models can opt out of text generation by setting this to
700 + `True`.
701 + """
702 +
703 @classmethod
704 def get_generation_prompt(cls, audio: np.ndarray,
705 stt_config: SpeechToTextConfig, language: str,
vllm/model_executor/models/registry.py
+10
@@ -284,6 +284,7 @@ class _ModelInfo:
284 is_hybrid: bool
285 has_noops: bool
286 supports_transcription: bool
287 + supports_transcription_only: bool
288 supports_v0_only: bool
289
290 @staticmethod
@@ -299,6 +300,8 @@ class _ModelInfo:
300 is_attention_free=is_attention_free(model),
301 is_hybrid=is_hybrid(model),
302 supports_transcription=supports_transcription(model),
303 + supports_transcription_only=(supports_transcription(model) and
304 + model.supports_transcription_only),
305 supports_v0_only=supports_v0_only(model),
306 has_noops=has_noops(model),
307 )
@@ -573,6 +576,13 @@ class _ModelRegistry:
576 model_cls, _ = self.inspect_model_cls(architectures)
577 return model_cls.supports_transcription
578
579 + def is_transcription_only_model(
580 + self,
581 + architectures: Union[str, list[str]],
582 + ) -> bool:
583 + model_cls, _ = self.inspect_model_cls(architectures)
584 + return model_cls.supports_transcription_only
585 +
586 def is_v1_compatible(
587 self,
588 architectures: Union[str, list[str]],
vllm/model_executor/models/whisper.py
+3
@@ -772,6 +772,9 @@ class WhisperForConditionalGeneration(nn.Module, SupportsTranscription,
772 ".fc2.": ".mlp.fc2."
773 })
774
775 + # Whisper only supports audio-conditioned generation.
776 + supports_transcription_only = True
777 +
778 @classmethod
779 def validate_language(cls, language: str) -> bool:
780 if language in ISO639_1_SUPPORTED_LANGS: