add qwen3

vllm-v0.6.2/vllm_mlu/vllm_mlu/engine/__init__.py

@@ -0,0 +1,2 @@
+import vllm_mlu.engine.arg_utils
+import vllm_mlu.engine.llm_engine

vllm-v0.6.2/vllm_mlu/vllm_mlu/engine/arg_utils.py

@@ -0,0 +1,120 @@
+from vllm.config import ModelConfig, VllmConfig
+from vllm.engine.arg_utils import EngineArgs, AsyncEngineArgs
+from vllm_mlu._mlu_utils import (BlockSizeInfo, USE_PAGED, get_device_name)
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm.logger import init_logger
+from vllm.utils import FlexibleArgumentParser
+
+logger = init_logger(__name__)
+
+
+vllm__engine__arg_utils__EngineArgs__create_model_config_org = EngineArgs.create_model_config
+vllm__engine__arg_utils__EngineArgs__create_engine_config_org = EngineArgs.create_engine_config
+vllm__engine__arg_utils__EngineArgs__add_cli_args_org = EngineArgs.add_cli_args
+vllm_engine__arg_utils__EngineArgs____post_init__org = EngineArgs.__post_init__
+
+
+def vllm_engine__arg_utils__EngineArgs____post_init__(self,) -> None:
+    '''
+    =============================
+    Add by vllm_mlu
+    =============================
+    @brief: 1. In MLU3XX device, when the tensor_parallel_size > 1, the enforce_eager is forced to set False.
+            2. For unpaged mode, set default block_size=2048.
+    '''
+    unsupport_graph_device = "3" in get_device_name()
+    if unsupport_graph_device and self.tensor_parallel_size > 1 and self.enforce_eager != True:
+        self.enforce_eager = True
+        logger.warning("The current device only support eager mode, when the tensor_parallel_size > 1. "
+                       "The param enforce_eager is forced to set True")
+
+    if not USE_PAGED and self.block_size == 16:
+        self.block_size = 2048
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    vllm_engine__arg_utils__EngineArgs____post_init__org(self)
+
+
+def vllm__engine__arg_utils__EngineArgs__create_model_config(self) -> ModelConfig:
+    model_config = vllm__engine__arg_utils__EngineArgs__create_model_config_org(self)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: set context mlugraph info for model config
+    '''
+    model_config.set_context_mlugraph_info(
+        getattr(self, "enable_context_mlugraph", False),
+        getattr(self, "context_batch_size_to_capture", None),
+        getattr(self, "context_seq_len_to_capture", None))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return model_config
+
+
+def vllm__engine__arg_utils__EngineArgs__create_engine_config(self) -> VllmConfig:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: disable custom_all_reduce, re-set block_size to support paged and unpaged mode.
+    '''
+    # MLU not support custom all reduce
+    self.disable_custom_all_reduce = True
+    BlockSizeInfo.set_block_size(self.block_size)
+    if not USE_PAGED and self.enable_chunked_prefill:
+        raise ValueError("Not support chunked_prefill in unpaged mode.")
+
+    engine_config = vllm__engine__arg_utils__EngineArgs__create_engine_config_org(self)
+    engine_config.cache_config.block_size = BlockSizeInfo.BLOCK_SIZE
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return engine_config
+
+
+@staticmethod
+def vllm__engine__arg_utils__EngineArgs__add_cli_args(
+        parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
+    parser = vllm__engine__arg_utils__EngineArgs__add_cli_args_org(parser)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: 1. remove block_size choices, set default value to -1
+            2. add kv_cache_dtype choices of 'int8'
+    '''
+    for action in parser._actions:
+        if action.dest == "block_size":
+            action.choices = None
+            action.default = -1
+        elif action.dest == "kv_cache_dtype":
+            action.choices += ['int8']
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return parser
+
+
+MluHijackObject.apply_hijack(EngineArgs,
+                             EngineArgs.__post_init__,
+                             vllm_engine__arg_utils__EngineArgs____post_init__)
+MluHijackObject.apply_hijack(EngineArgs,
+                             EngineArgs.create_model_config,
+                             vllm__engine__arg_utils__EngineArgs__create_model_config)
+MluHijackObject.apply_hijack(EngineArgs,
+                             EngineArgs.create_engine_config,
+                             vllm__engine__arg_utils__EngineArgs__create_engine_config)
+MluHijackObject.apply_hijack(EngineArgs,
+                             EngineArgs.add_cli_args,
+                             vllm__engine__arg_utils__EngineArgs__add_cli_args)

vllm-v0.6.2/vllm_mlu/vllm_mlu/engine/async_llm_engine.py

@@ -0,0 +1,35 @@
+
+
+from vllm.engine.async_llm_engine import AsyncLLMEngine
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+# for client init/reset server scheduler profile data
+async def vllm__engine__async_llm_engine__AsyncLLMEngine__init_scheduler_view(self):
+    for scheduler in self.engine.scheduler:
+        if hasattr(scheduler, "init_scheduler_view"):
+            scheduler.init_scheduler_view()
+        else:
+            logger.warning("Can not find any scheduler view, " + 
+                           "please 'export VLLM_SCHEDULER_PROFILE=true' first.")
+
+
+# for client pulling server scheduler profile data
+async def vllm__engine__async_llm_engine__AsyncLLMEngine__save_scheduler_view(self):
+    for idx, scheduler in enumerate(self.engine.scheduler):
+        if hasattr(scheduler, "save_scheduler_view"):
+            scheduler.save_scheduler_view(idx)
+        else:
+            logger.warning("Can not find any scheduler view, " + 
+                           "please 'export VLLM_SCHEDULER_PROFILE=true' first.")
+
+
+MluHijackObject.apply_hijack(AsyncLLMEngine,
+                             "init_scheduler_view",
+                             vllm__engine__async_llm_engine__AsyncLLMEngine__init_scheduler_view)
+MluHijackObject.apply_hijack(AsyncLLMEngine,
+                             "save_scheduler_view",
+                             vllm__engine__async_llm_engine__AsyncLLMEngine__save_scheduler_view)

vllm-v0.6.2/vllm_mlu/vllm_mlu/engine/llm_engine.py

@@ -0,0 +1,209 @@
+import time
+from typing import Optional, List, Union, Mapping
+from vllm.engine.llm_engine import LLMEngine
+from vllm_mlu._mlu_utils import USE_PAGED, BlockSizeInfo
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm.sampling_params import SamplingParams
+from vllm.lora.request import LoRARequest
+from vllm.logger import init_logger
+from vllm.inputs import PromptType
+from vllm.pooling_params import PoolingParams
+from vllm.prompt_adapter.request import PromptAdapterRequest
+from vllm.utils import deprecate_kwargs
+
+logger = init_logger(__name__)
+
+
+@deprecate_kwargs(
+    "inputs",
+    additional_message="Please use the 'prompt' parameter instead.",
+)
+def vllm_engine__llm_engine__LLMEngine__add_request(
+        self,
+        request_id: str,
+        prompt: Optional[PromptType] = None,
+        params: Optional[Union[SamplingParams, PoolingParams]] = None,
+        arrival_time: Optional[float] = None,
+        lora_request: Optional[LoRARequest] = None,
+        trace_headers: Optional[Mapping[str, str]] = None,
+        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        priority: int = 0,
+        *,
+        inputs: Optional[PromptType] = None,  # DEPRECATED
+) -> None:
+    """Add a request to the engine's request pool.
+
+    The request is added to the request pool and will be processed by the
+    scheduler as `engine.step()` is called. The exact scheduling policy is
+    determined by the scheduler.
+
+    Args:
+        request_id: The unique ID of the request.
+        prompt: The prompt to the LLM. See :class:`~vllm.inputs.PromptType`
+            for more details about the format of each input.
+        params: Parameters for sampling or pooling.
+            :class:`~vllm.SamplingParams` for text generation.
+            :class:`~vllm.PoolingParams` for pooling.
+        arrival_time: The arrival time of the request. If None, we use
+            the current monotonic time.
+        trace_headers: OpenTelemetry trace headers.
+        priority: The priority of the request.
+            Only applicable with priority scheduling.
+
+    Details:
+        - Set arrival_time to the current time if it is None.
+        - Set prompt_token_ids to the encoded prompt if it is None.
+        - Create `n` number of :class:`~vllm.Sequence` objects.
+        - Create a :class:`~vllm.SequenceGroup` object
+            from the list of :class:`~vllm.Sequence`.
+        - Add the :class:`~vllm.SequenceGroup` object to the scheduler.
+
+    Example:
+        >>> # initialize engine
+        >>> engine = LLMEngine.from_engine_args(engine_args)
+        >>> # set request arguments
+        >>> example_prompt = "Who is the president of the United States?"
+        >>> sampling_params = SamplingParams(temperature=0.0)
+        >>> request_id = 0
+        >>>
+        >>> # add the request to the engine
+        >>> engine.add_request(
+        >>>    str(request_id),
+        >>>    example_prompt,
+        >>>    SamplingParams(temperature=0.0))
+        >>> # continue the request processing
+        >>> ...
+    """
+    if inputs is not None:
+        prompt = inputs
+    assert prompt is not None and params is not None
+
+    if lora_request is not None and not self.lora_config:
+        raise ValueError(f"Got lora_request {lora_request} but LoRA is "
+                            "not enabled!")
+
+    if priority != 0 and not self.scheduler_config.policy == "priority":
+        raise ValueError(f"Got priority {priority} but "
+                            "Priority scheduling is not enabled.")
+
+    if isinstance(params, SamplingParams) \
+        and (params.guided_decoding or params.logits_processors) \
+        and self.scheduler_config.num_scheduler_steps > 1:
+        raise ValueError(
+            "Guided decoding and logits processors are not supported "
+            "in multi-step decoding")
+
+    if arrival_time is None:
+        arrival_time = time.time()
+
+    if self.tokenizer is not None:
+        self._validate_token_prompt(
+            prompt,
+            tokenizer=self.get_tokenizer(lora_request=lora_request))
+
+    preprocessed_inputs = self.input_preprocessor.preprocess(
+        prompt,
+        request_id=request_id,
+        lora_request=lora_request,
+        prompt_adapter_request=prompt_adapter_request,
+    )
+    processed_inputs = self.input_processor(preprocessed_inputs)
+
+    '''
+    =============================
+    Added by vllm_mlu
+    =============================
+    @brief: check input_len + output_len <= block_size
+    '''
+    def check_block_size_valid(input_len, output_len):
+        if BlockSizeInfo.BLOCK_SIZE < input_len + output_len:
+            raise ValueError(f"BLOCK_SIZE({BlockSizeInfo.BLOCK_SIZE}) can't smaller than " +
+                             f"input_len({input_len}) + output_len({output_len}).")
+
+    if isinstance(params, SamplingParams):
+        output_len = params.max_tokens
+
+        # Check for 'prompt_token_ids' in different levels of processed_inputs
+        if not USE_PAGED:
+            for key in ['prompt_token_ids', 'encoder', 'decoder']:
+                if key in processed_inputs:
+                    if key == 'prompt_token_ids':
+                        input_len = len(processed_inputs[key])
+                    elif isinstance(processed_inputs[key], dict) and 'prompt_token_ids' in processed_inputs[key]:
+                        input_len = len(processed_inputs[key]['prompt_token_ids'])
+                    else:
+                        continue
+
+                    check_block_size_valid(input_len, output_len)
+    '''
+    ==================
+    End of modification
+    ==================
+    '''
+
+    self._add_processed_request(
+        request_id=request_id,
+        processed_inputs=processed_inputs,
+        params=params,
+        arrival_time=arrival_time,
+        lora_request=lora_request,
+        prompt_adapter_request=prompt_adapter_request,
+        trace_headers=trace_headers,
+        priority=priority,
+    )
+
+
+def vllm__engine__llm_engine__LLMEngine__get_latency(self):
+    latency = self.model_executor.get_latency()
+    return latency
+
+
+def vllm__engine__llm_engine__LLMEngine__get_memory_usage(self):
+    return self.model_executor.get_memory_usage()
+
+
+def vllm__engine__llm_engine__LLMEngine__get_block_usage(self):
+    assert len(self.scheduler) == 1, f"Only support pipeline_parallel_size=1."
+    num_free_gpu_blocks = self.scheduler[0].block_manager.get_num_free_gpu_blocks()
+    num_free_cpu_blocks = self.scheduler[0].block_manager.get_num_free_cpu_blocks()
+    return (num_free_gpu_blocks, num_free_cpu_blocks)
+
+
+# for client init/reset server scheduler profile data
+def vllm__engine__llm_engine__LLMEngine__init_scheduler_view(self):
+    for scheduler in self.scheduler:
+        if hasattr(scheduler, "init_scheduler_view"):
+            scheduler.init_scheduler_view()
+        else:
+            logger.warning("Can not find any scheduler view, " +
+                           "please 'export VLLM_SCHEDULER_PROFILE=true' first.")
+
+
+# for client pulling server scheduler profile data
+def vllm__engine__llm_engine__LLMEngine__save_scheduler_view(self):
+    for idx, scheduler in enumerate(self.scheduler):
+        if hasattr(scheduler, "save_scheduler_view"):
+            scheduler.save_scheduler_view(idx)
+        else:
+            logger.warning("Can not find any scheduler view, " +
+                           "please 'export VLLM_SCHEDULER_PROFILE=true' first.")
+
+
+MluHijackObject.apply_hijack(LLMEngine,
+                             "init_scheduler_view",
+                             vllm__engine__llm_engine__LLMEngine__init_scheduler_view)
+MluHijackObject.apply_hijack(LLMEngine,
+                             "save_scheduler_view",
+                             vllm__engine__llm_engine__LLMEngine__save_scheduler_view)
+MluHijackObject.apply_hijack(LLMEngine,
+                             LLMEngine.add_request,
+                             vllm_engine__llm_engine__LLMEngine__add_request)
+MluHijackObject.apply_hijack(LLMEngine,
+                             "get_latency",
+                             vllm__engine__llm_engine__LLMEngine__get_latency)
+MluHijackObject.apply_hijack(LLMEngine,
+                             "get_memory_usage",
+                             vllm__engine__llm_engine__LLMEngine__get_memory_usage)
+MluHijackObject.apply_hijack(LLMEngine,
+                             "get_block_usage",
+                             vllm__engine__llm_engine__LLMEngine__get_block_usage)

vllm-v0.6.2/vllm_mlu/vllm_mlu/engine/multiprocessing/__init__.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class RPCSchedulerProfileRequest(Enum):
+    INIT_SCHEDULER_VIEW = 1
+    SAVE_SCHEDULER_VIEW = 2

vllm-v0.6.2/vllm_mlu/vllm_mlu/engine/multiprocessing/client.py

@@ -0,0 +1,32 @@
+from vllm.engine.multiprocessing.client import MQLLMEngineClient
+from vllm.logger import init_logger
+
+from vllm_mlu.engine.multiprocessing import RPCSchedulerProfileRequest
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+class MQLLMEngineClient_V2(MQLLMEngineClient):
+
+    async def init_scheduler_view(self):
+        """Send INIT_SCHEDULER_VIEW request to RPC Server."""
+
+        await self._send_one_way_rpc_request(
+            request=RPCSchedulerProfileRequest.INIT_SCHEDULER_VIEW,
+            socket=self.input_socket)
+
+
+    async def save_scheduler_view(self):
+        """Send SAVE_SCHEDULER_VIEW request to RPC Server."""
+
+        await self._send_one_way_rpc_request(
+            request=RPCSchedulerProfileRequest.SAVE_SCHEDULER_VIEW,
+            socket=self.input_socket)
+
+
+MluHijackObject.apply_hijack(MQLLMEngineClient,
+                             "init_scheduler_view",
+                             MQLLMEngineClient_V2.init_scheduler_view)
+MluHijackObject.apply_hijack(MQLLMEngineClient,
+                             "save_scheduler_view",
+                             MQLLMEngineClient_V2.save_scheduler_view)

vllm-v0.6.2/vllm_mlu/vllm_mlu/engine/multiprocessing/engine.py

@@ -0,0 +1,183 @@
+import pickle
+from typing import Iterator, List, Optional, Union
+
+import cloudpickle
+import zmq
+
+from vllm import SamplingParams
+# yapf conflicts with isort for this block
+# yapf: disable
+from vllm.engine.multiprocessing import (RPCAbortRequest, RPCProcessRequest,
+                                         RPCUProfileRequest)
+from vllm.engine.llm_engine import LLMEngine
+# yapf conflicts with isort for this block
+# yapf: disable
+from vllm.engine.multiprocessing import (IPC_DATA_EXT, IPC_HEALTH_EXT,
+                                         IPC_INPUT_EXT, IPC_OUTPUT_EXT,
+                                         RPCAbortRequest, RPCProcessRequest,
+                                         RPCUProfileRequest)
+from vllm.engine.multiprocessing.engine import (MQLLMEngine,
+                                                POLLING_TIMEOUT_MS)
+from vllm.logger import init_logger
+
+from vllm_mlu.engine.multiprocessing import RPCSchedulerProfileRequest
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+vllm__engine__multiprocessing__engine__MQLLMEngine____init____org = MQLLMEngine.__init__
+
+class MQLLMEngine_V2(MQLLMEngine):
+
+    def __init__(self,
+                 ipc_path: str,
+                 use_async_sockets: bool,
+                 *args,
+                 log_requests: bool = True,
+                 **kwargs) -> None:
+        # For MQLLMEngine, we can use cached outputs, since each new request
+        # output is immediately pickled and send over the socket, which frees
+        # the python object to be reused again.
+        kwargs['use_cached_outputs'] = True
+
+        self.engine = LLMEngine(*args, **kwargs)
+        self.log_requests = log_requests
+
+        self.use_async_sockets = use_async_sockets
+        if self.use_async_sockets:
+            self.engine.process_request_outputs_callback = \
+                self._async_socket_engine_callback
+
+        self.ctx = zmq.Context()  # type: ignore[attr-defined]
+
+        # Receive input from the client.
+        self.input_socket = self.ctx.socket(zmq.constants.PULL)
+        self.input_socket.bind(f"{ipc_path}{IPC_INPUT_EXT}")
+
+        # Send output stream back to client.
+        self.output_socket = self.ctx.socket(zmq.constants.PUSH)
+        self.output_socket.bind(f"{ipc_path}{IPC_OUTPUT_EXT}")
+
+        # Send heartbeats back to client.
+        self.heartbeat_socket = self.ctx.socket(zmq.constants.PUSH)
+        self.heartbeat_socket.bind(f"{ipc_path}{IPC_HEALTH_EXT}")
+
+        # IPC path for the data socket.
+        self.data_ipc_path = f"{ipc_path}{IPC_DATA_EXT}"
+
+        # Error state.
+        self._errored_with: Optional[BaseException] = None
+
+        self.collect_scheduler_view = False
+
+    def run_engine_loop(self):
+        """Core busy loop of the LLMEngine."""
+
+        while True:
+            if not self.engine.has_unfinished_requests():
+                # Poll until there is work to do.
+                while self.input_socket.poll(timeout=POLLING_TIMEOUT_MS) == 0:
+                    # When there's no work, check on engine health and send
+                    # health status back to client
+                    self._health_check()
+                    self.engine.do_log_stats()
+                    logger.debug("Waiting for new requests in engine loop.")
+
+            # Handle any input from the client.
+            self.handle_new_input()
+
+            '''
+            =============================
+            Add by vllm_mlu
+            =============================
+            @brief: support scheduler view
+            '''
+            if self.collect_scheduler_view:
+                self.collect_scheduler_view = False
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+
+            # Engine step.
+            request_outputs = self.engine_step()
+
+            # Send request outputs (if async, done in engine_step callback).
+            if not self.use_async_sockets:
+                self._send_outputs(request_outputs)
+
+    def handle_new_input(self):
+        """Handle new input from the socket"""
+        try:
+            while self.input_socket.poll(timeout=0) != 0:
+                frames = self.input_socket.recv_multipart(copy=False)
+                request = pickle.loads(frames[0].buffer)
+
+                '''
+                =============================
+                Add by vllm_mlu
+                =============================
+                @brief: support scheduler view
+                '''
+
+                if isinstance(request, RPCProcessRequest):
+                    if len(frames) > 1:
+                        # Use cloudpickle for logits processors
+                        assert isinstance(request.params, SamplingParams)
+                        lprocs = cloudpickle.loads(frames[1].buffer)
+                        request.params.logits_processors = lprocs
+                    self._handle_process_request(request)
+                elif isinstance(request, RPCAbortRequest):
+                    self._handle_abort_request(request)
+                elif isinstance(request, RPCUProfileRequest):
+                    if request == RPCUProfileRequest.START_PROFILE:
+                        self.start_profile()
+                    else:
+                        self.stop_profile()
+                elif isinstance(request, RPCSchedulerProfileRequest):
+                    self.collect_scheduler_view = True
+                    if request == RPCSchedulerProfileRequest.INIT_SCHEDULER_VIEW:
+                        self.init_scheduler_view()
+                    elif request == RPCSchedulerProfileRequest.SAVE_SCHEDULER_VIEW:
+                        self.save_scheduler_view()
+                else:
+                    raise ValueError("Unknown RPCRequest Type: "
+                                     f"{type(request)}")
+                '''
+                ==================
+                End of MLU Hijack
+                ==================
+                '''
+
+        except Exception as e:
+            self._set_errored(e)
+            self._send_unhealthy(e)
+            raise e
+
+    def init_scheduler_view(self):
+        """Init scheduler view."""
+        self.engine.init_scheduler_view()
+
+    def save_scheduler_view(self):
+        """Save scheduler view."""
+        self.engine.save_scheduler_view()
+
+
+MluHijackObject.apply_hijack(MQLLMEngine,
+                             MQLLMEngine.__init__,
+                             MQLLMEngine_V2.__init__)
+MluHijackObject.apply_hijack(MQLLMEngine,
+                             MQLLMEngine.run_engine_loop,
+                             MQLLMEngine_V2.run_engine_loop)
+MluHijackObject.apply_hijack(MQLLMEngine,
+                             MQLLMEngine.handle_new_input,
+                             MQLLMEngine_V2.handle_new_input)
+MluHijackObject.apply_hijack(MQLLMEngine,
+                             "init_scheduler_view",
+                             MQLLMEngine_V2.init_scheduler_view)
+MluHijackObject.apply_hijack(MQLLMEngine,
+                             "save_scheduler_view",
+                             MQLLMEngine_V2.save_scheduler_view)