chore: import upstream snapshot with attribution

swift/infer_engine/__init__.py

@@ -0,0 +1,39 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import TYPE_CHECKING
+
+from swift.utils.import_utils import _LazyModule
+
+if TYPE_CHECKING:
+    from .base import BaseInferEngine
+    from .grpo_vllm_engine import GRPOVllmEngine
+    from .infer_client import InferClient
+    from .infer_engine import InferEngine
+    from .lmdeploy_engine import LmdeployEngine
+    from .protocol import ChatCompletionResponse, Function, InferRequest, RequestConfig
+    from .sglang_engine import SglangEngine
+    from .transformers_engine import TransformersEngine
+    from .utils import AdapterRequest, patch_vllm_memory_leak, prepare_generation_config
+    from .vllm_engine import VllmEngine
+else:
+    _import_structure = {
+        'vllm_engine': ['VllmEngine'],
+        'grpo_vllm_engine': ['GRPOVllmEngine'],
+        'lmdeploy_engine': ['LmdeployEngine'],
+        'sglang_engine': ['SglangEngine'],
+        'transformers_engine': ['TransformersEngine'],
+        'infer_client': ['InferClient'],
+        'infer_engine': ['InferEngine'],
+        'base': ['BaseInferEngine'],
+        'utils': ['prepare_generation_config', 'AdapterRequest', 'patch_vllm_memory_leak'],
+        'protocol': ['InferRequest', 'RequestConfig', 'Function', 'ChatCompletionResponse'],
+    }
+
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

swift/infer_engine/base.py

@@ -0,0 +1,58 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from abc import ABC, abstractmethod
+from typing import AsyncIterator, Iterator, List, Optional, Union
+
+from swift.metrics import Metric
+from .protocol import ChatCompletionResponse, ChatCompletionStreamResponse, InferRequest, RequestConfig
+
+
+class BaseInferEngine(ABC):
+
+    @abstractmethod
+    def infer(self,
+              infer_requests: List[InferRequest],
+              request_config: Optional[RequestConfig] = None,
+              metrics: Optional[List[Metric]] = None,
+              *,
+              use_tqdm: Optional[bool] = None,
+              **kwargs) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
+        """
+        This method performs inference on a list of inference requests.
+
+        The method takes a list of inference requests and processes them according to the provided configuration.
+        It can optionally use tqdm for progress visualization and accept additional keyword arguments.
+
+        Args:
+            infer_requests (List[InferRequest]): A list of inference requests to be processed.
+            request_config (Optional[RequestConfig]): Configuration for the request, if any.
+            metrics (Optional[List[Metric]]): A list of usage information to return.
+            use_tqdm (Optional[bool]): Whether to use tqdm for progress visualization.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
+                The result of the inference.
+        """
+        pass
+
+    @abstractmethod
+    async def infer_async(self,
+                          infer_request: InferRequest,
+                          request_config: Optional[RequestConfig] = None,
+                          **kwargs) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]:
+        """
+        This method performs asynchronous inference on a single inference request.
+
+        The method takes an inference request and processes it according to the provided configuration.
+        It can accept additional keyword arguments.
+
+        Args:
+            infer_request (InferRequest): An inference request to be processed.
+            request_config (Optional[RequestConfig]): Configuration for the request, if any.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]: The result of
+                the asynchronous inference.
+        """
+        pass

swift/infer_engine/grpo_vllm_engine.py

@@ -0,0 +1,147 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+import torch
+from PIL import Image
+from tqdm.asyncio import tqdm_asyncio
+from typing import Any, Dict, List, Optional, Union
+
+from swift.metrics import Metric
+from swift.rlhf_trainers.utils import VLLM_LORA_INT_ID, VLLM_LORA_NAME, VLLM_LORA_PATH
+from swift.template import Template
+from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatMessage, InferRequest, RequestConfig,
+                       RolloutOutput)
+from .utils import AdapterRequest
+from .vllm_engine import VllmEngine
+
+try:
+    os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'
+    os.environ['VLLM_ENGINE_ITERATION_TIMEOUT_S'] = '86400'
+    from vllm.lora.request import LoRARequest
+except Exception:
+    raise
+
+
+class GRPOVllmEngine(VllmEngine):
+
+    def infer(
+        self,
+        infer_requests: List[Union[InferRequest, Dict[str, Any]]],
+        request_config: Optional[RequestConfig] = None,
+        metrics: Optional[List[Metric]] = None,
+        *,
+        use_tqdm: Optional[bool] = None,
+        adapter_request: Optional[AdapterRequest] = None,
+    ) -> List[RolloutOutput]:
+        if not adapter_request and self.enable_lora:
+            lora_loaded = VLLM_LORA_INT_ID in self.engine.list_loras()
+            if lora_loaded:
+                adapter_request = LoRARequest(
+                    lora_name=VLLM_LORA_NAME, lora_int_id=VLLM_LORA_INT_ID, lora_path=VLLM_LORA_PATH)
+
+        res = super().infer(
+            infer_requests,
+            request_config,
+            metrics,
+            use_tqdm=use_tqdm,
+            adapter_request=adapter_request,
+        )
+        if not isinstance(res, list):
+            res = [res]
+        for i, result in enumerate(res):
+            if not isinstance(result, RolloutOutput):
+                if not isinstance(result, ChatCompletionResponse):
+                    raise TypeError('Result must be a ChatCompletionResponse or RolloutOutput instance.')
+                res[i] = RolloutOutput(response=result)
+
+        return res
+
+    async def async_infer(self,
+                          infer_requests: List[InferRequest],
+                          request_config: Optional[RequestConfig] = None,
+                          metrics: Optional[List[Metric]] = None,
+                          *,
+                          use_tqdm: Optional[bool] = None,
+                          **kwargs) -> List[RolloutOutput]:
+        if request_config is None:
+            request_config = RequestConfig()
+        assert request_config.n == 1
+
+        tasks = [self.infer_async(infer_request, request_config, **kwargs) for infer_request in infer_requests]
+        if use_tqdm is None:
+            use_tqdm = len(infer_requests) > 1
+        res = await self._batch_infer_stream(tasks, request_config.stream, use_tqdm, metrics)
+
+        for i, result in enumerate(res):
+            if not isinstance(result, RolloutOutput):
+                if not isinstance(result, ChatCompletionResponse):
+                    raise TypeError('Result must be a ChatCompletionResponse or RolloutOutput instance.')
+                res[i] = RolloutOutput(response=result)
+
+        return res
+
+    async def _batch_infer_stream(self,
+                                  tasks,
+                                  stream: bool = True,
+                                  use_tqdm: bool = True,
+                                  metrics: Optional[List[Metric]] = None):
+        assert not stream
+        prog_bar = tqdm_asyncio(total=len(tasks), dynamic_ncols=True, disable=not use_tqdm)
+
+        async def _new_run(task):
+            try:
+                res = await task
+            except Exception as e:
+                if getattr(self, 'strict', True):
+                    raise
+                res = e
+            prog_bar.update()
+            self._update_metrics(res, metrics)
+            return res
+
+        new_tasks = [_new_run(task) for task in tasks]
+        return await self.batch_run(new_tasks)
+
+    def _create_chat_completion_response(self, result, inputs, request_config, request_id) -> ChatCompletionResponse:
+        assert result is not None
+        num_generated_tokens = sum(len(output.token_ids) for output in result.outputs)
+        usage_info = self._get_usage_info(len(result.prompt_token_ids), num_generated_tokens)
+        choices = []
+        for output in result.outputs:
+            output.token_ids = list(output.token_ids)
+            response = self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs'])
+            logprobs = self._get_logprobs(output.logprobs, output.token_ids, request_config.top_logprobs)
+            toolcall = self._get_toolcall(response)
+
+            token_ids = output.token_ids if request_config.return_details else None
+            choice = ChatCompletionResponseChoice(
+                index=output.index,
+                message=ChatMessage(role='assistant', content=response, tool_calls=toolcall),
+                finish_reason=output.finish_reason,
+                logprobs=logprobs,
+                token_ids=token_ids,
+                routed_experts=getattr(output, 'routed_experts', None))
+            choices.append(choice)
+        prompt_token_ids = None
+        images_size = None
+        if request_config.return_details:
+            prompt_token_ids = result.prompt_token_ids
+            images = inputs['template_inputs'].images
+            if all(isinstance(image, Image.Image) for image in images):
+                images_size = [image.size for image in images]
+        return ChatCompletionResponse(
+            model=self.model_name,
+            choices=choices,
+            usage=usage_info,
+            id=request_id,
+            prompt_token_ids=prompt_token_ids,
+            images_size=images_size)
+
+    def _add_adapter(self, adapter_request: Optional[Union[AdapterRequest, LoRARequest]] = None):
+        assert self.enable_lora, f'adapter_request: {adapter_request}, self.enable_lora: {self.enable_lora}'
+        from vllm.lora.request import LoRARequest
+        if isinstance(adapter_request, AdapterRequest):
+            return super()._add_adapter(adapter_request)
+        elif isinstance(adapter_request, LoRARequest):
+            return adapter_request
+        else:
+            raise ValueError(f'Invalid adapter request: {adapter_request}')

swift/infer_engine/infer_client.py

@@ -0,0 +1,158 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import aiohttp
+import json
+from copy import deepcopy
+from dacite import from_dict
+from dataclasses import asdict
+from requests.exceptions import HTTPError
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
+
+from swift.metrics import Metric
+from .infer_engine import InferEngine
+from .protocol import (ChatCompletionRequest, ChatCompletionResponse, ChatCompletionStreamResponse, InferRequest,
+                       ModelList, RequestConfig)
+
+
+class InferClient(InferEngine):
+
+    def __init__(self,
+                 host: str = '127.0.0.1',
+                 port: int = 8000,
+                 api_key: str = 'EMPTY',
+                 *,
+                 base_url: Optional[str] = None,
+                 timeout: Optional[int] = 86400) -> None:
+        """
+        Initialize the InferClient.
+
+        Args:
+            host (str): The hostname of the inference server. Defaults to '127.0.0.1'.
+            port (str): The port of the inference server. Defaults to '8000'.
+            api_key (str): The API key for authentication. Defaults to 'EMPTY'.
+            timeout (Optional[int]): The timeout for requests in seconds. Defaults to None.
+        """
+        self.api_key = api_key
+        self.host = host
+        self.port = port
+        self.timeout = timeout
+        if base_url is None:
+            base_url = f'http://{self.host}:{self.port}/v1'
+        self.base_url = base_url
+        self._models = None
+
+    @property
+    def models(self):
+        if self._models is None:
+            models = []
+            for model in self.get_model_list().data:
+                models.append(model.id)
+            assert len(models) > 0, f'models: {models}'
+            self._models = models
+        return self._models
+
+    def get_model_list(self) -> ModelList:
+        """Get model list from the inference server.
+        """
+        coro = self.get_model_list_async()
+        return self.safe_asyncio_run(coro)
+
+    def _get_request_kwargs(self) -> Dict[str, Any]:
+        request_kwargs = {}
+        if isinstance(self.timeout, (int, float)) and self.timeout > 0:
+            request_kwargs['timeout'] = self.timeout
+        if self.api_key is not None:
+            request_kwargs['headers'] = {'Authorization': f'Bearer {self.api_key}'}
+        return request_kwargs
+
+    async def get_model_list_async(self) -> ModelList:
+        url = f"{self.base_url.rstrip('/')}/models"
+        async with aiohttp.ClientSession() as session:
+            async with session.get(url, **self._get_request_kwargs()) as resp:
+                resp_obj = await resp.json()
+        return from_dict(ModelList, resp_obj)
+
+    def infer(
+            self,
+            infer_requests: List[InferRequest],
+            request_config: Optional[RequestConfig] = None,
+            metrics: Optional[List[Metric]] = None,
+            *,
+            model: Optional[str] = None,
+            use_tqdm: Optional[bool] = None
+    ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
+        """
+        Perform inference using the specified model.
+
+        Args:
+            infer_requests (List[InferRequest]): A list of inference requests.
+            request_config (Optional[RequestConfig]): Configuration for the request. Defaults to None.
+            metrics (Optional[List[Metric]]): The usage information to return. Defaults to None.
+            model (Optional[str]): The model name to be used for inference. Defaults to None.
+            use_tqdm (Optional[bool]): Whether to use tqdm for progress tracking. Defaults to None.
+
+        Returns:
+            List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
+            The inference responses or an iterator of streaming responses.
+        """
+        return super().infer(infer_requests, request_config, metrics, model=model, use_tqdm=use_tqdm)
+
+    @staticmethod
+    def _prepare_request_data(model: str, infer_request: InferRequest, request_config: RequestConfig) -> Dict[str, Any]:
+        if not isinstance(infer_request, dict):
+            infer_request = asdict(infer_request)
+        res = asdict(ChatCompletionRequest(model, **infer_request, **asdict(request_config)))
+        # ignore empty
+        empty_request = ChatCompletionRequest('', [])
+        for k in list(res.keys()):
+            if res[k] == getattr(empty_request, k):
+                res.pop(k)
+        return res
+
+    @staticmethod
+    def _parse_stream_data(data: bytes) -> Optional[str]:
+        data = data.decode(encoding='utf-8')
+        data = data.strip()
+        if len(data) == 0:
+            return
+        assert data.startswith('data:'), f'data: {data}'
+        return data[5:].strip()
+
+    async def infer_async(
+        self,
+        infer_request: InferRequest,
+        request_config: Optional[RequestConfig] = None,
+        *,
+        model: Optional[str] = None,
+    ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]:
+        request_config = deepcopy(request_config or RequestConfig())
+        if model is None:
+            if len(self.models) == 1:
+                model = self.models[0]
+            else:
+                raise ValueError(f'Please explicitly specify the model. Available models: {self.models}.')
+        url = f"{self.base_url.rstrip('/')}/chat/completions"
+
+        request_data = self._prepare_request_data(model, infer_request, request_config)
+        if request_config.stream:
+
+            async def _gen_stream() -> AsyncIterator[ChatCompletionStreamResponse]:
+                async with aiohttp.ClientSession() as session:
+                    async with session.post(url, json=request_data, **self._get_request_kwargs()) as resp:
+                        async for data in resp.content:
+                            data = self._parse_stream_data(data)
+                            if data == '[DONE]':
+                                break
+                            if data is not None:
+                                resp_obj = json.loads(data)
+                                if resp_obj['object'] == 'error':
+                                    raise HTTPError(resp_obj['message'])
+                                yield from_dict(ChatCompletionStreamResponse, resp_obj)
+
+            return _gen_stream()
+        else:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(url, json=request_data, **self._get_request_kwargs()) as resp:
+                    resp_obj = await resp.json()
+                    if resp_obj['object'] == 'error':
+                        raise HTTPError(resp_obj['message'])
+                    return from_dict(ChatCompletionResponse, resp_obj)

swift/infer_engine/infer_engine.py

@@ -0,0 +1,314 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import asyncio
+import concurrent.futures
+import os
+from queue import Queue
+from threading import Thread
+from tqdm import tqdm
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+from swift.metrics import Metric
+from swift.model import get_ckpt_dir
+from swift.template import Template, get_template
+from swift.utils import Processor, ProcessorMixin, get_logger
+from .base import BaseInferEngine
+from .protocol import (ChatCompletionMessageToolCall, ChatCompletionResponse, ChatCompletionStreamResponse,
+                       InferRequest, RequestConfig, UsageInfo)
+
+logger = get_logger()
+
+
+class InferEngine(BaseInferEngine, ProcessorMixin):
+
+    def __init__(self, template: Template):
+        processor = template.processor
+        self.template = template
+        self.template_type = template.template_meta.template_type
+        self.processor = processor
+        self.model_info = processor.model_info
+        self.model_meta = processor.model_meta
+        self.model_dir = self.model_info.model_dir
+        self.model_name = self.model_info.model_name
+        self.max_model_len = self.model_info.max_model_len
+        self.task_type = self.model_info.task_type
+        self.config = self.model_info.config
+        self.max_tokens_offset = 0
+
+    def _get_template(self, processor: Processor, template_type: Optional[str] = None):
+        ckpt_dir = get_ckpt_dir(processor.model_info.model_dir, getattr(self, 'adapters', None))
+        logger.info('Create the template for the infer_engine')
+        if ckpt_dir:
+            from swift.arguments import BaseArguments
+            args = BaseArguments.from_pretrained(ckpt_dir)
+            template = args.get_template(processor)
+        else:
+            template = get_template(processor, template_type=template_type)
+        return template
+
+    def _get_stop_words(self, stop_words: List[Union[str, List[int], None]]) -> List[str]:
+        stop: List[str] = []
+        for stop_word in stop_words:
+            if stop_word is None:
+                continue
+            elif isinstance(stop_word, list):
+                stop_word = self.tokenizer.decode(stop_word)
+            assert isinstance(stop_word, str)
+            if stop_word not in stop:
+                stop.append(stop_word)
+        return stop
+
+    def _get_stop_token_ids(self, stop_words: List[Union[str, List[int], None]]) -> List[int]:
+        stop_token_ids: List[int] = []
+        for stop_word in stop_words:
+            if stop_word is None:
+                continue
+            if isinstance(stop_word, str):
+                stop_word = self.tokenizer.encode(stop_word, add_special_tokens=False)
+            if isinstance(stop_word, list):
+                if len(stop_word) != 1:
+                    continue
+                else:
+                    stop_token = stop_word[0]
+            elif isinstance(stop_word, int):
+                stop_token = stop_word
+            assert isinstance(stop_token, int)
+            if stop_token not in stop_token_ids:
+                stop_token_ids.append(stop_token)
+        return stop_token_ids
+
+    def async_iter_to_iter(self, async_iter, prog_bar, metrics) -> Iterator:
+        queue = Queue()
+
+        async def _run_async_iter():
+            try:
+                async for item in await async_iter:
+                    queue.put(item)
+            except Exception as e:
+                if getattr(self, 'strict', True):
+                    raise
+                queue.put(e)
+            else:
+                queue.put(None)
+
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+        thread = Thread(target=lambda: loop.run_until_complete(_run_async_iter()))
+        thread.start()
+        pre_output = None
+        while True:
+            output = queue.get()
+            if output is None or isinstance(output, Exception):
+                prog_bar.update()
+                self._update_metrics(pre_output, metrics)
+                return
+            pre_output = output
+            yield output
+
+    @staticmethod
+    async def batch_run(tasks):
+        return await asyncio.gather(*tasks)
+
+    def _batch_infer_stream(
+        self,
+        tasks,
+        stream: bool = True,
+        use_tqdm: bool = True,
+        metrics: Optional[List[Metric]] = None
+    ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
+
+        prog_bar = tqdm(total=len(tasks), dynamic_ncols=True, disable=not use_tqdm)
+        if stream:
+            return [self.async_iter_to_iter(task, prog_bar, metrics) for task in tasks]
+        else:
+
+            async def _new_run(task):
+                try:
+                    res = await task
+                except Exception as e:
+                    if getattr(self, 'strict', True):
+                        raise
+                    res = e
+                prog_bar.update()
+                self._update_metrics(res, metrics)
+                return res
+
+            new_tasks = [_new_run(task) for task in tasks]
+            try:
+                loop = asyncio.get_event_loop()
+            except RuntimeError:
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+            return loop.run_until_complete(self.batch_run(new_tasks))
+
+    @staticmethod
+    def _get_usage_info(num_prompt_tokens: int, num_generated_tokens: int) -> UsageInfo:
+        return UsageInfo(
+            prompt_tokens=num_prompt_tokens,
+            completion_tokens=num_generated_tokens,
+            total_tokens=num_prompt_tokens + num_generated_tokens,
+        )
+
+    @staticmethod
+    def _update_usage_info(origin_use_info: UsageInfo, num_generated_tokens: int) -> UsageInfo:
+        return UsageInfo(
+            prompt_tokens=origin_use_info.prompt_tokens,
+            completion_tokens=origin_use_info.completion_tokens + num_generated_tokens,
+            total_tokens=origin_use_info.total_tokens + num_generated_tokens,
+        )
+
+    @staticmethod
+    def _update_metrics(result, metrics: Optional[List[Metric]] = None):
+        if metrics is None:
+            return result
+        result_origin = result
+        if not isinstance(result, (list, tuple)):
+            result = [result]
+        for response in result:
+            if response is None or isinstance(response, Exception):
+                continue
+            for metric in metrics:
+                metric.update(response)
+        return result_origin
+
+    def infer(self,
+              infer_requests: List[InferRequest],
+              request_config: Optional[RequestConfig] = None,
+              metrics: Optional[List[Metric]] = None,
+              *,
+              use_tqdm: Optional[bool] = None,
+              **kwargs) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
+        if request_config is None:
+            request_config = RequestConfig()
+        tasks = [self.infer_async(infer_request, request_config, **kwargs) for infer_request in infer_requests]
+        if use_tqdm is None:
+            use_tqdm = not request_config.stream and len(infer_requests) > 1
+        return self._batch_infer_stream(tasks, request_config.stream, use_tqdm, metrics)
+
+    def _get_toolcall(self, response: str) -> Optional[List[ChatCompletionMessageToolCall]]:
+        try:
+            functions = self.template.agent_template.get_toolcall(response)
+        except Exception:
+            functions = None
+        if functions:
+            return [ChatCompletionMessageToolCall(function=function) for function in functions]
+
+    @staticmethod
+    def _get_num_tokens(inputs: Dict[str, Any]) -> int:
+        if 'input_ids' in inputs:  # 1d or 2d
+            input_ids = inputs['input_ids']
+            if isinstance(input_ids, list):
+                return len(input_ids)
+            else:
+                return input_ids.shape[-1]
+        elif 'inputs_embeds' in inputs:  # 2d or 3d
+            return inputs['inputs_embeds'].shape[-2]
+        raise ValueError(f'Unable to retrieve input_ids and inputs_embeds. inputs: {inputs}')
+
+    def set_default_max_tokens(self, request_config: RequestConfig, inputs: Dict[str, Any]) -> None:
+        max_model_len = self.max_model_len
+        assert isinstance(inputs, dict)
+        # The num_tokens takes the maximum value from inputs_list.
+        num_tokens = self._get_num_tokens(inputs)
+        max_tokens = request_config.max_tokens
+        if max_model_len is None:
+            max_model_len = 8192
+            logger.warning(
+                'The current model is unable to retrieve `max_model_len`. It is set to the default value of 8192.')
+        max_max_tokens = max_model_len - num_tokens + self.max_tokens_offset
+        if max_tokens is None:
+            request_config.max_tokens = max_max_tokens
+        elif max_max_tokens < request_config.max_tokens:
+            logger.warning(f'max_model_len({max_model_len}) - num_tokens({num_tokens}) < max_tokens({max_tokens}). '
+                           f'Setting max_tokens: {max_model_len - num_tokens}')
+            request_config.max_tokens = max_max_tokens
+
+    def _get_logprobs(self,
+                      logprobs_list: Optional[List[Dict[int, float]]],
+                      token_ids: List[int],
+                      top_logprobs: Optional[int] = None) -> Optional[Dict[str, Any]]:
+        if logprobs_list is None or len(token_ids) == 0:
+            return None
+        if len(token_ids) > 0:
+            logprobs_list = logprobs_list[-len(token_ids):]
+        res = []
+        for logprobs, token_id in zip(logprobs_list, token_ids):
+            token = self.tokenizer.decode(token_id)
+            _res = {'token': token, 'logprob': logprobs[token_id], 'bytes': list(token.encode('utf8'))}
+            if top_logprobs is not None:
+                logprobs = {k: logprobs[k] for k in sorted(logprobs, key=lambda k: -logprobs[k])[:top_logprobs]}
+                res_top_logprobs = []
+                for k, logprob in logprobs.items():
+                    if logprob == float('-inf'):
+                        continue
+                    token = self.tokenizer.decode(k)
+                    res_top_logprobs.append({'token': token, 'logprob': logprob, 'bytes': list(token.encode('utf8'))})
+                _res['top_logprobs'] = res_top_logprobs
+            res.append(_res)
+        return {'content': res}
+
+    @staticmethod
+    def _get_finish_reason(max_tokens: int, completion_tokens: int, is_finished: bool):
+        if is_finished:
+            if completion_tokens >= max_tokens:
+                finish_reason = 'length'
+            else:
+                finish_reason = 'stop'
+        else:
+            finish_reason = None
+        return finish_reason
+
+    @staticmethod
+    def thread_run(target, args=(), kwargs=None):
+        kwargs = kwargs or {}
+
+        def func(target, queue, args, kwargs):
+            try:
+                queue.put(target(*args, **kwargs))
+            except Exception as e:
+                queue.put(e)
+
+        queue = Queue()
+        thread = Thread(target=func, args=(target, queue, args, kwargs))
+        thread.start()
+        thread.join()
+        result = queue.get()
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    @staticmethod
+    def safe_asyncio_run(coro):
+
+        def asyncio_run(core):
+            return asyncio.run(core)
+
+        return InferEngine.thread_run(asyncio_run, args=(coro, ))
+
+    def _batch_encode(self, infer_requests: List[InferRequest], strict: bool):
+        max_workers = max(min(32, os.cpu_count(), len(infer_requests)), 1)
+        error_list = []
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = [
+                executor.submit(self.template.encode, infer_request, return_template_inputs=True)
+                for infer_request in infer_requests
+            ]
+            concurrent.futures.wait(futures)
+            batched_inputs = []
+            for i, future in enumerate(futures):
+                try:
+                    batched_inputs.append(future.result())
+                except Exception as e:
+                    if strict:
+                        raise
+                    error_list.append((i, e))
+                    continue
+        return batched_inputs, error_list
+
+    @staticmethod
+    def _add_error_list(outputs, error_list):
+        for i, error in error_list:
+            outputs.insert(i, error)
+        return outputs

swift/infer_engine/lmdeploy_engine.py

@@ -0,0 +1,350 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import asyncio
+import inspect
+import lmdeploy
+import os
+import time
+import torch
+from contextlib import contextmanager
+from copy import deepcopy
+from lmdeploy import PytorchEngineConfig, TurbomindEngineConfig, VisionConfig, pipeline
+from lmdeploy.api import autoget_backend_config
+from lmdeploy.serve import async_engine
+from packaging import version
+from PIL import Image
+from transformers import GenerationConfig
+from transformers.utils.versions import require_version
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
+
+from swift.metrics import Metric
+from swift.model import get_processor
+from swift.template import Template
+from swift.utils import get_logger, get_seed, safe_snapshot_download
+from .infer_engine import InferEngine
+from .patch import patch_auto_config, patch_auto_tokenizer
+from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
+                       ChatCompletionStreamResponse, ChatMessage, DeltaMessage, InferRequest, RequestConfig)
+from .utils import InferStreamer
+
+try:
+    from lmdeploy import EngineGenerationConfig as LmdeployGenerationConfig
+except ImportError:
+    # compat lmdeploy >= 0.6.*
+    from lmdeploy import GenerationConfig as LmdeployGenerationConfig
+
+logger = get_logger()
+
+
+class LmdeployEngine(InferEngine):
+
+    def __init__(
+        self,
+        model_id_or_path: str,
+        *,
+        template: Optional[Template] = None,
+        torch_dtype: Optional[torch.dtype] = None,
+        model_type: Optional[str] = None,
+        template_type: Optional[str] = None,
+        use_hf: Optional[bool] = None,
+        hub_token: Optional[str] = None,
+        revision: Optional[str] = None,
+        # engine_kwargs
+        tp: int = 1,
+        session_len: Optional[int] = None,
+        cache_max_entry_count: float = 0.8,
+        quant_policy: int = 0,  # e.g. 4, 8
+        vision_batch_size: int = 1,  # max_batch_size in VisionConfig
+        engine_kwargs: Optional[Dict[str, Any]] = None,
+        devices: Optional[List[int]] = None,
+    ) -> None:
+        self.model_id_or_path = model_id_or_path
+        self.torch_dtype = torch_dtype
+        self.model_type = model_type
+        self.use_hf = use_hf
+        self.hub_token = hub_token
+        self.revision = revision
+        self.tp = tp
+        self.session_len = session_len
+        self.cache_max_entry_count = cache_max_entry_count
+        self.quant_policy = quant_policy
+        self.vision_batch_size = vision_batch_size
+        self.devices = devices
+        if template is None:
+            processor = self._get_processor()
+            template = self._get_template(processor, template_type=template_type)
+        else:
+            safe_snapshot_download(
+                model_id_or_path,
+                revision=revision,
+                download_model=True,
+                use_hf=use_hf,
+                ignore_patterns=getattr(template.model_meta, 'ignore_patterns', None),
+                hub_token=hub_token)
+        super().__init__(template)
+
+        if self.max_model_len is not None:
+            self.max_model_len -= 1
+        self._prepare_engine_kwargs(engine_kwargs)
+        self.config.torch_dtype = self.torch_dtype = self.torch_dtype or self.model_info.torch_dtype
+        self._prepare_engine()
+        self._load_generation_config()
+
+    def _get_processor(self):
+        return get_processor(
+            model_id_or_path=self.model_id_or_path,
+            torch_dtype=self.torch_dtype,
+            download_model=True,
+            model_type=self.model_type,
+            use_hf=self.use_hf,
+            hub_token=self.hub_token,
+            revision=self.revision)
+
+    def _prepare_engine_kwargs(self, engine_kwargs):
+        if engine_kwargs is None:
+            engine_kwargs = {}
+        engine_kwargs['tp'] = self.tp
+        engine_kwargs['session_len'] = self.session_len
+        engine_kwargs['cache_max_entry_count'] = self.cache_max_entry_count
+        engine_kwargs['quant_policy'] = self.quant_policy
+        if 'devices' in inspect.signature(TurbomindEngineConfig).parameters:
+            engine_kwargs['devices'] = self.devices
+        backend_config = TurbomindEngineConfig(**engine_kwargs)
+        backend_config = autoget_backend_config(self.model_dir, backend_config)
+        self.backend_config = backend_config
+        logger.info(f'backend_config: {backend_config}')
+
+        pipeline_kwargs = {}
+        is_multimodal = self.model_meta.is_multimodal
+        if is_multimodal:
+            require_version(
+                'lmdeploy<0.9', 'LmdeployEngine will no longer maintain inference for '
+                'multimodal models in lmdeploy>=0.9.')
+            vision_config = VisionConfig(max_batch_size=self.vision_batch_size)
+            pipeline_kwargs['vision_config'] = vision_config
+            logger.info(f'vision_config: {vision_config}')
+        self.pipeline_kwargs = pipeline_kwargs
+
+    @contextmanager
+    def _patch_pipeline(self):
+        _old_best_match_model = async_engine.best_match_model
+
+        def _best_match_model(*args, **kwargs) -> Optional[str]:
+            return self.model_info.model_type
+
+        async_engine.best_match_model = _best_match_model
+        try:
+            yield
+        finally:
+            async_engine.best_match_model = _old_best_match_model
+
+    def _prepare_engine(self):
+        with patch_auto_tokenizer(self.tokenizer), patch_auto_config(self.config), self._patch_pipeline():
+            engine = pipeline(self.model_dir, backend_config=self.backend_config, **self.pipeline_kwargs)
+        self.engine = engine
+
+    def _load_generation_config(self):
+        generation_config_path = os.path.join(self.model_dir, 'generation_config.json')
+        if os.path.isfile(generation_config_path):
+            generation_config = GenerationConfig.from_pretrained(self.model_dir)
+            kwargs = generation_config.to_dict()
+            max_new_tokens = kwargs.get('max_new_tokens')
+            if max_new_tokens is None:
+                kwargs.pop('max_new_tokens', None)
+            parameters = inspect.signature(LmdeployGenerationConfig).parameters
+            for k, v in kwargs.copy().items():
+                if k not in parameters or v is None:
+                    kwargs.pop(k)
+            self.generation_config = LmdeployGenerationConfig(**kwargs)
+        else:
+            self.generation_config = LmdeployGenerationConfig()
+
+    def _add_stop_words(self, generation_config: LmdeployGenerationConfig, request_config: RequestConfig) -> None:
+        template_meta = self.template.template_meta
+        stop_words = (request_config.stop or []) + (self.generation_config.stop_words or []) + template_meta.stop_words
+        generation_config.stop_words = self._get_stop_token_ids(stop_words)
+        # compat lmdeploy >= 0.6.*
+        generation_config.stop_token_ids = generation_config.stop_words
+
+    def _prepare_generation_config(self, request_config: RequestConfig) -> LmdeployGenerationConfig:
+        kwargs = {'max_new_tokens': request_config.max_tokens}
+        for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty']:
+            new_value = getattr(request_config, key)
+            if new_value is None:
+                kwargs[key] = getattr(self.generation_config, key)
+            else:
+                kwargs[key] = new_value
+        if request_config.seed is None:
+            request_config.seed = get_seed()
+        kwargs['random_seed'] = request_config.seed
+        if request_config.temperature == 0:
+            kwargs['temperature'] = 1  # avoid unnecessary process
+            kwargs['top_k'] = 1
+
+        if request_config.logprobs:
+            kwargs['logprobs'] = 1
+            if request_config.top_logprobs is not None:
+                kwargs['logprobs'] = max(1, request_config.top_logprobs)
+
+        res = LmdeployGenerationConfig(**kwargs)
+        return res
+
+    async def _infer_stream_async(
+        self,
+        inputs: Dict[str, Any],
+        generation_config: LmdeployGenerationConfig,
+        request_config: RequestConfig,
+    ) -> AsyncIterator[ChatCompletionStreamResponse]:
+        session_id = time.time_ns()
+        kwargs = {'stream_output': True, 'gen_config': generation_config, 'sequence_start': True, 'sequence_end': True}
+        if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'):
+            async with self.engine.model_inst(session_id) as inst:
+                context = self.engine.safe_run(inst, session_id, **inputs, **kwargs)
+        else:
+            context = self.engine.safe_run(session_id)
+
+        infer_streamer = InferStreamer(self.template, template_inputs=inputs['template_inputs'])
+        token_idx = 0
+        async with context as gen:
+            if version.parse(lmdeploy.__version__) < version.parse('0.6.5'):
+                generator = await self.engine.get_generator(False, session_id)
+                gen = generator.async_stream_infer(session_id=session_id, **inputs, **kwargs)
+            is_finished = False
+            while not is_finished:
+                try:
+                    output = await gen.__anext__()
+                except StopAsyncIteration:
+                    is_finished = True
+                delta_text = infer_streamer.get_printable_text(output.token_ids, is_finished)
+                if not delta_text and not is_finished:
+                    continue
+
+                logprobs = self._get_logprobs(output.logprobs, output.token_ids[token_idx:],
+                                              request_config.top_logprobs)
+                token_idx = len(output.token_ids)
+
+                usage_info = self._get_usage_info(len(inputs['input_ids']), output.num_token)
+                toolcall = None
+                if is_finished:
+                    toolcall = self._get_toolcall(
+                        self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs']))
+                finish_reason = self._get_finish_reason(generation_config.max_new_tokens, output.num_token,
+                                                        output.status.name == 'FINISH')
+                choices = [
+                    ChatCompletionResponseStreamChoice(
+                        index=0,
+                        delta=DeltaMessage(role='assistant', content=delta_text, tool_calls=toolcall),
+                        finish_reason=finish_reason,
+                        logprobs=logprobs)
+                ]
+                yield ChatCompletionStreamResponse(model=self.model_name, choices=choices, usage=usage_info)
+
+    async def _infer_full_async(
+        self,
+        inputs: Dict[str, Any],
+        generation_config: LmdeployGenerationConfig,
+        request_config: RequestConfig,
+    ) -> ChatCompletionResponse:
+        session_id = time.time_ns()
+        kwargs = {'stream_output': False, 'gen_config': generation_config, 'sequence_start': True, 'sequence_end': True}
+        if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'):
+            async with self.engine.model_inst(session_id) as inst:
+                async with self.engine.safe_run(inst, session_id, **inputs, **kwargs) as gen:
+                    async for output in gen:
+                        pass
+                if self.engine.backend == 'pytorch':
+                    # manually end pytorch session
+                    await inst.async_end(session_id)
+
+        else:
+            async with self.engine.safe_run(session_id):
+                generator = await self.engine.get_generator(False, session_id)
+                async for output in generator.async_stream_infer(session_id=session_id, **inputs, **kwargs):
+                    pass
+
+        response = self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs'])
+        logprobs = self._get_logprobs(output.logprobs, output.token_ids, request_config.top_logprobs)
+
+        usage_info = self._get_usage_info(len(inputs['input_ids']), output.num_token)
+        toolcall = self._get_toolcall(response)
+        finish_reason = self._get_finish_reason(generation_config.max_new_tokens, output.num_token,
+                                                output.status.name == 'FINISH')
+        token_ids = output.token_ids if request_config.return_details else None
+        choices = [
+            ChatCompletionResponseChoice(
+                index=0,
+                message=ChatMessage(role='assistant', content=response, tool_calls=toolcall),
+                finish_reason=finish_reason,
+                logprobs=logprobs,
+                token_ids=token_ids)
+        ]
+        prompt_token_ids = None
+        images_size = None
+        if request_config.return_details:
+            prompt_token_ids = inputs['input_ids']
+            images = inputs['template_inputs'].images
+            if all(isinstance(image, Image.Image) for image in images):
+                images_size = [image.size for image in images]
+        return ChatCompletionResponse(
+            model=self.model_name,
+            choices=choices,
+            usage=usage_info,
+            prompt_token_ids=prompt_token_ids,
+            images_size=images_size)
+
+    async def infer_async(self,
+                          infer_request: InferRequest,
+                          request_config: Optional[RequestConfig] = None,
+                          *,
+                          pre_infer_hook=None,
+                          **kwargs) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]:
+        request_config = deepcopy(request_config or RequestConfig())
+        self.template.set_mode('lmdeploy')
+
+        loop = asyncio.get_running_loop()
+        with torch.inference_mode():
+            inputs = await loop.run_in_executor(None, self.template.encode, infer_request, True)
+        images = inputs.pop('images', None)
+        if images:
+            if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'):
+                messages = self.engine._convert_prompts(('', images))
+                messages = await self.engine.async_convert_to_pil_images(messages)
+                results = await self.engine.vl_encoder.preprocess(messages)
+                if self.engine.backend == 'turbomind':
+                    results = await self.engine.vl_encoder.async_infer(results)
+                    inputs['images'] = [result['content'] for result in results if result['role'] == 'forward'][0]
+                    await self.template.prepare_lmdeploy_turbomind_inputs(inputs)
+                else:
+                    inputs['images'] = results[1]['content']
+                    await self.template.prepare_lmdeploy_pytorch_inputs(inputs)
+            else:
+                inputs['images'] = await self.engine.vl_encoder.async_infer(images)
+                await self.template.prepare_lmdeploy_turbomind_inputs(inputs)
+
+        self.set_default_max_tokens(request_config, inputs)
+        generation_config = self._prepare_generation_config(request_config)
+        self._add_stop_words(generation_config, request_config)
+        kwargs.update({'inputs': inputs, 'generation_config': generation_config, 'request_config': request_config})
+        if pre_infer_hook:
+            kwargs = pre_infer_hook(kwargs)
+        if request_config.stream:
+            return self._infer_stream_async(**kwargs)
+        else:
+            return await self._infer_full_async(**kwargs)
+
+    def _batch_infer_stream(self, *args, **kwargs):
+        if hasattr(self.engine, 'vl_encoder'):
+            self.engine.vl_encoder._loop_task = None
+        if hasattr(self.engine, 'free_insts'):
+            self.engine.free_insts = None
+        return super()._batch_infer_stream(*args, **kwargs)
+
+    def infer(
+        self,
+        infer_requests: List[InferRequest],
+        request_config: Optional[RequestConfig] = None,
+        metrics: Optional[List[Metric]] = None,
+        *,
+        use_tqdm: Optional[bool] = None,
+        **kwargs,
+    ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
+        return super().infer(infer_requests, request_config, metrics, use_tqdm=use_tqdm, **kwargs)

swift/infer_engine/patch.py

@@ -0,0 +1,34 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from contextlib import contextmanager
+from functools import wraps
+from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedTokenizerBase
+
+
+@contextmanager
+def patch_auto_tokenizer(tokenizer: PreTrainedTokenizerBase):
+    _old_from_pretrained = AutoTokenizer.from_pretrained
+
+    @wraps(_old_from_pretrained)
+    def _from_pretrained(*args, **kwargs):
+        return tokenizer
+
+    AutoTokenizer.from_pretrained = _from_pretrained
+    try:
+        yield
+    finally:
+        AutoTokenizer.from_pretrained = _old_from_pretrained
+
+
+@contextmanager
+def patch_auto_config(config: PretrainedConfig):
+    _old_from_pretrained = AutoConfig.from_pretrained
+
+    @wraps(_old_from_pretrained)
+    def _from_pretrained(*args, **kwargs):
+        return (config, {}) if 'return_unused_kwargs' in kwargs else config
+
+    AutoConfig.from_pretrained = _from_pretrained
+    try:
+        yield
+    finally:
+        AutoConfig.from_pretrained = _old_from_pretrained

swift/infer_engine/protocol.py

@@ -0,0 +1,622 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import base64
+import io
+import json
+import numpy as np
+import os
+import time
+import uuid
+from copy import deepcopy
+from dataclasses import asdict, dataclass, field, fields
+from PIL import Image
+from pydantic import AfterValidator, BaseModel, Field, PlainSerializer, field_validator
+from typing import Annotated, Any, Dict, List, Literal, Optional, Tuple, Union
+
+from swift.template import Messages, Tool
+from swift.utils import remove_response
+
+
+def serialize_ndarray(value):
+    if value is None:
+        return None
+    if isinstance(value, np.ndarray):
+        return {
+            'data': base64.b64encode(value.tobytes()).decode('ascii'),
+            'shape': value.shape,
+            'dtype': str(value.dtype),
+            '__ndarray__': True
+        }
+    return value
+
+
+def deserialize_ndarray(value):
+    if value is None:
+        return None
+    if isinstance(value, dict) and value.get('__ndarray__'):
+        data = base64.b64decode(value['data'])
+        return np.frombuffer(data, dtype=value['dtype']).reshape(value['shape'])
+    return value
+
+
+NumpyArray = Annotated[Any, PlainSerializer(serialize_ndarray, return_type=Dict), AfterValidator(deserialize_ndarray)]
+
+
+@dataclass
+class InferRequest:
+    """
+    Data structure for inference requests.
+
+    Attributes:
+        messages (Messages):
+            The input conversation in messages format. Each message is a dict containing at least
+            a "role" field (e.g., "user", "assistant", "system") and a "content" field.
+            Example:
+                [{
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image",  # can also be audio/video
+                            "image": "<url/path/base64/PIL.Image>",
+                        },
+                        {"type": "text", "text": "Please describe the picture."},
+                    ],
+                }]
+            The above is equivalent to:
+                [{"role": "user", "content": "<image>Please describe the picture."}]
+            with an additional argument:
+                images = ["<url/path/base64/PIL.Image>"]
+
+        images (List[Union[str, Image.Image]]):
+            Optional, a list of images associated with the request.
+            Each image can be a URL, local path, base64 string, or PIL.Image object.
+
+        audios (List[str]):
+            Optional, a list of audio resources associated with the request.
+
+        videos (List[str]):
+            Optional, a list of video resources associated with the request.
+
+        tools (Optional[List[Tool]]):
+            An optional list of tools. These should be organized in the agent_template format for
+            tools requested by the system, for example 'react_en'.
+
+        objects (Dict[str, Any]):
+            Container for additional multimodal objects, grouped by type (key).
+    """
+    messages: Messages
+
+    images: List[Union[str, Image.Image]] = field(default_factory=list)
+    audios: List[str] = field(default_factory=list)
+    videos: List[str] = field(default_factory=list)
+
+    tools: Optional[List[Tool]] = None
+    objects: Dict[str, Any] = field(default_factory=dict)
+    chat_template_kwargs: Dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self):
+        for key in ['images', 'audios', 'videos']:
+            val = getattr(self, key)
+            if isinstance(val, str):
+                setattr(self, key, [val])
+        assert isinstance(self.messages, list), f'messages: {self.messages}'
+
+    @staticmethod
+    def remove_response(messages) -> Optional[str]:
+        return remove_response(messages)
+
+    @staticmethod
+    def _to_printable(obj, key: Optional[str] = None):
+        if isinstance(obj, str) and key not in {'content', 'text'} and len(obj) >= 1000:
+            return f'<<<base64:{obj[:50]}..>>>'
+        elif isinstance(obj, list):
+            res = []
+            for item in obj:
+                res.append(InferRequest._to_printable(item))
+            return res
+        elif isinstance(obj, dict):
+            res = {}
+            for k, v in obj.items():
+                res[k] = InferRequest._to_printable(v, key=k)
+            return res
+        return obj
+
+    def to_printable(self):
+        return InferRequest._to_printable(asdict(self))
+
+
+@dataclass
+class RolloutInferRequest(InferRequest):
+    """
+    An inference request class for rollout scenarios.
+
+    This class extends `InferRequest` and specifically overrides the `images` attribute
+    to be a list of strings for compatibility with POST requests. Each string may
+    represent an image URL or a Base64-encoded image.
+
+    Inherits all fields from `InferRequest`:
+        messages (Messages):
+            Input conversation messages, supporting multimodal content.
+        audios (List[str]):
+            List of audio resources associated with the request.
+        videos (List[str]):
+            List of video resources associated with the request.
+        tools (Optional[List[Tool]]):
+            List of tools, organized by the agent template (e.g. 'react_en').
+        objects (Dict[str, Any]):
+            Optional container for additional multimodal objects.
+
+    Additional / Overridden fields:
+        images (List[str]):
+            List of image resources, each as a string (URL or base64).
+        data_dict (Dict):
+            Optional dictionary for extra request data.
+        uuid (Optional[str]):
+            Optional unique identifier for this request instance.
+    """
+    images: List[str] = field(default_factory=list)
+    data_dict: Dict = field(default_factory=dict)
+    uuid: Optional[str] = None
+
+
+def random_uuid() -> str:
+    return str(uuid.uuid4().hex)
+
+
+@dataclass
+class Model:
+    id: str  # model_type
+
+    object: str = 'model'
+    created: int = field(default_factory=lambda: int(time.time()))
+    owned_by: str = 'ms-swift'
+
+
+@dataclass
+class ModelList:
+    data: List[Model]
+    object: str = 'list'
+
+
+@dataclass
+class RequestConfig:
+    """NOTE: The following behavior is inconsistent with the OpenAI API.
+    Default values for OpenAI:
+        temperature = 1.
+        top_k = -1
+        top_p = 1.
+        repetition_penalty = 1.
+    """
+    max_tokens: Optional[int] = None  # None: max_model_len - num_tokens
+    # None: use deploy_args
+    temperature: Optional[float] = None
+    top_k: Optional[int] = None
+    top_p: Optional[float] = None
+    repetition_penalty: Optional[float] = None
+    num_beams: int = 1
+    stop: Optional[List[str]] = field(default_factory=list)
+
+    seed: Optional[int] = None
+    stream: bool = False
+    logprobs: bool = False
+    top_logprobs: Optional[int] = None
+    prompt_logprobs: Optional[int] = None
+
+    n: int = 1
+    best_of: Optional[int] = None
+    presence_penalty: float = 0.
+    frequency_penalty: float = 0.
+    length_penalty: float = 1.
+    # Return token_ids additionally (non-stream)
+    return_details: bool = False
+    # vLLM structured outputs (guided decoding)
+    structured_outputs_regex: Optional[str] = None
+
+    def __post_init__(self):
+        if self.stop is None:
+            self.stop = []
+
+
+@dataclass
+class CompletionRequestMixin:
+    model: str
+    prompt: str
+
+
+@dataclass
+class EmbeddingRequestMixin:
+    input: str
+    model: str
+    encoding_format: Literal['float', 'base64'] = 'float'
+
+
+@dataclass
+class ChatCompletionRequestMixin:
+    model: str
+    messages: Messages
+    tools: Optional[List[Tool]] = None
+    tool_choice: Optional[Union[str, Dict]] = None
+    chat_template_kwargs: Dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self):
+        if self.tool_choice is None:
+            self.tool_choice = 'none' if self.tools is None else 'auto'
+
+        if self.tools:
+            if self.tool_choice == 'none':
+                self.tools = None
+            elif isinstance(self.tool_choice, dict):
+                name = self.tool_choice['function']['name']
+                tool = next(tool for tool in self.tools if tool['function']['name'] == name)
+                if tool is None:
+                    raise ValueError(f"Tool choice '{name}' not found in tools.")
+                self.tools = [tool]
+
+
+@dataclass
+class MultiModalRequestMixin:
+    images: List[str] = field(default_factory=list)
+    audios: List[str] = field(default_factory=list)
+    videos: List[str] = field(default_factory=list)
+    objects: Dict[str, Any] = field(default_factory=dict)
+
+    @staticmethod
+    def to_base64(mm_data: Union[str, Image.Image, bytes]) -> str:
+        if isinstance(mm_data, dict) and 'bytes' in mm_data:
+            mm_data = mm_data['bytes'] or mm_data['path']
+        if isinstance(mm_data, str) and not os.path.isfile(mm_data):
+            # base64 or url
+            return mm_data
+        if isinstance(mm_data, str):
+            # local_path
+            with open(mm_data, 'rb') as f:
+                bytes_ = f.read()
+        elif isinstance(mm_data, Image.Image):
+            bytes_io = io.BytesIO()
+            mm_data.save(bytes_io, format='png')
+            bytes_ = bytes_io.getvalue()
+        else:
+            bytes_ = mm_data
+        img_base64: str = base64.b64encode(bytes_).decode('utf-8')
+        return img_base64
+
+    def __post_init__(self):
+        for key in ['images', 'audios', 'videos']:
+            values = getattr(self, key)
+            if isinstance(values, str):
+                values = [values]
+                setattr(self, key, values)
+            for i, val in enumerate(values):
+                values[i] = self.to_base64(val)
+
+
+@dataclass
+class CompletionRequest(RequestConfig, MultiModalRequestMixin, CompletionRequestMixin):
+
+    def __post_init__(self):
+        RequestConfig.__post_init__(self)
+        MultiModalRequestMixin.__post_init__(self)
+
+
+@dataclass
+class EmbeddingRequest(RequestConfig, MultiModalRequestMixin, EmbeddingRequestMixin):
+
+    def __post_init__(self):
+        RequestConfig.__post_init__(self)
+        MultiModalRequestMixin.__post_init__(self)
+
+    def parse(self) -> Tuple['InferRequest', 'RequestConfig']:
+        data = asdict(self)
+        res = []
+        for cls_type in [InferRequest, RequestConfig]:
+            parameters = set(f.name for f in fields(cls_type))
+            _data = {k: v for k, v in data.items() if k in parameters}
+            res.append(cls_type(**_data))
+        return tuple(res)
+
+
+@dataclass
+class ChatCompletionRequest(RequestConfig, MultiModalRequestMixin, ChatCompletionRequestMixin):
+
+    def __post_init__(self):
+        RequestConfig.__post_init__(self)
+        MultiModalRequestMixin.__post_init__(self)
+        ChatCompletionRequestMixin.__post_init__(self)
+        self.convert_to_base64()
+
+    def convert_to_base64(self):
+        for message in self.messages:
+            content = message['content']
+            if isinstance(content, str):
+                continue
+            for item in content:
+                key: str = item['type']
+                if key == 'text':
+                    continue
+
+                key_origin = key
+                value = item[key]
+                if key.endswith('_url'):
+                    key = key[:-len('_url')]
+                is_dict = False
+                if isinstance(value, dict):
+                    is_dict = True
+                    value = value['url']
+                if isinstance(value, str) and (value.startswith('data:') or value.startswith('http')
+                                               or len(value) > 200):
+                    continue
+
+                # local_path / PIL.Image
+                if isinstance(value, str) and os.path.isfile(value):
+                    suffix = os.path.splitext(value)[1][1:].lower()
+                elif isinstance(value, Image.Image):
+                    suffix = 'jpeg'
+                else:
+                    raise ValueError(f'value: {value}')
+                mm_data_base64 = self.to_base64(value)
+                new_value = f'data:{key}/{suffix};base64,{mm_data_base64}'
+                if is_dict:
+                    new_value = {'url': new_value}
+                item[key_origin] = new_value
+
+    def parse(self) -> Tuple['InferRequest', 'RequestConfig']:
+        data = asdict(self)
+        res = []
+        for cls_type in [InferRequest, RequestConfig]:
+            parameters = set(f.name for f in fields(cls_type))
+            _data = {k: v for k, v in data.items() if k in parameters}
+            res.append(cls_type(**_data))
+        return tuple(res)
+
+    @classmethod
+    def from_cmpl_request(cls, cmpl_request: Union[CompletionRequest, EmbeddingRequest]) -> 'ChatCompletionRequest':
+        cmpl_request = asdict(cmpl_request)
+        if 'prompt' in cmpl_request:
+            prompt = cmpl_request.pop('prompt')
+        else:
+            prompt = cmpl_request.pop('input')
+        cmpl_request['messages'] = [{'role': 'user', 'content': prompt}]
+        if 'encoding_format' in cmpl_request:
+            cmpl_request.pop('encoding_format')
+        return cls(**cmpl_request)
+
+
+@dataclass
+class UsageInfo:
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+@dataclass
+class Function:
+    name: str
+    arguments: Optional[Union[str, Any]]
+
+    def __post_init__(self):
+        if not isinstance(self.arguments, str):
+            self.arguments = json.dumps(self.arguments, ensure_ascii=False)
+        self.name = self.name.strip()
+        self.arguments = self.arguments.strip()
+
+
+@dataclass
+class ChatCompletionMessageToolCall:
+    function: Function
+    type: str = 'function'
+    id: str = field(default_factory=lambda: f'toolcall-{random_uuid()}')
+
+
+@dataclass
+class ChatMessage:
+    role: Literal['system', 'user', 'assistant']
+    content: Union[str, List[Dict[str, Any]], int, float, List[float]]
+    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
+    reasoning_content: Optional[str] = None
+
+
+@dataclass
+class ChatCompletionResponseChoice:
+    index: int
+    message: ChatMessage
+    finish_reason: Literal['stop', 'length', None]
+    logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None
+    token_ids: Optional[List[int]] = None
+    routed_experts: Optional[NumpyArray] = None
+
+    def to_cmpl_choice(self) -> 'CompletionResponseChoice':
+        self = deepcopy(self)
+        assert not self.message.tool_calls, f'message: {self.message}'
+        return CompletionResponseChoice(self.index, self.message.content, self.finish_reason, self.logprobs)
+
+
+@dataclass
+class EmbeddingResponseData:
+    object: str = 'embedding'
+    index: int = 0
+    embedding: List[str] = field(default_factory=lambda: [])
+
+
+@dataclass
+class EmbeddingResponse:
+    model: str
+    data: List[EmbeddingResponseData]
+    usage: UsageInfo
+    id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}')
+    object: str = 'list'
+    created: int = field(default_factory=lambda: int(time.time()))
+
+
+@dataclass
+class CompletionResponseChoice:
+    index: int
+    text: str
+    finish_reason: Literal['stop', 'length', None]
+    logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None
+
+
+@dataclass
+class ChatCompletionResponse:
+    model: str
+    choices: List[ChatCompletionResponseChoice]
+    usage: UsageInfo
+    id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}')
+    object: str = 'chat.completion'
+    created: int = field(default_factory=lambda: int(time.time()))
+    prompt_token_ids: Optional[List[int]] = None
+    prompt_logprobs: Optional[List] = None
+    images_size: Optional[List[Tuple[int, int]]] = None
+
+    def to_cmpl_response(self) -> 'CompletionResponse':
+        self = deepcopy(self)
+        choices = [choice.to_cmpl_choice() for choice in self.choices]
+        id_ = f'cmpl{self.id[len("chatcmpl"):]}'
+        return CompletionResponse(
+            self.model,
+            choices,
+            self.usage,
+            id_,
+            created=self.created,
+            prompt_token_ids=self.prompt_token_ids,
+            prompt_logprobs=self.prompt_logprobs,
+        )
+
+
+class RolloutOutput(BaseModel):
+    """
+    Output structure for rollout.
+
+    Attributes:
+        response (ChatCompletionResponse):
+            The model's response
+
+        messages (Optional[Messages]):
+            (Optional) Conversation history for the final rollout; required for multi-turn scenarios.
+            NOTE:
+                - If provided, this messages sequence will overwrite the original messages.
+                - If not provided, 'response' will be appended as the latest turn in the original messages.
+                - For multi-turn training, you need to manually return the updated messages, including the full history.
+                - The messages should include the latest assistant response as the final message.
+
+        response_token_ids (Optional[List[List[int]]]):
+            (Optional) Token IDs generated at each rollout turn.
+            If provided, the training process will skip tokenizing the response.
+
+        response_loss_mask (Optional[List[List[int]]]):
+            (Optional) Loss masks corresponding to each rollout turn.
+            If provided, the training process will skip computing loss masks for the response (as controlled by the `loss_scale` parameter). # noqa
+
+        rollout_infos (Dict[str, Any]):
+            (Optional) Additional rollout information. This must be JSON-serializable.
+    """
+    response: ChatCompletionResponse
+    # multi turn
+    messages: Optional[Messages] = None
+    response_token_ids: List[List[int]] = Field(default_factory=list)
+    response_loss_mask: List[List[int]] = Field(default_factory=list)
+    rollout_infos: Dict[str, Any] = Field(default_factory=dict)
+    # rollout logprobs for each turn (used for rollout importance sampling correction in multi-turn scenarios)
+    rollout_logprobs: List[List[float]] = Field(default_factory=list)
+
+    prompt_logprobs: Optional[List] = None
+
+    @field_validator('response_token_ids', 'response_loss_mask', 'rollout_logprobs', mode='before')
+    @classmethod
+    def _wrap_flat_list(cls, v):
+        if isinstance(v, list) and v and isinstance(v[0], (int, float)):
+            return [v]
+        return v
+
+    def model_post_init(self, __context):
+        # Ensure multimodal data in rollout_infos is serializable (e.g., images to base64)
+        super().model_post_init(__context)
+        self.mminfo_to_serializable()
+
+    def mminfo_to_serializable(self):
+        mm_keys = ['images', 'audios', 'videos']
+
+        for key, values in self.rollout_infos.items():
+            if key in mm_keys:
+                if not isinstance(values, list):
+                    values = [values]
+                for i, value in enumerate(values):
+                    values[i] = MultiModalRequestMixin.to_base64(value)
+                self.rollout_infos[key] = values
+
+
+@dataclass
+class CompletionResponse:
+    model: str
+    choices: List[CompletionResponseChoice]
+    usage: UsageInfo
+    id: str = field(default_factory=lambda: f'cmpl-{random_uuid()}')
+    object: str = 'text_completion'
+    created: int = field(default_factory=lambda: int(time.time()))
+    prompt_token_ids: Optional[List[int]] = None
+    prompt_logprobs: Optional[List] = None
+
+
+@dataclass
+class DeltaMessage:
+    role: Literal['system', 'user', 'assistant', None] = None
+    content: Optional[str] = None
+    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
+    reasoning_content: Optional[str] = None
+
+
+@dataclass
+class ChatCompletionResponseStreamChoice:
+    index: int
+    delta: DeltaMessage
+    finish_reason: Literal['stop', 'length', None]
+    logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None
+
+    def to_cmpl_choice(self) -> 'CompletionResponseStreamChoice':
+        self = deepcopy(self)
+        assert not self.delta.tool_calls
+        return CompletionResponseStreamChoice(self.index, self.delta.content, self.finish_reason, self.logprobs)
+
+
+@dataclass
+class CompletionResponseStreamChoice:
+    index: int
+    text: str
+    finish_reason: Literal['stop', 'length', None]
+    logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None
+
+
+@dataclass
+class ChatCompletionStreamResponse:
+    model: str
+    choices: List[ChatCompletionResponseStreamChoice]
+    usage: Optional[UsageInfo] = None
+    id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}')
+    object: str = 'chat.completion.chunk'
+    created: int = field(default_factory=lambda: int(time.time()))
+
+    def to_cmpl_response(self) -> 'CompletionStreamResponse':
+        self = deepcopy(self)
+        choices = [choice.to_cmpl_choice() for choice in self.choices]
+        id_ = f'cmpl{self.id[len("chatcmpl"):]}'
+        return CompletionStreamResponse(self.model, choices, self.usage, id_, created=self.created)
+
+
+@dataclass
+class CompletionStreamResponse:
+    model: str
+    choices: List[CompletionResponseStreamChoice]
+    usage: Optional[UsageInfo] = None
+    id: str = field(default_factory=lambda: f'cmpl-{random_uuid()}')
+    object: str = 'text_completion.chunk'
+    created: int = field(default_factory=lambda: int(time.time()))
+
+
+class InitCommunicatorRequest(BaseModel):
+    host: str
+    port: int
+    world_size: int
+
+
+class UpdateWeightsRequest(BaseModel):
+    name: str
+    dtype: str
+    shape: list[int]

swift/infer_engine/sglang_engine.py

@@ -0,0 +1,302 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import asyncio
+import inspect
+import os
+import sglang as sgl
+import torch
+from copy import deepcopy
+from PIL import Image
+from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.srt.server_args import ServerArgs
+from transformers import GenerationConfig
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
+
+from swift.metrics import Metric
+from swift.model import get_processor
+from swift.template import Template
+from swift.utils import get_logger, safe_snapshot_download
+from .infer_engine import InferEngine
+from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
+                       ChatCompletionStreamResponse, ChatMessage, DeltaMessage, EmbeddingResponse,
+                       EmbeddingResponseData, InferRequest, RequestConfig, random_uuid)
+from .utils import InferStreamer
+
+logger = get_logger()
+
+
+class SglangEngine(InferEngine):
+
+    def __init__(
+        self,
+        model_id_or_path: str,
+        *,
+        template: Optional[Template] = None,
+        torch_dtype: Optional[torch.dtype] = None,
+        model_type: Optional[str] = None,
+        template_type: Optional[str] = None,
+        use_hf: Optional[bool] = None,
+        hub_token: Optional[str] = None,
+        revision: Optional[str] = None,
+        # engine kwargs
+        tp_size: int = 1,
+        pp_size: int = 1,
+        dp_size: int = 1,
+        ep_size: int = 1,
+        enable_ep_moe: bool = False,
+        mem_fraction_static: Optional[float] = None,
+        context_length: Optional[int] = None,
+        disable_cuda_graph: bool = False,
+        quantization: Optional[str] = None,
+        task_type: Optional[str] = None,
+        kv_cache_dtype: str = 'auto',
+        enable_dp_attention: bool = False,
+        disable_custom_all_reduce: bool = True,
+        speculative_algorithm: Optional[str] = None,
+        speculative_num_steps: Optional[int] = None,
+        speculative_eagle_topk: Optional[int] = None,
+        speculative_num_draft_tokens: Optional[int] = None,
+        log_level='error',
+        engine_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        self.model_id_or_path = model_id_or_path
+        self.torch_dtype = torch_dtype
+        self.model_type = model_type
+        self.use_hf = use_hf
+        self.hub_token = hub_token
+        self.revision = revision
+        self.tp_size = tp_size
+        self.pp_size = pp_size
+        self.dp_size = dp_size
+        self.ep_size = ep_size
+        self.enable_ep_moe = enable_ep_moe
+        self.mem_fraction_static = mem_fraction_static
+        self.context_length = context_length
+        self.disable_cuda_graph = disable_cuda_graph
+        self.quantization = quantization
+        self.task_type = task_type
+        self.kv_cache_dtype = kv_cache_dtype
+        self.enable_dp_attention = enable_dp_attention
+        self.disable_custom_all_reduce = disable_custom_all_reduce
+        self.speculative_algorithm = speculative_algorithm
+        self.speculative_num_steps = speculative_num_steps
+        self.speculative_eagle_topk = speculative_eagle_topk
+        self.speculative_num_draft_tokens = speculative_num_draft_tokens
+        self.log_level = log_level
+        if template is None:
+            processor = self._get_processor()
+            template = self._get_template(processor, template_type=template_type)
+        else:
+            safe_snapshot_download(
+                model_id_or_path,
+                revision=revision,
+                download_model=True,
+                use_hf=use_hf,
+                ignore_patterns=getattr(template.model_meta, 'ignore_patterns', None),
+                hub_token=hub_token)
+        super().__init__(template)
+        self._prepare_server_args(engine_kwargs)
+        self.engine = sgl.Engine(server_args=self.server_args)
+        self._load_generation_config()
+        if speculative_num_draft_tokens is not None:
+            self.max_tokens_offset = -speculative_num_draft_tokens
+
+    def _get_processor(self):
+        return get_processor(
+            model_id_or_path=self.model_id_or_path,
+            torch_dtype=self.torch_dtype,
+            download_model=True,
+            model_type=self.model_type,
+            use_hf=self.use_hf,
+            hub_token=self.hub_token,
+            revision=self.revision,
+            task_type=self.task_type)
+
+    def _prepare_server_args(self, engine_kwargs):
+        if engine_kwargs is None:
+            engine_kwargs = {}
+        if self.context_length is not None:
+            self.max_model_len = self.context_length
+            logger.info(f'Setting max_model_len: {self.context_length}')
+        if self.max_model_len is not None:
+            self.max_model_len -= 1
+        parameters = inspect.signature(ServerArgs).parameters
+        if 'pp_size' in parameters:
+            engine_kwargs['pp_size'] = self.pp_size
+        if 'enable_ep_moe' in parameters:
+            engine_kwargs['enable_ep_moe'] = self.enable_ep_moe
+        self.server_args = ServerArgs(
+            model_path=self.model_dir,
+            dtype=self.model_info.torch_dtype,
+            tp_size=self.tp_size,
+            dp_size=self.dp_size,
+            ep_size=self.ep_size,
+            mem_fraction_static=self.mem_fraction_static,
+            context_length=self.context_length,
+            disable_cuda_graph=self.disable_cuda_graph,
+            quantization=self.quantization,
+            kv_cache_dtype=self.kv_cache_dtype,
+            enable_dp_attention=self.enable_dp_attention,
+            disable_custom_all_reduce=self.disable_custom_all_reduce,
+            speculative_algorithm=self.speculative_algorithm,
+            speculative_num_steps=self.speculative_num_steps,
+            speculative_eagle_topk=self.speculative_eagle_topk,
+            speculative_num_draft_tokens=self.speculative_num_draft_tokens,
+            log_level=self.log_level,
+            skip_tokenizer_init=True,
+            trust_remote_code=True,
+            **engine_kwargs,
+        )
+        if self.task_type == 'embedding':
+            self.server_args.is_embedding = True
+
+    def _load_generation_config(self) -> None:
+        generation_config_path = os.path.join(self.model_dir, 'generation_config.json')
+        if os.path.isfile(generation_config_path):
+            generation_config = GenerationConfig.from_pretrained(self.model_dir)
+        else:
+            generation_config = GenerationConfig()
+        kwargs = generation_config.to_dict()
+        top_k = kwargs.get('top_k')
+        if top_k == 0:
+            kwargs['top_k'] = -1
+
+        parameters = inspect.signature(SamplingParams).parameters
+        self.generation_config = {k: v for k, v in kwargs.items() if k in parameters and v is not None}
+
+    def _prepare_generation_config(self, request_config: RequestConfig) -> Dict[str, Any]:
+        kwargs = {'max_new_tokens': request_config.max_tokens}
+        for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty']:
+            new_value = getattr(request_config, key)
+            if new_value is None:
+                kwargs[key] = self.generation_config.get(key)
+            else:
+                kwargs[key] = new_value
+        for key in ['n', 'frequency_penalty', 'presence_penalty']:
+            kwargs[key] = getattr(request_config, key)
+
+        return kwargs
+
+    def _add_stop_words(self, generation_config: Dict[str, Any], request_config: RequestConfig) -> None:
+        template_meta = self.template.template_meta
+        stop_words = (request_config.stop or []) + (self.generation_config.get('stop') or []) + template_meta.stop_words
+        generation_config['stop_token_ids'] = self._get_stop_token_ids(stop_words)
+
+    def _create_chat_completion_response(self, output, inputs, return_details: bool = False):
+        assert output is not None
+        meta_info = output['meta_info']
+        usage_info = self._get_usage_info(meta_info['prompt_tokens'], meta_info['completion_tokens'])
+        response = self.template.decode_generate_ids(output['output_ids'], template_inputs=inputs['template_inputs'])
+        toolcall = self._get_toolcall(response)
+        token_ids = output['output_ids'] if return_details else None
+        choice = ChatCompletionResponseChoice(
+            index=0,
+            message=ChatMessage(role='assistant', content=response, tool_calls=toolcall),
+            finish_reason=meta_info['finish_reason']['type'],
+            logprobs=None,
+            token_ids=token_ids)
+        prompt_token_ids = None
+        images_size = None
+        if return_details:
+            prompt_token_ids = output.get('prompt_token_ids')
+            images = inputs['template_inputs'].images
+            if all(isinstance(image, Image.Image) for image in images):
+                images_size = [image.size for image in images]
+        return ChatCompletionResponse(
+            model=self.model_name,
+            choices=[choice],
+            usage=usage_info,
+            id=random_uuid(),
+            prompt_token_ids=prompt_token_ids,
+            images_size=images_size)
+
+    def infer(
+        self,
+        infer_requests: List[InferRequest],
+        request_config: Optional[RequestConfig] = None,
+        metrics: Optional[List[Metric]] = None,
+        *,
+        use_tqdm: Optional[bool] = None,
+    ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
+        return super().infer(infer_requests, request_config, metrics, use_tqdm=use_tqdm)
+
+    async def infer_async(self,
+                          infer_request: InferRequest,
+                          request_config: Optional[RequestConfig] = None,
+                          *,
+                          pre_infer_hook=None,
+                          **kwargs) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]:
+        request_config = deepcopy(request_config or RequestConfig())
+        self.template.set_mode('sglang')
+        loop = asyncio.get_running_loop()
+        with torch.inference_mode():
+            inputs = await loop.run_in_executor(None, self.template.encode, infer_request, True)
+        if self.task_type == 'embedding':
+            inputs.pop('length', None)
+        self.set_default_max_tokens(request_config, inputs)
+        generation_config = self._prepare_generation_config(request_config)
+        self._add_stop_words(generation_config, request_config)
+        kwargs.update({'inputs': inputs, 'generation_config': generation_config, 'request_config': request_config})
+        if pre_infer_hook:
+            kwargs = pre_infer_hook(kwargs)
+        if request_config.stream:
+            return self._infer_stream_async(**kwargs)
+        elif self.task_type == 'embedding':
+            kwargs.pop('generation_config', None)
+            return await self._infer_embedding_async(**kwargs)
+        else:
+            return await self._infer_full_async(**kwargs)
+
+    async def _infer_embedding_async(self, inputs: Dict[str, Any], **kwargs) -> EmbeddingResponse:
+        from sglang.srt.managers.io_struct import EmbeddingReqInput
+        obj = EmbeddingReqInput(
+            input_ids=inputs['input_ids'], image_data=inputs.get('images'), audio_data=inputs.get('audios'))
+        generator = self.engine.tokenizer_manager.generate_request(obj, None)
+        output = await generator.__anext__()
+        usage_info = self._get_usage_info(output['meta_info']['prompt_tokens'], 0)
+        return EmbeddingResponse(
+            model=self.model_name,
+            data=[EmbeddingResponseData(embedding=output['embedding'])],
+            usage=usage_info,
+            id=random_uuid())
+
+    async def _infer_full_async(self, inputs: Dict[str, Any], generation_config: Dict[str, Any],
+                                request_config: RequestConfig) -> ChatCompletionResponse:
+        engine_inputs = {k: v for k, v in inputs.items() if k != 'template_inputs'}
+        output = await self.engine.async_generate(**engine_inputs, sampling_params=generation_config)
+        output['prompt_token_ids'] = inputs['input_ids']
+        return self._create_chat_completion_response(output, inputs, request_config.return_details)
+
+    async def _infer_stream_async(self, inputs: Dict[str, Any], generation_config: Dict[str, Any],
+                                  **kwargs) -> AsyncIterator[ChatCompletionStreamResponse]:
+        engine_inputs = {k: v for k, v in inputs.items() if k != 'template_inputs'}
+        result_generator = await self.engine.async_generate(
+            **engine_inputs, sampling_params=generation_config, stream=True)
+        infer_streamer = InferStreamer(self.template, template_inputs=inputs['template_inputs'])
+        async for output in result_generator:
+            res = self._create_chat_completion_stream_response(output, infer_streamer)
+            if res is None:
+                continue
+            yield res
+
+    def _create_chat_completion_stream_response(self, output, infer_streamer) -> Optional[ChatCompletionStreamResponse]:
+        assert output is not None
+        meta_info = output['meta_info']
+        finish_reason = meta_info['finish_reason']
+        is_finished = finish_reason is not None
+        delta_text = infer_streamer.get_printable_text(output['output_ids'], is_finished)
+        if not delta_text and not is_finished:
+            return
+        toolcall = None
+        if is_finished:
+            finish_reason = finish_reason['type']
+            toolcall = self._get_toolcall(
+                self.template.decode_generate_ids(output['output_ids'], **infer_streamer.decode_kwargs))
+        meta_info = output['meta_info']
+        usage_info = self._get_usage_info(meta_info['prompt_tokens'], meta_info['completion_tokens'])
+        # TODO: logprobs
+        choice = ChatCompletionResponseStreamChoice(
+            index=0,
+            delta=DeltaMessage(role='assistant', content=delta_text, tool_calls=toolcall),
+            finish_reason=finish_reason,
+            logprobs=None)
+        return ChatCompletionStreamResponse(model=self.model_name, choices=[choice], usage=usage_info)

swift/infer_engine/transformers_engine.py

@@ -0,0 +1,594 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import asyncio
+import hashlib
+import inspect
+import json
+import os
+import pickle
+import time
+import torch
+import torch.nn.functional as F
+import transformers
+from copy import deepcopy
+from packaging import version
+from PIL import Image
+from queue import Queue
+from threading import Thread
+from torch import nn
+from tqdm import tqdm
+from transformers import GenerationConfig, LogitsProcessorList
+from transformers.utils import is_torch_npu_available
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
+
+from swift.metrics import Metric
+from swift.model import get_model_processor
+from swift.template import Template
+from swift.tuners import Swift
+from swift.utils import get_last_valid_indices, patch_kernels, safe_snapshot_download, to_device
+from .infer_engine import InferEngine
+from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
+                       ChatCompletionStreamResponse, ChatMessage, DeltaMessage, EmbeddingResponse,
+                       EmbeddingResponseData, InferRequest, RequestConfig, random_uuid)
+from .utils import AdapterRequest, InferStreamer, LogitsStreamer, TokensIteratorStreamer, prepare_generation_config
+
+_TRANSFORMERS_GE_5_2 = version.parse(transformers.__version__) >= version.parse('5.2.0')
+_kernels_patched = False
+
+
+class _GenerationConfig(GenerationConfig):
+
+    def __repr__(self) -> str:
+        parameters = inspect.signature(self.to_json_string).parameters
+        kwargs = {}
+        if 'ignore_metadata' in parameters:
+            kwargs['ignore_metadata'] = True
+        gen_kwargs = json.loads(self.to_json_string(**kwargs))
+        gen_kwargs.pop('transformers_version', None)
+        return f'GenerationConfig({gen_kwargs})'
+
+
+class TransformersEngine(InferEngine):
+
+    def __init__(
+            self,
+            model: Union[str, nn.Module],
+            *,
+            template: Optional[Template] = None,
+            adapters: Optional[List[str]] = None,
+            max_batch_size: int = 1,  # 0/1: no limit
+            reranker_use_activation: bool = True,
+            # model kwargs
+            torch_dtype: Optional[torch.dtype] = None,
+            model_type: Optional[str] = None,
+            attn_impl: Optional[str] = None,
+            experts_impl: Optional[str] = None,
+            device_map: Optional[Union[str, Dict[str, Any]]] = None,
+            task_type: Optional[str] = None,
+            quantization_config=None,
+            model_kwargs: Optional[Dict[str, Any]] = None,
+            template_type: Optional[str] = None,
+            # hub kwargs
+            use_hf: Optional[bool] = None,
+            revision: Optional[str] = None,
+            hub_token: Optional[str] = None,
+            **kwargs):
+        if isinstance(adapters, str):
+            adapters = [adapters]
+        self.adapters = adapters or []
+        self.max_batch_size = max_batch_size
+        self.reranker_use_activation = reranker_use_activation
+
+        self.torch_dtype = torch_dtype
+        self.model_type = model_type
+        self.attn_impl = attn_impl
+        self.experts_impl = experts_impl
+        self.device_map = device_map
+        self.task_type = task_type
+        self.quantization_config = quantization_config
+        self.model_kwargs = model_kwargs
+
+        self.use_hf = use_hf
+        self.revision = revision
+        self.hub_token = hub_token
+        global _kernels_patched
+        if _TRANSFORMERS_GE_5_2 and not _kernels_patched:
+            if use_hf is not None and 'USE_HF' not in os.environ:
+                os.environ['USE_HF'] = str(use_hf)
+            patch_kernels()
+            _kernels_patched = True
+        if isinstance(model, str):
+            self.model, processor = self._get_model_processor(model, **kwargs)
+            template = self._get_template(processor, template_type=template_type)
+        elif isinstance(model, nn.Module):
+            self.model = model
+            if template is None:
+                raise ValueError('`template` is required when `model` is a nn.Module')
+        super().__init__(template)
+        for adapter in self.adapters:
+            self._add_adapter(safe_snapshot_download(adapter, use_hf=self.use_hf, hub_token=self.hub_token))
+        self.engine = self.model  # dummy
+        self.generation_config = getattr(self.model, 'generation_config', None)
+        self._queue = Queue()
+        self._task_pool = {}
+        self._adapters_pool = {}
+        self._task_thread = None
+
+    def _get_model_processor(self, model_id_or_path, **kwargs):
+        return get_model_processor(
+            model_id_or_path,
+            torch_dtype=self.torch_dtype,
+            model_type=self.model_type,
+            use_hf=self.use_hf,
+            hub_token=self.hub_token,
+            revision=self.revision,
+            device_map=self.device_map,
+            quantization_config=self.quantization_config,
+            attn_impl=self.attn_impl,
+            experts_impl=self.experts_impl,
+            task_type=self.task_type,
+            model_kwargs=self.model_kwargs,
+            **kwargs)
+
+    def _start_infer_worker(self):
+        self._task_thread = Thread(target=self._infer_worker, daemon=True)
+        self._task_thread.start()
+
+    def _fetch_infer_requests(self):
+        while not self._queue.empty():
+            infer_request, kwargs, queue = self._queue.get()
+            info = hashlib.sha256(pickle.dumps((kwargs['request_config']))).hexdigest()
+            if info not in self._task_pool:
+                self._task_pool[info] = kwargs, []
+            self._task_pool[info][1].append((infer_request, queue))
+        if len(self._task_pool) == 0:
+            return
+        key, (kwargs, data) = next(iter(self._task_pool.items()))
+        max_batch_size = self.max_batch_size
+        if max_batch_size <= 0:
+            max_batch_size = len(data)
+        data, remain_data = data[:max_batch_size], data[max_batch_size:]
+        if remain_data:
+            self._task_pool[key] = kwargs, remain_data
+        else:
+            self._task_pool.pop(key)
+        kwargs = kwargs.copy()
+        kwargs['infer_requests'] = [d[0] for d in data]
+        queue_list = [d[1] for d in data]
+        return kwargs, queue_list
+
+    def _infer_worker(self):
+        while True:
+            time.sleep(0.01)
+            item = self._fetch_infer_requests()
+            if item is not None:
+                kwargs, queue_list = item
+                request_config = kwargs['request_config']
+                res_list_or_gen = self._infer(**kwargs)
+                if request_config.stream:
+                    finished = False
+                    while not finished:
+                        try:
+                            res_list = next(res_list_or_gen)
+                        except StopIteration:
+                            finished = True
+                            res_list = [None] * len(queue_list)
+                        for (queue, loop), res in zip(queue_list, res_list):
+                            asyncio.run_coroutine_threadsafe(queue.put(res), loop)
+                else:
+                    for (queue, loop), res in zip(queue_list, res_list_or_gen):
+                        asyncio.run_coroutine_threadsafe(queue.put(res), loop)
+
+    def _add_adapter(self, adapter_path: str, adapter_name: Optional[str] = None) -> None:
+        self.model = Swift.from_pretrained(self.model, adapter_path, adapter_name)
+
+    def _prepare_generation_config(self, request_config: RequestConfig) -> _GenerationConfig:
+        generation_config = prepare_generation_config(self.generation_config, request_config, self.tokenizer)
+        generation_config.return_dict_in_generate = True
+        if request_config.logprobs:
+            generation_config.output_logits = True
+        generation_config.num_return_sequences = request_config.n
+        return _GenerationConfig(**generation_config.to_dict())
+
+    def _add_stop_words(self, generation_config: _GenerationConfig, request_config: RequestConfig) -> None:
+        template_meta = self.template.template_meta
+        stop_words = (request_config.stop or []) + template_meta.stop_words
+        generation_config.stop_words = self._get_stop_words(stop_words)
+
+    @staticmethod
+    def preprocess_logits(batched_logits: Optional[List[torch.Tensor]], batched_generate_ids: torch.Tensor,
+                          top_logprobs: Optional[int]):
+        top_logprobs = top_logprobs or 1
+        batch_size = batched_generate_ids.shape[0]
+        if batched_logits is None:
+            return None
+        batched_logprobs = []
+        for i in range(batch_size):
+            logprobs_list = []
+            generate_ids = batched_generate_ids[i]
+            for j, logits in enumerate(batched_logits):
+                token = generate_ids[j].item()
+                logprobs = torch.log_softmax(logits[i], -1)
+                tokens = [token] + logprobs.argsort(descending=True, dim=-1)[:top_logprobs].tolist()
+                logprobs_list.append({token: logprobs[token].item() for token in tokens})
+            batched_logprobs.append(logprobs_list)
+        return batched_logprobs
+
+    @staticmethod
+    def _update_batched_logprobs(batched_logprobs: List[torch.Tensor], logits_streamer: Optional[LogitsStreamer],
+                                 generate_ids: torch.Tensor, top_logprobs: int) -> None:
+        seq_len = generate_ids.shape[1] - len(batched_logprobs[0])
+        if logits_streamer is None or seq_len == 0:
+            return
+
+        res = []
+        for i in range(seq_len):
+            res.append(logits_streamer.queue.get())
+        new_batched_logprobs = TransformersEngine.preprocess_logits(res, generate_ids[:, -seq_len:], top_logprobs)
+        for logprobs, new_logprobs in zip(batched_logprobs, new_batched_logprobs):
+            logprobs += new_logprobs
+
+    def _infer_stream(self, inputs: Dict[str, Any], *, generation_config: GenerationConfig,
+                      adapter_request: Optional[AdapterRequest], request_config: RequestConfig,
+                      template_inputs) -> Iterator[List[Optional[ChatCompletionStreamResponse]]]:
+        if generation_config.num_beams != 1:
+            error_msg = 'Streaming generation does not support beam search.'
+            raise ValueError(error_msg)
+        streamer = TokensIteratorStreamer()
+        generate_kwargs = {
+            'generation_config': generation_config,
+            'streamer': streamer,
+            **inputs,
+        }
+        adapter_names = self._get_adapter_names(adapter_request)
+        if adapter_names is not None:
+            generate_kwargs['adapter_names'] = adapter_names
+        num_prompt_tokens = self._get_num_tokens(inputs)
+
+        logits_streamer = None
+        if generation_config.output_logits:
+            generate_kwargs['logits_processor'] = LogitsProcessorList([LogitsStreamer()])
+
+        def _model_generate(**kwargs):
+            if is_torch_npu_available():
+                torch.npu.set_device(self.model.device)
+            self.template.generate(self.model, **kwargs)
+
+        generate_kwargs = self.template.prepare_generate_kwargs(generate_kwargs, model=self.model)
+        thread = Thread(target=_model_generate, kwargs=generate_kwargs)
+        thread.start()
+        batch_size = inputs['attention_mask'].shape[0]
+        all_is_finished = False
+        is_finished = [False] * batch_size
+        infer_streamers = [InferStreamer(self.template, template_inputs=template_inputs[i]) for i in range(batch_size)]
+        request_id_list = [f'chatcmpl-{random_uuid()}' for _ in range(batch_size)]
+        token_idxs = [0] * batch_size
+
+        raw_batched_generate_ids = None  # or torch.Tensor: [batch_size, seq_len]
+        batched_logprobs = [[] for _ in range(batch_size)]
+        while not all_is_finished:
+            try:
+                batched_tokens = next(streamer)
+                if batched_tokens.ndim == 1:
+                    batched_tokens = batched_tokens[:, None]
+
+                raw_batched_generate_ids = torch.concat(
+                    [batched_tokens]
+                    if raw_batched_generate_ids is None else [raw_batched_generate_ids, batched_tokens],
+                    dim=1)
+            except StopIteration:
+                all_is_finished = True
+
+            batched_generate_ids = self.template.get_generate_ids(raw_batched_generate_ids, num_prompt_tokens)
+            self._update_batched_logprobs(batched_logprobs, logits_streamer, batched_generate_ids,
+                                          request_config.top_logprobs)
+
+            res = []
+            for i in range(batched_generate_ids.shape[0]):
+                if is_finished[i]:
+                    res.append(None)
+                    continue
+                generate_ids = batched_generate_ids[i]
+
+                # ignore pad_token
+                masks = generate_ids != self.tokenizer.pad_token_id
+                generate_ids = generate_ids[masks].tolist()
+                logprobs_list = None
+                if batched_logprobs[i]:
+                    logprobs_list = [logprobs for m, logprobs in zip(masks, batched_logprobs[i]) if m.item()]
+
+                is_finished[i] = (
+                    all_is_finished or is_finished[i]
+                    or len(generate_ids) > 0 and generate_ids[-1] == self.tokenizer.pad_token_id)
+                delta_text = infer_streamers[i].get_printable_text(generate_ids, is_finished[i])
+                if not delta_text and not is_finished[i]:
+                    res.append(None)
+                    continue
+                logprobs = self._get_logprobs(logprobs_list, generate_ids[token_idxs[i]:], request_config.top_logprobs)
+                token_idxs[i] = len(generate_ids)
+
+                usage_info = self._get_usage_info(num_prompt_tokens, len(generate_ids))
+                toolcall = None
+                if is_finished[i]:
+                    toolcall = self._get_toolcall(
+                        self.template.decode_generate_ids(generate_ids, template_inputs=template_inputs[i]))
+                finish_reason = self._get_finish_reason(generation_config.max_new_tokens, usage_info.completion_tokens,
+                                                        is_finished[i])
+
+                choices = [
+                    ChatCompletionResponseStreamChoice(
+                        index=0,
+                        delta=DeltaMessage(role='assistant', content=delta_text, tool_calls=toolcall),
+                        finish_reason=finish_reason,
+                        logprobs=logprobs)
+                ]
+                res.append(
+                    ChatCompletionStreamResponse(
+                        model=self.model_name, choices=choices, usage=usage_info, id=request_id_list[i]))
+            if any(res):
+                yield res
+
+    def _get_adapter_names(self, adapter_request: Optional[AdapterRequest]) -> Optional[List[str]]:
+        if adapter_request is None:
+            if self._adapters_pool:
+                return ['__base__']
+            return
+        adapter_name = adapter_request.name
+        if adapter_name not in self._adapters_pool:
+            self._adapters_pool[adapter_name] = adapter_request
+            self._add_adapter(adapter_request.path, adapter_name)
+        return [adapter_name]
+
+    def _infer_forward(self, inputs: Dict[str, Any], adapter_request: Optional[AdapterRequest],
+                       request_config: RequestConfig, **kwargs):
+        call_kwargs = {}
+        top_logprobs = request_config.top_logprobs or 20
+        adapter_names = self._get_adapter_names(adapter_request)
+        if adapter_names is not None:
+            call_kwargs['adapter_names'] = adapter_names
+        num_prompt_tokens = self._get_num_tokens(inputs)
+        inputs.pop('labels', None)
+        output = self.model(**inputs, **call_kwargs)
+        if hasattr(output, 'logits'):
+            logits = output.logits
+        elif 'last_hidden_state' in output:
+            # embeddings
+            logits = output['last_hidden_state']
+        else:
+            raise NotImplementedError('Only support `logits` or `hidden_state` in output.')
+        task_type = self.template.task_type
+        if task_type == 'seq_cls':
+            preds, logprobs = self.template.decode_seq_cls(logits, top_logprobs)
+        elif task_type == 'prm':
+            preds = self.template.decode_prm(inputs['input_ids'], logits)
+            logprobs = [None] * len(preds)
+        elif task_type == 'embedding':
+            preds = logits
+            logprobs = [None] * len(preds)
+        elif task_type in ('reranker', 'generative_reranker'):
+            if task_type == 'generative_reranker':
+                attention_mask = inputs.get('attention_mask')
+                last_valid_indices = -1 if attention_mask is None else get_last_valid_indices(attention_mask)
+                batch_indices = torch.arange(logits.shape[0], device=logits.device)
+                logits = logits[batch_indices, last_valid_indices]
+            preds = logits.float()
+            if self.reranker_use_activation:
+                preds = F.sigmoid(preds)
+            preds = preds.tolist()
+            logprobs = [None] * len(preds)
+        else:
+            raise ValueError(f'Unsupported task_type: {task_type}')
+
+        res = []
+        for i, pred in enumerate(preds):
+            usage_info = self._get_usage_info(num_prompt_tokens, 1)
+            if task_type == 'embedding':
+                res.append(
+                    EmbeddingResponse(
+                        model=self.model_name, usage=usage_info, data=[EmbeddingResponseData(embedding=pred.tolist())]))
+            else:
+                choices = [
+                    ChatCompletionResponseChoice(
+                        index=0,
+                        message=ChatMessage(role='assistant', content=pred, tool_calls=None),
+                        finish_reason='stop',
+                        logprobs=logprobs[i])
+                ]
+                res.append(ChatCompletionResponse(model=self.model_name, choices=choices, usage=usage_info))
+        return res
+
+    def _infer_full(self, inputs: Dict[str, Any], *, generation_config: GenerationConfig,
+                    adapter_request: Optional[AdapterRequest], request_config: RequestConfig,
+                    template_inputs) -> List[ChatCompletionResponse]:
+        # bos_token TODO: encoder-decoder
+        generate_kwargs = {'generation_config': generation_config, **inputs}
+        adapter_names = self._get_adapter_names(adapter_request)
+        if adapter_names is not None:
+            generate_kwargs['adapter_names'] = adapter_names
+        num_prompt_tokens = self._get_num_tokens(inputs)
+        generate_kwargs = self.template.prepare_generate_kwargs(generate_kwargs, model=self.model)
+        output = dict(self.template.generate(self.model, **generate_kwargs))
+        output.pop('past_key_values', None)
+        batched_generate_ids = output['sequences']
+        batched_generate_ids = self.template.get_generate_ids(batched_generate_ids, num_prompt_tokens)
+        self.template.debug_logger({'generate_ids': batched_generate_ids})  # debug
+        batched_logprobs = self.preprocess_logits(
+            output.get('logits'), batched_generate_ids, request_config.top_logprobs)
+
+        res = []
+        num_return_sequences = generation_config.num_return_sequences
+        for i in range(inputs['attention_mask'].shape[0]):
+            choices = []
+            usage_info = self._get_usage_info(num_prompt_tokens, 0)
+            for j in range(num_return_sequences):
+                batched_index = i * num_return_sequences + j
+                generate_ids = batched_generate_ids[batched_index]
+
+                # ignore pad_token
+                masks = generate_ids != self.tokenizer.pad_token_id
+                generate_ids = generate_ids[masks].tolist()
+                logprobs_list = None
+                if batched_logprobs is not None:
+                    logprobs_list = [
+                        logprobs for m, logprobs in zip(masks, batched_logprobs[batched_index]) if m.item()
+                    ]
+
+                logprobs = self._get_logprobs(logprobs_list, generate_ids, request_config.top_logprobs)
+                usage_info = self._update_usage_info(usage_info, len(generate_ids))
+                response = self.template.decode_generate_ids(generate_ids, template_inputs=template_inputs[i])
+                finish_reason = self._get_finish_reason(generation_config.max_new_tokens, len(generate_ids), True)
+                toolcall = self._get_toolcall(response)
+                token_ids = generate_ids if request_config.return_details else None
+                choices.append(
+                    ChatCompletionResponseChoice(
+                        index=j,
+                        message=ChatMessage(role='assistant', content=response, tool_calls=toolcall),
+                        finish_reason=finish_reason,
+                        logprobs=logprobs,
+                        token_ids=token_ids))
+            prompt_token_ids = None
+            images_size = None
+            if request_config.return_details:
+                if 'input_ids' in inputs:
+                    non_pad_indices = (inputs['input_ids'][i] != self.tokenizer.pad_token_id).nonzero()
+                    if non_pad_indices.numel() > 0:
+                        idx = non_pad_indices.min().item()
+                        prompt_token_ids = inputs['input_ids'][i][idx:].tolist()
+                if all(isinstance(image, Image.Image) for image in template_inputs[i].images):
+                    images_size = [image.size for image in template_inputs[i].images]
+            res.append(
+                ChatCompletionResponse(
+                    model=self.model_name,
+                    choices=choices,
+                    usage=usage_info,
+                    prompt_token_ids=prompt_token_ids,
+                    images_size=images_size))
+        return res
+
+    async def infer_async(
+        self,
+        infer_request: InferRequest,
+        request_config: Optional[RequestConfig] = None,
+        *,
+        adapter_request: Optional[AdapterRequest] = None,
+        pre_infer_hook=None,
+    ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]:
+        if request_config is None:
+            request_config = RequestConfig()
+        queue = asyncio.Queue()
+        self._queue.put((infer_request, {
+            'request_config': request_config,
+            'adapter_request': adapter_request,
+            'pre_infer_hook': pre_infer_hook
+        }, (queue, asyncio.get_event_loop())))
+        await asyncio.sleep(0)
+        if self._task_thread is None:
+            self._start_infer_worker()
+        if request_config.stream:
+
+            async def _gen_wrapper():
+                while True:
+                    item = await queue.get()
+                    await asyncio.sleep(0)
+                    if item is None:
+                        break
+                    yield item
+
+            return _gen_wrapper()
+        else:
+            return await queue.get()
+
+    # Ensure `template._post_encode` has no gradient.
+    @torch.inference_mode()
+    def _infer(
+        self,
+        infer_requests: List[InferRequest],
+        request_config: RequestConfig,
+        *,
+        adapter_request: Optional[AdapterRequest] = None,
+        pre_infer_hook=None,
+    ) -> Union[List[ChatCompletionResponse], Iterator[List[Optional[ChatCompletionStreamResponse]]]]:
+        self.model.eval()
+        request_config = deepcopy(request_config)
+        if self.template.use_model:
+            self.template.model = self.model
+
+        if self.model_info.task_type == 'causal_lm':
+            self.template.set_mode('transformers')
+
+        batched_inputs, error_list = self._batch_encode(infer_requests, strict=getattr(self, 'strict', True))
+        if len(batched_inputs) > 0:
+            template_inputs = [inputs.pop('template_inputs') for inputs in batched_inputs]
+            inputs = to_device(self.template.data_collator(batched_inputs), self.model.device)
+            self.template.debug_logger(inputs)  # debug
+            if self.model_meta.is_multimodal:
+                _, inputs = self.template.pre_forward_hook(self.model, None, inputs)
+            if self.model_info.task_type == 'causal_lm':
+                self.set_default_max_tokens(request_config, inputs)
+                generation_config = self._prepare_generation_config(request_config)
+                self._add_stop_words(generation_config, request_config)
+            else:
+                generation_config = request_config
+
+            kwargs = {
+                'inputs': inputs,
+                'generation_config': generation_config,
+                'adapter_request': adapter_request,
+                'request_config': request_config,
+                'template_inputs': template_inputs,
+            }
+            if pre_infer_hook:
+                kwargs = pre_infer_hook(kwargs)
+        else:
+            kwargs = {}
+        if request_config.stream:
+
+            def _gen_wrapper():
+                if len(kwargs) > 0:
+                    for res in self._infer_stream(**kwargs):
+                        yield self._add_error_list(res, error_list)
+                else:
+                    yield self._add_error_list([], error_list)
+
+            return _gen_wrapper()
+        else:
+            if len(kwargs) > 0:
+                infer_func = self._infer_forward if self.template.task_type in {
+                    'seq_cls', 'prm', 'embedding', 'reranker', 'generative_reranker'
+                } else self._infer_full
+                res = infer_func(**kwargs)
+            else:
+                res = []
+            return self._add_error_list(res, error_list)
+
+    def infer(
+        self,
+        infer_requests: List[InferRequest],
+        request_config: Optional[RequestConfig] = None,
+        metrics: Optional[List[Metric]] = None,
+        *,
+        use_tqdm: Optional[bool] = None,
+        adapter_request: Optional[AdapterRequest] = None
+    ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
+        if request_config is None:
+            request_config = RequestConfig()
+        if request_config.stream:
+            return super().infer(
+                infer_requests, request_config, metrics, use_tqdm=use_tqdm, adapter_request=adapter_request)
+        # Has higher stability than calling super().infer
+        if use_tqdm is None:
+            use_tqdm = not request_config.stream and len(infer_requests) > 1
+        prog_bar = tqdm(total=len(infer_requests), dynamic_ncols=True, disable=not use_tqdm)
+        # If self.max_batch_size <= 0, then process all infer_requests at once.
+        max_batch_size = self.max_batch_size
+        if max_batch_size <= 0:
+            max_batch_size = len(infer_requests)
+        res = []
+        i = 0
+        while i < len(infer_requests):
+            infer_requests_samples = infer_requests[i:i + max_batch_size]
+            res += self._infer(infer_requests_samples, request_config, adapter_request=adapter_request)
+            i += max_batch_size
+            prog_bar.update(len(infer_requests_samples))
+        prog_bar.close()
+        self._update_metrics(res, metrics)
+        return res

swift/infer_engine/utils.py

@@ -0,0 +1,633 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+import os
+import re
+import torch
+from collections import OrderedDict
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager, nullcontext
+from dataclasses import dataclass
+from itertools import repeat
+from packaging import version
+from queue import Queue
+from transformers import GenerationConfig, LogitsProcessor
+from transformers.generation.streamers import BaseStreamer
+from typing import List, Optional, Union
+
+from swift.model.register import fix_do_sample_warning
+from swift.utils import get_device, synchronize
+from .protocol import RequestConfig
+
+
+@dataclass
+class AdapterRequest:
+    name: str
+    path: str
+
+
+class InferTools:
+
+    @staticmethod
+    def _is_chinese_char(cp: int) -> bool:
+        """Checks whether CP is the codepoint of a CJK character."""
+        # copy from transformers.generation.streamers.TextStreamer
+        if ((0x4E00 <= cp <= 0x9FFF) or (0x3400 <= cp <= 0x4DBF) or (0x20000 <= cp <= 0x2A6DF)
+                or (0x2A700 <= cp <= 0x2B73F) or (0x2B740 <= cp <= 0x2B81F) or (0x2B820 <= cp <= 0x2CEAF)
+                or (0xF900 <= cp <= 0xFAFF) or (0x2F800 <= cp <= 0x2FA1F)):
+            return True
+
+        return False
+
+
+class InferStreamer(InferTools):
+
+    def __init__(self, template, **decode_kwargs):
+        self.template = template
+        self.tokenizer = template.tokenizer
+
+        self.cache_idx = 0  # token idx
+        self.print_idx = 0
+        self.decode_kwargs = decode_kwargs
+        self.first_num_space = -1  # The number of whitespace characters before the first token.
+        self.first_token = True
+
+    def _align_blank_suffix(self, response: str) -> str:
+        # Avoid the occurrence of repeated words in sentence.
+        cur_num_space = len(response) - len(response.lstrip(' '))
+        if self.first_num_space == -1:
+            self.first_num_space = cur_num_space
+        elif cur_num_space < self.first_num_space:
+            response = ' ' * (self.first_num_space - cur_num_space) + response
+        elif cur_num_space > self.first_num_space:
+            response = response[cur_num_space - self.first_num_space:]
+        return response
+
+    def _get_response(self, response: str, is_finished: bool, token_len: int) -> str:
+        # After the symbol for a new line, we flush the cache.
+        if self.first_token:
+            printable_text = response
+            self.first_token = False
+        elif response.endswith('\n') or is_finished:
+            printable_text = response[self.print_idx:]
+            self.cache_idx += token_len
+            self.first_num_space = -1
+            self.print_idx = 0
+        # If the last token is a CJK character, we print the characters.
+        elif len(response) > 0 and self._is_chinese_char(ord(response[-1])):
+            printable_text = response[self.print_idx:]
+            self.print_idx += len(printable_text)
+        # Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words,
+        # which may change with the subsequent token -- there are probably smarter ways to do this!)
+        else:
+            printable_text = response[self.print_idx:response.rfind(' ') + 1]
+            self.print_idx += len(printable_text)
+        return printable_text
+
+    def get_printable_text(self, raw_tokens: List[int], is_finished: bool) -> str:
+        raw_tokens = raw_tokens[self.cache_idx:]
+        if self.first_token:
+            raw_tokens = []
+        response = self.template.decode_generate_ids(
+            raw_tokens, is_finished=is_finished, first_token=self.first_token, **self.decode_kwargs)
+        response = self._align_blank_suffix(response)
+        return self._get_response(response, is_finished, len(raw_tokens))
+
+
+class StreamerMixin:
+
+    def __init__(self):
+        self.queue = Queue()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> torch.Tensor:
+        value = self.queue.get()
+        if value is None:
+            raise StopIteration()
+        else:
+            return value
+
+
+class TokensIteratorStreamer(StreamerMixin, BaseStreamer):
+
+    def put(self, value: torch.Tensor) -> None:
+        self.queue.put(value)
+
+    def end(self) -> None:
+        self.queue.put(None)
+
+
+class LogitsStreamer(LogitsProcessor):
+
+    def __init__(self):
+        self.queue = Queue()
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        self.queue.put(scores)
+        return scores
+
+
+def _set_generation_config_default_value(model_generation_config: GenerationConfig,
+                                         generation_config: GenerationConfig) -> GenerationConfig:
+    for k, v in model_generation_config.to_dict().items():
+        new_v = getattr(generation_config, k, None)
+        if k in ['max_length']:
+            continue
+        if k in ['no_repeat_ngram_size'] or v is not None and new_v is None:
+            setattr(generation_config, k, v)
+    return generation_config
+
+
+def prepare_generation_config(model_generation_config: Optional[GenerationConfig], request_config: RequestConfig,
+                              tokenizer) -> Optional[GenerationConfig]:
+    if model_generation_config is None or request_config is None:
+        return model_generation_config
+    kwargs = {'max_new_tokens': request_config.max_tokens}
+    # not use: 'n', 'best_of', 'frequency_penalty', 'presence_penalty'
+    for key in ['length_penalty']:
+        kwargs[key] = getattr(request_config, key)
+    for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty', 'num_beams']:
+        new_value = getattr(request_config, key)
+        if new_value is None:
+            kwargs[key] = getattr(model_generation_config, key, None)
+        else:
+            kwargs[key] = new_value
+
+    if kwargs.get('top_k') is not None and kwargs['top_k'] <= 0:
+        kwargs['top_k'] = None
+
+    if not getattr(model_generation_config, 'do_sample', False) and request_config.temperature in {0, None}:
+        kwargs['temperature'] = 0
+    if kwargs['temperature'] == 0:
+        kwargs['do_sample'] = False
+        kwargs['temperature'] = 1
+        kwargs['top_p'] = 1
+        kwargs['top_k'] = 50
+    else:
+        kwargs['do_sample'] = True
+    generation_config = GenerationConfig(**kwargs)
+    generation_config = _set_generation_config_default_value(model_generation_config, generation_config)
+    fix_do_sample_warning(generation_config)
+
+    if generation_config.eos_token_id is None:
+        generation_config.eos_token_id = tokenizer.eos_token_id
+    generation_config.pad_token_id = tokenizer.pad_token_id
+    return generation_config
+
+
+def patch_lmdeploy(load_weights=False):
+    """This patch allows lmdeploy selects device and reload state_dict"""
+    import lmdeploy
+    assert version.parse(lmdeploy.__version__) >= version.parse('0.7.0')
+    from lmdeploy.messages import TurbomindEngineConfig
+    from lmdeploy.turbomind.deploy import loader
+    from lmdeploy.turbomind.deploy.loader import create_loader
+    from lmdeploy.turbomind.deploy.source_model import llama
+
+    def _create_loader(model_path: str, pattern: str):
+        if not isinstance(model_path, (str, os.PathLike)):
+
+            def generate():
+                generator = OrderedDict()
+                model_dict = {}
+                if not isinstance(model_path, dict):
+                    for key, value in list(model_path):
+                        model_dict[key] = value
+                else:
+                    model_dict = model_path
+                for key, value in model_dict.items():
+                    match = re.findall(pattern, key)
+                    if not match:
+                        if -1 not in generator:
+                            generator[-1] = {}
+                        generator[-1][key] = value
+                    else:
+                        layer = int(match[0])
+                        if layer not in generator:
+                            generator[layer] = {}
+                        generator[layer][key] = value
+                return generator
+
+            return generate()
+        else:
+            return create_loader(model_path, pattern)
+
+    loader.create_loader = _create_loader
+    llama.create_loader = _create_loader
+
+    TurbomindEngineConfig.devices = [0]
+
+    from lmdeploy.turbomind.turbomind import TurboMind
+    from lmdeploy.turbomind.utils import ModelSource
+
+    @contextmanager
+    def patch_threadpool_map():
+        ThreadPoolExecutor.map_origin = ThreadPoolExecutor.map
+        ThreadPoolExecutor.map = lambda *args, **kwargs: []
+        yield
+        ThreadPoolExecutor.map = ThreadPoolExecutor.map_origin
+        del ThreadPoolExecutor.map_origin
+
+    @contextmanager
+    def tm_model_context(self):
+
+        def _get_tm_model(model_path,
+                          model_name,
+                          chat_template_name,
+                          engine_config: TurbomindEngineConfig,
+                          group_size: int = None,
+                          out_dir: str = None):
+            from lmdeploy.turbomind.deploy.converter import get_tm_model_origin
+            tm_model = get_tm_model_origin(model_path, model_name, chat_template_name, engine_config, group_size,
+                                           out_dir)
+            self.tm_model = tm_model
+            return tm_model
+
+        from lmdeploy.turbomind.deploy import converter
+        converter.get_tm_model_origin = converter.get_tm_model
+        converter.get_tm_model = _get_tm_model
+        yield
+        converter.get_tm_model = converter.get_tm_model_origin
+        del converter.get_tm_model_origin
+
+    def __init__(self,
+                 model_path: str,
+                 tokenizer: object,
+                 model_name: str = None,
+                 chat_template_name: str = None,
+                 engine_config: TurbomindEngineConfig = None,
+                 model_source: ModelSource = ModelSource.WORKSPACE,
+                 **kwargs):
+        self.gpu_list = engine_config.devices
+        with patch_threadpool_map(), tm_model_context(self):
+            self.__origin_init__(model_path, tokenizer, model_name, chat_template_name, engine_config, model_source,
+                                 **kwargs)
+
+        with ThreadPoolExecutor(max_workers=self.gpu_count) as e:
+            ranks = [self.node_id * self.gpu_count + device_id for device_id in range(self.gpu_count)]
+            if not load_weights:
+                for _ in e.map(self.model_comm.process_weight, self.gpu_list, ranks):
+                    pass
+            if version.parse(lmdeploy.__version__) < version.parse('0.7.2'):
+                for _ in e.map(self.model_comm.create_engine, self.gpu_list, ranks, repeat(self.nccl_params)):
+                    pass
+            else:
+                for _ in e.map(self.model_comm.create_engine, self.gpu_list, ranks):
+                    pass
+
+    def _create_weight(self, model_comm):
+        """Allocate weight buffer, load params if from_workspace."""
+
+        # TODO: support mpi
+        self.node_id = 0
+        self.node_num = 1
+        if version.parse(lmdeploy.__version__) < version.parse('0.7.2'):
+            self.nccl_params = model_comm.create_nccl_params(self.node_id)
+        synchronize()
+
+        # create weight
+        def _create_weight_func(index, device_id):
+            rank = self.node_id * self.gpu_count + index
+            model_comm.create_shared_weights(device_id, rank)
+
+        with ThreadPoolExecutor(max_workers=self.gpu_count) as executor:
+            futures = []
+            for idx, device_id in enumerate(self.gpu_list):
+                futures.append(executor.submit(_create_weight_func, idx, device_id))
+            for future in futures:
+                future.result()
+
+    def _get_model_params(self, model_comm, tm_params):
+        """Get turbomind model params when loading from hf."""
+
+        def _get_params(idx, device_id, que):
+            rank = self.node_id * self.gpu_count + idx
+            out = model_comm.get_params(device_id, rank)
+            que.put(out)
+
+        que = Queue()
+        with ThreadPoolExecutor(max_workers=self.gpu_count) as executor:
+            futures = []
+            for idx, device_id in enumerate(self.gpu_list):
+                futures.append(executor.submit(_get_params, idx, device_id, que))
+            for future in futures:
+                future.result()
+
+        for _ in range(self.gpu_count):
+            tensor_map = que.get()
+            for k, v in tensor_map.items():
+                if k not in tm_params:
+                    tm_params[k] = []
+                tm_params[k].append(v)
+
+    def _load_weights(self, state_dict):
+        tm_params = self.tm_model.tm_params
+        self._get_model_params(self.model_comm, tm_params)
+        input_model = self.tm_model.input_model
+        model_path = input_model.model_path
+        input_model.model_path = state_dict
+        self.tm_model.export()
+        input_model.model_path = model_path
+
+    from lmdeploy.turbomind.turbomind import TurboMindInstance
+
+    def create_instance(self, cuda_stream_id=0):
+        return TurboMindInstance(self, self.config, cuda_stream_id, self.gpu_list)
+
+    TurboMind.__origin_init__ = TurboMind.__init__
+    TurboMind.__init__ = __init__
+    TurboMind._create_weight = _create_weight
+    TurboMind._get_model_params = _get_model_params
+    TurboMind.create_instance = create_instance
+    if load_weights:
+        TurboMind.load_weights = _load_weights
+
+    def __init_ins__(self, tm_model, config, cuda_stream_id=0, gpu_list=None):
+        if gpu_list is None:
+            gpu_list = [0]
+        self.gpu_list = gpu_list
+        self.__origin_init__(tm_model, config, cuda_stream_id)
+
+    def _create_model_instance(self, device_id):
+        model_inst = self.tm_model.model_comm.create_model_instance(self.gpu_list[0])
+        return model_inst
+
+    TurboMindInstance.__origin_init__ = TurboMindInstance.__init__
+    TurboMindInstance.__init__ = __init_ins__
+    TurboMindInstance._create_model_instance = _create_model_instance
+
+
+def patch_npu_vllm(vllm_device: str, *, colocate: bool = False):
+    if isinstance(vllm_device, int):
+        vllm_device = get_device(vllm_device)
+    device_type = vllm_device.split(':')[0]
+    if device_type == 'npu':
+        from swift.model.npu_patch.vllm_ascend import patch_vllm_ascend_runtime
+        from swift.model.npu_patch.vllm_ascend_memory import vllm_ascend_mem_get_info_context
+        patch_vllm_ascend_runtime(colocate=colocate)
+        return vllm_ascend_mem_get_info_context(vllm_device)
+
+    return nullcontext()
+
+
+def patch_vllm_triton_device_guard():
+    import functools
+    try:
+        from vllm.v1.worker import gpu_worker as _gw
+        _orig_fn = _gw.init_worker_distributed_environment
+    except (ImportError, AttributeError):
+        return
+
+    if getattr(_gw, '_swift_dist_env_patched', False):
+        return
+
+    @functools.wraps(_orig_fn)
+    def _patched_init_worker_distributed_environment(*args, **kwargs):
+        if not torch.cuda.is_available():
+            return _orig_fn(*args, **kwargs)
+        expected_device = torch.cuda.current_device()
+        result = _orig_fn(*args, **kwargs)
+        actual_device = torch.cuda.current_device()
+        if actual_device != expected_device:
+            torch.cuda.set_device(expected_device)
+        return result
+
+    _gw.init_worker_distributed_environment = _patched_init_worker_distributed_environment
+    _gw._swift_dist_env_patched = True
+
+
+def patch_vllm_memory_leak():
+    # fix vllm 0.7.3 memory leak
+    # https://github.com/vllm-project/vllm/pull/14326
+    import vllm
+    try:
+        vllm_version = version.parse(vllm.__version__)
+        needs_patch = (vllm_version == version.parse('0.7.3'))
+    except version.InvalidVersion:
+        needs_patch = False
+
+    if not needs_patch:
+        return
+
+    def patch_vllm_abort_seq_group():
+        from typing import Dict, Iterable
+        from vllm.core.scheduler import Scheduler
+        from vllm.sequence import SequenceGroup, SequenceGroupBase, SequenceStatus
+
+        def new_abort_seq_group(
+            self,
+            request_id: Union[str, Iterable[str]],
+            seq_id_to_seq_group: Optional[Dict[str, SequenceGroupBase]] = None,
+        ) -> None:
+            if isinstance(request_id, str):
+                request_id = (request_id, )
+            request_ids = set(request_id)
+            seq_id_to_seq_group = seq_id_to_seq_group or {}
+            for state_queue in [self.waiting, self.running, self.swapped]:
+                aborted_groups: List[SequenceGroup] = []
+                for seq_group in state_queue:
+                    # When n>1, seq_group.request_id looks like
+                    # foo_parallel_sample_0, while request_ids is just foo, and we
+                    # should resolve it as real_request_id to match.
+                    if seq_group.request_id in seq_id_to_seq_group:
+                        real_request_id = seq_id_to_seq_group[seq_group.request_id].group_id
+                    else:
+                        real_request_id = seq_group.request_id
+                    if real_request_id in request_ids:
+                        # Appending aborted group into pending list.
+                        aborted_groups.append(seq_group)
+                        # We can't remove real_request_id in request_ids here,
+                        # because there may be other seq groups sharing the same
+                        # real_request_id
+                for aborted_group in aborted_groups:
+                    # Remove the sequence group from the state queue.
+                    state_queue.remove(aborted_group)
+                    # Remove the aborted request from the Mamba cache.
+                    self._finished_requests_ids.append(aborted_group.request_id)
+                    for seq in aborted_group.get_seqs():
+                        if seq.is_finished():
+                            continue
+                        seq.status = SequenceStatus.FINISHED_ABORTED
+                        self.free_seq(seq)
+                    if aborted_group.request_id in seq_id_to_seq_group:
+                        del seq_id_to_seq_group[aborted_group.request_id]
+
+                    self._free_seq_group_cross_attn_blocks(aborted_group)
+
+        origin_method = Scheduler.abort_seq_group
+        Scheduler._old_abort_seq_group = origin_method
+        Scheduler.abort_seq_group = new_abort_seq_group
+
+    def patch_vllm_engine():
+        from vllm.engine.llm_engine import LLMEngine, SchedulerOutputState
+        from vllm.outputs import PoolingRequestOutput, RequestOutput
+        from vllm.sequence import ExecuteModelRequest
+
+        def new_abort_request(self, request_id) -> None:
+            for scheduler in self.scheduler:
+                scheduler.abort_seq_group(request_id, seq_id_to_seq_group=self.seq_id_to_seq_group)
+
+        origin_method = LLMEngine.abort_request
+        LLMEngine._old_abort_request = origin_method
+        LLMEngine.abort_request = new_abort_request
+
+        def new_step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]:
+            if self.parallel_config.pipeline_parallel_size > 1:
+                raise NotImplementedError('Pipeline parallelism is only supported through AsyncLLMEngine '
+                                          'as performance will be severely degraded otherwise.')
+
+            # For llm_engine, there is no pipeline parallel support, so the engine
+            # used is always 0.
+            virtual_engine = 0
+
+            # These are cached outputs from previous iterations. None if on first
+            # iteration
+            cached_outputs = self.cached_scheduler_outputs[virtual_engine]
+            seq_group_metadata_list = cached_outputs.seq_group_metadata_list
+            scheduler_outputs = cached_outputs.scheduler_outputs
+            allow_async_output_proc = cached_outputs.allow_async_output_proc
+
+            ctx = self.scheduler_contexts[virtual_engine]
+
+            # Clear outputs for each new scheduler iteration
+            ctx.request_outputs.clear()
+
+            # Skip the scheduler if there are any remaining steps in the seq groups.
+            # This ensures that the scheduler is only called again when the current
+            # batch has completed.
+            # The scheduler is also skipped if a single request caused the last
+            # engine step to fail, and the previous schedule needs to be rerun.
+            if not self._has_remaining_steps(seq_group_metadata_list):
+                # Schedule iteration
+                (seq_group_metadata_list, scheduler_outputs,
+                 allow_async_output_proc) = self.scheduler[virtual_engine].schedule()
+
+                ctx.seq_group_metadata_list = seq_group_metadata_list
+                ctx.scheduler_outputs = scheduler_outputs
+
+                finished_requests_ids = self.scheduler[virtual_engine].get_and_reset_finished_requests_ids()
+                # When n>1, elements in self.seq_id_to_seq_group should be deleted
+                # here, otherwise memory leaks.
+                for finished_request_id in finished_requests_ids:
+                    if finished_request_id in self.seq_id_to_seq_group:
+                        del self.seq_id_to_seq_group[finished_request_id]
+
+                # Maybe switch from async mode to sync mode
+                if not allow_async_output_proc and len(ctx.output_queue) > 0:
+                    self._process_model_outputs(ctx=ctx)
+
+                if (self.scheduler_config.is_multi_step and scheduler_outputs.num_lookahead_slots > 0):
+                    # cache the scheduler outputs for the next iteration if we have
+                    # lookahead slots
+                    self._cache_scheduler_outputs_for_multi_step(virtual_engine, seq_group_metadata_list,
+                                                                 scheduler_outputs, allow_async_output_proc)
+            else:
+                finished_requests_ids = list()
+
+            assert seq_group_metadata_list is not None
+            assert scheduler_outputs is not None
+
+            if not scheduler_outputs.is_empty():
+
+                # Check if we have a cached last_output from the previous iteration.
+                # For supporting PP this is probably the best way to pass the
+                # sampled_token_ids, as a separate broadcast over all the PP stages
+                # will cause one virtual engine's microbatch to block the pipeline.
+                last_sampled_token_ids = \
+                    self._get_last_sampled_token_ids(virtual_engine)
+
+                execute_model_req = ExecuteModelRequest(
+                    seq_group_metadata_list=seq_group_metadata_list,
+                    blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
+                    blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
+                    blocks_to_copy=scheduler_outputs.blocks_to_copy,
+                    num_lookahead_slots=scheduler_outputs.num_lookahead_slots,
+                    running_queue_size=scheduler_outputs.running_queue_size,
+                    finished_requests_ids=finished_requests_ids,
+                    # We use ExecuteModelRequest to pass the last sampled_token_ids
+                    # to each of the non-last PP stages for in-place prepare_input.
+                    last_sampled_token_ids=last_sampled_token_ids)
+
+                if allow_async_output_proc:
+                    execute_model_req.async_callback = self.async_callbacks[virtual_engine]
+
+                outputs = self.model_executor.execute_model(execute_model_req=execute_model_req)
+
+                # We need to do this here so that last step's sampled_token_ids can
+                # be passed to the next iteration for PP.
+                if self.scheduler_config.is_multi_step:
+                    self._update_cached_scheduler_output(virtual_engine, outputs)
+            else:
+                # Nothing scheduled => If there is pending async postprocessor,
+                # then finish it here.
+                if len(ctx.output_queue) > 0:
+                    self._process_model_outputs(ctx=ctx)
+                # No outputs in this case
+                outputs = []
+
+            # Finish the current step for all the sequence groups.
+            if self.scheduler_config.is_multi_step:
+                for seq_group in seq_group_metadata_list:
+                    seq_group.finish_step()
+
+            if not self._has_remaining_steps(seq_group_metadata_list):
+                # clear the cache if we have finished all the steps.
+                if self.scheduler_config.is_multi_step:
+                    self.cached_scheduler_outputs[0] = SchedulerOutputState()
+
+                # is_first_step_output is True only when the num_steps of all
+                # the sequences are 1. When the num_steps > 1,
+                # multi_step_model_runner does the first-step output append.
+                is_first_step_output: bool = False if not seq_group_metadata_list \
+                    else seq_group_metadata_list[0].state.num_steps == 1
+
+                # Add results to the output_queue
+                ctx.append_output(
+                    outputs=outputs,
+                    seq_group_metadata_list=seq_group_metadata_list,
+                    scheduler_outputs=scheduler_outputs,
+                    is_async=allow_async_output_proc,
+                    is_last_step=True,
+                    is_first_step_output=is_first_step_output)
+
+                if outputs and allow_async_output_proc:
+                    assert len(outputs) == 1, ('Async postprocessor expects only a single output set')
+
+                    self._advance_to_next_step(outputs[0], seq_group_metadata_list,
+                                               scheduler_outputs.scheduled_seq_groups)
+
+                # Check if need to run the usual non-async path
+                if not allow_async_output_proc:
+                    self._process_model_outputs(ctx=ctx)
+
+                    # Log stats.
+                    self.do_log_stats(scheduler_outputs, outputs)
+
+                    # Tracing
+                    self.do_tracing(scheduler_outputs)
+            else:
+                # Multi-step case
+                return ctx.request_outputs
+
+            if not self.has_unfinished_requests():
+                # Drain async postprocessor (if exists)
+                if len(ctx.output_queue) > 0:
+                    self._process_model_outputs(ctx=ctx)
+                assert len(ctx.output_queue) == 0
+
+                # Stop the execute model loop in parallel workers until there are
+                # more requests to process. This avoids waiting indefinitely in
+                # torch.distributed ops which may otherwise timeout, and unblocks
+                # the RPC thread in the workers so that they can process any other
+                # queued control plane messages, such as add/remove lora adapters.
+                self.model_executor.stop_remote_worker_execution_loop()
+
+            return ctx.request_outputs
+
+        origin_method = LLMEngine.step
+        LLMEngine._old_step = origin_method
+        LLMEngine.step = new_step
+
+    patch_vllm_abort_seq_group()
+    patch_vllm_engine()

swift/infer_engine/vllm_engine.py

@@ -0,0 +1,938 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import asyncio
+import inspect
+import multiprocessing
+import os
+import time
+import torch
+from contextlib import contextmanager, nullcontext
+from copy import copy, deepcopy
+from packaging import version
+from PIL import Image
+from tqdm import tqdm
+from transformers import AutoConfig, GenerationConfig
+from transformers.utils import is_torch_npu_available
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
+
+from swift.metrics import Metric
+from swift.model import get_processor
+from swift.template import Template
+from swift.utils import (disable_deepspeed_zero3, get_device, get_dist_setting, get_logger, is_dist,
+                         safe_snapshot_download)
+from .infer_engine import InferEngine
+from .patch import patch_auto_tokenizer
+from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
+                       ChatCompletionStreamResponse, ChatMessage, DeltaMessage, EmbeddingResponse,
+                       EmbeddingResponseData, InferRequest, RequestConfig, random_uuid)
+from .utils import AdapterRequest, InferStreamer, patch_npu_vllm, patch_vllm_memory_leak, patch_vllm_triton_device_guard
+
+logger = get_logger()
+try:
+    # After setting the environment variables, import vllm. This way of writing allows lint to pass.
+    os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'
+    os.environ['VLLM_ENGINE_ITERATION_TIMEOUT_S'] = '86400'
+    import vllm
+    from vllm import AsyncEngineArgs, AsyncLLMEngine, EngineArgs, LLMEngine, SamplingParams
+    from vllm.pooling_params import PoolingParams
+    try:
+        # vLLM v0.12+ uses StructuredOutputsParams
+        from vllm.sampling_params import StructuredOutputsParams
+    except ImportError:
+        # Fallback for older vLLM versions
+        from vllm.sampling_params import GuidedDecodingParams as StructuredOutputsParams
+except Exception:
+    raise
+
+try:
+    from vllm.reasoning import ReasoningParserManager
+except ImportError:
+    ReasoningParserManager = None
+
+dtype_mapping = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}
+
+
+def _patch_vllm_dp_coordinator_timeout():
+    # https://github.com/vllm-project/vllm/pull/37452 introduced a 30-second default timeout,
+    # which is prone to timing out in spawn scenarios. Patch it to 180 seconds here.
+    try:
+        from vllm.v1.engine import coordinator as coordinator_module
+    except ImportError:
+        return
+
+    coordinator_cls = coordinator_module.DPCoordinator
+    if not hasattr(coordinator_cls, '_wait_for_zmq_addrs'):
+        return
+
+    if getattr(coordinator_cls, '_swift_timeout_patched', False):
+        return
+
+    def _wait_for_zmq_addrs(self, zmq_addr_pipe):
+        t0 = time.monotonic()
+        try:
+            ready = multiprocessing.connection.wait([zmq_addr_pipe, self.proc.sentinel], timeout=180)
+            elapsed = time.monotonic() - t0
+            if not ready:
+                raise RuntimeError(f'DP Coordinator process failed to report ZMQ addresses '
+                                   f'within 180s (elapsed={elapsed:.1f}s).')
+            try:
+                return zmq_addr_pipe.recv()
+            except EOFError:
+                raise RuntimeError('DP Coordinator process failed during startup.') from None
+        finally:
+            zmq_addr_pipe.close()
+
+    coordinator_cls._wait_for_zmq_addrs = _wait_for_zmq_addrs
+    coordinator_cls._swift_timeout_patched = True
+
+
+_patch_vllm_dp_coordinator_timeout()
+
+
+@contextmanager
+def _patch_rope_validation_ignore_keys():
+    """Accept list-style RoPE validation ignore keys from older vLLM configs.
+
+    vLLM 0.18.x Qwen3.5 configs may pass ``ignore_keys_at_rope_validation``
+    as a list, while Transformers 5.x treats it as a set and performs a set
+    union during RoPE validation. vLLM release tags from 0.19.0 onward changed
+    the Qwen3.5 configs to set literals, but 0.18-based vLLM/vLLM-Ascend stacks
+    still need this compatibility layer. See vLLM PR:
+    https://github.com/vllm-project/vllm/pull/37338
+    """
+    from transformers import PretrainedConfig
+
+    origin_convert = getattr(PretrainedConfig, 'convert_rope_params_to_dict', None)
+    if origin_convert is None:
+        yield
+        return
+
+    def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation=None, **kwargs):
+        if isinstance(ignore_keys_at_rope_validation, list):
+            ignore_keys_at_rope_validation = set(ignore_keys_at_rope_validation)
+        return origin_convert(self, ignore_keys_at_rope_validation=ignore_keys_at_rope_validation, **kwargs)
+
+    PretrainedConfig.convert_rope_params_to_dict = convert_rope_params_to_dict
+    try:
+        yield
+    finally:
+        PretrainedConfig.convert_rope_params_to_dict = origin_convert
+
+
+class VllmEngine(InferEngine):
+
+    def __init__(
+        self,
+        model_id_or_path: str,
+        *,
+        template: Optional[Template] = None,
+        torch_dtype: Optional[torch.dtype] = None,
+        adapters: Optional[List[str]] = None,
+        use_async_engine: bool = False,
+        model_type: Optional[str] = None,
+        template_type: Optional[str] = None,
+        use_hf: Optional[bool] = None,
+        hub_token: Optional[str] = None,
+        revision: Optional[str] = None,
+        # engine_kwargs
+        gpu_memory_utilization: float = 0.9,
+        tensor_parallel_size: int = 1,
+        pipeline_parallel_size: int = 1,
+        enable_expert_parallel: bool = False,
+        max_model_len: Optional[int] = None,
+        max_num_seqs: int = 256,
+        disable_custom_all_reduce: bool = True,
+        enforce_eager: bool = False,
+        limit_mm_per_prompt: Optional[Dict[str, Any]] = None,
+        seed: Optional[int] = None,
+        task_type: Optional[str] = None,  # embedding
+        disable_cascade_attn: bool = False,
+        load_format: str = 'auto',
+        mm_processor_cache_gb: Optional[float] = None,
+        logprobs_mode: Optional[str] = None,
+        speculative_config: Optional[Union[str, dict]] = None,
+        # lora
+        enable_lora: bool = False,
+        max_loras: int = 1,
+        max_lora_rank: int = 16,
+        enable_prefix_caching: Optional[bool] = None,
+        enable_sleep_mode: bool = False,
+        distributed_executor_backend: Optional[str] = None,
+        quantization: Optional[str] = None,
+        # reasoning parser
+        reasoning_parser: Optional[str] = None,
+        engine_kwargs: Optional[Dict[str, Any]] = None,
+        num_labels: Optional[int] = None,
+        reranker_use_activation: bool = True,
+    ) -> None:
+        self.model_id_or_path = model_id_or_path
+        self.torch_dtype = torch_dtype
+        if isinstance(adapters, str):
+            adapters = [adapters]
+        self.default_adapter_request = None
+        if isinstance(adapters, list) and adapters:
+            assert len(adapters) == 1, 'Only one adapter is supported for now.'
+            enable_lora = True
+            self.default_adapter_request = AdapterRequest('default', adapters[0])
+        self.adapters = adapters or []
+        self.use_async_engine = use_async_engine
+        self.model_type = model_type
+        self.use_hf = use_hf
+        self.hub_token = hub_token
+        self.revision = revision
+
+        self.gpu_memory_utilization = gpu_memory_utilization
+        self.tensor_parallel_size = tensor_parallel_size
+        self.pipeline_parallel_size = pipeline_parallel_size
+        self.enable_expert_parallel = enable_expert_parallel
+        self.max_num_seqs = max_num_seqs
+        self.disable_custom_all_reduce = disable_custom_all_reduce
+        self.enforce_eager = enforce_eager
+        self.limit_mm_per_prompt = limit_mm_per_prompt
+        self.seed = seed
+        self.task_type = task_type
+        self.disable_cascade_attn = disable_cascade_attn
+        self.load_format = load_format
+        self.mm_processor_cache_gb = mm_processor_cache_gb
+        self.logprobs_mode = logprobs_mode
+        self.speculative_config = speculative_config
+
+        self.enable_lora = enable_lora
+        self.max_loras = max_loras
+        self.max_lora_rank = max_lora_rank
+        self.enable_prefix_caching = enable_prefix_caching
+        self.enable_sleep_mode = enable_sleep_mode
+        self.distributed_executor_backend = distributed_executor_backend
+        self.quantization = quantization
+        self.num_labels = num_labels
+        self.reranker_use_activation = reranker_use_activation
+        self._config_cls = None
+
+        patch_vllm_memory_leak()
+        patch_vllm_triton_device_guard()
+        self._adapters_pool = {}
+        if template is None:
+            processor = self._get_processor()
+            template = self._get_template(processor, template_type=template_type)
+        else:
+            safe_snapshot_download(
+                model_id_or_path,
+                revision=revision,
+                download_model=True,
+                use_hf=use_hf,
+                ignore_patterns=getattr(template.model_meta, 'ignore_patterns', None),
+                hub_token=hub_token)
+        super().__init__(template)
+        if max_model_len is not None:
+            self.max_model_len = max_model_len
+            logger.info(f'Setting max_model_len: {max_model_len}')
+        self._prepare_engine_kwargs(max_model_len, engine_kwargs)
+        context = nullcontext()
+        if is_torch_npu_available() and (tensor_parallel_size == 1 or pipeline_parallel_size == 1):
+            colocate = (
+                getattr(self, '_swift_vllm_colocate_runtime', False)
+                or self.distributed_executor_backend == 'external_launcher')
+            context = patch_npu_vllm(get_device(), colocate=colocate)
+        with context:
+            self._prepare_engine()
+        self._load_generation_config()
+        self._fix_vllm_bug()
+        self.patch_remove_log()
+        self._request_count = 0
+        self._prepare_reasoning_parser(reasoning_parser)
+
+    def _get_processor(self):
+        return get_processor(
+            model_id_or_path=self.model_id_or_path,
+            torch_dtype=self.torch_dtype,
+            download_model=True,
+            model_type=self.model_type,
+            use_hf=self.use_hf,
+            hub_token=self.hub_token,
+            revision=self.revision,
+            num_labels=self.num_labels,
+            task_type=self.task_type)
+
+    def _prepare_engine(self) -> None:
+        with patch_auto_tokenizer(self.tokenizer), self._patch_auto_config(), \
+                _patch_rope_validation_ignore_keys(), disable_deepspeed_zero3():
+            llm_engine_cls = AsyncLLMEngine if self.use_async_engine else LLMEngine
+            engine = llm_engine_cls.from_engine_args(self.engine_args)
+        self.engine = engine
+
+    @contextmanager
+    def _patch_auto_config(self):
+        _old_from_pretrained = AutoConfig.from_pretrained
+
+        def _from_pretrained(*args, **kwargs):
+            config = deepcopy(self.config)
+            if self._version_ge('0.19'):
+                if self.model_type == 'deepseek_v4':
+                    return _old_from_pretrained(*args, **kwargs)
+                if self._config_cls is None:
+                    hf_config = _old_from_pretrained(*args, **kwargs)
+                    self._config_cls = hf_config.__class__
+                if not isinstance(config, self._config_cls):
+                    config.__class__ = self._config_cls
+            return config
+
+        AutoConfig.from_pretrained = _from_pretrained
+        try:
+            yield
+        finally:
+            AutoConfig.from_pretrained = _old_from_pretrained
+
+    def _prepare_engine_kwargs(self, max_model_len, engine_kwargs) -> None:
+        if engine_kwargs is None:
+            engine_kwargs = {}
+        if self.task_type == 'embedding':
+            self.task = 'embed'
+        elif self.task_type == 'seq_cls':
+            self.task = 'classify'
+        elif self.task_type in ('reranker', 'generative_reranker'):
+            self.task = 'score'
+        disable_log_stats = engine_kwargs.pop('disable_log_stats', True)
+        if self.use_async_engine:
+            engine_cls = AsyncEngineArgs
+        else:
+            engine_cls = EngineArgs
+        parameters = inspect.signature(engine_cls).parameters
+        if self.use_async_engine and 'disable_log_requests' in parameters:
+            engine_kwargs['disable_log_requests'] = True
+        if 'enable_lora' in parameters and self.enable_lora:
+            engine_kwargs['enable_lora'] = self.enable_lora
+            engine_kwargs['max_loras'] = self.max_loras
+            engine_kwargs['max_lora_rank'] = self.max_lora_rank
+        else:
+            assert not self.enable_lora, (
+                'The current version of vLLM does not support `enable_lora`. Please upgrade vLLM.')
+
+        if 'limit_mm_per_prompt' in parameters and self.limit_mm_per_prompt:
+            engine_kwargs['limit_mm_per_prompt'] = self.limit_mm_per_prompt
+        else:
+            assert not self.limit_mm_per_prompt, (
+                'The current version of vLLM does not support `limit_mm_per_prompt`. Please upgrade vLLM.')
+        for key in [
+                'enable_expert_parallel', 'enable_sleep_mode', 'disable_cascade_attn', 'load_format',
+                'mm_processor_cache_gb', 'speculative_config', 'logprobs_mode', 'quantization'
+        ]:
+            if key in parameters:
+                value = getattr(self, key, None)
+                if value is not None:
+                    engine_kwargs[key] = value
+            else:
+                logger.warning(f'The current version of vLLM does not support `{key}`. Ignored.')
+        for key in ['task', 'seed']:
+            val = getattr(self, key, None)
+            if val is not None:
+                engine_kwargs[key] = val
+
+        model_info = self.model_info
+        arch_mapping = {'deepseek_vl2': ['DeepseekVLV2ForCausalLM'], 'chatglm4v': ['GLM4VForCausalLM']}
+        if self.model_meta.model_type in arch_mapping:
+            architectures = arch_mapping[self.model_meta.model_type]
+            engine_kwargs['hf_overrides'] = {'architectures': architectures}
+        self.template.set_mode('vllm')
+        engine_kwargs.update(self.template.prepare_engine_kwargs())
+        if self.enable_prefix_caching is not None:
+            engine_kwargs['enable_prefix_caching'] = self.enable_prefix_caching
+        engine_args = engine_cls(
+            model=self.model_dir,
+            dtype=dtype_mapping[model_info.torch_dtype],
+            gpu_memory_utilization=self.gpu_memory_utilization,
+            tensor_parallel_size=self.tensor_parallel_size,
+            pipeline_parallel_size=self.pipeline_parallel_size,
+            max_model_len=max_model_len,
+            max_num_seqs=self.max_num_seqs,
+            disable_log_stats=disable_log_stats,
+            disable_custom_all_reduce=self.disable_custom_all_reduce,
+            enforce_eager=self.enforce_eager,
+            trust_remote_code=True,
+            distributed_executor_backend=self.distributed_executor_backend,
+            **engine_kwargs,
+        )
+        self.engine_args = engine_args
+
+    def _prepare_reasoning_parser(self, reasoning_parser: Optional[str]) -> None:
+        self.reasoning_parser = None
+        if not reasoning_parser:
+            return
+
+        # Validate reasoning_parser if provided
+        if ReasoningParserManager is None:
+            raise ImportError('the version of vLLM is too old, please upgrade vLLM')
+
+        valid_reasoning_parsers = list(ReasoningParserManager.reasoning_parsers.keys())
+        if reasoning_parser not in valid_reasoning_parsers:
+            raise ValueError(f'Invalid reasoning_parser: {reasoning_parser}. '
+                             f'Available parsers: {valid_reasoning_parsers}')
+        logger.info(f'Using reasoning_parser: {reasoning_parser}')
+
+        reasoning_parser_cls = ReasoningParserManager.get_reasoning_parser(reasoning_parser)
+        self.reasoning_parser = reasoning_parser_cls(self.tokenizer)
+
+    def _fix_vllm_bug(self) -> None:
+        # fix vllm==0.4 bug (very slow)
+        tokenizer = self.tokenizer
+        if self._version_ge(
+                '0.4') and not self._version_ge('0.6') and not tokenizer.__class__.__name__.startswith('Cached'):
+            _tokenizer_len = len(tokenizer)
+            __old_len__ = tokenizer.__class__.__len__
+
+            def __len__(self) -> int:
+                if self is tokenizer:
+                    return _tokenizer_len
+                else:
+                    return __old_len__(self)
+
+            tokenizer.__class__.__len__ = __len__
+
+    def _load_generation_config(self) -> None:
+        generation_config_path = os.path.join(self.model_dir, 'generation_config.json')
+        if os.path.isfile(generation_config_path):
+            generation_config = GenerationConfig.from_pretrained(self.model_dir)
+            kwargs = generation_config.to_dict()
+            max_new_tokens = kwargs.get('max_new_tokens')
+            if max_new_tokens is not None:
+                kwargs['max_tokens'] = max_new_tokens
+            top_k = kwargs.get('top_k')
+            if top_k == 0:
+                kwargs['top_k'] = -1
+            parameters = inspect.signature(SamplingParams).parameters
+            for k, v in kwargs.copy().items():
+                if k not in parameters or v is None:
+                    kwargs.pop(k)
+            self.generation_config = SamplingParams(**kwargs)
+        else:
+            self.generation_config = SamplingParams()
+
+    def _add_stop_words(self, generation_config: SamplingParams, request_config: RequestConfig) -> None:
+        template_meta = self.template.template_meta
+        stop_words = (request_config.stop or []) + (self.generation_config.stop or []) + template_meta.stop_words
+        generation_config.stop = self._get_stop_words(stop_words)
+        # stop parameter is not effective in v1 engine (test version: vllm 0.8.5.post)
+        generation_config.stop_token_ids = self._get_stop_token_ids(stop_words)
+
+    @staticmethod
+    def _version_ge(base_version: str):
+        vllm_version = vllm.__version__
+        if vllm_version is None or 'dev' in vllm_version:
+            return True
+        return version.parse(vllm_version) >= version.parse(base_version)
+
+    def _add_adapter(self, adapter_request: Optional[AdapterRequest] = None):
+        assert self.enable_lora, f'adapter_request: {adapter_request}, self.enable_lora: {self.enable_lora}'
+        from vllm.lora.request import LoRARequest
+        adapter_name = adapter_request.name
+        adapter_path = adapter_request.path
+        if adapter_name in self._adapters_pool:
+            lora_request = self._adapters_pool[adapter_name]
+        else:
+            lora_request = LoRARequest(
+                lora_name=adapter_name, lora_path=adapter_path, lora_int_id=len(self._adapters_pool) + 1)
+            self._adapters_pool[adapter_name] = lora_request
+        return lora_request
+
+    def _add_request(self,
+                     inputs: Dict[str, Any],
+                     generation_config: SamplingParams,
+                     request_id: str,
+                     adapter_request: Optional[AdapterRequest] = None):
+        kwargs = {}
+        adapter_request = adapter_request or self.default_adapter_request
+        if adapter_request:
+            kwargs['lora_request'] = self._add_adapter(adapter_request)
+
+        input_ids = inputs['input_ids']
+        if self._version_ge('0.4.3'):
+            llm_inputs = {'prompt_token_ids': input_ids}
+            mm_data = {}
+            for key in ['images', 'audios', 'videos']:
+                media_data = inputs.get(key) or []
+                if media_data:
+                    if self._version_ge('0.6'):
+
+                        mm_data[key.rstrip('s')] = media_data[0] if (
+                            len(media_data) == 1 and
+                            # compat qwen3_vl
+                            not isinstance(media_data[0], tuple)) else media_data
+                    else:
+                        assert len(media_data) == 1, (
+                            f'The current version of vllm only supports single {key}. Please upgrade to vllm >= 0.6.0')
+                        mm_data[key.rstrip('s')] = media_data[0]
+            if mm_data:
+                llm_inputs['multi_modal_data'] = mm_data
+            mm_processor_kwargs = inputs.get('mm_processor_kwargs')
+            if mm_processor_kwargs:
+                llm_inputs['mm_processor_kwargs'] = mm_processor_kwargs
+
+            has_task_arg = 'task' in inspect.signature(PoolingParams).parameters
+            has_activation_arg = 'activation' in inspect.signature(PoolingParams).parameters
+            task_mapping = {
+                'embedding': 'embed',
+                'seq_cls': 'classify',
+                'reranker': 'score',
+                'generative_reranker': 'score',
+            }
+            if self.task_type in task_mapping:
+                pooling_kwargs = {}
+                if has_task_arg:
+                    pooling_kwargs['task'] = task_mapping[self.task_type]
+                if self.task_type in ('reranker', 'generative_reranker') and \
+                        has_activation_arg and self.reranker_use_activation:
+                    pooling_kwargs['activation'] = True
+                pooling_params = PoolingParams(**pooling_kwargs)
+                return self.engine.encode(llm_inputs, pooling_params, request_id)
+            elif self.use_async_engine:
+                return self.engine.generate(llm_inputs, generation_config, request_id, **kwargs)
+            else:
+                return self.engine.add_request(request_id, llm_inputs, generation_config, **kwargs)
+        else:
+            if self.use_async_engine:
+                return self.engine.generate(None, generation_config, request_id, input_ids, **kwargs)
+            else:
+                return self.engine.add_request(request_id, None, generation_config, input_ids, **kwargs)
+
+    def _get_logprobs(self,
+                      logprobs_list: Optional[List[Dict[int, float]]],
+                      token_ids: List[int],
+                      top_logprobs: Optional[int] = None) -> Optional[Dict[str, Any]]:
+        if logprobs_list is None or len(token_ids) == 0:
+            return None
+        if len(token_ids) > 0:
+            logprobs_list = logprobs_list[-len(token_ids):]
+        for logprobs in logprobs_list:
+            for token_id, logprob in logprobs.items():
+                logprobs[token_id] = logprob.logprob
+        return super()._get_logprobs(logprobs_list, token_ids, top_logprobs)
+
+    def _prepare_generation_config(self, request_config: RequestConfig) -> SamplingParams:
+        kwargs = {'max_tokens': request_config.max_tokens}
+        for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty']:
+            new_value = getattr(request_config, key)
+            if new_value is None:
+                kwargs[key] = getattr(self.generation_config, key)
+            else:
+                kwargs[key] = new_value
+
+        # Convert Swift's Chat Completions API style (logprobs: bool, top_logprobs: int)
+        # to vLLM's SamplingParams style (logprobs: int)
+        # vLLM semantics:
+        #   - logprobs=None: no logprobs returned
+        #   - logprobs=0: only sampled token's logprob
+        #   - logprobs=N: top-N tokens + sampled token (up to N+1 total)
+        if request_config.logprobs:
+            # If logprobs=True, return log probabilities
+            if request_config.top_logprobs is not None and request_config.top_logprobs > 0:
+                # Return top_logprobs most likely tokens at each position
+                # (plus sampled token if not in top-N)
+                kwargs['logprobs'] = request_config.top_logprobs
+            else:
+                # Return only the sampled token's logprob
+                kwargs['logprobs'] = 0
+
+        if request_config.prompt_logprobs is not None:
+            kwargs['prompt_logprobs'] = request_config.prompt_logprobs
+
+        # TODO: beam search
+        for key in ['n', 'best_of', 'frequency_penalty', 'presence_penalty', 'seed']:
+            if hasattr(SamplingParams, key):
+                kwargs[key] = getattr(request_config, key)
+
+        # Handle structured outputs (guided decoding)
+        # vLLM v0.12+ uses 'structured_outputs' parameter, older versions use 'guided_decoding'
+        if request_config.structured_outputs_regex:
+            structured_outputs_param = StructuredOutputsParams(regex=request_config.structured_outputs_regex)
+            if hasattr(SamplingParams, 'structured_outputs'):
+                kwargs['structured_outputs'] = structured_outputs_param
+            else:
+                # Fallback for older vLLM versions
+                kwargs['guided_decoding'] = structured_outputs_param
+
+        res = SamplingParams(**kwargs)
+
+        if hasattr(res, 'output_kind') and res.n > 1:
+            # fix n > 1 in V1 Engine
+            from vllm.sampling_params import RequestOutputKind
+            res.output_kind = RequestOutputKind.FINAL_ONLY
+        return res
+
+    @property
+    def inner_model(self):
+        return self.engine.model_executor.driver_worker.worker.model_runner.model
+
+    @property
+    def inner_model_executor(self):
+        return self.engine.model_executor
+
+    async def _infer_stream_async(
+        self,
+        inputs: Dict[str, Any],
+        generation_config: SamplingParams,
+        adapter_request: Optional[AdapterRequest],
+        request_config: RequestConfig,
+    ) -> AsyncIterator[ChatCompletionStreamResponse]:
+        request_id = random_uuid()
+        result_generator = self._add_request(inputs, generation_config, request_id, adapter_request=adapter_request)
+        infer_streamers = [
+            InferStreamer(self.template, template_inputs=inputs['template_inputs']) for _ in range(generation_config.n)
+        ]
+        token_idxs = [0 for _ in range(generation_config.n)]
+        async for result in result_generator:
+            res = self._create_chat_completion_stream_response(result, request_config, request_id, infer_streamers,
+                                                               token_idxs)
+            if res is None:
+                continue
+            yield res
+
+    def _create_chat_completion_stream_response(self, result, request_config, request_id, infer_streamers,
+                                                token_idxs) -> Optional[ChatCompletionStreamResponse]:
+        is_diff = False
+        is_finished = False
+        for output in result.outputs:
+            output.token_ids = list(output.token_ids)
+            output.delta_text = infer_streamers[output.index].get_printable_text(output.token_ids, output.finished())
+            output.is_finished = output.finish_reason is not None
+            is_diff |= bool(output.delta_text)
+            is_finished |= output.is_finished
+        if not is_diff and not is_finished:
+            return
+
+        num_generated_tokens = sum(len(output.token_ids) for output in result.outputs)
+        usage_info = self._get_usage_info(len(result.prompt_token_ids), num_generated_tokens)
+        choices = []
+        previous_texts = [''] * len(result.outputs)
+        for output in result.outputs:
+            i = output.index
+            logprobs = self._get_logprobs(output.logprobs, output.token_ids[token_idxs[i]:],
+                                          request_config.top_logprobs)
+
+            # Handle reasoning content in streaming
+            delta_content = output.delta_text
+            delta_reasoning_content = None
+
+            if self.reasoning_parser and output.delta_text:
+                try:
+                    # Get token IDs for the delta (new tokens in this step)
+                    delta_token_ids = output.token_ids[token_idxs[i]:]
+                    previous_token_ids = output.token_ids[:token_idxs[i]]
+
+                    # Get current accumulated text for this output
+                    previous_text = previous_texts[i]
+                    current_text = previous_text + output.delta_text
+                    previous_texts[i] = current_text
+                    # Extract reasoning content from the delta
+                    delta_message = self.reasoning_parser.extract_reasoning_content_streaming(
+                        previous_text, current_text, output.delta_text, previous_token_ids, output.token_ids,
+                        delta_token_ids)
+
+                    if delta_message:
+                        delta_reasoning_content = delta_message.reasoning_content
+                        if delta_message.content:
+                            delta_content = delta_message.content
+                        else:
+                            delta_content = None
+
+                except Exception as e:
+                    logger.warning(f'Failed to extract reasoning content in streaming: {e}')
+                    # Fallback to original delta_text
+                    delta_content = output.delta_text
+            token_idxs[i] = len(output.token_ids)
+
+            toolcall = None
+            if output.is_finished:
+                toolcall = self._get_toolcall(
+                    self.template.decode_generate_ids(output.token_ids, **infer_streamers[i].decode_kwargs))
+
+            choice = ChatCompletionResponseStreamChoice(
+                index=i,
+                delta=DeltaMessage(
+                    role='assistant',
+                    content=delta_content,
+                    reasoning_content=delta_reasoning_content,
+                    tool_calls=toolcall),
+                finish_reason=output.finish_reason,
+                logprobs=logprobs)
+            choices.append(choice)
+        return ChatCompletionStreamResponse(model=self.model_name, choices=choices, usage=usage_info, id=request_id)
+
+    @staticmethod
+    def _format_prompt_logprobs(prompt_logprobs):
+        if prompt_logprobs is None:
+            return None
+        result = []
+        for pos_lps in prompt_logprobs:
+            if pos_lps is None:
+                result.append(None)
+            else:
+                pos_dict = {}
+                for token_id, lp_obj in pos_lps.items():
+                    pos_dict[str(token_id)] = {
+                        'logprob': lp_obj.logprob,
+                        'rank': getattr(lp_obj, 'rank', None),
+                        'decoded_token': getattr(lp_obj, 'decoded_token', ''),
+                    }
+                result.append(pos_dict)
+        return result
+
+    def _create_embedding_response(self, result, generation_config, request_id) -> EmbeddingResponse:
+        assert result is not None
+        embedding = result.outputs.data.cpu().numpy().tolist()
+        usage_info = self._get_usage_info(len(result.prompt_token_ids), 0)
+        return EmbeddingResponse(
+            model=self.model_name, data=[EmbeddingResponseData(embedding=embedding)], usage=usage_info, id=request_id)
+
+    def _create_chat_completion_response(
+        self,
+        result,
+        inputs,
+        request_config,
+        request_id,
+    ) -> ChatCompletionResponse:
+        assert result is not None
+        num_generated_tokens = sum(len(output.token_ids) for output in result.outputs)
+        usage_info = self._get_usage_info(len(result.prompt_token_ids), num_generated_tokens)
+        choices = []
+        for output in result.outputs:
+            output.token_ids = list(output.token_ids)
+            response = self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs'])
+
+            # Extract reasoning content if reasoning_parser is enabled
+            reasoning_content = None
+            content = response
+            if self.reasoning_parser:
+                try:
+                    reasoning_content, content = self.reasoning_parser.extract_reasoning_content(
+                        response,
+                        request=None  # We don't have the original request here
+                    )
+                except Exception as e:
+                    logger.warning(f'Failed to extract reasoning content: {e}')
+                    # Fallback to original response
+                    content = response
+
+            logprobs = self._get_logprobs(output.logprobs, output.token_ids, request_config.top_logprobs)
+            toolcall = self._get_toolcall(content)  # Use content instead of response for tool calls
+            token_ids = output.token_ids if request_config.return_details else None
+            choice = ChatCompletionResponseChoice(
+                index=output.index,
+                message=ChatMessage(
+                    role='assistant', content=content, reasoning_content=reasoning_content, tool_calls=toolcall),
+                finish_reason=output.finish_reason,
+                logprobs=logprobs,
+                token_ids=token_ids)
+            choices.append(choice)
+        prompt_token_ids = None
+        images_size = None
+        if request_config.return_details:
+            prompt_token_ids = result.prompt_token_ids
+            images = inputs['template_inputs'].images
+            if all(isinstance(image, Image.Image) for image in images):
+                images_size = [image.size for image in images]
+        formatted_prompt_logprobs = None
+        if request_config.prompt_logprobs is not None:
+            formatted_prompt_logprobs = self._format_prompt_logprobs(result.prompt_logprobs)
+        return ChatCompletionResponse(
+            model=self.model_name,
+            choices=choices,
+            usage=usage_info,
+            id=request_id,
+            prompt_token_ids=prompt_token_ids,
+            prompt_logprobs=formatted_prompt_logprobs,
+            images_size=images_size)
+
+    def _create_seq_cls_response(
+        self,
+        result,
+        request_config,
+        request_id,
+    ) -> ChatCompletionResponse:
+        assert result is not None
+        choices = []
+        preds = result.outputs.data
+        if preds.dim() == 1:
+            preds = preds.unsqueeze(0)
+        if self.task_type == 'seq_cls':
+            top_logprobs = request_config.top_logprobs or 20
+            preds, logprobs = self.template.decode_seq_cls(preds, top_logprobs)
+        else:
+            logprobs = [None] * len(preds)
+        num_prompt_token_ids = 0
+        num_generated_tokens = 0
+        for i, pred in enumerate(preds):
+            num_prompt_token_ids += len(result.prompt_token_ids)
+            num_generated_tokens += 1
+            if isinstance(pred, torch.Tensor):
+                pred = pred.tolist()
+            choices.append(
+                ChatCompletionResponseChoice(
+                    index=0,
+                    message=ChatMessage(role='assistant', content=pred, tool_calls=None),
+                    finish_reason='stop',
+                    logprobs=logprobs[i]))
+        usage_info = self._get_usage_info(num_prompt_token_ids, num_generated_tokens)
+        return ChatCompletionResponse(
+            model=self.model_name,
+            choices=choices,
+            usage=usage_info,
+            id=request_id,
+            prompt_token_ids=result.prompt_token_ids)
+
+    async def _infer_full_async(
+        self,
+        inputs: Dict[str, Any],
+        generation_config: SamplingParams,
+        adapter_request: Optional[AdapterRequest],
+        request_config: RequestConfig,
+        request_id: Optional[str] = None,
+    ) -> Union[ChatCompletionResponse, EmbeddingResponse]:
+        if request_id is None:
+            request_id = random_uuid()
+        result_generator = self._add_request(inputs, generation_config, request_id, adapter_request=adapter_request)
+        result = None
+        async for result in result_generator:
+            pass
+        if self.task_type == 'embedding':
+            return self._create_embedding_response(result, generation_config, request_id)
+        elif self.task_type in ('seq_cls', 'reranker', 'generative_reranker'):
+            return self._create_seq_cls_response(result, request_config, request_id)
+        else:
+            return self._create_chat_completion_response(result, inputs, request_config, request_id)
+
+    def _batch_infer_stream(self, *args, **kwargs):
+        if hasattr(self.engine, 'engine'):
+            self.engine.engine.model_executor.parallel_worker_tasks = None
+        return super()._batch_infer_stream(*args, **kwargs)
+
+    def infer(
+        self,
+        infer_requests: List[InferRequest],
+        request_config: Optional[RequestConfig] = None,
+        metrics: Optional[List[Metric]] = None,
+        *,
+        use_tqdm: Optional[bool] = None,
+        adapter_request: Optional[AdapterRequest] = None,
+    ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
+        if self.use_async_engine:
+            return super().infer(
+                infer_requests,
+                request_config,
+                metrics,
+                use_tqdm=use_tqdm,
+                adapter_request=adapter_request,
+            )
+        else:
+            request_config = deepcopy(request_config or RequestConfig())
+            if request_config.stream and len(infer_requests) > 1:
+                raise ValueError('If you want to use stream batch inference, you need to set use_async_engine to True.')
+            if use_tqdm is None:
+                use_tqdm = len(infer_requests) > 1
+            rank = get_dist_setting()[0]
+            if is_dist() and rank % self.engine_args.tensor_parallel_size != 0:
+                use_tqdm = False
+            self.template.set_mode('vllm')
+            batched_inputs, error_list = self._batch_encode(infer_requests, strict=getattr(self, 'strict', True))
+            request_id_list = []
+            for i, inputs in enumerate(batched_inputs):
+                request_id = str(self._request_count)
+                request_id_list.append(request_id)
+                self._request_count += 1
+                _request_config = deepcopy(request_config)
+                self.set_default_max_tokens(_request_config, inputs)
+                generation_config = self._prepare_generation_config(_request_config)
+                if generation_config.seed is not None:
+                    generation_config.seed += i
+                self._add_stop_words(generation_config, _request_config)
+                self._add_request(inputs, generation_config, request_id, adapter_request=adapter_request)
+            prog_bar = tqdm(total=len(batched_inputs), dynamic_ncols=True, disable=not use_tqdm)
+            outputs = {}
+            if request_config.stream:
+
+                def _gen_wrapper():
+                    infer_streamers = [
+                        InferStreamer(self.template, template_inputs=inputs['template_inputs'])
+                        for _ in range(generation_config.n)
+                    ]
+                    token_idxs = [0 for _ in range(generation_config.n)]
+                    while self.engine.has_unfinished_requests():
+                        result = self.engine.step()
+                        if not result:
+                            continue
+                        result = result[0]
+                        res = self._create_chat_completion_stream_response(result, request_config, request_id,
+                                                                           infer_streamers, token_idxs)
+                        if res is None:
+                            continue
+                        yield res
+                        if result.finished:
+                            break
+
+                    self._update_metrics(res, metrics)
+
+                return [_gen_wrapper()]
+            else:
+                while self.engine.has_unfinished_requests():
+                    step_outputs = self.engine.step()
+                    for output in step_outputs:
+                        if output.finished:
+                            outputs[output.request_id] = output
+                            prog_bar.update()
+                prog_bar.close()
+                outputs = [outputs[request_id] for request_id in request_id_list]
+                res = [
+                    self._create_chat_completion_response(result, inputs, request_config, request_id)
+                    for request_id, inputs, result in zip(request_id_list, batched_inputs, outputs)
+                ]
+                self._update_metrics(res, metrics)
+                return self._add_error_list(res, error_list)
+
+    async def infer_async(
+        self,
+        infer_request: InferRequest,
+        request_config: Optional[RequestConfig] = None,
+        *,
+        adapter_request: Optional[AdapterRequest] = None,
+        pre_infer_hook=None,
+    ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]:
+        if not self.use_async_engine:
+            raise ValueError('If you want to use `infer_async`, you need to pass `use_async_engine` as True.')
+        request_config = deepcopy(request_config or RequestConfig())
+        self.template.set_mode('vllm')
+        loop = asyncio.get_running_loop()
+        with torch.inference_mode():
+            inputs = await loop.run_in_executor(None, self.template.encode, infer_request, True)
+        self.set_default_max_tokens(request_config, inputs)
+        generation_config = self._prepare_generation_config(request_config)
+        self._add_stop_words(generation_config, request_config)
+        kwargs = {
+            'inputs': inputs,
+            'generation_config': generation_config,
+            'adapter_request': adapter_request,
+            'request_config': request_config,
+        }
+        if hasattr(infer_request, 'uuid') and infer_request.uuid:
+            # RolloutInferRequest
+            kwargs.update({'request_id': infer_request.uuid})
+        if pre_infer_hook:
+            kwargs = pre_infer_hook(kwargs)
+        if request_config.stream:
+            return self._infer_stream_async(**kwargs)
+        else:
+            return await self._infer_full_async(**kwargs)
+
+    @staticmethod
+    def patch_remove_log():
+        from vllm.engine import async_llm_engine
+        if not hasattr(async_llm_engine, '_log_task_completion'):
+            return
+
+        async_llm_engine._origin_log_task_completion = async_llm_engine._log_task_completion
+
+        def new_log_task_completion(task, error_callback) -> None:
+            try:
+                return_value = task.result()
+                raise AssertionError(f'The engine background task should never finish without an '
+                                     f'exception. {return_value}')
+            except asyncio.exceptions.CancelledError:
+                pass
+
+        async_llm_engine._log_task_completion = new_log_task_completion