xinference 1.0.1__py3-none-any.whl → 1.2.1__py3-none-any.whl

xinference/model/llm/vllm/xavier/engine.py

@@ -0,0 +1,247 @@
+# Copyright 2022-2025 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+from typing import Dict, List, Optional, Type, Union
+
+from vllm import AsyncEngineArgs, EmbeddingRequestOutput, RequestOutput
+from vllm.config import VllmConfig
+from vllm.engine.async_llm_engine import AsyncLLMEngine, _AsyncLLMEngine
+from vllm.engine.llm_engine import SchedulerOutputState
+from vllm.engine.metrics_types import StatLoggerBase
+from vllm.executor.executor_base import ExecutorBase
+from vllm.sequence import ExecuteModelRequest
+from vllm.usage.usage_lib import UsageContext
+
+from .executor import XavierExecutor
+from .scheduler import XavierScheduler
+
+logger = logging.getLogger(__name__)
+
+
+class XavierInternalEngine(_AsyncLLMEngine):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._xavier_config = kwargs["vllm_config"].xavier_config
+        self.scheduler = [
+            XavierScheduler(
+                self.scheduler_config,
+                self.cache_config,
+                self.lora_config,
+                self.parallel_config.pipeline_parallel_size,
+                self.async_callbacks[v_id]
+                if self.model_config.use_async_output_proc
+                else None,
+                xavier_config=self._xavier_config,
+                virtual_engine=v_id,
+            )
+            for v_id in range(self.parallel_config.pipeline_parallel_size)
+        ]
+        self.output_processor.scheduler = self.scheduler
+        self.model_executor.scheduler = self.scheduler
+
+    async def step_async(
+        self, virtual_engine: int
+    ) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
+        """Performs one decoding iteration and returns newly generated results.
+        The workers are ran asynchronously if possible.
+
+        This function performs one decoding iteration of the engine. It first
+        schedules the sequences to be executed in the next iteration and the
+        token blocks to be swapped in/out/copy. Then, it executes the model
+        and updates the scheduler with the model outputs. Finally, it decodes
+        the sequences and returns the newly generated results.
+        """
+        # 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.
+        if not self._has_remaining_steps(seq_group_metadata_list):
+            # Schedule iteration
+            """Xinference Change!!!
+            Why copy the entire function code of vllm:
+            The purpose here is to modify the way the `schedule` function is invoked to asynchronous calling.
+            No other modifications were made elsewhere.
+            """
+            (
+                seq_group_metadata_list,
+                scheduler_outputs,
+                allow_async_output_proc,
+            ) = await self.scheduler[virtual_engine].schedule()
+
+            ctx.seq_group_metadata_list = seq_group_metadata_list
+            ctx.scheduler_outputs = scheduler_outputs
+
+            # 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,
+                )
+
+        assert seq_group_metadata_list is not None
+        assert scheduler_outputs is not None
+
+        if not scheduler_outputs.is_empty():
+            finished_requests_ids = self.scheduler[
+                virtual_engine
+            ].get_and_reset_finished_requests_ids()
+
+            # 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,
+                virtual_engine=virtual_engine,
+                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]
+
+            # Execute the model.
+            outputs = await self.model_executor.execute_model_async(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:
+            if len(ctx.output_queue) > 0:
+                self._process_model_outputs(ctx=ctx)
+            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[virtual_engine] = 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
+            )
+
+            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,
+                )
+
+            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
+
+        return ctx.request_outputs
+
+
+class XavierEngine(AsyncLLMEngine):
+    _engine_class: Type[_AsyncLLMEngine] = XavierInternalEngine
+    _xavier_config: Optional[Dict] = None
+
+    @classmethod
+    def _get_executor_cls(cls, engine_config: VllmConfig) -> Type[ExecutorBase]:
+        logger.debug(f"Initializing Xavier executor.")
+        return XavierExecutor
+
+    @classmethod
+    def from_engine_args(
+        cls,
+        engine_args: AsyncEngineArgs,
+        engine_config: Optional[VllmConfig] = None,
+        start_engine_loop: bool = True,
+        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
+        stat_loggers: Optional[Dict[str, StatLoggerBase]] = None,
+        xavier_config: Optional[Dict] = None,
+    ) -> "AsyncLLMEngine":
+        cls._xavier_config = xavier_config
+        return super().from_engine_args(
+            engine_args, engine_config, start_engine_loop, usage_context, stat_loggers
+        )
+
+    def __init__(self, *args, **kwargs):
+        # set xavier_config to `vllm_config`,
+        # because it may be needed everywhere in the vllm internal components
+        kwargs["vllm_config"].xavier_config = self._xavier_config
+        super().__init__(*args, **kwargs)
+
+    async def init_xavier(self):
+        await self.engine.model_executor.init_transfer()

xinference/model/llm/vllm/xavier/executor.py

@@ -0,0 +1,134 @@
+# Copyright 2022-2025 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING, List, Optional, Set, Tuple, Union
+
+import xoscar as xo
+from vllm.executor.gpu_executor import GPUExecutorAsync
+from vllm.model_executor.layers.sampler import SamplerOutput
+from vllm.sequence import ExecuteModelRequest, PoolerOutput
+from vllm.utils import is_pin_memory_available
+from vllm.worker.cache_engine import CacheEngine
+
+if TYPE_CHECKING:
+    from .scheduler import XavierScheduler
+
+
+class XavierExecutor(GPUExecutorAsync):
+    scheduler: Optional[List["XavierScheduler"]] = None
+
+    def _init_executor(self) -> None:
+        super()._init_executor()
+        self._transfer_ref = None
+        self._block_tracker_ref = None
+
+    async def init_transfer(self):
+        """
+        In vllm, the `cache_engine` is the entity that truly manages the KV cache tensors.
+        Retrieve the necessary transmission information from the `cache_engine`.
+        """
+        transfer_ref = await self._get_transfer_ref()
+        ref_cache_engine: CacheEngine = self.driver_worker.cache_engine[0]
+        buffer_dtype = ref_cache_engine.dtype
+        buffer_device = "cpu"
+        buffer_pin_memory = is_pin_memory_available()
+        num_attn_layers = ref_cache_engine.num_attention_layers
+        kv_cache_shape = ref_cache_engine.gpu_cache[0].shape
+        assert kv_cache_shape[0] == 2
+        buffer_num = 2
+        transfer_block_num = self.vllm_config.xavier_config.get("transfer_block_num")
+        buffer_shape = (
+            transfer_block_num,
+            num_attn_layers,
+            kv_cache_shape[0],
+            *kv_cache_shape[2:],
+        )
+        await transfer_ref.setup(
+            self.driver_worker.cache_engine,
+            self.scheduler,
+            num_buffer=buffer_num,
+            buffer_shape=buffer_shape,
+            buffer_dtype=buffer_dtype,
+            buffer_device=buffer_device,
+            pin_memory=buffer_pin_memory,
+        )
+
+    async def _get_block_tracker_ref(self):
+        if self._block_tracker_ref is None:
+            block_tracker_address = self.vllm_config.xavier_config.get(
+                "block_tracker_address"
+            )
+            block_tracker_uid = self.vllm_config.xavier_config.get("block_tracker_uid")
+            self._block_tracker_ref = await xo.actor_ref(
+                address=block_tracker_address, uid=block_tracker_uid
+            )
+        return self._block_tracker_ref
+
+    async def _get_transfer_ref(self):
+        from .transfer import TransferActor
+
+        if self._transfer_ref is None:
+            transfer_address = self.vllm_config.xavier_config.get("rank_address")
+            rank = self.vllm_config.xavier_config.get("rank")
+            self._transfer_ref = await xo.actor_ref(
+                address=transfer_address, uid=f"{TransferActor.default_uid()}-{rank}"
+            )
+        return self._transfer_ref
+
+    def get_rank(self) -> int:
+        return self.vllm_config.xavier_config.get("rank")
+
+    async def execute_model_async(
+        self,
+        execute_model_req: ExecuteModelRequest,
+    ) -> List[Union[SamplerOutput, PoolerOutput]]:
+        """
+        Collect information about the blocks involved in the execution before the vllm `ModelRunner` executes.
+        This information will be used by the tracker after execution to register the locally computed blocks.
+        """
+        virtual_engine = execute_model_req.virtual_engine
+        block_tracker_ref = await self._get_block_tracker_ref()
+        scheduler = self.scheduler[virtual_engine]  # type: ignore
+        rank = self.get_rank()
+        executed_blocks_details: Set[Tuple[int, int]] = set()
+        for meta in execute_model_req.seq_group_metadata_list:
+            block_tables = meta.block_tables
+            for seq_id, block_ids in block_tables.items():
+                for _id in block_ids:
+                    b = scheduler.block_manager.get_block_by_block_id(seq_id, _id)
+                    # The `executed` attribute is used to prevent duplicate registration of the block.
+                    executed = scheduler.block_manager.get_block_status_by_block_id(
+                        "executed", _id
+                    )
+                    detail = (b.content_hash, b.block_id)
+                    if (b.content_hash is not None) and (not executed):
+                        executed_blocks_details.add(detail)
+
+        res = await super().execute_model_async(execute_model_req)
+
+        if executed_blocks_details:
+            """
+            Why not collect and register the information after execution?
+            Because after execution, the model's execution callback hook will release the block_id,
+            causing the block manager to lose access to the correct information.
+            """
+            await block_tracker_ref.register_blocks(
+                virtual_engine, list(executed_blocks_details), rank
+            )
+
+            for _, _id in executed_blocks_details:
+                scheduler.block_manager.set_block_status_by_block_id(
+                    "executed", _id, True
+                )
+
+        return res