@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 46d81d6951

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 14, 23:45 UTC

[V1] Get supported tasks from model runner instead of model config (#21585)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Jul 25, 2025 at 20:36 UTC 46d81d69511ac11f27d19af80d19dd7b2cce8613

19 files changed +200 -54

vllm/entrypoints/llm.py

+17 -7

@@ -14,6 +14,7 @@ from pydantic import ValidationError

14 from tqdm.auto import tqdm

15 from typing_extensions import TypeVar, deprecated

17 +import vllm.envs as envs

18 from vllm.beam_search import (BeamSearchInstance, BeamSearchOutput,

19 BeamSearchSequence,

20 create_sort_beams_key_function)

@@ -44,9 +45,10 @@ from vllm.model_executor.layers.quantization import QuantizationMethods

45 from vllm.outputs import (ClassificationRequestOutput, EmbeddingRequestOutput,

46 PoolingRequestOutput, RequestOutput,

47 ScoringRequestOutput)

47 -from vllm.pooling_params import PoolingParams, PoolingTask

48 +from vllm.pooling_params import PoolingParams

49 from vllm.sampling_params import (BeamSearchParams, GuidedDecodingParams,

50 RequestOutputKind, SamplingParams)

51 +from vllm.tasks import PoolingTask

52 from vllm.transformers_utils.tokenizer import (AnyTokenizer, MistralTokenizer,

53 get_cached_tokenizer)

54 from vllm.usage.usage_lib import UsageContext

@@ -277,6 +279,16 @@ class LLM:

279 self.request_counter = Counter()

280 self.default_sampling_params: Union[dict[str, Any], None] = None

281

282 + if envs.VLLM_USE_V1:

283 + supported_tasks = self.llm_engine \

284 + .get_supported_tasks() # type: ignore

285 + else:

286 + supported_tasks = self.llm_engine.model_config.supported_tasks

287 +

288 + logger.info("Supported_tasks: %s", supported_tasks)

289 +

290 + self.supported_tasks = supported_tasks

291 +

292 def get_tokenizer(

293 self,

294 lora_request: Optional[LoRARequest] = None,

@@ -1170,8 +1182,7 @@ class LLM:

1182 A list of `EmbeddingRequestOutput` objects containing the

1183 embedding vectors in the same order as the input prompts.

1184 """

1173 - model_config = self.llm_engine.model_config

1174 - if "embed" not in model_config.supported_tasks:

1185 + if "embed" not in self.supported_tasks:

1186 raise ValueError("Embedding API is not supported by this model. "

1187 "Please set `--task embed`.")

1188

@@ -1215,8 +1226,7 @@ class LLM:

1226 A list of `ClassificationRequestOutput` objects containing the

1227 embedding vectors in the same order as the input prompts.

1228 """

1218 - model_config = self.llm_engine.model_config

1219 - if "classify" not in model_config.supported_tasks:

1229 + if "classify" not in self.supported_tasks:

1230 raise ValueError(

1231 "Classification API is not supported by this model. "

1232 "Please set `--task classify`.")

@@ -1397,8 +1407,8 @@ class LLM:

1407

1408 raise ValueError(" ".join(messages))

1409

1400 - if all(t not in model_config.supported_tasks

1401 - for t in ("embed", "classify")):

1410 + supported_tasks = self.supported_tasks

1411 + if all(t not in supported_tasks 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

vllm/entrypoints/openai/api_server.py

+19 -13

@@ -1586,6 +1586,14 @@ async def init_app_state(

1586 state.vllm_config = vllm_config

1587 model_config = vllm_config.model_config

1588

1589 + if envs.VLLM_USE_V1:

1590 + supported_tasks = await engine_client \

1591 + .get_supported_tasks() # type: ignore

1592 + else:

1593 + supported_tasks = model_config.supported_tasks

1594 +

1595 + logger.info("Supported_tasks: %s", supported_tasks)

1596 +

1597 resolved_chat_template = load_chat_template(args.chat_template)

1598 if resolved_chat_template is not None:

1599 # Get the tokenizer to check official template

@@ -1647,7 +1655,7 @@ async def init_app_state(

1655 reasoning_parser=args.reasoning_parser,

1656 enable_prompt_tokens_details=args.enable_prompt_tokens_details,

1657 enable_force_include_usage=args.enable_force_include_usage,

1650 - ) if "generate" in model_config.supported_tasks else None

1658 + ) if "generate" in supported_tasks else None

1659 state.openai_serving_chat = OpenAIServingChat(

1660 engine_client,

1661 model_config,

@@ -1664,7 +1672,7 @@ async def init_app_state(

1672 reasoning_parser=args.reasoning_parser,

1673 enable_prompt_tokens_details=args.enable_prompt_tokens_details,

1674 enable_force_include_usage=args.enable_force_include_usage,

1667 - ) if "generate" in model_config.supported_tasks else None

1675 + ) if "generate" in supported_tasks else None

1676 state.openai_serving_completion = OpenAIServingCompletion(

1677 engine_client,

1678 model_config,

@@ -1673,7 +1681,7 @@ async def init_app_state(

1681 return_tokens_as_token_ids=args.return_tokens_as_token_ids,

1682 enable_prompt_tokens_details=args.enable_prompt_tokens_details,

1683 enable_force_include_usage=args.enable_force_include_usage,

1676 - ) if "generate" in model_config.supported_tasks else None

1684 + ) if "generate" in supported_tasks else None

1685 state.openai_serving_pooling = OpenAIServingPooling(

1686 engine_client,

1687 model_config,

@@ -1681,7 +1689,7 @@ async def init_app_state(

1689 request_logger=request_logger,

1690 chat_template=resolved_chat_template,

1691 chat_template_content_format=args.chat_template_content_format,

1684 - ) if "encode" in model_config.supported_tasks else None

1692 + ) if "encode" in supported_tasks else None

1693 state.openai_serving_embedding = OpenAIServingEmbedding(

1694 engine_client,

1695 model_config,

@@ -1689,24 +1697,22 @@ async def init_app_state(

1697 request_logger=request_logger,

1698 chat_template=resolved_chat_template,

1699 chat_template_content_format=args.chat_template_content_format,

1692 - ) if "embed" in model_config.supported_tasks else None

1700 + ) if "embed" in supported_tasks else None

1701 state.openai_serving_classification = ServingClassification(

1702 engine_client,

1703 model_config,

1704 state.openai_serving_models,

1705 request_logger=request_logger,

1698 - ) if "classify" in model_config.supported_tasks else None

1706 + ) if "classify" in supported_tasks else None

1707

1700 - enable_serving_reranking = ("classify" in model_config.supported_tasks

1701 - and getattr(model_config.hf_config,

1702 - "num_labels", 0) == 1)

1708 + enable_serving_reranking = ("classify" in supported_tasks and getattr(

1709 + model_config.hf_config, "num_labels", 0) == 1)

1710 state.openai_serving_scores = ServingScores(

1711 engine_client,

1712 model_config,

1713 state.openai_serving_models,

1714 request_logger=request_logger,

1708 - ) if ("embed" in model_config.supported_tasks

1709 - or enable_serving_reranking) else None

1715 + ) if ("embed" in supported_tasks or enable_serving_reranking) else None

1716

1717 state.openai_serving_tokenization = OpenAIServingTokenization(

1718 engine_client,

@@ -1721,13 +1727,13 @@ async def init_app_state(

1727 model_config,

1728 state.openai_serving_models,

1729 request_logger=request_logger,

1724 - ) if "transcription" in model_config.supported_tasks else None

1730 + ) if "transcription" in supported_tasks else None

1731 state.openai_serving_translation = OpenAIServingTranslation(

1732 engine_client,

1733 model_config,

1734 state.openai_serving_models,

1735 request_logger=request_logger,

1730 - ) if "transcription" in model_config.supported_tasks else None

1736 + ) if "transcription" in supported_tasks else None

1737 state.task = model_config.task

1738

1739 state.enable_server_load_tracking = args.enable_server_load_tracking

vllm/entrypoints/openai/run_batch.py

+14 -7

@@ -14,6 +14,7 @@ import torch

14 from prometheus_client import start_http_server

15 from tqdm import tqdm

17 +import vllm.envs as envs

18 from vllm.config import VllmConfig

19 from vllm.engine.arg_utils import AsyncEngineArgs, optional_type

20 from vllm.engine.protocol import EngineClient

@@ -335,6 +336,14 @@ async def run_batch(

336

337 model_config = vllm_config.model_config

338

339 + if envs.VLLM_USE_V1:

340 + supported_tasks = await engine_client \

341 + .get_supported_tasks() # type: ignore

342 + else:

343 + supported_tasks = model_config.supported_tasks

344 +

345 + logger.info("Supported_tasks: %s", supported_tasks)

346 +

347 # Create the openai serving objects.

348 openai_serving_models = OpenAIServingModels(

349 engine_client=engine_client,

@@ -351,7 +360,7 @@ async def run_batch(

360 chat_template=None,

361 chat_template_content_format="auto",

362 enable_prompt_tokens_details=args.enable_prompt_tokens_details,

354 - ) if "generate" in model_config.supported_tasks else None

363 + ) if "generate" in supported_tasks else None

364 openai_serving_embedding = OpenAIServingEmbedding(

365 engine_client,

366 model_config,

@@ -359,19 +368,17 @@ async def run_batch(

368 request_logger=request_logger,

369 chat_template=None,

370 chat_template_content_format="auto",

362 - ) if "embed" in model_config.supported_tasks else None

371 + ) if "embed" in supported_tasks else None

372

364 - enable_serving_reranking = ("classify" in model_config.supported_tasks

365 - and getattr(model_config.hf_config,

366 - "num_labels", 0) == 1)

373 + enable_serving_reranking = ("classify" in supported_tasks and getattr(

374 + model_config.hf_config, "num_labels", 0) == 1)

375

376 openai_serving_scores = ServingScores(

377 engine_client,

378 model_config,

379 openai_serving_models,

380 request_logger=request_logger,

373 - ) if ("embed" in model_config.supported_tasks

374 - or enable_serving_reranking) else None

381 + ) if ("embed" in supported_tasks or enable_serving_reranking) else None

382

383 tracker = BatchProgressTracker()

384 logger.info("Reading batch from %s...", args.input_file)

vllm/executor/executor_base.py

+4 -4

@@ -16,8 +16,8 @@ from vllm.config import VllmConfig

16 from vllm.logger import init_logger

17 from vllm.lora.request import LoRARequest

18 from vllm.model_executor.layers.sampler import SamplerOutput

19 -from vllm.pooling_params import PoolingTask

19 from vllm.sequence import ExecuteModelRequest, PoolerOutput

20 +from vllm.tasks import SupportedTask

21 from vllm.utils import make_async

22 from vllm.worker.worker_base import WorkerBase

@@ -136,9 +136,9 @@ class ExecutorBase(ABC):

136 return self.collective_rpc(rpc_func)

137

138 @cached_property # Avoid unnecessary RPC calls

139 - def supported_pooling_tasks(self) -> tuple[PoolingTask, ...]:

140 - output = self.collective_rpc("get_supported_pooling_tasks")

141 - return tuple({task for tasks in output for task in tasks})

139 + def supported_tasks(self) -> tuple[SupportedTask, ...]:

140 + output = self.collective_rpc("get_supported_tasks")

141 + return output[0]

142

143 def execute_model(

144 self, execute_model_req: ExecuteModelRequest

vllm/model_executor/layers/pooler.py

+2 -1

@@ -16,8 +16,9 @@ from vllm.config import ModelConfig, PoolerConfig

16 from vllm.model_executor.pooling_metadata import ( # noqa: E501

17 PoolingMetadata as V0PoolingMetadata)

18 from vllm.model_executor.pooling_metadata import PoolingTensors

19 -from vllm.pooling_params import PoolingParams, PoolingTask

19 +from vllm.pooling_params import PoolingParams

20 from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput

21 +from vllm.tasks import PoolingTask

22 from vllm.utils import resolve_obj_by_qualname

23 from vllm.v1.pool.metadata import PoolingMetadata as V1PoolingMetadata

vllm/model_executor/models/bert.py

+1 -1

@@ -26,8 +26,8 @@ from vllm.model_executor.layers.quantization import QuantizationConfig

26 from vllm.model_executor.layers.vocab_parallel_embedding import (

27 VocabParallelEmbedding)

28 from vllm.model_executor.pooling_metadata import PoolingMetadata

29 -from vllm.pooling_params import PoolingTask

29 from vllm.sequence import IntermediateTensors

30 +from vllm.tasks import PoolingTask

32 from .interfaces import SupportsCrossEncoding, SupportsQuant, SupportsV0Only

33 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix

vllm/model_executor/models/gritlm.py

+1 -1

@@ -16,8 +16,8 @@ from vllm.model_executor.layers.pooler import (DispatchPooler, Pooler,

16 get_prompt_token_ids)

17 from vllm.model_executor.models.llama import LlamaForCausalLM

18 from vllm.model_executor.pooling_metadata import PoolingMetadata

19 -from vllm.pooling_params import PoolingTask

19 from vllm.sequence import PoolerOutput

20 +from vllm.tasks import PoolingTask

21 from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config

23 from .interfaces import SupportsV0Only

vllm/model_executor/models/modernbert.py

+1 -1

@@ -23,8 +23,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

23 VocabParallelEmbedding)

24 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

25 from vllm.model_executor.pooling_metadata import PoolingMetadata

26 -from vllm.pooling_params import PoolingTask

26 from vllm.sequence import IntermediateTensors

27 +from vllm.tasks import PoolingTask

29 from .interfaces import SupportsCrossEncoding, SupportsV0Only

30 from .utils import WeightsMapper, maybe_prefix

vllm/pooling_params.py

+2 -3

@@ -1,17 +1,16 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from typing import TYPE_CHECKING, Literal, Optional

4 +from typing import TYPE_CHECKING, Optional

6 import msgspec

8 from vllm.sampling_params import RequestOutputKind

9 +from vllm.tasks import PoolingTask

11 if TYPE_CHECKING:

12 from vllm.config import ModelConfig

13 -PoolingTask = Literal["encode", "embed", "classify", "score"]

14 -

15 class PoolingParams(

16 msgspec.Struct,

vllm/tasks.py new

+11

@@ -0,0 +1,11 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import Literal, get_args

4 +

5 +GenerationTask = Literal["generate", "transcription"]

6 +GENERATION_TASKS = get_args(GenerationTask)

7 +

8 +PoolingTask = Literal["encode", "embed", "classify", "score"]

9 +POOLING_TASKS = get_args(PoolingTask)

10 +

11 +SupportedTask = Literal[GenerationTask, PoolingTask]

vllm/v1/engine/async_llm.py

@@ -21,6 +21,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

21 from vllm.outputs import PoolingRequestOutput, RequestOutput

22 from vllm.pooling_params import PoolingParams

23 from vllm.sampling_params import SamplingParams

24 +from vllm.tasks import SupportedTask

25 from vllm.transformers_utils.config import (

26 maybe_register_config_serialize_by_value)

27 from vllm.transformers_utils.tokenizer import AnyTokenizer

@@ -211,6 +212,9 @@ class AsyncLLM(EngineClient):

212 if handler := getattr(self, "output_handler", None):

213 handler.cancel()

214

215 + async def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

216 + return await self.engine_core.get_supported_tasks_async()

217 +

218 async def add_request(

219 self,

220 request_id: str,

vllm/v1/engine/core.py

+9 -2

@@ -23,6 +23,7 @@ from vllm.executor.multiproc_worker_utils import _add_prefix

23 from vllm.logger import init_logger

24 from vllm.logging_utils.dump_input import dump_engine_exception

25 from vllm.lora.request import LoRARequest

26 +from vllm.tasks import POOLING_TASKS, SupportedTask

27 from vllm.transformers_utils.config import (

28 maybe_register_config_serialize_by_value)

29 from vllm.utils import (bind_process_name, make_zmq_socket,

@@ -195,11 +196,17 @@ class EngineCore:

196 "warmup model) took %.2f seconds"), elapsed)

197 return num_gpu_blocks, num_cpu_blocks, scheduler_kv_cache_config

198

199 + def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

200 + return self.model_executor.supported_tasks

201 +

202 def add_request(self, request: EngineCoreRequest):

203 """Add request to the scheduler."""

204 if pooling_params := request.pooling_params:

201 - supported_pooling_tasks = (

202 - self.model_executor.supported_pooling_tasks)

205 + supported_pooling_tasks = [

206 + task for task in self.get_supported_tasks()

207 + if task in POOLING_TASKS

208 + ]

209 +

210 if pooling_params.task not in supported_pooling_tasks:

211 raise ValueError(f"Unsupported task: {pooling_params.task!r} "

212 f"Supported tasks: {supported_pooling_tasks}")

vllm/v1/engine/core_client.py

+16

@@ -21,6 +21,7 @@ import zmq.asyncio

21 from vllm.config import VllmConfig

22 from vllm.logger import init_logger

23 from vllm.lora.request import LoRARequest

24 +from vllm.tasks import SupportedTask

25 from vllm.utils import get_open_port, get_open_zmq_inproc_path, make_zmq_socket

26 from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,

27 EngineCoreRequestType,

@@ -104,6 +105,9 @@ class EngineCoreClient(ABC):

105 def get_output(self) -> EngineCoreOutputs:

106 raise NotImplementedError

107

108 + def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

109 + raise NotImplementedError

110 +

111 def add_request(self, request: EngineCoreRequest) -> None:

112 raise NotImplementedError

113

@@ -170,6 +174,9 @@ class EngineCoreClient(ABC):

174 async def get_output_async(self) -> EngineCoreOutputs:

175 raise NotImplementedError

176

177 + async def get_supported_tasks_async(self) -> tuple[SupportedTask, ...]:

178 + raise NotImplementedError

179 +

180 async def add_request_async(self, request: EngineCoreRequest) -> None:

181 raise NotImplementedError

182

@@ -238,6 +245,9 @@ class InprocClient(EngineCoreClient):

245 outputs, _ = self.engine_core.step()

246 return outputs.get(0) or EngineCoreOutputs()

247

248 + def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

249 + return self.engine_core.get_supported_tasks()

250 +

251 def add_request(self, request: EngineCoreRequest) -> None:

252 self.engine_core.add_request(request)

253

@@ -608,6 +618,9 @@ class SyncMPClient(MPClient):

618

619 return future.result()

620

621 + def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

622 + return self.call_utility("get_supported_tasks")

623 +

624 def add_request(self, request: EngineCoreRequest) -> None:

625 if self.is_dp:

626 self.engines_running = True

@@ -802,6 +815,9 @@ class AsyncMPClient(MPClient):

815 self._ensure_output_queue_task()

816 return await future

817

818 + async def get_supported_tasks_async(self) -> tuple[SupportedTask, ...]:

819 + return await self.call_utility_async("get_supported_tasks")

820 +

821 async def add_request_async(self, request: EngineCoreRequest) -> None:

822 request.client_index = self.client_index

823 await self._send_input(EngineCoreRequestType.ADD, request)

vllm/v1/engine/llm_engine.py

@@ -18,6 +18,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

18 from vllm.outputs import PoolingRequestOutput, RequestOutput

19 from vllm.pooling_params import PoolingParams

20 from vllm.sampling_params import SamplingParams

21 +from vllm.tasks import SupportedTask

22 from vllm.transformers_utils.tokenizer_group import (

23 TokenizerGroup, init_tokenizer_from_configs)

24 from vllm.usage.usage_lib import UsageContext

@@ -176,6 +177,9 @@ class LLMEngine:

177 def validate_outputs(cls, outputs, output_type):

178 return outputs

179

180 + def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

181 + return self.engine_core.get_supported_tasks()

182 +

183 def abort_request(self, request_ids: list[str]) -> None:

184 """Remove request_ids from EngineCore and Detokenizer."""

185

vllm/v1/worker/gpu_model_runner.py

+31 -4

@@ -30,15 +30,17 @@ from vllm.logger import init_logger

30 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaBase

31 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding

32 from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader

33 -from vllm.model_executor.models.interfaces import is_mixture_of_experts

34 -from vllm.model_executor.models.interfaces_base import (VllmModelForPooling,

35 - is_pooling_model)

33 +from vllm.model_executor.models.interfaces import (is_mixture_of_experts,

34 + supports_transcription)

35 +from vllm.model_executor.models.interfaces_base import (

36 + VllmModelForPooling, is_pooling_model, is_text_generation_model)

37 from vllm.multimodal import MULTIMODAL_REGISTRY

38 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange

39 from vllm.multimodal.utils import group_mm_inputs_by_modality

39 -from vllm.pooling_params import PoolingParams, PoolingTask

40 +from vllm.pooling_params import PoolingParams

41 from vllm.sampling_params import SamplingType

42 from vllm.sequence import IntermediateTensors, PoolerOutput

43 +from vllm.tasks import GenerationTask, PoolingTask, SupportedTask

44 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,

45 GiB_bytes, LazyLoader, check_use_alibi, get_dtype_size,

46 is_pin_memory_available, round_up)

@@ -1153,6 +1155,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1155 def get_model(self) -> nn.Module:

1156 return self.model

1157

1158 + def get_supported_generation_tasks(self) -> list[GenerationTask]:

1159 + model = self.get_model()

1160 + supported_tasks = list[GenerationTask]()

1161 +

1162 + if is_text_generation_model(model):

1163 + supported_tasks.append("generate")

1164 +

1165 + if supports_transcription(model):

1166 + if model.supports_transcription_only:

1167 + return ["transcription"]

1168 +

1169 + supported_tasks.append("transcription")

1170 +

1171 + return supported_tasks

1172 +

1173 def get_supported_pooling_tasks(self) -> list[PoolingTask]:

1174 model = self.get_model()

1175 if not is_pooling_model(model):

@@ -1160,6 +1177,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1177

1178 return list(model.pooler.get_supported_tasks())

1179

1180 + def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

1181 + tasks = list[SupportedTask]()

1182 +

1183 + if self.model_config.runner_type == "generate":

1184 + tasks.extend(self.get_supported_generation_tasks())

1185 + if self.model_config.runner_type == "pooling":

1186 + tasks.extend(self.get_supported_pooling_tasks())

1187 +

1188 + return tuple(tasks)

1189 +

1190 def apply_grammar_bitmask(

1191 self,

1192 scheduler_output: "SchedulerOutput",

vllm/v1/worker/gpu_worker.py

+3 -3

@@ -23,8 +23,8 @@ from vllm.logger import init_logger

23 from vllm.lora.request import LoRARequest

24 from vllm.model_executor import set_random_seed

25 from vllm.platforms import current_platform

26 -from vllm.pooling_params import PoolingTask

26 from vllm.sequence import IntermediateTensors

27 +from vllm.tasks import SupportedTask

28 from vllm.utils import GiB_bytes, MemorySnapshot, memory_profiling

29 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType

30 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec

@@ -320,8 +320,8 @@ class Worker(WorkerBase):

320 def get_model(self) -> nn.Module:

321 return self.model_runner.get_model()

322

323 - def get_supported_pooling_tasks(self) -> list[PoolingTask]:

324 - return self.model_runner.get_supported_pooling_tasks()

323 + def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

324 + return self.model_runner.get_supported_tasks()

325

326 @torch.inference_mode()

327 def execute_model(

vllm/v1/worker/tpu_model_runner.py

+29 -2

@@ -27,13 +27,15 @@ from vllm.logger import init_logger

27 from vllm.lora.layers import BaseLayerWithLoRA

28 from vllm.model_executor.model_loader import get_model_loader

29 from vllm.model_executor.model_loader.tpu import TPUModelLoader

30 -from vllm.model_executor.models.interfaces_base import is_pooling_model

30 +from vllm.model_executor.models.interfaces import supports_transcription

31 +from vllm.model_executor.models.interfaces_base import (

32 + is_pooling_model, is_text_generation_model)

33 from vllm.multimodal import MULTIMODAL_REGISTRY

34 from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargs,

35 PlaceholderRange)

36 from vllm.multimodal.utils import group_mm_inputs_by_modality

35 -from vllm.pooling_params import PoolingTask

37 from vllm.sequence import IntermediateTensors

38 +from vllm.tasks import GenerationTask, PoolingTask, SupportedTask

39 from vllm.utils import (LayerBlockType, cdiv, is_pin_memory_available,

40 prev_power_of_2)

41 from vllm.v1.attention.backends.pallas import (TPU_STR_DTYPE_TO_TORCH_DTYPE,

@@ -489,6 +491,21 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

491 def get_model(self) -> nn.Module:

492 return self.model

493

494 + def get_supported_generation_tasks(self) -> list[GenerationTask]:

495 + model = self.get_model()

496 + supported_tasks = list[GenerationTask]()

497 +

498 + if is_text_generation_model(model):

499 + supported_tasks.append("generate")

500 +

501 + if supports_transcription(model):

502 + if model.supports_transcription_only:

503 + return ["transcription"]

504 +

505 + supported_tasks.append("transcription")

506 +

507 + return supported_tasks

508 +

509 def get_supported_pooling_tasks(self) -> list[PoolingTask]:

510 model = self.get_model()

511 if not is_pooling_model(model):

@@ -496,6 +513,16 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

513

514 return list(model.pooler.get_supported_tasks())

515

516 + def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

517 + tasks = list[SupportedTask]()

518 +

519 + if self.model_config.runner_type == "generate":

520 + tasks.extend(self.get_supported_generation_tasks())

521 + if self.model_config.runner_type == "pooling":

522 + tasks.extend(self.get_supported_pooling_tasks())

523 +

524 + return tuple(tasks)

525 +

526 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:

527 """

528 Generates the KVCacheSpec by parsing the kv cache format from each

vllm/v1/worker/tpu_worker.py

+3 -3

@@ -21,7 +21,7 @@ from vllm.logger import init_logger

21 from vllm.lora.request import LoRARequest

22 from vllm.model_executor import set_random_seed

23 from vllm.platforms import current_platform

24 -from vllm.pooling_params import PoolingTask

24 +from vllm.tasks import SupportedTask

25 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv

26 from vllm.v1.attention.backends.pallas import TPU_HEAD_SIZE_ALIGNMENT

27 from vllm.v1.core.sched.output import SchedulerOutput

@@ -282,8 +282,8 @@ class TPUWorker:

282 def get_model(self) -> nn.Module:

283 return self.model_runner.get_model()

284

285 - def get_supported_pooling_tasks(self) -> list[PoolingTask]:

286 - return self.model_runner.get_supported_pooling_tasks()

285 + def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

286 + return self.model_runner.get_supported_tasks()

287

288 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:

289 return self.model_runner.get_kv_cache_spec()

vllm/worker/model_runner_base.py

+29 -2

@@ -12,9 +12,11 @@ import torch.nn as nn

12 from vllm.config import VllmConfig

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.sampler import SamplerOutput

15 -from vllm.model_executor.models.interfaces_base import is_pooling_model

16 -from vllm.pooling_params import PoolingTask

15 +from vllm.model_executor.models.interfaces import supports_transcription

16 +from vllm.model_executor.models.interfaces_base import (

17 + is_pooling_model, is_text_generation_model)

18 from vllm.sequence import IntermediateTensors, SequenceGroupMetadata

19 +from vllm.tasks import GenerationTask, PoolingTask, SupportedTask

21 if TYPE_CHECKING:

22 from vllm.attention import AttentionMetadata

@@ -224,6 +226,21 @@ class ModelRunnerBase(ABC, Generic[T]):

226 def get_model(self) -> nn.Module:

227 raise NotImplementedError

228

229 + def get_supported_generation_tasks(self) -> list[GenerationTask]:

230 + model = self.get_model()

231 + supported_tasks = list[GenerationTask]()

232 +

233 + if is_text_generation_model(model):

234 + supported_tasks.append("generate")

235 +

236 + if supports_transcription(model):

237 + if model.supports_transcription_only:

238 + return ["transcription"]

239 +

240 + supported_tasks.append("transcription")

241 +

242 + return supported_tasks

243 +

244 def get_supported_pooling_tasks(self) -> list[PoolingTask]:

245 model = self.get_model()

246 if not is_pooling_model(model):

@@ -231,6 +248,16 @@ class ModelRunnerBase(ABC, Generic[T]):

248

249 return list(model.pooler.get_supported_tasks())

250

251 + def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

252 + tasks = list[SupportedTask]()

253 +

254 + if self.model_config.runner_type == "generate":

255 + tasks.extend(self.get_supported_generation_tasks())

256 + if self.model_config.runner_type == "pooling":

257 + tasks.extend(self.get_supported_pooling_tasks())

258 +

259 + return tuple(tasks)

260 +

261 def execute_model(

262 self,

263 model_input: T,