From d2146120303d1b5cd4196072b9231611499d0f7d Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 23:14:02 +0800 Subject: [PATCH 1/2] REF: extract shared Xavier core and Torch backend --- xinference/core/supervisor.py | 18 +- .../model/llm/vllm/xavier/actor_loop.py | 62 +- .../model/llm/vllm/xavier/block_tracker.py | 153 +---- .../model/llm/vllm/xavier/collective.py | 96 +-- .../llm/vllm/xavier/collective_manager.py | 94 +-- .../model/llm/vllm/xavier/direct_handoff.py | 271 +-------- .../model/llm/vllm/xavier/direct_history.py | 365 +----------- .../model/llm/vllm/xavier/gpu_transfer.py | 551 +----------------- xinference/model/llm/vllm/xavier/profiling.py | 58 +- .../model/llm/vllm/xavier/request_transfer.py | 100 +--- xinference/model/llm/vllm/xavier/scheduler.py | 2 +- xinference/model/llm/vllm/xavier/snapshot.py | 89 +-- .../vllm/xavier/test/test_direct_handoff.py | 13 +- .../llm/vllm/xavier/test/test_gpu_transfer.py | 14 +- .../model/llm/vllm/xavier/tiered_snapshot.py | 252 +------- xinference/model/llm/vllm/xavier/transfer.py | 12 +- xinference/model/llm/vllm/xavier/transport.py | 20 +- xinference/model/llm/vllm/xavier/utils.py | 63 +- .../model/llm/vllm/xavier/v1_connector.py | 16 +- .../xavier/xavier_remote_kvcache_manager.py | 2 +- .../llm/vllm/xavier/xavier_scheduler_hook.py | 2 +- xinference/model/llm/xavier/__init__.py | 3 + xinference/model/llm/xavier/actor_loop.py | 60 ++ .../model/llm/xavier/backends/__init__.py | 3 + .../llm/xavier/backends/torch/__init__.py | 3 + .../xavier/backends/torch/direct_handoff.py | 265 +++++++++ .../xavier/backends/torch/direct_history.py | 361 ++++++++++++ .../llm/xavier/backends/torch/gpu_transfer.py | 542 +++++++++++++++++ .../xavier/backends/torch/request_transfer.py | 96 +++ .../llm/xavier/backends/torch/snapshot.py | 87 +++ .../xavier/backends/torch/tiered_snapshot.py | 249 ++++++++ xinference/model/llm/xavier/block_tracker.py | 153 +++++ xinference/model/llm/xavier/collective.py | 93 +++ .../model/llm/xavier/collective_manager.py | 99 ++++ xinference/model/llm/xavier/profiling.py | 57 ++ xinference/model/llm/xavier/tests/__init__.py | 2 + .../llm/xavier/tests/test_shared_core.py | 178 ++++++ xinference/model/llm/xavier/transport.py | 21 + xinference/model/llm/xavier/utils.py | 60 ++ 39 files changed, 2487 insertions(+), 2098 deletions(-) create mode 100644 xinference/model/llm/xavier/__init__.py create mode 100644 xinference/model/llm/xavier/actor_loop.py create mode 100644 xinference/model/llm/xavier/backends/__init__.py create mode 100644 xinference/model/llm/xavier/backends/torch/__init__.py create mode 100644 xinference/model/llm/xavier/backends/torch/direct_handoff.py create mode 100644 xinference/model/llm/xavier/backends/torch/direct_history.py create mode 100644 xinference/model/llm/xavier/backends/torch/gpu_transfer.py create mode 100644 xinference/model/llm/xavier/backends/torch/request_transfer.py create mode 100644 xinference/model/llm/xavier/backends/torch/snapshot.py create mode 100644 xinference/model/llm/xavier/backends/torch/tiered_snapshot.py create mode 100644 xinference/model/llm/xavier/block_tracker.py create mode 100644 xinference/model/llm/xavier/collective.py create mode 100644 xinference/model/llm/xavier/collective_manager.py create mode 100644 xinference/model/llm/xavier/profiling.py create mode 100644 xinference/model/llm/xavier/tests/__init__.py create mode 100644 xinference/model/llm/xavier/tests/test_shared_core.py create mode 100644 xinference/model/llm/xavier/transport.py create mode 100644 xinference/model/llm/xavier/utils.py diff --git a/xinference/core/supervisor.py b/xinference/core/supervisor.py index 1687d695f9..2cfabd6901 100644 --- a/xinference/core/supervisor.py +++ b/xinference/core/supervisor.py @@ -1059,10 +1059,10 @@ async def signal_handler(): signal.SIGTERM, lambda: asyncio.create_task(signal_handler()) ) - from ..model.llm.vllm.xavier.block_tracker import VLLMBlockTracker - from ..model.llm.vllm.xavier.collective_manager import CollectiveManager + from ..model.llm.xavier.block_tracker import BlockTracker + from ..model.llm.xavier.collective_manager import CollectiveManager - self._block_tracker_mapping: Dict[str, xo.ActorRefType[VLLMBlockTracker]] = {} # type: ignore + self._block_tracker_mapping: Dict[str, xo.ActorRefType[BlockTracker]] = {} # type: ignore self._collective_manager_mapping: Dict[ # type: ignore str, xo.ActorRefType[CollectiveManager] ] = {} @@ -3099,7 +3099,7 @@ async def launch_builtin_model( and model_engine is not None and model_engine.lower() == "vllm" ) - from ..model.llm.vllm.xavier.transport import validate_gpu_cache_budget + from ..model.llm.xavier.transport import validate_gpu_cache_budget gpu_cache_bytes = validate_gpu_cache_budget( kwargs.pop("xavier_gpu_cache_bytes", None), enable_xavier, replica @@ -3266,15 +3266,13 @@ async def _launch_model(): nonlocal download_hub try: if enable_xavier: - from ..model.llm.vllm.xavier.block_tracker import VLLMBlockTracker - from ..model.llm.vllm.xavier.collective_manager import ( - CollectiveManager, - ) + from ..model.llm.xavier.block_tracker import BlockTracker + from ..model.llm.xavier.collective_manager import CollectiveManager self._block_tracker_mapping[model_uid] = await xo.create_actor( - VLLMBlockTracker, + BlockTracker, address=self.address, - uid=f"{VLLMBlockTracker.default_uid()}-{model_uid}", + uid=f"{BlockTracker.default_uid()}-{model_uid}", ) self._collective_manager_mapping[model_uid] = await xo.create_actor( CollectiveManager, diff --git a/xinference/model/llm/vllm/xavier/actor_loop.py b/xinference/model/llm/vllm/xavier/actor_loop.py index b778e654b7..4d0f953e23 100644 --- a/xinference/model/llm/vllm/xavier/actor_loop.py +++ b/xinference/model/llm/vllm/xavier/actor_loop.py @@ -1,60 +1,14 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Share the synchronous actor RPC loop among connectors in an EngineCore thread.""" +"""Compatibility alias for the shared Xavier actor_loop module.""" -import asyncio -import logging -import os -import threading +import sys +from typing import TYPE_CHECKING -logger = logging.getLogger(__name__) -_SHUTDOWN_TIMEOUT = 5.0 -_local = threading.local() +from ...xavier import actor_loop as _implementation +if TYPE_CHECKING: + from ...xavier.actor_loop import acquire_actor_loop as acquire_actor_loop + from ...xavier.actor_loop import release_actor_loop as release_actor_loop -def acquire_actor_loop() -> asyncio.AbstractEventLoop: - # Never reuse a parent's event loop after fork, or a different thread's loop. - if getattr(_local, "pid", None) != os.getpid() or not getattr(_local, "users", 0): - _local.loop = asyncio.new_event_loop() - _local.pid = os.getpid() - _local.users = 0 - _local.users += 1 - return _local.loop - - -def release_actor_loop(loop: asyncio.AbstractEventLoop) -> None: - if ( - getattr(_local, "pid", None) != os.getpid() - or getattr(_local, "loop", None) is not loop - ): - return - _local.users -= 1 - if _local.users: - return - try: - # xoscar listener tasks close their clients in cancellation cleanup. - pending = asyncio.all_tasks(loop) - for task in pending: - task.cancel() - - async def drain(): - if pending: - await asyncio.gather(*pending, return_exceptions=True) - await loop.shutdown_asyncgens() - - cleanup = loop.create_task(drain()) - _, unfinished = loop.run_until_complete( - asyncio.wait({cleanup}, timeout=_SHUTDOWN_TIMEOUT) - ) - if unfinished: - logger.warning("Timed out draining Xavier actor loop during shutdown") - cleanup.cancel() - # Give cancellation one turn, without waiting on unresponsive peers. - loop.run_until_complete(asyncio.sleep(0)) - elif not cleanup.cancelled(): - cleanup.result() - finally: - try: - loop.close() - finally: - del _local.loop +sys.modules[__name__] = _implementation diff --git a/xinference/model/llm/vllm/xavier/block_tracker.py b/xinference/model/llm/vllm/xavier/block_tracker.py index 639645f795..fa01e6f0b5 100644 --- a/xinference/model/llm/vllm/xavier/block_tracker.py +++ b/xinference/model/llm/vllm/xavier/block_tracker.py @@ -1,149 +1,14 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd -# -# 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 random -from typing import Dict, List, Optional, Set, Tuple +# Licensed under the Apache License, Version 2.0. +"""Compatibility alias for the shared Xavier block_tracker module.""" -import xoscar as xo +import sys +from typing import TYPE_CHECKING +from ...xavier import block_tracker as _implementation -class VLLMBlockTracker(xo.StatelessActor): - @classmethod - def default_uid(cls): - return f"vllm-block-tracker-actor" +if TYPE_CHECKING: + from ...xavier.block_tracker import BlockTracker as BlockTracker + from ...xavier.block_tracker import VLLMBlockTracker as VLLMBlockTracker - def __init__(self): - super().__init__() - # engine -> hash -> (rank, block_id) - self._hash_to_rank_and_block_id: Dict[int, Dict[int, Set[Tuple[int, int]]]] = {} # type: ignore - # engine -> rank -> (hash, block_id) - self._rank_to_hash_and_block_id: Dict[int, Dict[int, Set[Tuple[int, int]]]] = {} # type: ignore - self._unavailable_ranks: Set[int] = set() # type: ignore - - def register_blocks( - self, virtual_engine: int, block_infos: List[Tuple[int, int]], rank: int - ): - if virtual_engine not in self._hash_to_rank_and_block_id: - self._hash_to_rank_and_block_id[virtual_engine] = {} - hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine] - - if virtual_engine not in self._rank_to_hash_and_block_id: - self._rank_to_hash_and_block_id[virtual_engine] = {} - rank_to_hash_and_block_id = self._rank_to_hash_and_block_id[virtual_engine] - if rank not in rank_to_hash_and_block_id: - rank_to_hash_and_block_id[rank] = set() - hash_and_block_id = rank_to_hash_and_block_id[rank] - - # Replace-on-register: a vLLM block id can be reused for new content - # (its KV cache slot is reallocated), so before registering the new - # (hash, block_id) mappings, drop any stale (old_hash, block_id) this - # rank still holds for the same block ids. Without this a stale hash - # keeps pointing at a block that now holds different KV, and a later - # query on that hash would load the wrong KV. This also bounds the - # tracker to one entry per (rank, block_id). - new_block_ids = {block_id for _, block_id in block_infos} - new_pairs = set(block_infos) - stale = [ - (stale_hash, stale_block_id) - for stale_hash, stale_block_id in hash_and_block_id - if stale_block_id in new_block_ids - and (stale_hash, stale_block_id) not in new_pairs - ] - for stale_hash, stale_block_id in stale: - hash_and_block_id.discard((stale_hash, stale_block_id)) - stale_entries = hash_to_rank_and_block_id.get(stale_hash) - if stale_entries is not None: - stale_entries.discard((rank, stale_block_id)) - if not stale_entries: - hash_to_rank_and_block_id.pop(stale_hash, None) - - for hash_content, block_id in block_infos: - hash_to_rank_and_block_id.setdefault(hash_content, set()).add( - (rank, block_id) - ) - hash_and_block_id.add((hash_content, block_id)) - - def query_blocks( - self, - virtual_engine: int, - hash_contents: List[Tuple[int, int]], - exclude_rank: Optional[int] = None, - ) -> Dict[int, Set[Tuple[int, int, int]]]: - if virtual_engine not in self._hash_to_rank_and_block_id: - return {} - hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine] - remote: Dict[int, Set[Tuple[int, int, int]]] = {} - for hash_content, _id in hash_contents: - if ( - hash_content in hash_to_rank_and_block_id - ) and hash_to_rank_and_block_id[hash_content]: - # exclude ranks that are in the recovery process or the caller's own rank - rank_and_block_id = [ - (r, b) - for r, b in hash_to_rank_and_block_id[hash_content] - if r not in self._unavailable_ranks and r != exclude_rank - ] - if rank_and_block_id: - # TODO: Randomly select here, and try to distribute requests as evenly as possible. - # There may be better methods in the future. - rank, block_id = random.choice(rank_and_block_id) - if rank not in remote: - remote[rank] = { - (hash_content, block_id, _id), - } - else: - remote[rank].add((hash_content, block_id, _id)) - return remote - - def unregister_block(self, virtual_engine: int, rank: int, block_id: int): - self.unregister_blocks(virtual_engine, rank, [block_id]) - - def unregister_blocks(self, virtual_engine: int, rank: int, block_ids: List[int]): - entries = self._rank_to_hash_and_block_id.get(virtual_engine, {}).get(rank) - hashes = self._hash_to_rank_and_block_id.get(virtual_engine, {}) - if not entries: - return - removed_ids = set(block_ids) - removed = { - (content_hash, block_id) - for content_hash, block_id in entries - if block_id in removed_ids - } - entries.difference_update(removed) - for content_hash, block_id in removed: - locations = hashes.get(content_hash) - if locations is not None: - locations.discard((rank, block_id)) - if not locations: - hashes.pop(content_hash, None) - - def unregister_rank(self, rank: int): - """ - This rank is in the recovery process, and its query results will be excluded. - """ - self._unavailable_ranks.add(rank) - - def register_rank(self, rank: int): - """ - After recovery is successful, clear all stale data of the rank and mark the rank as available. - """ - for _, rank_to_hash_and_block_id in self._rank_to_hash_and_block_id.items(): - rank_to_hash_and_block_id.pop(rank, None) - - for _, hash_to_rank_and_block_id in self._hash_to_rank_and_block_id.items(): - for _, rank_and_block_id in hash_to_rank_and_block_id.items(): - to_delete = [(r, b) for r, b in rank_and_block_id if r == rank] - if to_delete: - rank_and_block_id.difference_update(to_delete) - - self._unavailable_ranks.discard(rank) +sys.modules[__name__] = _implementation diff --git a/xinference/model/llm/vllm/xavier/collective.py b/xinference/model/llm/vllm/xavier/collective.py index 3098f5305f..9e7426305a 100644 --- a/xinference/model/llm/vllm/xavier/collective.py +++ b/xinference/model/llm/vllm/xavier/collective.py @@ -1,93 +1,13 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd -# -# 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 List, Optional +# Licensed under the Apache License, Version 2.0. +"""Compatibility alias for the shared Xavier collective module.""" -from .transport import get_transport_host +import sys +from typing import TYPE_CHECKING -logger = logging.getLogger(__name__) +from ...xavier import collective as _implementation +if TYPE_CHECKING: + from ...xavier.collective import CollectiveRank as CollectiveRank -class CollectiveRank: - def __init__( - self, - rank: int, - world_size: int, - rank_address: str, - store_address: str, - store_port: int, - world_addresses: List[str], - ): - self._rank = rank - self._world_size = world_size - self._rank_address = rank_address - self._world_addresses = world_addresses - self._store_address = store_address - self._store_port = store_port - self._device = None - self._tcp_store = None - self._context = None - - def init_rank(self): - from xoscar.collective import xoscar_pygloo as xp - - logger.debug( - f"init_rank called: rank={self._rank}, world_size={self._world_size}, " - f"rank_address={self._rank_address}, store_address={self._store_address}, " - f"store_port={self._store_port}, world_addresses={self._world_addresses}" - ) - - self._context = xp.rendezvous.Context(rank=self._rank, size=self._world_size) - - attr = xp.transport.tcp.attr(get_transport_host(self._rank_address)) - self._device = xp.transport.tcp.CreateDevice(attr) - - opt = xp.rendezvous.TCPStoreOptions() - opt.port = self._store_port - opt.numWorkers = self._world_size - opt.isServer = self._rank == 0 - opt.waitWorkers = False - - if self._store_address is None: - raise ValueError( - f"store_address is None for rank {self._rank}. " - f"This usually means the rank 0 model was not properly initialized before other ranks." - ) - if self._store_port is None: - raise ValueError( - f"store_port is None for rank {self._rank}. " - f"This usually means the rank 0 model was not properly initialized before other ranks." - ) - - self._tcp_store = xp.rendezvous.TCPStore(self._store_address, opt) - if self._world_addresses: - self.connect_full_mesh() - - def connect_full_mesh( - self, prefix: Optional[str] = None, world_addresses: Optional[List[str]] = None - ): - from xoscar.collective import xoscar_pygloo as xp - - assert self._device is not None - assert self._tcp_store is not None - assert self._context is not None - if world_addresses is not None: - self._world_addresses = world_addresses - prefix_store = xp.rendezvous.PrefixStore( - prefix or str(self._world_size), self._tcp_store - ) - self._context.connectFullMesh(prefix_store, self._device) - logger.debug( - f"Rank {self._rank} arrives successfully, world addresses: {self._world_addresses}" - ) +sys.modules[__name__] = _implementation diff --git a/xinference/model/llm/vllm/xavier/collective_manager.py b/xinference/model/llm/vllm/xavier/collective_manager.py index 2b89d10700..f294adf122 100644 --- a/xinference/model/llm/vllm/xavier/collective_manager.py +++ b/xinference/model/llm/vllm/xavier/collective_manager.py @@ -11,17 +11,16 @@ # 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 asyncio import logging -import traceback from typing import TYPE_CHECKING, Any, Dict, List, Optional, no_type_check import xoscar as xo -from .block_tracker import VLLMBlockTracker +from ...xavier.collective_manager import CollectiveManager as CollectiveManager +from ...xavier.collective_manager import with_lock as with_lock if TYPE_CHECKING: - from .transfer import Rank0TransferActor, TransferActor + from .transfer import Rank0TransferActor logger = logging.getLogger(__name__) @@ -66,90 +65,3 @@ async def start_transfer_for_vllm(self, rank_addresses: List[str]): logger.debug( f"Init transfer actor: {self._transfer_ref.address}, rank: {self._rank} done for vllm." # type: ignore ) - - -def with_lock(method): - async def wrapper(self, *args, **kwargs): - async with self._lock: - return await method(self, *args, **kwargs) - - return wrapper - - -class CollectiveManager(xo.StatelessActor): - @classmethod - def default_uid(cls): - return f"xavier-collective-manager" - - def __init__(self, model_uid: str): - super().__init__() - self._model_uid = model_uid - self._tracker_ref: Optional[xo.ActorRefType["VLLMBlockTracker"]] = None - self._rank_to_ref: Dict[int, xo.ActorRefType["TransferActor"]] = {} - self._lock = asyncio.Lock() - - async def __post_create__(self): - self._tracker_ref = await xo.actor_ref( - address=self.address, - uid=f"{VLLMBlockTracker.default_uid()}-{self._model_uid}", - ) - - async def unregister_rank(self, rank: int): - self._rank_to_ref.pop(rank, None) - await self._tracker_ref.unregister_rank(rank) # type: ignore - await asyncio.gather( - *( - ref.release_consumer_leases_v1(rank) - for other, ref in self._rank_to_ref.items() - if other != 0 - ), - return_exceptions=True, - ) - logger.debug(f"Unregister rank: {rank}") - - async def register_rank(self, rank: int, address: str, update: bool = False): - from .transfer import TransferActor - - rank_ref = await xo.actor_ref( - address=address, uid=f"{TransferActor.default_uid()}-{rank}" - ) - self._rank_to_ref[rank] = rank_ref - logger.debug(f"Register rank: {rank}, address: {address}") - if update: - await self._update_world() - await self._tracker_ref.register_rank(rank) # type: ignore - - @with_lock - async def _update_world(self): - """ - Locking is used to prevent chaos when multiple replicas trigger recovery simultaneously. - """ - from .....core.utils import gen_random_string - - prefix = gen_random_string(6) - tasks = [] - rank_to_ref = self._rank_to_ref.copy() - world_addresses = [ref.address for _, ref in sorted(rank_to_ref.items())] - for rank, ref in rank_to_ref.items(): - tasks.append(ref.connect_full_mesh(prefix, world_addresses)) - try: - logger.debug( - f"Rebuild collective communication with world_addresses: {world_addresses}, prefix: {prefix}" - ) - await asyncio.gather(*tasks) - logger.debug( - f"Rebuild collective communication with world_addresses: {world_addresses}, prefix: {prefix} done." - ) - except Exception as e: - """ - The exception here is most likely due to another replica triggering recovery during the recovery process, - causing `connect_full_mesh` to time out. - Simply log the exception and - let the subsequent update process handle the reconstruction of the collective communication world. - """ - logger.error( - f"Rebuild collective communication with world_addresses: {world_addresses} failed. " - f"Exception: {e}" - ) - # Print the complete error stack - traceback.print_exception(type(e), e, e.__traceback__) diff --git a/xinference/model/llm/vllm/xavier/direct_handoff.py b/xinference/model/llm/vllm/xavier/direct_handoff.py index a365bac9c4..2d6f970233 100644 --- a/xinference/model/llm/vllm/xavier/direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/direct_handoff.py @@ -1,267 +1,16 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Request-scoped GPU handoff with independent tiered history.""" +"""Compatibility alias for the shared Xavier direct_handoff module.""" -import asyncio -import logging -import time -from dataclasses import dataclass, field +import sys +from typing import TYPE_CHECKING -import torch -import xoscar as xo +from ...xavier.backends.torch import direct_handoff as _implementation -from .direct_history import DirectHistoryMixin -from .gpu_transfer import GPUTransfer -from .request_transfer import LayerRead, batch_reads, unpack_reads +if TYPE_CHECKING: + from ...xavier.backends.torch.direct_handoff import ( + DirectGPUTransfer as DirectGPUTransfer, + ) + from ...xavier.backends.torch.direct_handoff import DirectRequest as DirectRequest -logger = logging.getLogger(__name__) - - -@dataclass -class DirectRequest: - request_id: str - blocks: set[int] - deadline: float - claimed: bool = False - reading: bool = False - released: bool = False - indices: dict[tuple[int, ...], torch.Tensor] = field(default_factory=dict) - hashes: dict[int, int] = field(default_factory=dict) - retaining: bool = False - retention_attempted: bool = False - - -class DirectGPUTransfer(DirectHistoryMixin, GPUTransfer): - def __init__(self, actor, caches, budget): - super().__init__(actor, caches, 0) - self.direct_requests: dict[str, DirectRequest] = {} - self.finished_sending: set[str] = set() - self.metrics.update( - direct_registered=0, direct_finished=0, direct_expired=0, index_uploads=0 - ) - - self._init_history(budget) - - def _index_tensor( - self, indices: dict[tuple[int, ...], torch.Tensor], values: list[int] - ) -> torch.Tensor: - # Reuse across layers and slabs within one request. Constructing a CUDA - # tensor from a Python list otherwise synchronizes an H2D copy each time. - key = tuple(values) - if key not in indices: - indices[key] = torch.tensor(values, dtype=torch.long, device=self.device) - self.metrics["index_uploads"] += 1 - return indices[key] - - async def _close(self): - self._history_closing = True - if self._history_task is not None: - await asyncio.gather(self._history_task, return_exceptions=True) - if self.history is not None: - self.metrics["history"] = self.history.stats() - await super()._close() - if self.history is not None: - for key in list(self.history.blocks): - self.history._drop(key) - self.history.leases.clear() - self.history.evicted.clear() - self._history_leases.clear() - self._history_probation.clear() - for state in self.direct_requests.values(): - state.indices.clear() - self.direct_requests.clear() - self.finished_sending.clear() - - def register_direct(self, ticket, request_id, blocks, hashes=None): - if self.closing or ticket in self.direct_requests: - raise ValueError("Invalid direct handoff registration") - if not blocks or any( - block < 0 or block >= len(cache) - for cache in self.caches.values() - for block in blocks - ): - raise ValueError("Invalid direct KV block IDs") - if hashes is not None and len(hashes) != len(blocks): - raise ValueError("History hashes and engine blocks differ") - self.direct_requests[ticket] = DirectRequest( - request_id, - set(blocks), - time.monotonic() + 120, - hashes=dict(zip(blocks, hashes or [])), - ) - self.metrics["direct_registered"] += 1 - - def _live_direct(self, ticket): - state = self.direct_requests.get(ticket) - if state is None or state.released: - return None - if not state.reading and time.monotonic() >= state.deadline: - self.metrics["direct_expired"] += 1 - state.retention_attempted = True - self.release_direct(ticket) - return None - return state - - def claim_direct(self, ticket): - state = self._live_direct(ticket) - if state is None: - return False - # Reclaim abandoned claims even when D dies before submitting a load. - # A scheduling retry or a completed slab refreshes this idle lease. - state.claimed = True - state.deadline = time.monotonic() + 600 - return True - - def abandon_direct(self, ticket): - state = self.direct_requests.get(ticket) - if state is not None and not state.claimed: - self.release_direct(ticket) - - def release_direct(self, ticket): - state = self.direct_requests.get(ticket) - if state is None: - return - state.released = True - if not state.reading: - if self._schedule_history(ticket, state): - return - state.indices.clear() - self.finished_sending.add(state.request_id) - del self.direct_requests[ticket] - self.metrics["direct_finished"] += 1 - - def poll_direct(self): - for ticket in list(self.direct_requests): - self._live_direct(ticket) - result = self.finished_sending - self.finished_sending = set() - return result - - async def send_direct(self, ticket, reads, remote_ref, slab_bytes): - async with self.send_lock: - state = self._live_direct(ticket) - if state is None: - return False - if slab_bytes not in self.send_buffers: - raise ValueError("Direct peer transfer slab sizes differ") - if sum(read.nbytes for read in reads) > slab_bytes: - raise ValueError("Direct transfer exceeds slab capacity") - for read in reads: - cache = self.caches[read.layer] - if ( - not set(read.keys) <= state.blocks - or cache.dtype != read.dtype - or tuple(cache.shape[1:]) != read.block_shape - ): - raise ValueError("Direct KV source layout or block IDs differ") - state.reading = True - state.deadline = time.monotonic() + 600 - try: - offset = 0 - for read in reads: - end = offset + read.nbytes - target = ( - self.send_buffer[offset:end] - .view(read.dtype) - .reshape(len(read.keys), *read.block_shape) - ) - cache = self.caches[read.layer] - torch.index_select( - cache, - 0, - self._index_tensor(state.indices, read.keys), - out=target, - ) - offset = end - await asyncio.to_thread(torch.cuda.synchronize, self.device) - await xo.copy_to([self.send_buffers[slab_bytes]], [remote_ref]) - self.metrics["wire_bytes"] += slab_bytes - self.metrics["useful_bytes"] += offset - finally: - # Even a failed gather may have queued reads of engine slots. - await asyncio.to_thread(torch.cuda.synchronize, self.device) - state.reading = False - state.deadline = time.monotonic() + 600 - if state.released: - self.release_direct(ticket) - return True - - async def load_direct(self, requests, tickets): - failed = False - try: - return await self._load_direct(requests, tickets) - except BaseException: - failed = True - raise - finally: - - results = await asyncio.gather( - *( - self.actor.release_remote_direct_gpu_v1(rank, ticket) - for ranks, ticket in zip(requests, tickets) - for rank in ranks - ), - return_exceptions=True, - ) - for result in results: - if isinstance(result, BaseException): - if not failed: - raise result - logger.warning("Direct ticket cleanup failed", exc_info=result) - - async def _load_direct(self, requests, tickets): - from .transfer import TransferActor - - self.metrics["load_calls"] += 1 - self.metrics["load_requests"] += len(requests) - invalid_blocks = set() - async with self.recv_lock: - for ranks, ticket in zip(requests, tickets): - if len(ranks) != 1: - raise ValueError("Direct handoff requires one producer") - rank, layers = next(iter(ranks.items())) - sender = await xo.actor_ref( - address=self.actor._world_addresses[rank], - uid=f"{TransferActor.default_uid()}-{rank}", - ) - indices: dict[tuple[int, ...], torch.Tensor] = {} - try: - reads = [ - LayerRead( - name, - list(mapping), - list(mapping.values()), - tuple(self.caches[name].shape[1:]), - self.caches[name].dtype, - ) - for name, mapping in layers.items() - if mapping - ] - for batch in batch_reads(reads, max_bytes=self.slab_bytes): - size = sum(read.nbytes for read in batch) - slab_bytes = min(n for n in self.recv_refs if n >= size) - available = await sender.send_direct_gpu_v1( - ticket, batch, self.recv_refs[slab_bytes], slab_bytes - ) - if not available: - # Never expose a partial request as valid KV. vLLM's - # load-error callback recomputes all its destinations. - invalid_blocks.update( - dest - for mapping in layers.values() - for dest in mapping.values() - ) - break - for read, blocks in unpack_reads( - self.recv_buffer[:size], batch - ): - cache = self.caches[read.layer] - cache[self._index_tensor(indices, read.destinations)] = ( - blocks - ) - await asyncio.to_thread(torch.cuda.synchronize, self.device) - self.metrics["gpu_batches"] += 1 - finally: - await asyncio.to_thread(torch.cuda.synchronize, self.device) - indices.clear() - return invalid_blocks +sys.modules[__name__] = _implementation diff --git a/xinference/model/llm/vllm/xavier/direct_history.py b/xinference/model/llm/vllm/xavier/direct_history.py index cf58f83545..aaa4b212ed 100644 --- a/xinference/model/llm/vllm/xavier/direct_history.py +++ b/xinference/model/llm/vllm/xavier/direct_history.py @@ -1,361 +1,16 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Bounded GPU/CPU history retained after request-scoped handoff.""" +"""Compatibility alias for the shared Xavier direct_history module.""" -import asyncio -import logging -import time -from collections import OrderedDict +import sys +from typing import TYPE_CHECKING -import torch +from ...xavier.backends.torch import direct_history as _implementation -from .request_transfer import LayerRead, batch_reads -from .tiered_snapshot import TieredKVSnapshotStore +if TYPE_CHECKING: + from ...xavier.backends.torch.direct_history import ( + DirectHistoryMixin as DirectHistoryMixin, + ) + from ...xavier.backends.torch.direct_history import HistoryStore as HistoryStore -logger = logging.getLogger(__name__) - - -class HistoryStore(TieredKVSnapshotStore): - """Independent GPU-first history with bounded CPU overflow.""" - - def __init__(self, budget, block_bytes, device, cpu_capacity=0): - super().__init__(max(1, cpu_capacity), budget, block_bytes, device) - self.cpu_capacity = cpu_capacity - self.capacity = self.gpu_capacity + cpu_capacity - - def reserve(self, lease, keys): - if not set(keys).issubset(self.ready): - return False - self.leases.setdefault(lease, set()).update(keys) - # Prefix heads outlive their tails. Reservation is not a cache hit: - # scheduling can retry or abort without ever restoring any blocks. - for key in reversed(keys): - self.touch(key) - return True - - def read(self, layer, keys, device=None): - # Batch same-tier blocks before H2D instead of synchronizing once per - # block. load_history partitions reads by tier and holds their leases. - values = [self.blocks[key][layer] for key in keys] - target = torch.device("cpu") if device is None else device - if all(value.device == values[0].device for value in values): - return torch.stack(values).to(target).contiguous() - return super().read(layer, keys, device) - - -class DirectHistoryMixin: - def _init_history(self, budget): - block_bytes = sum(c[0].numel() * c.element_size() for c in self.caches.values()) - cpu_capacity = self.store.cpu_capacity - self.history = ( - HistoryStore(budget, block_bytes, self.device, cpu_capacity) - if budget - else None - ) - self._history_task = None - self._history_closing = False - self._history_leases = {} - self._history_reading = set() - # Metadata only. A second completed request demonstrates reuse before - # new content can displace existing GPU history. Empty slots need no - # probation, preserving first-request retention while there is room. - self._history_probation: OrderedDict[int, None] = OrderedDict() - self._history_probation_limit = min( - 8192, 2 * self.history.capacity if self.history else 0 - ) - self.history_chunk_bytes = 2 * 1024 * 1024 - self.history_fill_chunk_bytes = 8 * 1024 * 1024 - self.history_pending_bytes = min( - self.history.capacity * block_bytes if self.history else 0, 64 * 1024 * 1024 - ) - # A soft admission deadline, not a cancellation timeout: an in-flight - # CUDA chunk always drains before the engine can reuse its blocks. - self.history_retention_seconds = 0.01 - self.metrics.update( - history_saved_blocks=0, - history_skipped_requests=0, - history_failures=0, - history_hit_blocks=0, - history_cpu_hit_blocks=0, - history_gpu_hit_blocks=0, - history_loaded_requests=0, - history_admission_rejected_blocks=0, - history_admission_reused_blocks=0, - history_save_chunks=0, - history_capacity_limited_blocks=0, - history_deadline_dropped_blocks=0, - history_closing_dropped_blocks=0, - ) - - def _expire_history_leases(self): - now = time.monotonic() - for lease, deadline in list(self._history_leases.items()): - if deadline <= now and lease not in self._history_reading: - self.release_history(lease) - - def reserve_history(self, lease, keys): - if self.closing or self.history is None: - return [] - self._expire_history_leases() - if lease in self._history_leases: - raise ValueError("Duplicate history lease") - matched = [] - for key in keys: - if key not in self.history.ready: - break - matched.append(key) - if matched: - if not self.history.reserve(lease, matched): - return [] - self._history_leases[lease] = time.monotonic() + 120 - return matched - - def release_history(self, lease): - if lease not in self._history_reading: - self._history_leases.pop(lease, None) - if self.history is not None: - self.history.release(lease) - - def _schedule_history(self, ticket, state): - if state.retaining: - return True - if state.retention_attempted: - return False - state.retention_attempted = True - if self._history_closing or self.history is None or not state.hashes: - return False - limit = min( - self.history.capacity, - self.history_pending_bytes // self.history.block_bytes, - ) - candidates = [] - free = max(0, self.history.capacity - len(self.history.blocks)) - prefix = [] - blocked = None - for block, key in state.hashes.items(): - if blocked is not None: - if key not in self.history.ready: - self.metrics[blocked] += 1 - continue - if len(prefix) >= self.history.capacity or ( - key not in self.history.ready and len(candidates) >= limit - ): - blocked = "history_capacity_limited_blocks" - if key not in self.history.ready: - self.metrics[blocked] += 1 - elif key in self.history.ready: - prefix.append(key) - elif free or key in self._history_probation: - candidates.append((key, block)) - prefix.append(key) - free = max(0, free - 1) - if key in self._history_probation: - self.metrics["history_admission_reused_blocks"] += 1 - else: - blocked = "history_admission_rejected_blocks" - self.metrics[blocked] += 1 - # Observe the full request, but keep heads newest in both probation and - # cache LRU. A rejected head must not admit unreachable suffix blocks. - for key in reversed(list(state.hashes.values())): - if key in self.history.ready: - self.history.touch(key) - self._history_probation.pop(key, None) - else: - self._history_probation[key] = None - self._history_probation.move_to_end(key) - while len(self._history_probation) > self._history_probation_limit: - self._history_probation.popitem(last=False) - if not candidates: - return False - # Observe completed requests even when the single writer is busy, but - # never queue another writer or extend engine block ownership for it. - if self._history_task is not None and not self._history_task.done(): - self.metrics["history_skipped_requests"] += 1 - return False - state.retaining = True - deadline = time.monotonic() + self.history_retention_seconds - task = self._history_task = asyncio.create_task( - self._retain_history(ticket, state, candidates, deadline) - ) - self.tasks.add(task) - - def done(future): - self.tasks.discard(future) - if not future.cancelled() and future.exception() is not None: - exc = future.exception() - logger.error( - "GPU history retention failed", - exc_info=exc, - ) - - task.add_done_callback(done) - return True - - async def _retain_history(self, ticket, state, candidates, deadline): - pending = [] - prefix_keys = list(state.hashes.values())[: self.history.capacity] - try: - # At most one writer and one bounded chunk. Yield between chunks; - # never queue the writer ahead of an already active handoff or load. - start = 0 - while start < len(candidates): - # Filling unused capacity can amortize per-layer gathers; - # replacements use smaller chunks to limit lock hold time. - free = self.history.gpu_capacity - self.history.counts["gpu"] - chunk_bytes = ( - self.history_fill_chunk_bytes - if free > 0 - else self.history_chunk_bytes - ) - count = max( - 1, - min( - 32, - min(self.slab_bytes, chunk_bytes) // self.history.block_bytes, - ), - ) - if free > 0: - count = min(count, free) - chunk = candidates[start : start + count] - while True: - if self._history_closing or time.monotonic() >= deadline: - reason = "closing" if self._history_closing else "deadline" - self.metrics[f"history_{reason}_dropped_blocks"] += ( - len(candidates) - start - ) - return - if not self.send_lock.locked() and not self.recv_lock.locked(): - break - await asyncio.sleep(0.001) - start += len(chunk) - async with self.send_lock: - self._expire_history_leases() - pairs = [(k, b) for k, b in chunk if k not in self.history.ready] - if not pairs: - continue - pending, ids = map(list, zip(*pairs)) - index = self._index_tensor(state.indices, ids) - layers = { - name: cache.index_select(0, index) - for name, cache in self.caches.items() - } - self.history.stage_blocks(pending, layers) - await asyncio.to_thread(torch.cuda.synchronize, self.device) - available = self.history.publish(pending, set(self.caches)) - self.metrics["history_saved_blocks"] += len(available) - self.metrics["history_save_chunks"] += 1 - for key in available: - self._history_probation.pop(key, None) - for key in reversed(prefix_keys): - if key in self.history.ready: - self.history.touch(key) - complete = len(available) == len(pending) - if not complete: - self.metrics["history_capacity_limited_blocks"] += ( - len(pending) - len(available) + len(candidates) - start - ) - pending = [] - del layers, index - if not complete: - # Capacity may be pinned by another reader. Do not - # continue past a prefix we could not publish. - return - await asyncio.sleep(0) - except Exception: - self.metrics["history_failures"] += 1 - logger.warning("Skipping failed GPU history retention", exc_info=True) - finally: - # Original engine slots stay owned even after D finishes, until all - # optional snapshot reads have drained. Never publish partial data. - await asyncio.to_thread(torch.cuda.synchronize, self.device) - for key in pending: - if key in self.history.blocks and key not in self.history.ready: - self.history._drop(key) - self.history.evicted.clear() - state.retaining = False - self.release_direct(ticket) - - async def load_history(self, requests, leases): - failed = False - try: - await self._load_history(requests, leases) - except BaseException: - failed = True - raise - finally: - errors = [] - for lease in leases: - try: - self.release_history(lease) - except BaseException as error: - errors.append(error) - if errors and not failed: - raise errors[0] - for cleanup_error in errors: - logger.warning("History lease cleanup failed", exc_info=cleanup_error) - - async def _load_history(self, requests, leases): - async with self.recv_lock: - for ranks, lease in zip(requests, leases): - indices = {} - try: - if self.history is None or lease not in self._history_leases: - raise RuntimeError("GPU history lease expired") - self._history_reading.add(lease) - reads = [] - written = set() - for layers in ranks.values(): - for name, mapping in layers.items(): - cache = self.caches[name] - if not set(mapping).issubset(self.history.leases[lease]): - raise ValueError("History read outside lease") - for tier in ("gpu", "cpu"): - pairs = [ - (key, dest) - for key, dest in mapping.items() - if self.history.tiers[key] == tier - ] - if pairs: - keys, destinations = map(list, zip(*pairs)) - reads.append( - LayerRead( - name, - keys, - destinations, - tuple(cache.shape[1:]), - cache.dtype, - ) - ) - for batch in batch_reads(reads, max_bytes=self.slab_bytes): - for read in batch: - blocks = self.history.read( - read.layer, read.keys, self.device - ) - if ( - blocks.dtype != read.dtype - or tuple(blocks.shape[1:]) != read.block_shape - ): - raise ValueError("GPU history layout mismatch") - self.caches[read.layer][ - self._index_tensor(indices, read.destinations) - ] = blocks - written.update(read.keys) - await asyncio.to_thread(torch.cuda.synchronize, self.device) - hits = {"gpu": 0, "cpu": 0} - for key in written: - tier = self.history.tiers[key] - hits[tier] += 1 - self.metrics[f"history_{tier}_hit_blocks"] += 1 - self.history.metrics[f"{tier}_hits"] += 1 - self.metrics["history_hit_blocks"] += len(written) - self.metrics["history_loaded_requests"] += 1 - logger.debug( - "Restored Xavier history: blocks=%s gpu=%s cpu=%s", - len(written), - hits["gpu"], - hits["cpu"], - ) - finally: - await asyncio.to_thread(torch.cuda.synchronize, self.device) - indices.clear() - self._history_reading.discard(lease) +sys.modules[__name__] = _implementation diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index a4712d8f7d..4f3245fb87 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -1,544 +1,19 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Local CUDA IPC caches and persistent xoscar NIXL transfer buffers.""" +"""Compatibility alias for the shared Xavier gpu_transfer module.""" -import asyncio -import logging -from importlib.metadata import version -from typing import Dict +import sys +from typing import TYPE_CHECKING -import torch -import xoscar as xo -from packaging.version import Version +from ...xavier.backends.torch import gpu_transfer as _implementation -from .request_transfer import ( - MAX_REQUEST_BLOCKS, - MAX_REQUEST_BYTES, - LayerRead, - batch_reads, - unpack_reads, -) -from .tiered_snapshot import TieredKVSnapshotStore +if TYPE_CHECKING: + from ...xavier.backends.torch.gpu_transfer import GPUTransfer as GPUTransfer + from ...xavier.backends.torch.gpu_transfer import ( + GPUTransferMixin as GPUTransferMixin, + ) + from ...xavier.backends.torch.gpu_transfer import ( + finish_before_cancel as finish_before_cancel, + ) -logger = logging.getLogger(__name__) - - -async def finish_before_cancel(task): - """Keep ownership until completion, including after repeated cancellation.""" - try: - return await asyncio.shield(task) - except asyncio.CancelledError: - while not task.done(): - try: - await asyncio.shield(task) - except asyncio.CancelledError: - continue - except Exception: - break - # Observe the result without replacing the caller's cancellation. - if not task.cancelled(): - task.exception() - raise - - -class GPUTransfer: - def __init__(self, actor, caches: Dict[str, torch.Tensor], budget: int): - if Version(version("xoscar")) < Version("0.11.1"): - raise RuntimeError("Xavier GPU transfer requires xoscar[nixl]>=0.11.1") - if not actor.address.startswith("nixl://"): - raise RuntimeError("Xavier GPU transfer requires a NIXL actor pool") - if not caches or any(not cache.is_cuda for cache in caches.values()): - raise ValueError("Xavier GPU transfer requires registered CUDA KV caches") - devices = {cache.device for cache in caches.values()} - if len(devices) != 1: - raise ValueError("Xavier GPU transfer requires one CUDA device per replica") - self.actor, self.caches = actor, caches - self.device = next(iter(devices)) - sizes = [cache[0].numel() * cache.element_size() for cache in caches.values()] - self.store = TieredKVSnapshotStore( - actor._snapshot_store.capacity, budget, sum(sizes), self.device - ) - # Fixed slabs avoid registration churn. Separate directions allow hybrids - # to serve a peer while they themselves are waiting for incoming data. - self.slab_bytes = max( - max(sizes), min(MAX_REQUEST_BYTES, sum(sizes) * MAX_REQUEST_BLOCKS) - ) - self.send_buffer = torch.zeros( - self.slab_bytes, dtype=torch.uint8, device=self.device - ) - self.recv_buffer = torch.zeros_like(self.send_buffer) - self.recv_ref = xo.buffer_ref(actor.address, self.recv_buffer) - # Keep stable views for small warm-cache transfers. Reuse their buffer - # identities, rather than creating and registering a slice per request. - small_bytes = min(self.slab_bytes, 256 * 1024) - self.send_buffers = {self.slab_bytes: self.send_buffer} - self.recv_buffers = {self.slab_bytes: self.recv_buffer} - self.recv_refs = {self.slab_bytes: self.recv_ref} - if small_bytes < self.slab_bytes: - self.send_buffers[small_bytes] = self.send_buffer[:small_bytes] - self.recv_buffers[small_bytes] = self.recv_buffer[:small_bytes] - self.recv_refs[small_bytes] = xo.buffer_ref( - actor.address, self.recv_buffers[small_bytes] - ) - self._small_slab_streak: Dict[int, int] = {} - self.send_lock, self.recv_lock = asyncio.Lock(), asyncio.Lock() - self.tasks: set[asyncio.Task] = set() - self.closing = False - self.metrics = dict( - gpu_batches=0, - cpu_batches=0, - wire_bytes=0, - useful_bytes=0, - load_calls=0, - load_requests=0, - ) - - async def run(self, function, *args): - if self.closing: - raise RuntimeError("Xavier GPU transfer is shutting down") - task = asyncio.create_task(function(*args)) - self.tasks.add(task) - - def completed(future): - self.tasks.discard(future) - if not future.cancelled(): - error = future.exception() - if error is not None: - logger.warning( - "Xavier GPU transfer task failed", - exc_info=(type(error), error, error.__traceback__), - ) - - task.add_done_callback(completed) - return await finish_before_cancel(task) - - async def load_requests_with_leases(self, requests, leases): - load_failed = False - try: - await self.load_requests(requests) - except BaseException: - load_failed = True - raise - finally: - results = await asyncio.gather( - *( - self.actor.release_remote_blocks_v1(lease, ranks) - for lease, ranks in leases - if lease - ), - return_exceptions=True, - ) - for result in results: - if isinstance(result, BaseException): - if not load_failed: - raise result - logger.warning( - "Failed to release Xavier snapshot lease after load failure", - exc_info=(type(result), result, result.__traceback__), - ) - - async def close(self): - task = getattr(self, "_close_task", None) - if task is None: - self.closing = True - task = self._close_task = asyncio.create_task(self._close()) - return await finish_before_cancel(task) - - async def _close(self): - if self.tasks: - await asyncio.gather(*list(self.tasks), return_exceptions=True) - try: - await asyncio.to_thread(torch.cuda.synchronize, self.device) - logger.info( - "Xavier GPU cache stats: %s", - dict(self.metrics, cache=self.store.stats()), - ) - finally: - self.recv_refs.clear() - self.recv_buffers.clear() - self.send_buffers.clear() - self.recv_ref = None - self.caches.clear() - self.send_buffer = self.recv_buffer = None - self.store._gpu_lru.clear() - self.store._block_sizes.clear() - self.store.blocks.clear() - self.store.ready.clear() - self.store.tiers.clear() - self.store.logical_dtypes.clear() - self.store.leases.clear() - self.store.evicted.clear() - self.store.counts.update(gpu=0, cpu=0) - self._small_slab_streak.clear() - - async def stage(self, entries): - if self.closing: - raise RuntimeError("Xavier GPU transfer is shutting down") - keys = {key for layers in entries for ids, _ in layers.values() for key in ids} - failed = False - copied = False - last_reused_keys = None - # These complete snapshots were copied on this call's stream. They may - # be reused before our final fence, unlike unpublished work of a peer call. - staged_keys: set[int] = set() - try: - for layers in entries: - if self._can_stage_blocks(layers, staged_keys): - copied = self._stage_blocks(layers, staged_keys) or copied - last_reused_keys = None - continue - for layer, (block_keys, ids) in layers.items(): - # Unpublished layers may still be copying in another stage - # awaiting its fence; only published data can skip the copy. - if all( - key in self.store.ready - and layer in self.store.blocks.get(key, {}) - for key in block_keys - ): - if block_keys != last_reused_keys: - for key in block_keys: - self.store.touch(key) - last_reused_keys = block_keys - continue - last_reused_keys = None - cache = self.caches[layer] - blocks = cache.index_select( - 0, torch.tensor(ids, dtype=torch.long, device=cache.device) - ) - self.store.stage(layer, block_keys, blocks) - copied = True - except Exception: - failed = True - logger.warning( - "Failed to stage Xavier GPU snapshots; skipping new snapshots", - exc_info=True, - ) - finally: - # Fence only if this call copied or failed: even a failed gather - # may have queued reads of EngineCore slots. Pure published reuse - # needs no fence before returning ownership to EngineCore. - try: - if copied or failed: - await asyncio.to_thread(torch.cuda.synchronize, self.device) - except BaseException: - failed = True - raise - finally: - if failed: - for key in keys: - if key in self.store.blocks and key not in self.store.ready: - self.store._drop(key) - - def _can_stage_blocks(self, layers, staged_keys): - if not layers or set(layers) != set(self.caches): - return False - keys, ids = next(iter(layers.values())) - return ( - bool(keys) - and len(keys) == len(ids) - and len(keys) == len(set(keys)) - and all(value == (keys, ids) for value in layers.values()) - and all( - key not in self.store.blocks - or ( - (key in self.store.ready or key in staged_keys) - and self.caches.keys() <= self.store.blocks[key].keys() - ) - for key in keys - ) - ) - - def _stage_blocks(self, layers, staged_keys): - keys, ids = next(iter(layers.values())) - count = max( - 1, min(MAX_REQUEST_BLOCKS, MAX_REQUEST_BYTES // self.store.block_bytes) - ) - pending_keys, pending_ids = [], [] - copied = False - - def flush(): - nonlocal copied - if not pending_keys: - return - index = torch.tensor(pending_ids, dtype=torch.long, device=self.device) - values = { - name: cache.index_select(0, index) - for name, cache in self.caches.items() - } - self.store.stage_blocks(pending_keys, values) - staged_keys.update(key for key in pending_keys if key in self.store.blocks) - copied = True - pending_keys.clear() - pending_ids.clear() - - for key, source in zip(keys, ids): - if key in self.store.ready or key in staged_keys: - flush() - # Admission of preceding blocks may have evicted this hit. - if key in self.store.blocks and ( - key in self.store.ready or key in staged_keys - ): - self.store.touch(key) - continue - pending_keys.append(key) - pending_ids.append(source) - if len(pending_keys) == count: - flush() - flush() - return copied - - def locations(self, reads): - result = {} - for read in reads: - for key in read.keys: - if ( - key not in self.store.ready - or read.layer not in self.store.blocks[key] - ): - raise KeyError("Requested Xavier snapshot is not published") - tensor = self.store.blocks[key][read.layer] - if ( - tensor.dtype != read.dtype - or tuple(tensor.shape) != read.block_shape - ): - raise ValueError("Xavier source and destination KV layouts differ") - result[key] = self.store.tiers[key] - return result - - async def send(self, reads, remote_ref, slab_bytes): - async with self.send_lock: - if slab_bytes not in self.send_buffers: - raise ValueError("Xavier peer transfer slab sizes differ") - locations = self.locations(reads) - if any(tier != "gpu" for tier in locations.values()): - raise ValueError("GPU transfer requested for a CPU snapshot") - size = sum(read.nbytes for read in reads) - if size > slab_bytes: - raise ValueError("Xavier GPU batch exceeds transfer slab") - offset = 0 - for read in reads: - end = offset + read.nbytes - target = ( - self.send_buffer[offset:end] - .view(read.dtype) - .reshape(len(read.keys), *read.block_shape) - ) - torch.stack( - [self.store.blocks[key][read.layer] for key in read.keys], - out=target, - ) - offset = end - await asyncio.to_thread(torch.cuda.synchronize, self.device) - await xo.copy_to([self.send_buffers[slab_bytes]], [remote_ref]) - self.metrics["wire_bytes"] += slab_bytes - self.metrics["useful_bytes"] += size - - def _select_slab_bytes(self, rank: int, size: int) -> int: - candidate = min(n for n in self.recv_refs if n >= size) - if candidate == self.slab_bytes: - self._small_slab_streak.pop(rank, None) - return candidate - # xoscar caches only the latest registration on each peer channel. - # Require consecutive small batches so full/small tails do not churn - # registrations. Sustained small workloads still reuse the small view. - streak = min(2, self._small_slab_streak.get(rank, 0) + 1) - self._small_slab_streak[rank] = streak - return candidate if streak == 2 else self.slab_bytes - - async def load(self, ranks): - return await self.load_requests([ranks]) - - async def load_requests(self, requests): - from .transfer import TransferActor - - # Index by destination, not hash: one cached block may feed multiple - # requests' slots. Reject conflicting writes before any transfer starts. - ranks = {} - owners = {} - request_bytes = {} - for request in requests: - for rank, layers in request.items(): - size = sum( - len(mapping) - * self.caches[name][0].numel() - * self.caches[name].element_size() - for name, mapping in layers.items() - ) - request_bytes[rank] = max(request_bytes.get(rank, 0), size) - for name, mapping in layers.items(): - for key, destination in mapping.items(): - target = (name, destination) - if target in owners: - if owners[target] != key: - raise ValueError("Conflicting Xavier KV destinations") - continue - owners[target] = key - ranks.setdefault(rank, {}).setdefault(name, {})[ - destination - ] = key - - self.metrics["load_calls"] += 1 - self.metrics["load_requests"] += len(requests) - async with self.recv_lock: - for rank, layers in ranks.items(): - sender = await xo.actor_ref( - address=self.actor._world_addresses[rank], - uid=f"{TransferActor.default_uid()}-{rank}", - ) - reads = [ - LayerRead( - name, - list(mapping.values()), - list(mapping), - tuple(self.caches[name].shape[1:]), - self.caches[name].dtype, - ) - for name, mapping in layers.items() - if mapping - ] - locations = await sender.gpu_snapshot_locations_v1(reads) - for tier in ("gpu", "cpu"): - selected = [] - for read in reads: - pairs = [ - (key, dest) - for key, dest in zip(read.keys, read.destinations) - if locations[key] == tier - ] - if pairs: - selected.append( - LayerRead( - read.layer, - [p[0] for p in pairs], - [p[1] for p in pairs], - read.block_shape, - read.dtype, - ) - ) - # Combining small requests must not turn each transfer back - # into a padded full-slab copy. Large requests retain the - # original batch bound to avoid fragmenting cold prefills. - limit = self.slab_bytes - if tier == "gpu": - limit = min( - (n for n in self.recv_refs if n >= request_bytes[rank]), - default=self.slab_bytes, - ) - for batch in batch_reads(selected, max_bytes=limit): - size = sum(read.nbytes for read in batch) - if tier == "gpu": - slab_bytes = self._select_slab_bytes(rank, size) - await sender.send_gpu_request_v1( - batch, self.recv_refs[slab_bytes], slab_bytes - ) - payload = self.recv_buffer[:size] - else: - payload = await self.actor.read_request_blocks_v1( - rank, batch - ) - for read, blocks in unpack_reads(payload, batch): - cache = self.caches[read.layer] - cache[ - torch.tensor( - read.destinations, - dtype=torch.long, - device=cache.device, - ) - ] = blocks.to(cache.device, non_blocking=True) - # Both producer's IPC ownership and receiver slab reuse - # require the cache writes to finish before acknowledging. - await asyncio.to_thread(torch.cuda.synchronize, self.device) - self.metrics[tier + "_batches"] += 1 - - -class GPUTransferMixin: - def map_gpu_caches_v1(self, descriptors, budget, direct_handoff=False): - from torch.multiprocessing.reductions import rebuild_cuda_tensor - - if getattr(self, "_gpu_transfer", None) is not None: - raise RuntimeError("Xavier GPU caches already registered") - caches = { - name: rebuild_cuda_tensor(*desc) for name, desc in descriptors.items() - } - runtime_type = GPUTransfer - if direct_handoff: - from .direct_handoff import DirectGPUTransfer - - runtime_type = DirectGPUTransfer - runtime = runtime_type(self, caches, budget) - self._gpu_transfer = runtime - self._snapshot_store = runtime.store - - async def stage_gpu_requests_v1(self, entries): - runtime = self._gpu_transfer - return await runtime.run(runtime.stage, entries) - - def gpu_snapshot_locations_v1(self, reads): - return self._gpu_transfer.locations(reads) - - async def send_gpu_request_v1(self, reads, remote_ref, slab_bytes): - runtime = self._gpu_transfer - return await runtime.run(runtime.send, reads, remote_ref, slab_bytes) - - async def load_gpu_requests_v1(self, requests, leases): - runtime = self._gpu_transfer - # The actor loop advances independently of EngineCore. Keep both writes - # and lease cleanup inside the protected task before returning readiness. - return await runtime.run(runtime.load_requests_with_leases, requests, leases) - - async def close_gpu_caches_v1(self): - runtime = getattr(self, "_gpu_transfer", None) - if runtime is not None: - # Keep the closed runtime as a tombstone: close() reuses its task, - # and map_gpu_caches_v1 must not re-register after shutdown/failure. - await runtime.close() - - def register_direct_gpu_v1(self, ticket, request_id, blocks, hashes=None): - self._gpu_transfer.register_direct(ticket, request_id, blocks, hashes) - - def claim_direct_gpu_v1(self, ticket): - return self._gpu_transfer.claim_direct(ticket) - - def abandon_direct_gpu_v1(self, ticket): - self._gpu_transfer.abandon_direct(ticket) - - async def claim_remote_direct_gpu_v1(self, rank, ticket): - ref = await xo.actor_ref( - address=self._world_addresses[rank], uid=f"{self.default_uid()}-{rank}" - ) - return await ref.claim_direct_gpu_v1(ticket) - - def release_direct_gpu_v1(self, ticket): - self._gpu_transfer.release_direct(ticket) - - def poll_direct_gpu_v1(self): - return self._gpu_transfer.poll_direct() - - async def send_direct_gpu_v1(self, ticket, reads, remote_ref, slab_bytes): - runtime = self._gpu_transfer - return await runtime.run( - runtime.send_direct, ticket, reads, remote_ref, slab_bytes - ) - - async def load_direct_gpu_v1(self, requests, tickets): - runtime = self._gpu_transfer - return await runtime.run(runtime.load_direct, requests, tickets) - - async def release_remote_direct_gpu_v1(self, rank, ticket): - ref = await xo.actor_ref( - address=self._world_addresses[rank], uid=f"{self.default_uid()}-{rank}" - ) - await ref.release_direct_gpu_v1(ticket) - - def reserve_direct_history_v1(self, lease, keys): - runtime = getattr(self, "_gpu_transfer", None) - return runtime.reserve_history(lease, keys) if runtime is not None else [] - - def release_direct_history_v1(self, lease): - runtime = getattr(self, "_gpu_transfer", None) - if runtime is not None: - runtime.release_history(lease) - - async def load_direct_history_v1(self, requests, leases): - runtime = self._gpu_transfer - return await runtime.run(runtime.load_history, requests, leases) +sys.modules[__name__] = _implementation diff --git a/xinference/model/llm/vllm/xavier/profiling.py b/xinference/model/llm/vllm/xavier/profiling.py index 2e70043119..23fb573692 100644 --- a/xinference/model/llm/vllm/xavier/profiling.py +++ b/xinference/model/llm/vllm/xavier/profiling.py @@ -1,57 +1,13 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Opt-in, synchronizing diagnostics; never enable for throughput measurements.""" +"""Compatibility alias for the shared Xavier profiling module.""" -import json -import logging -import os -import time -from contextlib import contextmanager -from typing import Any, Iterator +import sys +from typing import TYPE_CHECKING -logger = logging.getLogger(__name__) -_ENABLED = os.environ.get("XINFERENCE_XAVIER_PROFILE") == "1" -if _ENABLED: - # Keep configured actor/file logging; EngineCore may have no handlers. - if not logger.hasHandlers(): - logger.addHandler(logging.StreamHandler()) - logger.setLevel(logging.INFO) - logger.propagate = True +from ...xavier import profiling as _implementation +if TYPE_CHECKING: + from ...xavier.profiling import profile_stage as profile_stage -@contextmanager -def profile_stage(stage: str, *, device: Any = None, **fields: Any) -> Iterator[None]: - if not _ENABLED: - yield - return - - def synchronize() -> None: - if device is not None and device.type == "cuda": - import torch - - torch.cuda.synchronize(device) - - # Drain preceding GPU work outside the timed region. This intentionally - # serializes execution; profile results are separate from benchmark runs. - synchronize() - started = time.perf_counter() - succeeded = False - try: - yield - synchronize() - succeeded = True - finally: - elapsed = time.perf_counter() - started - logger.info( - "Xavier profile: %s", - json.dumps( - dict( - stage=stage, - elapsed_s=elapsed, - pid=os.getpid(), - succeeded=succeeded, - **fields, - ), - sort_keys=True, - ), - ) +sys.modules[__name__] = _implementation diff --git a/xinference/model/llm/vllm/xavier/request_transfer.py b/xinference/model/llm/vllm/xavier/request_transfer.py index 2402ce22af..d7b8a68c42 100644 --- a/xinference/model/llm/vllm/xavier/request_transfer.py +++ b/xinference/model/llm/vllm/xavier/request_transfer.py @@ -1,96 +1,16 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Bounded cross-layer Gloo messages for Xavier V1 snapshot reads.""" +"""Compatibility alias for the shared Xavier request_transfer module.""" -import math -from dataclasses import dataclass -from typing import List, Tuple +import sys +from typing import TYPE_CHECKING -import torch +from ...xavier.backends.torch import request_transfer as _implementation -MAX_REQUEST_BYTES = 16 * 1024 * 1024 -MAX_REQUEST_BLOCKS = 64 +if TYPE_CHECKING: + from ...xavier.backends.torch.request_transfer import LayerRead as LayerRead + from ...xavier.backends.torch.request_transfer import batch_reads as batch_reads + from ...xavier.backends.torch.request_transfer import pack_reads as pack_reads + from ...xavier.backends.torch.request_transfer import unpack_reads as unpack_reads - -@dataclass -class LayerRead: - layer: str - keys: List[int] - destinations: List[int] - block_shape: Tuple[int, ...] - dtype: torch.dtype - logical_dtype: torch.dtype | None = None - - @property - def nbytes(self) -> int: - return len(self.keys) * math.prod(self.block_shape) * self.dtype.itemsize - - -def batch_reads(reads, max_bytes=None): - max_bytes = MAX_REQUEST_BYTES if max_bytes is None else max_bytes - batch = [] - size = 0 - for read in reads: - if batch and read.dtype != batch[0].dtype: - yield batch - batch, size = [], 0 - block_bytes = math.prod(read.block_shape) * read.dtype.itemsize - if block_bytes <= 0: - raise ValueError("Invalid KV block size") - offset = 0 - while offset < len(read.keys): - limit = min(MAX_REQUEST_BLOCKS, len(read.keys) - offset) - count = max(1, min(limit, (max_bytes - size) // block_bytes)) - if batch and size + count * block_bytes > max_bytes: - yield batch - batch, size = [], 0 - continue - part = LayerRead( - read.layer, - read.keys[offset : offset + count], - read.destinations[offset : offset + count], - read.block_shape, - read.dtype, - read.logical_dtype, - ) - batch.append(part) - size += part.nbytes - offset += count - if size >= max_bytes: - yield batch - batch, size = [], 0 - if batch: - yield batch - - -def pack_reads(store, reads): - buffers = [] - for read in reads: - if not set(read.keys).issubset(store.ready): - raise KeyError("Requested KV snapshots are not published") - if any( - store.logical_dtypes[key][read.layer] != (read.logical_dtype or read.dtype) - for key in read.keys - ): - raise ValueError("Xavier logical KV dtype does not match the receiver") - tensor = store.read(read.layer, read.keys) - if tensor.dtype != read.dtype or tuple(tensor.shape) != ( - len(read.keys), - *read.block_shape, - ): - raise ValueError("KV snapshot layout does not match the receiver") - buffers.append(tensor.view(torch.uint8).reshape(-1)) - return torch.cat(buffers) - - -def unpack_reads(payload, reads): - expected = sum(read.nbytes for read in reads) - if payload.dtype != torch.uint8 or payload.numel() != expected: - raise ValueError("Invalid Xavier cross-layer payload size or dtype") - offset = 0 - for read in reads: - end = offset + read.nbytes - yield read, payload[offset:end].view(read.dtype).reshape( - len(read.keys), *read.block_shape - ) - offset = end +sys.modules[__name__] = _implementation diff --git a/xinference/model/llm/vllm/xavier/scheduler.py b/xinference/model/llm/vllm/xavier/scheduler.py index bf5408447a..50505432ae 100644 --- a/xinference/model/llm/vllm/xavier/scheduler.py +++ b/xinference/model/llm/vllm/xavier/scheduler.py @@ -31,9 +31,9 @@ SequenceStatus, ) +from ...xavier.utils import hash_block_tokens from .block_manager import XavierBlockManager from .cache_lifecycle import PDCacheLifecycleMixin -from .utils import hash_block_tokens from .xavier_remote_kvcache_manager import XavierRemoteKVCacheManager try: diff --git a/xinference/model/llm/vllm/xavier/snapshot.py b/xinference/model/llm/vllm/xavier/snapshot.py index 7d516a2a9f..5e6856631b 100644 --- a/xinference/model/llm/vllm/xavier/snapshot.py +++ b/xinference/model/llm/vllm/xavier/snapshot.py @@ -1,87 +1,14 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -from collections import OrderedDict -from typing import Dict, List, Set +"""Compatibility alias for the shared Xavier snapshot module.""" -import torch +import sys +from typing import TYPE_CHECKING +from ...xavier.backends.torch import snapshot as _implementation -class KVSnapshotStore: - """Bounded, content-addressed CPU blocks with atomic publication and read leases. +if TYPE_CHECKING: + from ...xavier.backends.torch.snapshot import KVSnapshotStore as KVSnapshotStore + from ...xavier.backends.torch.snapshot import block_major_view as block_major_view - All mutation runs on the TransferActor's event loop. Tensor storage is owned - by each block, never by a GPU slot or a view retaining an entire request. - """ - - def __init__(self, capacity: int): - self.capacity = capacity - self.blocks: OrderedDict[int, Dict[str, torch.Tensor]] = OrderedDict() - self.logical_dtypes: Dict[int, Dict[str, torch.dtype]] = {} - self.ready: Set[int] = set() - self.leases: Dict[str, Set[int]] = {} - self.evicted: Set[int] = set() - - def stage( - self, - layer: str, - keys: List[int], - tensors: torch.Tensor, - logical_dtype: torch.dtype | None = None, - ): - pinned = set().union(*self.leases.values()) if self.leases else set() - for key, tensor in zip(keys, tensors): - if key not in self.blocks: - if len(self.blocks) >= self.capacity: - victim = next((k for k in self.blocks if k not in pinned), None) - if victim is None: - continue # No room: leave this block to local recomputation. - del self.blocks[victim] - del self.logical_dtypes[victim] - self.ready.discard(victim) - self.evicted.add(victim) - self.blocks[key] = {} - self.logical_dtypes[key] = {} - self.blocks.move_to_end(key) - # Identical content is immutable, including while a reader holds it. - if layer not in self.blocks[key]: - self.blocks[key][layer] = tensor.clone().contiguous() - self.logical_dtypes[key][layer] = logical_dtype or tensors.dtype - - def publish(self, keys: List[int], layers: Set[str]) -> List[int]: - available = [] - for key in keys: - if layers and layers.issubset(self.blocks.get(key, {})): - self.ready.add(key) - available.append(key) - return available - - def reserve(self, lease: str, keys: List[int]) -> bool: - if not set(keys).issubset(self.ready): - return False - self.leases.setdefault(lease, set()).update(keys) - for key in keys: - self.blocks.move_to_end(key) - return True - - def release(self, lease: str): - self.leases.pop(lease, None) - - def release_consumer(self, rank: int): - for lease in list(self.leases): - if lease.startswith(f"{rank}:"): - self.release(lease) - - def read(self, layer: str, keys: List[int]) -> torch.Tensor: - return torch.stack([self.blocks[key][layer] for key in keys]).contiguous() - - -def block_major_view(tensor: torch.Tensor, num_blocks: int) -> torch.Tensor: - """Normalize supported attention layouts without guessing ambiguous axes.""" - if tensor.ndim == 5 and tensor.shape[0] == 2: - if tensor.shape[1] == num_blocks: - if num_blocks == 2: - raise ValueError("Ambiguous Xavier KV block axis with only two blocks") - return tensor.movedim(1, 0) - if tensor.shape[0] != num_blocks: - raise ValueError(f"Unsupported Xavier KV layout {tuple(tensor.shape)}") - return tensor +sys.modules[__name__] = _implementation diff --git a/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py index c676d16505..bd9cbdc2f7 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py @@ -42,8 +42,12 @@ async def send(*args): @pytest.mark.asyncio -async def test_direct_handoff_preserves_bits_and_uses_no_snapshots(monkeypatch): +@pytest.mark.parametrize("actor_uid", ["vllm-transfer-actor", "adapter-transfer"]) +async def test_direct_handoff_preserves_bits_and_uses_no_snapshots( + monkeypatch, actor_uid +): source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + dest.actor.default_uid = lambda: actor_uid source.caches["K"].view(torch.int16).copy_( torch.tensor( [ @@ -61,13 +65,18 @@ async def test_direct_handoff_preserves_bits_and_uses_no_snapshots(monkeypatch): ) source.register_direct("ticket", "producer", [1, 2]) peer = peer_for(source) - monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + actor_ref = AsyncMock(return_value=peer) + monkeypatch.setattr(xo, "actor_ref", actor_ref) async def copy(buffers, refs): refs[0].copy_(buffers[0]) monkeypatch.setattr(xo, "copy_to", copy) await dest.run(dest.load_direct, [{0: {"K": {1: 5, 2: 3}}}], ["ticket"]) + assert actor_ref.await_args_list[0].kwargs == { + "address": "source", + "uid": f"{actor_uid}-0", + } assert torch.equal( dest.caches["K"].view(torch.int16)[[5, 3]], source.caches["K"].view(torch.int16)[[1, 2]], diff --git a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py index d427560e2a..5963af8317 100644 --- a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py @@ -39,7 +39,9 @@ def runtime(monkeypatch, gpu_slots=1): load_calls=0, load_requests=0, ) - r.actor = SimpleNamespace(_world_addresses=["source"]) + r.actor = SimpleNamespace( + _world_addresses=["source"], default_uid=lambda: "vllm-transfer-actor" + ) return r @@ -50,8 +52,12 @@ def stage(r, key): @pytest.mark.asyncio @pytest.mark.parametrize("small_slab", [False, True]) -async def test_mixed_tier_load_preserves_bits_and_destinations(monkeypatch, small_slab): +@pytest.mark.parametrize("actor_uid", ["vllm-transfer-actor", "adapter-transfer"]) +async def test_mixed_tier_load_preserves_bits_and_destinations( + monkeypatch, small_slab, actor_uid +): source, dest = runtime(monkeypatch), runtime(monkeypatch) + dest.actor.default_uid = lambda: actor_uid if small_slab: source.send_buffers[4] = source.send_buffer[:4] dest.recv_buffers[4] = dest.recv_buffer[:4] @@ -68,11 +74,13 @@ async def copy(buffers, refs): gpu_snapshot_locations_v1=AsyncMock(side_effect=source.locations), send_gpu_request_v1=AsyncMock(side_effect=source.send), ) - monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + actor_ref = AsyncMock(return_value=peer) + monkeypatch.setattr(xo, "actor_ref", actor_ref) dest.actor.read_request_blocks_v1 = AsyncMock( side_effect=lambda rank, reads: pack_reads(source.store, reads) ) await dest.run(dest.load, {0: {"K": {1: 5, 2: 3}}}) + actor_ref.assert_awaited_once_with(address="source", uid=f"{actor_uid}-0") assert dest.caches["K"][[5, 3]].tolist() == [[1, -1], [2, -2]] assert dest.caches["K"][:3].count_nonzero() == 0 assert dest.metrics["gpu_batches"] == dest.metrics["cpu_batches"] == 1 diff --git a/xinference/model/llm/vllm/xavier/tiered_snapshot.py b/xinference/model/llm/vllm/xavier/tiered_snapshot.py index 457fd2a77f..dcd07a6db7 100644 --- a/xinference/model/llm/vllm/xavier/tiered_snapshot.py +++ b/xinference/model/llm/vllm/xavier/tiered_snapshot.py @@ -1,249 +1,15 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Budgeted hot snapshots with CPU overflow and shared publication/read leases.""" +"""Compatibility alias for the shared Xavier tiered_snapshot module.""" -from collections import OrderedDict -from typing import Dict, List, Tuple +import sys +from typing import TYPE_CHECKING -import torch +from ...xavier.backends.torch import tiered_snapshot as _implementation -from .snapshot import KVSnapshotStore +if TYPE_CHECKING: + from ...xavier.backends.torch.tiered_snapshot import ( + TieredKVSnapshotStore as TieredKVSnapshotStore, + ) - -class TieredKVSnapshotStore(KVSnapshotStore): - """GPU-first snapshot policy, independent of the actor transport. - - ``block_bytes`` reserves a complete block across all layers, including while - it is only partially staged. The GPU budget covers retained snapshots, not - temporary copy/read tensors or the CUDA allocator's reserved memory. - Leased blocks cannot migrate or be evicted. Reads default to CPU to retain - the base store's interface; GPU consumers must select their output device. - CUDA operations follow the caller's current stream. - """ - - def __init__( - self, - cpu_capacity: int, - gpu_budget_bytes: int, - block_bytes: int, - gpu_device: torch.device, - ): - if cpu_capacity <= 0 or gpu_budget_bytes < 0 or block_bytes <= 0: - raise ValueError("Invalid tiered snapshot capacity or budget") - super().__init__(cpu_capacity + gpu_budget_bytes // block_bytes) - self.cpu_capacity = cpu_capacity - self.gpu_capacity = gpu_budget_bytes // block_bytes - self.gpu_device = gpu_device - self.tiers: Dict[int, str] = {} - # GPU victim lookup must not walk an arbitrarily large CPU history. - # Keep the GPU subsequence in the same order as the global block LRU. - self._gpu_lru: OrderedDict[int, None] = OrderedDict() - self.counts = {"gpu": 0, "cpu": 0} - self.metrics = dict(demotions=0, gpu_hits=0, cpu_hits=0, skipped=0) - self.block_bytes = block_bytes - # Avoid rescanning every previous layer when validating the next one. - self._block_sizes: Dict[int, int] = {} - - def touch(self, key: int) -> None: - self.blocks.move_to_end(key) - if key in self._gpu_lru: - self._gpu_lru.move_to_end(key) - - def _drop(self, key: int) -> None: - self._gpu_lru.pop(key, None) - self.counts[self.tiers.pop(key)] -= 1 - del self.blocks[key] - del self.logical_dtypes[key] - del self._block_sizes[key] - self.ready.discard(key) - self.evicted.add(key) - - def _cpu_room(self, pinned: set) -> bool: - if self.counts["cpu"] < self.cpu_capacity: - return True - victim = next( - (k for k in self.blocks if self.tiers[k] == "cpu" and k not in pinned), - None, - ) - if victim is None: - return False - self._drop(victim) - return True - - @staticmethod - def _copy_to_cpu(layers: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - # A synchronous D2H copy per layer stalls the producer for every layer - # of every evicted block. Pack compatible layers and synchronize once - # per group. Each host allocation belongs to one block, so evicting a - # block releases its storage without retaining unrelated snapshots. - groups: Dict[ - Tuple[torch.device, torch.dtype], List[Tuple[str, torch.Tensor]] - ] = {} - host = {} - for layer, value in layers.items(): - if value.device.type != "cuda": - host[layer] = value.to("cpu", copy=True) - else: - groups.setdefault((value.device, value.dtype), []).append( - (layer, value) - ) - for group in groups.values(): - packed = torch.cat([value.reshape(-1) for _, value in group]).to("cpu") - values = packed.split([value.numel() for _, value in group]) - for (layer, original), value in zip(group, values): - host[layer] = value.reshape(original.shape) - return host - - def _admit(self, key: int, pinned: set) -> bool: - if self.gpu_capacity and self.counts["gpu"] >= self.gpu_capacity: - victim = next( - (k for k in self._gpu_lru if k not in pinned), - None, - ) - if victim is not None: - cpu_available = self.counts["cpu"] < self.cpu_capacity or any( - self.tiers[k] == "cpu" and k not in pinned for k in self.blocks - ) - if cpu_available: - # Copy before evicting CPU content or publishing the new tier. - host = self._copy_to_cpu(self.blocks[victim]) - self._cpu_room(pinned) - self.blocks[victim] = host - self.tiers[victim] = "cpu" - del self._gpu_lru[victim] - self.counts["gpu"] -= 1 - self.counts["cpu"] += 1 - self.metrics["demotions"] += 1 - else: - # CPU leases must not prevent reuse of an unleased GPU slot. - self._drop(victim) - tier = "gpu" if self.counts["gpu"] < self.gpu_capacity else "cpu" - if tier == "cpu" and not self._cpu_room(pinned): - self.metrics["skipped"] += 1 - return False - self.blocks[key] = {} - self.logical_dtypes[key] = {} - self._block_sizes[key] = 0 - self.tiers[key] = tier - if tier == "gpu": - self._gpu_lru[key] = None - self.counts[tier] += 1 - return True - - def stage( - self, - layer: str, - keys: List[int], - tensors: torch.Tensor, - logical_dtype: torch.dtype | None = None, - ): - # Reject the whole batch before changing content, placement, or LRU order. - for key, tensor in zip(keys, tensors): - existing = self.blocks.get(key, {}) - if layer in existing: - continue - used = self._block_sizes.get(key, 0) - if used + tensor.numel() * tensor.element_size() > self.block_bytes: - raise ValueError("Snapshot exceeds configured full-block size") - pinned = set().union(*self.leases.values()) if self.leases else set() - for key, tensor in zip(keys, tensors): - if key not in self.blocks and not self._admit(key, pinned): - continue - self.touch(key) - if layer not in self.blocks[key]: - device = ( - self.gpu_device if self.tiers[key] == "gpu" else torch.device("cpu") - ) - try: - self.blocks[key][layer] = ( - tensor.detach().to(device, copy=True).contiguous() - ) - self.logical_dtypes[key][layer] = logical_dtype or tensors.dtype - self._block_sizes[key] += tensor.numel() * tensor.element_size() - except Exception: - # A failed first layer must not consume an empty slot. - if not self.blocks[key]: - self._drop(key) - raise - - def stage_blocks(self, keys: List[int], layers: Dict[str, torch.Tensor]): - """Copy complete new blocks once per dtype, with independent ownership.""" - if not keys: - return - if len(keys) != len(set(keys)): - raise ValueError("Packed staging requires unique snapshot keys") - if not layers or any(len(value) != len(keys) for value in layers.values()): - raise ValueError("Snapshot keys and tensors differ in length") - groups: Dict[ - Tuple[torch.device, torch.dtype], List[Tuple[str, torch.Tensor]] - ] = {} - nbytes = sum( - value[0].numel() * value.element_size() for value in layers.values() - ) - if nbytes > self.block_bytes: - raise ValueError("Snapshot exceeds configured full-block size") - if any(key in self.blocks for key in keys): - raise ValueError("Packed staging requires new snapshot keys") - for name, value in layers.items(): - groups.setdefault((value.device, value.dtype), []).append((name, value)) - packed = [] - for group in groups.values(): - values = torch.cat( - [value.reshape(len(keys), -1) for _, value in group], dim=1 - ) - names = [name for name, _ in group] - shapes = [value.shape[1:] for _, value in group] - sizes = [value[0].numel() for _, value in group] - packed.append((names, shapes, sizes, values)) - dtypes = {name: value.dtype for name, value in layers.items()} - pinned = set().union(*self.leases.values()) if self.leases else set() - for index, key in enumerate(keys): - if not self._admit(key, pinned): - continue - try: - target = ( - self.gpu_device if self.tiers[key] == "gpu" else torch.device("cpu") - ) - content: Dict[str, torch.Tensor] = {} - for names, shapes, sizes, values in packed: - # Layers share only this block's allocation. Evicting one - # block cannot retain the rest of the staging batch. - block = values[index].detach().to(target, copy=True) - parts = ( - value.reshape(shape) - for value, shape in zip(block.split(sizes), shapes) - ) - content.update(zip(names, parts)) - self.blocks[key] = content - self.logical_dtypes[key] = dtypes.copy() - self._block_sizes[key] = nbytes - except Exception: - self._drop(key) - raise - - def reserve(self, lease: str, keys: List[int]) -> bool: - if not super().reserve(lease, keys): - return False - for key in keys: - if key in self._gpu_lru: - self._gpu_lru.move_to_end(key) - self.metrics[self.tiers[key] + "_hits"] += 1 - return True - - def read( - self, layer: str, keys: List[int], device: torch.device | None = None - ) -> torch.Tensor: - target = torch.device("cpu") if device is None else device - return torch.stack( - [self.blocks[key][layer].to(target) for key in keys] - ).contiguous() - - def stats(self) -> dict: - return dict( - self.metrics, - gpu_blocks=self.counts["gpu"], - cpu_blocks=self.counts["cpu"], - gpu_capacity=self.gpu_capacity, - cpu_capacity=self.cpu_capacity, - gpu_reserved_bytes=self.counts["gpu"] * self.block_bytes, - ) +sys.modules[__name__] = _implementation diff --git a/xinference/model/llm/vllm/xavier/transfer.py b/xinference/model/llm/vllm/xavier/transfer.py index ada4ae4f8c..65b5e9d053 100644 --- a/xinference/model/llm/vllm/xavier/transfer.py +++ b/xinference/model/llm/vllm/xavier/transfer.py @@ -21,9 +21,9 @@ import torch import xoscar as xo -from .collective import CollectiveRank -from .profiling import profile_stage -from .snapshot import KVSnapshotStore +from ...xavier.backends.torch.snapshot import KVSnapshotStore +from ...xavier.collective import CollectiveRank +from ...xavier.profiling import profile_stage try: from vllm.utils import TORCH_DTYPE_TO_NUMPY_DTYPE, Device @@ -117,7 +117,7 @@ def get_gloo_dtype(self, input_dtype: torch.dtype): return TypeMappingGloo[TORCH_DTYPE_TO_NUMPY_DTYPE[input_dtype]] -from .gpu_transfer import GPUTransferMixin +from ...xavier.backends.torch.gpu_transfer import GPUTransferMixin class TransferActor( @@ -189,7 +189,7 @@ def _get_cache_engine(self, virtual_engine: int) -> CacheEngine: return self._cache_engine[virtual_engine] # type: ignore def configure_snapshots_v1(self, capacity: int): - from .snapshot import KVSnapshotStore + from ...xavier.backends.torch.snapshot import KVSnapshotStore if capacity <= 0: raise ValueError("Xavier snapshot capacity must be positive") @@ -482,7 +482,7 @@ async def read_layer_blocks_v1( async def start_send_request_blocks_v1(self, to_rank, reads): from xoscar.collective import xoscar_pygloo as xp - from .request_transfer import pack_reads + from ...xavier.backends.torch.request_transfer import pack_reads # Validate and retain the complete payload before acknowledging the send. payload = pack_reads(self._snapshot_store, reads) diff --git a/xinference/model/llm/vllm/xavier/transport.py b/xinference/model/llm/vllm/xavier/transport.py index ac743f22cf..e13fd30640 100644 --- a/xinference/model/llm/vllm/xavier/transport.py +++ b/xinference/model/llm/vllm/xavier/transport.py @@ -13,7 +13,9 @@ # limitations under the License. from typing import Any, Dict, Optional -from urllib.parse import urlsplit + +from ...xavier.transport import get_transport_host as get_transport_host +from ...xavier.transport import validate_gpu_cache_budget as validate_gpu_cache_budget XAVIER_TRANSPORT_BACKEND_KEY = "vllm_transfer_backend_type" XAVIER_TRANSPORT_BACKEND_ALIAS_KEY = "transfer_backend_type" @@ -55,22 +57,6 @@ def set_xavier_transport_backend( return xavier_config -def validate_gpu_cache_budget(value, enabled: bool, replicas: int): - if value is None: - return None - if isinstance(value, bool) or not isinstance(value, int) or value < 0: - raise ValueError("xavier_gpu_cache_bytes must be a non-negative integer") - if not enabled or replicas <= 1: - raise ValueError( - "xavier_gpu_cache_bytes requires Xavier with multiple replicas" - ) - return value - - -def get_transport_host(address: str) -> Optional[str]: - return urlsplit(address if "://" in address else "tcp://" + address).hostname - - def gpu_pool_options(address: str, env: Dict[str, str]) -> Dict[str, str]: import os from importlib.metadata import version diff --git a/xinference/model/llm/vllm/xavier/utils.py b/xinference/model/llm/vllm/xavier/utils.py index adc963164a..b85e642e86 100644 --- a/xinference/model/llm/vllm/xavier/utils.py +++ b/xinference/model/llm/vllm/xavier/utils.py @@ -1,60 +1,13 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd -# -# 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 struct -from typing import List, Optional +# Licensed under the Apache License, Version 2.0. +"""Compatibility alias for the shared Xavier utils module.""" -import xxhash +import sys +from typing import TYPE_CHECKING +from ...xavier import utils as _implementation -def hash_block_tokens( - is_first_block: bool, - prev_block_hash: Optional[int], - cur_block_token_ids: List[int], - extra_hash: Optional[int] = None, - none_hash: int = -1, -) -> int: - """ - Computes a stable hash value corresponding to the contents of a block - and the contents of the preceding block(s), using xxhash instead of - Python hash() for cross-process stability. - """ +if TYPE_CHECKING: + from ...xavier.utils import hash_block_tokens as hash_block_tokens - if prev_block_hash is None: - prev_block_hash = none_hash - if extra_hash is None: - extra_hash = none_hash - - buf = bytearray() - - # 0. hash version - buf += b"v1" - - # 1. is_first_block: bool -> 1 byte - buf += struct.pack(" asyncio.AbstractEventLoop: + # Never reuse a parent's event loop after fork, or a different thread's loop. + if getattr(_local, "pid", None) != os.getpid() or not getattr(_local, "users", 0): + _local.loop = asyncio.new_event_loop() + _local.pid = os.getpid() + _local.users = 0 + _local.users += 1 + return _local.loop + + +def release_actor_loop(loop: asyncio.AbstractEventLoop) -> None: + if ( + getattr(_local, "pid", None) != os.getpid() + or getattr(_local, "loop", None) is not loop + ): + return + _local.users -= 1 + if _local.users: + return + try: + # xoscar listener tasks close their clients in cancellation cleanup. + pending = asyncio.all_tasks(loop) + for task in pending: + task.cancel() + + async def drain(): + if pending: + await asyncio.gather(*pending, return_exceptions=True) + await loop.shutdown_asyncgens() + + cleanup = loop.create_task(drain()) + _, unfinished = loop.run_until_complete( + asyncio.wait({cleanup}, timeout=_SHUTDOWN_TIMEOUT) + ) + if unfinished: + logger.warning("Timed out draining Xavier actor loop during shutdown") + cleanup.cancel() + # Give cancellation one turn, without waiting on unresponsive peers. + loop.run_until_complete(asyncio.sleep(0)) + elif not cleanup.cancelled(): + cleanup.result() + finally: + try: + loop.close() + finally: + del _local.loop diff --git a/xinference/model/llm/xavier/backends/__init__.py b/xinference/model/llm/xavier/backends/__init__.py new file mode 100644 index 0000000000..ec2f132353 --- /dev/null +++ b/xinference/model/llm/xavier/backends/__init__.py @@ -0,0 +1,3 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Optional Xavier storage and transfer implementations.""" diff --git a/xinference/model/llm/xavier/backends/torch/__init__.py b/xinference/model/llm/xavier/backends/torch/__init__.py new file mode 100644 index 0000000000..d4bdcb3be5 --- /dev/null +++ b/xinference/model/llm/xavier/backends/torch/__init__.py @@ -0,0 +1,3 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Torch snapshot storage and CUDA/NIXL transfer; loaded by engine adapters.""" diff --git a/xinference/model/llm/xavier/backends/torch/direct_handoff.py b/xinference/model/llm/xavier/backends/torch/direct_handoff.py new file mode 100644 index 0000000000..8fe2b9669d --- /dev/null +++ b/xinference/model/llm/xavier/backends/torch/direct_handoff.py @@ -0,0 +1,265 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Request-scoped GPU handoff with independent tiered history.""" + +import asyncio +import logging +import time +from dataclasses import dataclass, field + +import torch +import xoscar as xo + +from .direct_history import DirectHistoryMixin +from .gpu_transfer import GPUTransfer +from .request_transfer import LayerRead, batch_reads, unpack_reads + +logger = logging.getLogger(__name__) + + +@dataclass +class DirectRequest: + request_id: str + blocks: set[int] + deadline: float + claimed: bool = False + reading: bool = False + released: bool = False + indices: dict[tuple[int, ...], torch.Tensor] = field(default_factory=dict) + hashes: dict[int, int] = field(default_factory=dict) + retaining: bool = False + retention_attempted: bool = False + + +class DirectGPUTransfer(DirectHistoryMixin, GPUTransfer): + def __init__(self, actor, caches, budget): + super().__init__(actor, caches, 0) + self.direct_requests: dict[str, DirectRequest] = {} + self.finished_sending: set[str] = set() + self.metrics.update( + direct_registered=0, direct_finished=0, direct_expired=0, index_uploads=0 + ) + + self._init_history(budget) + + def _index_tensor( + self, indices: dict[tuple[int, ...], torch.Tensor], values: list[int] + ) -> torch.Tensor: + # Reuse across layers and slabs within one request. Constructing a CUDA + # tensor from a Python list otherwise synchronizes an H2D copy each time. + key = tuple(values) + if key not in indices: + indices[key] = torch.tensor(values, dtype=torch.long, device=self.device) + self.metrics["index_uploads"] += 1 + return indices[key] + + async def _close(self): + self._history_closing = True + if self._history_task is not None: + await asyncio.gather(self._history_task, return_exceptions=True) + if self.history is not None: + self.metrics["history"] = self.history.stats() + await super()._close() + if self.history is not None: + for key in list(self.history.blocks): + self.history._drop(key) + self.history.leases.clear() + self.history.evicted.clear() + self._history_leases.clear() + self._history_probation.clear() + for state in self.direct_requests.values(): + state.indices.clear() + self.direct_requests.clear() + self.finished_sending.clear() + + def register_direct(self, ticket, request_id, blocks, hashes=None): + if self.closing or ticket in self.direct_requests: + raise ValueError("Invalid direct handoff registration") + if not blocks or any( + block < 0 or block >= len(cache) + for cache in self.caches.values() + for block in blocks + ): + raise ValueError("Invalid direct KV block IDs") + if hashes is not None and len(hashes) != len(blocks): + raise ValueError("History hashes and engine blocks differ") + self.direct_requests[ticket] = DirectRequest( + request_id, + set(blocks), + time.monotonic() + 120, + hashes=dict(zip(blocks, hashes or [])), + ) + self.metrics["direct_registered"] += 1 + + def _live_direct(self, ticket): + state = self.direct_requests.get(ticket) + if state is None or state.released: + return None + if not state.reading and time.monotonic() >= state.deadline: + self.metrics["direct_expired"] += 1 + state.retention_attempted = True + self.release_direct(ticket) + return None + return state + + def claim_direct(self, ticket): + state = self._live_direct(ticket) + if state is None: + return False + # Reclaim abandoned claims even when D dies before submitting a load. + # A scheduling retry or a completed slab refreshes this idle lease. + state.claimed = True + state.deadline = time.monotonic() + 600 + return True + + def abandon_direct(self, ticket): + state = self.direct_requests.get(ticket) + if state is not None and not state.claimed: + self.release_direct(ticket) + + def release_direct(self, ticket): + state = self.direct_requests.get(ticket) + if state is None: + return + state.released = True + if not state.reading: + if self._schedule_history(ticket, state): + return + state.indices.clear() + self.finished_sending.add(state.request_id) + del self.direct_requests[ticket] + self.metrics["direct_finished"] += 1 + + def poll_direct(self): + for ticket in list(self.direct_requests): + self._live_direct(ticket) + result = self.finished_sending + self.finished_sending = set() + return result + + async def send_direct(self, ticket, reads, remote_ref, slab_bytes): + async with self.send_lock: + state = self._live_direct(ticket) + if state is None: + return False + if slab_bytes not in self.send_buffers: + raise ValueError("Direct peer transfer slab sizes differ") + if sum(read.nbytes for read in reads) > slab_bytes: + raise ValueError("Direct transfer exceeds slab capacity") + for read in reads: + cache = self.caches[read.layer] + if ( + not set(read.keys) <= state.blocks + or cache.dtype != read.dtype + or tuple(cache.shape[1:]) != read.block_shape + ): + raise ValueError("Direct KV source layout or block IDs differ") + state.reading = True + state.deadline = time.monotonic() + 600 + try: + offset = 0 + for read in reads: + end = offset + read.nbytes + target = ( + self.send_buffer[offset:end] + .view(read.dtype) + .reshape(len(read.keys), *read.block_shape) + ) + cache = self.caches[read.layer] + torch.index_select( + cache, + 0, + self._index_tensor(state.indices, read.keys), + out=target, + ) + offset = end + await asyncio.to_thread(torch.cuda.synchronize, self.device) + await xo.copy_to([self.send_buffers[slab_bytes]], [remote_ref]) + self.metrics["wire_bytes"] += slab_bytes + self.metrics["useful_bytes"] += offset + finally: + # Even a failed gather may have queued reads of engine slots. + await asyncio.to_thread(torch.cuda.synchronize, self.device) + state.reading = False + state.deadline = time.monotonic() + 600 + if state.released: + self.release_direct(ticket) + return True + + async def load_direct(self, requests, tickets): + failed = False + try: + return await self._load_direct(requests, tickets) + except BaseException: + failed = True + raise + finally: + + results = await asyncio.gather( + *( + self.actor.release_remote_direct_gpu_v1(rank, ticket) + for ranks, ticket in zip(requests, tickets) + for rank in ranks + ), + return_exceptions=True, + ) + for result in results: + if isinstance(result, BaseException): + if not failed: + raise result + logger.warning("Direct ticket cleanup failed", exc_info=result) + + async def _load_direct(self, requests, tickets): + self.metrics["load_calls"] += 1 + self.metrics["load_requests"] += len(requests) + invalid_blocks = set() + async with self.recv_lock: + for ranks, ticket in zip(requests, tickets): + if len(ranks) != 1: + raise ValueError("Direct handoff requires one producer") + rank, layers = next(iter(ranks.items())) + sender = await xo.actor_ref( + address=self.actor._world_addresses[rank], + uid=f"{self.actor.default_uid()}-{rank}", + ) + indices: dict[tuple[int, ...], torch.Tensor] = {} + try: + reads = [ + LayerRead( + name, + list(mapping), + list(mapping.values()), + tuple(self.caches[name].shape[1:]), + self.caches[name].dtype, + ) + for name, mapping in layers.items() + if mapping + ] + for batch in batch_reads(reads, max_bytes=self.slab_bytes): + size = sum(read.nbytes for read in batch) + slab_bytes = min(n for n in self.recv_refs if n >= size) + available = await sender.send_direct_gpu_v1( + ticket, batch, self.recv_refs[slab_bytes], slab_bytes + ) + if not available: + # Never expose a partial request as valid KV. vLLM's + # load-error callback recomputes all its destinations. + invalid_blocks.update( + dest + for mapping in layers.values() + for dest in mapping.values() + ) + break + for read, blocks in unpack_reads( + self.recv_buffer[:size], batch + ): + cache = self.caches[read.layer] + cache[self._index_tensor(indices, read.destinations)] = ( + blocks + ) + await asyncio.to_thread(torch.cuda.synchronize, self.device) + self.metrics["gpu_batches"] += 1 + finally: + await asyncio.to_thread(torch.cuda.synchronize, self.device) + indices.clear() + return invalid_blocks diff --git a/xinference/model/llm/xavier/backends/torch/direct_history.py b/xinference/model/llm/xavier/backends/torch/direct_history.py new file mode 100644 index 0000000000..cf58f83545 --- /dev/null +++ b/xinference/model/llm/xavier/backends/torch/direct_history.py @@ -0,0 +1,361 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Bounded GPU/CPU history retained after request-scoped handoff.""" + +import asyncio +import logging +import time +from collections import OrderedDict + +import torch + +from .request_transfer import LayerRead, batch_reads +from .tiered_snapshot import TieredKVSnapshotStore + +logger = logging.getLogger(__name__) + + +class HistoryStore(TieredKVSnapshotStore): + """Independent GPU-first history with bounded CPU overflow.""" + + def __init__(self, budget, block_bytes, device, cpu_capacity=0): + super().__init__(max(1, cpu_capacity), budget, block_bytes, device) + self.cpu_capacity = cpu_capacity + self.capacity = self.gpu_capacity + cpu_capacity + + def reserve(self, lease, keys): + if not set(keys).issubset(self.ready): + return False + self.leases.setdefault(lease, set()).update(keys) + # Prefix heads outlive their tails. Reservation is not a cache hit: + # scheduling can retry or abort without ever restoring any blocks. + for key in reversed(keys): + self.touch(key) + return True + + def read(self, layer, keys, device=None): + # Batch same-tier blocks before H2D instead of synchronizing once per + # block. load_history partitions reads by tier and holds their leases. + values = [self.blocks[key][layer] for key in keys] + target = torch.device("cpu") if device is None else device + if all(value.device == values[0].device for value in values): + return torch.stack(values).to(target).contiguous() + return super().read(layer, keys, device) + + +class DirectHistoryMixin: + def _init_history(self, budget): + block_bytes = sum(c[0].numel() * c.element_size() for c in self.caches.values()) + cpu_capacity = self.store.cpu_capacity + self.history = ( + HistoryStore(budget, block_bytes, self.device, cpu_capacity) + if budget + else None + ) + self._history_task = None + self._history_closing = False + self._history_leases = {} + self._history_reading = set() + # Metadata only. A second completed request demonstrates reuse before + # new content can displace existing GPU history. Empty slots need no + # probation, preserving first-request retention while there is room. + self._history_probation: OrderedDict[int, None] = OrderedDict() + self._history_probation_limit = min( + 8192, 2 * self.history.capacity if self.history else 0 + ) + self.history_chunk_bytes = 2 * 1024 * 1024 + self.history_fill_chunk_bytes = 8 * 1024 * 1024 + self.history_pending_bytes = min( + self.history.capacity * block_bytes if self.history else 0, 64 * 1024 * 1024 + ) + # A soft admission deadline, not a cancellation timeout: an in-flight + # CUDA chunk always drains before the engine can reuse its blocks. + self.history_retention_seconds = 0.01 + self.metrics.update( + history_saved_blocks=0, + history_skipped_requests=0, + history_failures=0, + history_hit_blocks=0, + history_cpu_hit_blocks=0, + history_gpu_hit_blocks=0, + history_loaded_requests=0, + history_admission_rejected_blocks=0, + history_admission_reused_blocks=0, + history_save_chunks=0, + history_capacity_limited_blocks=0, + history_deadline_dropped_blocks=0, + history_closing_dropped_blocks=0, + ) + + def _expire_history_leases(self): + now = time.monotonic() + for lease, deadline in list(self._history_leases.items()): + if deadline <= now and lease not in self._history_reading: + self.release_history(lease) + + def reserve_history(self, lease, keys): + if self.closing or self.history is None: + return [] + self._expire_history_leases() + if lease in self._history_leases: + raise ValueError("Duplicate history lease") + matched = [] + for key in keys: + if key not in self.history.ready: + break + matched.append(key) + if matched: + if not self.history.reserve(lease, matched): + return [] + self._history_leases[lease] = time.monotonic() + 120 + return matched + + def release_history(self, lease): + if lease not in self._history_reading: + self._history_leases.pop(lease, None) + if self.history is not None: + self.history.release(lease) + + def _schedule_history(self, ticket, state): + if state.retaining: + return True + if state.retention_attempted: + return False + state.retention_attempted = True + if self._history_closing or self.history is None or not state.hashes: + return False + limit = min( + self.history.capacity, + self.history_pending_bytes // self.history.block_bytes, + ) + candidates = [] + free = max(0, self.history.capacity - len(self.history.blocks)) + prefix = [] + blocked = None + for block, key in state.hashes.items(): + if blocked is not None: + if key not in self.history.ready: + self.metrics[blocked] += 1 + continue + if len(prefix) >= self.history.capacity or ( + key not in self.history.ready and len(candidates) >= limit + ): + blocked = "history_capacity_limited_blocks" + if key not in self.history.ready: + self.metrics[blocked] += 1 + elif key in self.history.ready: + prefix.append(key) + elif free or key in self._history_probation: + candidates.append((key, block)) + prefix.append(key) + free = max(0, free - 1) + if key in self._history_probation: + self.metrics["history_admission_reused_blocks"] += 1 + else: + blocked = "history_admission_rejected_blocks" + self.metrics[blocked] += 1 + # Observe the full request, but keep heads newest in both probation and + # cache LRU. A rejected head must not admit unreachable suffix blocks. + for key in reversed(list(state.hashes.values())): + if key in self.history.ready: + self.history.touch(key) + self._history_probation.pop(key, None) + else: + self._history_probation[key] = None + self._history_probation.move_to_end(key) + while len(self._history_probation) > self._history_probation_limit: + self._history_probation.popitem(last=False) + if not candidates: + return False + # Observe completed requests even when the single writer is busy, but + # never queue another writer or extend engine block ownership for it. + if self._history_task is not None and not self._history_task.done(): + self.metrics["history_skipped_requests"] += 1 + return False + state.retaining = True + deadline = time.monotonic() + self.history_retention_seconds + task = self._history_task = asyncio.create_task( + self._retain_history(ticket, state, candidates, deadline) + ) + self.tasks.add(task) + + def done(future): + self.tasks.discard(future) + if not future.cancelled() and future.exception() is not None: + exc = future.exception() + logger.error( + "GPU history retention failed", + exc_info=exc, + ) + + task.add_done_callback(done) + return True + + async def _retain_history(self, ticket, state, candidates, deadline): + pending = [] + prefix_keys = list(state.hashes.values())[: self.history.capacity] + try: + # At most one writer and one bounded chunk. Yield between chunks; + # never queue the writer ahead of an already active handoff or load. + start = 0 + while start < len(candidates): + # Filling unused capacity can amortize per-layer gathers; + # replacements use smaller chunks to limit lock hold time. + free = self.history.gpu_capacity - self.history.counts["gpu"] + chunk_bytes = ( + self.history_fill_chunk_bytes + if free > 0 + else self.history_chunk_bytes + ) + count = max( + 1, + min( + 32, + min(self.slab_bytes, chunk_bytes) // self.history.block_bytes, + ), + ) + if free > 0: + count = min(count, free) + chunk = candidates[start : start + count] + while True: + if self._history_closing or time.monotonic() >= deadline: + reason = "closing" if self._history_closing else "deadline" + self.metrics[f"history_{reason}_dropped_blocks"] += ( + len(candidates) - start + ) + return + if not self.send_lock.locked() and not self.recv_lock.locked(): + break + await asyncio.sleep(0.001) + start += len(chunk) + async with self.send_lock: + self._expire_history_leases() + pairs = [(k, b) for k, b in chunk if k not in self.history.ready] + if not pairs: + continue + pending, ids = map(list, zip(*pairs)) + index = self._index_tensor(state.indices, ids) + layers = { + name: cache.index_select(0, index) + for name, cache in self.caches.items() + } + self.history.stage_blocks(pending, layers) + await asyncio.to_thread(torch.cuda.synchronize, self.device) + available = self.history.publish(pending, set(self.caches)) + self.metrics["history_saved_blocks"] += len(available) + self.metrics["history_save_chunks"] += 1 + for key in available: + self._history_probation.pop(key, None) + for key in reversed(prefix_keys): + if key in self.history.ready: + self.history.touch(key) + complete = len(available) == len(pending) + if not complete: + self.metrics["history_capacity_limited_blocks"] += ( + len(pending) - len(available) + len(candidates) - start + ) + pending = [] + del layers, index + if not complete: + # Capacity may be pinned by another reader. Do not + # continue past a prefix we could not publish. + return + await asyncio.sleep(0) + except Exception: + self.metrics["history_failures"] += 1 + logger.warning("Skipping failed GPU history retention", exc_info=True) + finally: + # Original engine slots stay owned even after D finishes, until all + # optional snapshot reads have drained. Never publish partial data. + await asyncio.to_thread(torch.cuda.synchronize, self.device) + for key in pending: + if key in self.history.blocks and key not in self.history.ready: + self.history._drop(key) + self.history.evicted.clear() + state.retaining = False + self.release_direct(ticket) + + async def load_history(self, requests, leases): + failed = False + try: + await self._load_history(requests, leases) + except BaseException: + failed = True + raise + finally: + errors = [] + for lease in leases: + try: + self.release_history(lease) + except BaseException as error: + errors.append(error) + if errors and not failed: + raise errors[0] + for cleanup_error in errors: + logger.warning("History lease cleanup failed", exc_info=cleanup_error) + + async def _load_history(self, requests, leases): + async with self.recv_lock: + for ranks, lease in zip(requests, leases): + indices = {} + try: + if self.history is None or lease not in self._history_leases: + raise RuntimeError("GPU history lease expired") + self._history_reading.add(lease) + reads = [] + written = set() + for layers in ranks.values(): + for name, mapping in layers.items(): + cache = self.caches[name] + if not set(mapping).issubset(self.history.leases[lease]): + raise ValueError("History read outside lease") + for tier in ("gpu", "cpu"): + pairs = [ + (key, dest) + for key, dest in mapping.items() + if self.history.tiers[key] == tier + ] + if pairs: + keys, destinations = map(list, zip(*pairs)) + reads.append( + LayerRead( + name, + keys, + destinations, + tuple(cache.shape[1:]), + cache.dtype, + ) + ) + for batch in batch_reads(reads, max_bytes=self.slab_bytes): + for read in batch: + blocks = self.history.read( + read.layer, read.keys, self.device + ) + if ( + blocks.dtype != read.dtype + or tuple(blocks.shape[1:]) != read.block_shape + ): + raise ValueError("GPU history layout mismatch") + self.caches[read.layer][ + self._index_tensor(indices, read.destinations) + ] = blocks + written.update(read.keys) + await asyncio.to_thread(torch.cuda.synchronize, self.device) + hits = {"gpu": 0, "cpu": 0} + for key in written: + tier = self.history.tiers[key] + hits[tier] += 1 + self.metrics[f"history_{tier}_hit_blocks"] += 1 + self.history.metrics[f"{tier}_hits"] += 1 + self.metrics["history_hit_blocks"] += len(written) + self.metrics["history_loaded_requests"] += 1 + logger.debug( + "Restored Xavier history: blocks=%s gpu=%s cpu=%s", + len(written), + hits["gpu"], + hits["cpu"], + ) + finally: + await asyncio.to_thread(torch.cuda.synchronize, self.device) + indices.clear() + self._history_reading.discard(lease) diff --git a/xinference/model/llm/xavier/backends/torch/gpu_transfer.py b/xinference/model/llm/xavier/backends/torch/gpu_transfer.py new file mode 100644 index 0000000000..04769551e9 --- /dev/null +++ b/xinference/model/llm/xavier/backends/torch/gpu_transfer.py @@ -0,0 +1,542 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Local CUDA IPC caches and persistent xoscar NIXL transfer buffers.""" + +import asyncio +import logging +from importlib.metadata import version +from typing import Dict + +import torch +import xoscar as xo +from packaging.version import Version + +from .request_transfer import ( + MAX_REQUEST_BLOCKS, + MAX_REQUEST_BYTES, + LayerRead, + batch_reads, + unpack_reads, +) +from .tiered_snapshot import TieredKVSnapshotStore + +logger = logging.getLogger(__name__) + + +async def finish_before_cancel(task): + """Keep ownership until completion, including after repeated cancellation.""" + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError: + continue + except Exception: + break + # Observe the result without replacing the caller's cancellation. + if not task.cancelled(): + task.exception() + raise + + +class GPUTransfer: + def __init__(self, actor, caches: Dict[str, torch.Tensor], budget: int): + if Version(version("xoscar")) < Version("0.11.1"): + raise RuntimeError("Xavier GPU transfer requires xoscar[nixl]>=0.11.1") + if not actor.address.startswith("nixl://"): + raise RuntimeError("Xavier GPU transfer requires a NIXL actor pool") + if not caches or any(not cache.is_cuda for cache in caches.values()): + raise ValueError("Xavier GPU transfer requires registered CUDA KV caches") + devices = {cache.device for cache in caches.values()} + if len(devices) != 1: + raise ValueError("Xavier GPU transfer requires one CUDA device per replica") + self.actor, self.caches = actor, caches + self.device = next(iter(devices)) + sizes = [cache[0].numel() * cache.element_size() for cache in caches.values()] + self.store = TieredKVSnapshotStore( + actor._snapshot_store.capacity, budget, sum(sizes), self.device + ) + # Fixed slabs avoid registration churn. Separate directions allow hybrids + # to serve a peer while they themselves are waiting for incoming data. + self.slab_bytes = max( + max(sizes), min(MAX_REQUEST_BYTES, sum(sizes) * MAX_REQUEST_BLOCKS) + ) + self.send_buffer = torch.zeros( + self.slab_bytes, dtype=torch.uint8, device=self.device + ) + self.recv_buffer = torch.zeros_like(self.send_buffer) + self.recv_ref = xo.buffer_ref(actor.address, self.recv_buffer) + # Keep stable views for small warm-cache transfers. Reuse their buffer + # identities, rather than creating and registering a slice per request. + small_bytes = min(self.slab_bytes, 256 * 1024) + self.send_buffers = {self.slab_bytes: self.send_buffer} + self.recv_buffers = {self.slab_bytes: self.recv_buffer} + self.recv_refs = {self.slab_bytes: self.recv_ref} + if small_bytes < self.slab_bytes: + self.send_buffers[small_bytes] = self.send_buffer[:small_bytes] + self.recv_buffers[small_bytes] = self.recv_buffer[:small_bytes] + self.recv_refs[small_bytes] = xo.buffer_ref( + actor.address, self.recv_buffers[small_bytes] + ) + self._small_slab_streak: Dict[int, int] = {} + self.send_lock, self.recv_lock = asyncio.Lock(), asyncio.Lock() + self.tasks: set[asyncio.Task] = set() + self.closing = False + self.metrics = dict( + gpu_batches=0, + cpu_batches=0, + wire_bytes=0, + useful_bytes=0, + load_calls=0, + load_requests=0, + ) + + async def run(self, function, *args): + if self.closing: + raise RuntimeError("Xavier GPU transfer is shutting down") + task = asyncio.create_task(function(*args)) + self.tasks.add(task) + + def completed(future): + self.tasks.discard(future) + if not future.cancelled(): + error = future.exception() + if error is not None: + logger.warning( + "Xavier GPU transfer task failed", + exc_info=(type(error), error, error.__traceback__), + ) + + task.add_done_callback(completed) + return await finish_before_cancel(task) + + async def load_requests_with_leases(self, requests, leases): + load_failed = False + try: + await self.load_requests(requests) + except BaseException: + load_failed = True + raise + finally: + results = await asyncio.gather( + *( + self.actor.release_remote_blocks_v1(lease, ranks) + for lease, ranks in leases + if lease + ), + return_exceptions=True, + ) + for result in results: + if isinstance(result, BaseException): + if not load_failed: + raise result + logger.warning( + "Failed to release Xavier snapshot lease after load failure", + exc_info=(type(result), result, result.__traceback__), + ) + + async def close(self): + task = getattr(self, "_close_task", None) + if task is None: + self.closing = True + task = self._close_task = asyncio.create_task(self._close()) + return await finish_before_cancel(task) + + async def _close(self): + if self.tasks: + await asyncio.gather(*list(self.tasks), return_exceptions=True) + try: + await asyncio.to_thread(torch.cuda.synchronize, self.device) + logger.info( + "Xavier GPU cache stats: %s", + dict(self.metrics, cache=self.store.stats()), + ) + finally: + self.recv_refs.clear() + self.recv_buffers.clear() + self.send_buffers.clear() + self.recv_ref = None + self.caches.clear() + self.send_buffer = self.recv_buffer = None + self.store._gpu_lru.clear() + self.store._block_sizes.clear() + self.store.blocks.clear() + self.store.ready.clear() + self.store.tiers.clear() + self.store.logical_dtypes.clear() + self.store.leases.clear() + self.store.evicted.clear() + self.store.counts.update(gpu=0, cpu=0) + self._small_slab_streak.clear() + + async def stage(self, entries): + if self.closing: + raise RuntimeError("Xavier GPU transfer is shutting down") + keys = {key for layers in entries for ids, _ in layers.values() for key in ids} + failed = False + copied = False + last_reused_keys = None + # These complete snapshots were copied on this call's stream. They may + # be reused before our final fence, unlike unpublished work of a peer call. + staged_keys: set[int] = set() + try: + for layers in entries: + if self._can_stage_blocks(layers, staged_keys): + copied = self._stage_blocks(layers, staged_keys) or copied + last_reused_keys = None + continue + for layer, (block_keys, ids) in layers.items(): + # Unpublished layers may still be copying in another stage + # awaiting its fence; only published data can skip the copy. + if all( + key in self.store.ready + and layer in self.store.blocks.get(key, {}) + for key in block_keys + ): + if block_keys != last_reused_keys: + for key in block_keys: + self.store.touch(key) + last_reused_keys = block_keys + continue + last_reused_keys = None + cache = self.caches[layer] + blocks = cache.index_select( + 0, torch.tensor(ids, dtype=torch.long, device=cache.device) + ) + self.store.stage(layer, block_keys, blocks) + copied = True + except Exception: + failed = True + logger.warning( + "Failed to stage Xavier GPU snapshots; skipping new snapshots", + exc_info=True, + ) + finally: + # Fence only if this call copied or failed: even a failed gather + # may have queued reads of EngineCore slots. Pure published reuse + # needs no fence before returning ownership to EngineCore. + try: + if copied or failed: + await asyncio.to_thread(torch.cuda.synchronize, self.device) + except BaseException: + failed = True + raise + finally: + if failed: + for key in keys: + if key in self.store.blocks and key not in self.store.ready: + self.store._drop(key) + + def _can_stage_blocks(self, layers, staged_keys): + if not layers or set(layers) != set(self.caches): + return False + keys, ids = next(iter(layers.values())) + return ( + bool(keys) + and len(keys) == len(ids) + and len(keys) == len(set(keys)) + and all(value == (keys, ids) for value in layers.values()) + and all( + key not in self.store.blocks + or ( + (key in self.store.ready or key in staged_keys) + and self.caches.keys() <= self.store.blocks[key].keys() + ) + for key in keys + ) + ) + + def _stage_blocks(self, layers, staged_keys): + keys, ids = next(iter(layers.values())) + count = max( + 1, min(MAX_REQUEST_BLOCKS, MAX_REQUEST_BYTES // self.store.block_bytes) + ) + pending_keys, pending_ids = [], [] + copied = False + + def flush(): + nonlocal copied + if not pending_keys: + return + index = torch.tensor(pending_ids, dtype=torch.long, device=self.device) + values = { + name: cache.index_select(0, index) + for name, cache in self.caches.items() + } + self.store.stage_blocks(pending_keys, values) + staged_keys.update(key for key in pending_keys if key in self.store.blocks) + copied = True + pending_keys.clear() + pending_ids.clear() + + for key, source in zip(keys, ids): + if key in self.store.ready or key in staged_keys: + flush() + # Admission of preceding blocks may have evicted this hit. + if key in self.store.blocks and ( + key in self.store.ready or key in staged_keys + ): + self.store.touch(key) + continue + pending_keys.append(key) + pending_ids.append(source) + if len(pending_keys) == count: + flush() + flush() + return copied + + def locations(self, reads): + result = {} + for read in reads: + for key in read.keys: + if ( + key not in self.store.ready + or read.layer not in self.store.blocks[key] + ): + raise KeyError("Requested Xavier snapshot is not published") + tensor = self.store.blocks[key][read.layer] + if ( + tensor.dtype != read.dtype + or tuple(tensor.shape) != read.block_shape + ): + raise ValueError("Xavier source and destination KV layouts differ") + result[key] = self.store.tiers[key] + return result + + async def send(self, reads, remote_ref, slab_bytes): + async with self.send_lock: + if slab_bytes not in self.send_buffers: + raise ValueError("Xavier peer transfer slab sizes differ") + locations = self.locations(reads) + if any(tier != "gpu" for tier in locations.values()): + raise ValueError("GPU transfer requested for a CPU snapshot") + size = sum(read.nbytes for read in reads) + if size > slab_bytes: + raise ValueError("Xavier GPU batch exceeds transfer slab") + offset = 0 + for read in reads: + end = offset + read.nbytes + target = ( + self.send_buffer[offset:end] + .view(read.dtype) + .reshape(len(read.keys), *read.block_shape) + ) + torch.stack( + [self.store.blocks[key][read.layer] for key in read.keys], + out=target, + ) + offset = end + await asyncio.to_thread(torch.cuda.synchronize, self.device) + await xo.copy_to([self.send_buffers[slab_bytes]], [remote_ref]) + self.metrics["wire_bytes"] += slab_bytes + self.metrics["useful_bytes"] += size + + def _select_slab_bytes(self, rank: int, size: int) -> int: + candidate = min(n for n in self.recv_refs if n >= size) + if candidate == self.slab_bytes: + self._small_slab_streak.pop(rank, None) + return candidate + # xoscar caches only the latest registration on each peer channel. + # Require consecutive small batches so full/small tails do not churn + # registrations. Sustained small workloads still reuse the small view. + streak = min(2, self._small_slab_streak.get(rank, 0) + 1) + self._small_slab_streak[rank] = streak + return candidate if streak == 2 else self.slab_bytes + + async def load(self, ranks): + return await self.load_requests([ranks]) + + async def load_requests(self, requests): + # Index by destination, not hash: one cached block may feed multiple + # requests' slots. Reject conflicting writes before any transfer starts. + ranks = {} + owners = {} + request_bytes = {} + for request in requests: + for rank, layers in request.items(): + size = sum( + len(mapping) + * self.caches[name][0].numel() + * self.caches[name].element_size() + for name, mapping in layers.items() + ) + request_bytes[rank] = max(request_bytes.get(rank, 0), size) + for name, mapping in layers.items(): + for key, destination in mapping.items(): + target = (name, destination) + if target in owners: + if owners[target] != key: + raise ValueError("Conflicting Xavier KV destinations") + continue + owners[target] = key + ranks.setdefault(rank, {}).setdefault(name, {})[ + destination + ] = key + + self.metrics["load_calls"] += 1 + self.metrics["load_requests"] += len(requests) + async with self.recv_lock: + for rank, layers in ranks.items(): + sender = await xo.actor_ref( + address=self.actor._world_addresses[rank], + uid=f"{self.actor.default_uid()}-{rank}", + ) + reads = [ + LayerRead( + name, + list(mapping.values()), + list(mapping), + tuple(self.caches[name].shape[1:]), + self.caches[name].dtype, + ) + for name, mapping in layers.items() + if mapping + ] + locations = await sender.gpu_snapshot_locations_v1(reads) + for tier in ("gpu", "cpu"): + selected = [] + for read in reads: + pairs = [ + (key, dest) + for key, dest in zip(read.keys, read.destinations) + if locations[key] == tier + ] + if pairs: + selected.append( + LayerRead( + read.layer, + [p[0] for p in pairs], + [p[1] for p in pairs], + read.block_shape, + read.dtype, + ) + ) + # Combining small requests must not turn each transfer back + # into a padded full-slab copy. Large requests retain the + # original batch bound to avoid fragmenting cold prefills. + limit = self.slab_bytes + if tier == "gpu": + limit = min( + (n for n in self.recv_refs if n >= request_bytes[rank]), + default=self.slab_bytes, + ) + for batch in batch_reads(selected, max_bytes=limit): + size = sum(read.nbytes for read in batch) + if tier == "gpu": + slab_bytes = self._select_slab_bytes(rank, size) + await sender.send_gpu_request_v1( + batch, self.recv_refs[slab_bytes], slab_bytes + ) + payload = self.recv_buffer[:size] + else: + payload = await self.actor.read_request_blocks_v1( + rank, batch + ) + for read, blocks in unpack_reads(payload, batch): + cache = self.caches[read.layer] + cache[ + torch.tensor( + read.destinations, + dtype=torch.long, + device=cache.device, + ) + ] = blocks.to(cache.device, non_blocking=True) + # Both producer's IPC ownership and receiver slab reuse + # require the cache writes to finish before acknowledging. + await asyncio.to_thread(torch.cuda.synchronize, self.device) + self.metrics[tier + "_batches"] += 1 + + +class GPUTransferMixin: + def map_gpu_caches_v1(self, descriptors, budget, direct_handoff=False): + from torch.multiprocessing.reductions import rebuild_cuda_tensor + + if getattr(self, "_gpu_transfer", None) is not None: + raise RuntimeError("Xavier GPU caches already registered") + caches = { + name: rebuild_cuda_tensor(*desc) for name, desc in descriptors.items() + } + runtime_type = GPUTransfer + if direct_handoff: + from .direct_handoff import DirectGPUTransfer + + runtime_type = DirectGPUTransfer + runtime = runtime_type(self, caches, budget) + self._gpu_transfer = runtime + self._snapshot_store = runtime.store + + async def stage_gpu_requests_v1(self, entries): + runtime = self._gpu_transfer + return await runtime.run(runtime.stage, entries) + + def gpu_snapshot_locations_v1(self, reads): + return self._gpu_transfer.locations(reads) + + async def send_gpu_request_v1(self, reads, remote_ref, slab_bytes): + runtime = self._gpu_transfer + return await runtime.run(runtime.send, reads, remote_ref, slab_bytes) + + async def load_gpu_requests_v1(self, requests, leases): + runtime = self._gpu_transfer + # The actor loop advances independently of EngineCore. Keep both writes + # and lease cleanup inside the protected task before returning readiness. + return await runtime.run(runtime.load_requests_with_leases, requests, leases) + + async def close_gpu_caches_v1(self): + runtime = getattr(self, "_gpu_transfer", None) + if runtime is not None: + # Keep the closed runtime as a tombstone: close() reuses its task, + # and map_gpu_caches_v1 must not re-register after shutdown/failure. + await runtime.close() + + def register_direct_gpu_v1(self, ticket, request_id, blocks, hashes=None): + self._gpu_transfer.register_direct(ticket, request_id, blocks, hashes) + + def claim_direct_gpu_v1(self, ticket): + return self._gpu_transfer.claim_direct(ticket) + + def abandon_direct_gpu_v1(self, ticket): + self._gpu_transfer.abandon_direct(ticket) + + async def claim_remote_direct_gpu_v1(self, rank, ticket): + ref = await xo.actor_ref( + address=self._world_addresses[rank], uid=f"{self.default_uid()}-{rank}" + ) + return await ref.claim_direct_gpu_v1(ticket) + + def release_direct_gpu_v1(self, ticket): + self._gpu_transfer.release_direct(ticket) + + def poll_direct_gpu_v1(self): + return self._gpu_transfer.poll_direct() + + async def send_direct_gpu_v1(self, ticket, reads, remote_ref, slab_bytes): + runtime = self._gpu_transfer + return await runtime.run( + runtime.send_direct, ticket, reads, remote_ref, slab_bytes + ) + + async def load_direct_gpu_v1(self, requests, tickets): + runtime = self._gpu_transfer + return await runtime.run(runtime.load_direct, requests, tickets) + + async def release_remote_direct_gpu_v1(self, rank, ticket): + ref = await xo.actor_ref( + address=self._world_addresses[rank], uid=f"{self.default_uid()}-{rank}" + ) + await ref.release_direct_gpu_v1(ticket) + + def reserve_direct_history_v1(self, lease, keys): + runtime = getattr(self, "_gpu_transfer", None) + return runtime.reserve_history(lease, keys) if runtime is not None else [] + + def release_direct_history_v1(self, lease): + runtime = getattr(self, "_gpu_transfer", None) + if runtime is not None: + runtime.release_history(lease) + + async def load_direct_history_v1(self, requests, leases): + runtime = self._gpu_transfer + return await runtime.run(runtime.load_history, requests, leases) diff --git a/xinference/model/llm/xavier/backends/torch/request_transfer.py b/xinference/model/llm/xavier/backends/torch/request_transfer.py new file mode 100644 index 0000000000..2402ce22af --- /dev/null +++ b/xinference/model/llm/xavier/backends/torch/request_transfer.py @@ -0,0 +1,96 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Bounded cross-layer Gloo messages for Xavier V1 snapshot reads.""" + +import math +from dataclasses import dataclass +from typing import List, Tuple + +import torch + +MAX_REQUEST_BYTES = 16 * 1024 * 1024 +MAX_REQUEST_BLOCKS = 64 + + +@dataclass +class LayerRead: + layer: str + keys: List[int] + destinations: List[int] + block_shape: Tuple[int, ...] + dtype: torch.dtype + logical_dtype: torch.dtype | None = None + + @property + def nbytes(self) -> int: + return len(self.keys) * math.prod(self.block_shape) * self.dtype.itemsize + + +def batch_reads(reads, max_bytes=None): + max_bytes = MAX_REQUEST_BYTES if max_bytes is None else max_bytes + batch = [] + size = 0 + for read in reads: + if batch and read.dtype != batch[0].dtype: + yield batch + batch, size = [], 0 + block_bytes = math.prod(read.block_shape) * read.dtype.itemsize + if block_bytes <= 0: + raise ValueError("Invalid KV block size") + offset = 0 + while offset < len(read.keys): + limit = min(MAX_REQUEST_BLOCKS, len(read.keys) - offset) + count = max(1, min(limit, (max_bytes - size) // block_bytes)) + if batch and size + count * block_bytes > max_bytes: + yield batch + batch, size = [], 0 + continue + part = LayerRead( + read.layer, + read.keys[offset : offset + count], + read.destinations[offset : offset + count], + read.block_shape, + read.dtype, + read.logical_dtype, + ) + batch.append(part) + size += part.nbytes + offset += count + if size >= max_bytes: + yield batch + batch, size = [], 0 + if batch: + yield batch + + +def pack_reads(store, reads): + buffers = [] + for read in reads: + if not set(read.keys).issubset(store.ready): + raise KeyError("Requested KV snapshots are not published") + if any( + store.logical_dtypes[key][read.layer] != (read.logical_dtype or read.dtype) + for key in read.keys + ): + raise ValueError("Xavier logical KV dtype does not match the receiver") + tensor = store.read(read.layer, read.keys) + if tensor.dtype != read.dtype or tuple(tensor.shape) != ( + len(read.keys), + *read.block_shape, + ): + raise ValueError("KV snapshot layout does not match the receiver") + buffers.append(tensor.view(torch.uint8).reshape(-1)) + return torch.cat(buffers) + + +def unpack_reads(payload, reads): + expected = sum(read.nbytes for read in reads) + if payload.dtype != torch.uint8 or payload.numel() != expected: + raise ValueError("Invalid Xavier cross-layer payload size or dtype") + offset = 0 + for read in reads: + end = offset + read.nbytes + yield read, payload[offset:end].view(read.dtype).reshape( + len(read.keys), *read.block_shape + ) + offset = end diff --git a/xinference/model/llm/xavier/backends/torch/snapshot.py b/xinference/model/llm/xavier/backends/torch/snapshot.py new file mode 100644 index 0000000000..7d516a2a9f --- /dev/null +++ b/xinference/model/llm/xavier/backends/torch/snapshot.py @@ -0,0 +1,87 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +from collections import OrderedDict +from typing import Dict, List, Set + +import torch + + +class KVSnapshotStore: + """Bounded, content-addressed CPU blocks with atomic publication and read leases. + + All mutation runs on the TransferActor's event loop. Tensor storage is owned + by each block, never by a GPU slot or a view retaining an entire request. + """ + + def __init__(self, capacity: int): + self.capacity = capacity + self.blocks: OrderedDict[int, Dict[str, torch.Tensor]] = OrderedDict() + self.logical_dtypes: Dict[int, Dict[str, torch.dtype]] = {} + self.ready: Set[int] = set() + self.leases: Dict[str, Set[int]] = {} + self.evicted: Set[int] = set() + + def stage( + self, + layer: str, + keys: List[int], + tensors: torch.Tensor, + logical_dtype: torch.dtype | None = None, + ): + pinned = set().union(*self.leases.values()) if self.leases else set() + for key, tensor in zip(keys, tensors): + if key not in self.blocks: + if len(self.blocks) >= self.capacity: + victim = next((k for k in self.blocks if k not in pinned), None) + if victim is None: + continue # No room: leave this block to local recomputation. + del self.blocks[victim] + del self.logical_dtypes[victim] + self.ready.discard(victim) + self.evicted.add(victim) + self.blocks[key] = {} + self.logical_dtypes[key] = {} + self.blocks.move_to_end(key) + # Identical content is immutable, including while a reader holds it. + if layer not in self.blocks[key]: + self.blocks[key][layer] = tensor.clone().contiguous() + self.logical_dtypes[key][layer] = logical_dtype or tensors.dtype + + def publish(self, keys: List[int], layers: Set[str]) -> List[int]: + available = [] + for key in keys: + if layers and layers.issubset(self.blocks.get(key, {})): + self.ready.add(key) + available.append(key) + return available + + def reserve(self, lease: str, keys: List[int]) -> bool: + if not set(keys).issubset(self.ready): + return False + self.leases.setdefault(lease, set()).update(keys) + for key in keys: + self.blocks.move_to_end(key) + return True + + def release(self, lease: str): + self.leases.pop(lease, None) + + def release_consumer(self, rank: int): + for lease in list(self.leases): + if lease.startswith(f"{rank}:"): + self.release(lease) + + def read(self, layer: str, keys: List[int]) -> torch.Tensor: + return torch.stack([self.blocks[key][layer] for key in keys]).contiguous() + + +def block_major_view(tensor: torch.Tensor, num_blocks: int) -> torch.Tensor: + """Normalize supported attention layouts without guessing ambiguous axes.""" + if tensor.ndim == 5 and tensor.shape[0] == 2: + if tensor.shape[1] == num_blocks: + if num_blocks == 2: + raise ValueError("Ambiguous Xavier KV block axis with only two blocks") + return tensor.movedim(1, 0) + if tensor.shape[0] != num_blocks: + raise ValueError(f"Unsupported Xavier KV layout {tuple(tensor.shape)}") + return tensor diff --git a/xinference/model/llm/xavier/backends/torch/tiered_snapshot.py b/xinference/model/llm/xavier/backends/torch/tiered_snapshot.py new file mode 100644 index 0000000000..457fd2a77f --- /dev/null +++ b/xinference/model/llm/xavier/backends/torch/tiered_snapshot.py @@ -0,0 +1,249 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Budgeted hot snapshots with CPU overflow and shared publication/read leases.""" + +from collections import OrderedDict +from typing import Dict, List, Tuple + +import torch + +from .snapshot import KVSnapshotStore + + +class TieredKVSnapshotStore(KVSnapshotStore): + """GPU-first snapshot policy, independent of the actor transport. + + ``block_bytes`` reserves a complete block across all layers, including while + it is only partially staged. The GPU budget covers retained snapshots, not + temporary copy/read tensors or the CUDA allocator's reserved memory. + Leased blocks cannot migrate or be evicted. Reads default to CPU to retain + the base store's interface; GPU consumers must select their output device. + CUDA operations follow the caller's current stream. + """ + + def __init__( + self, + cpu_capacity: int, + gpu_budget_bytes: int, + block_bytes: int, + gpu_device: torch.device, + ): + if cpu_capacity <= 0 or gpu_budget_bytes < 0 or block_bytes <= 0: + raise ValueError("Invalid tiered snapshot capacity or budget") + super().__init__(cpu_capacity + gpu_budget_bytes // block_bytes) + self.cpu_capacity = cpu_capacity + self.gpu_capacity = gpu_budget_bytes // block_bytes + self.gpu_device = gpu_device + self.tiers: Dict[int, str] = {} + # GPU victim lookup must not walk an arbitrarily large CPU history. + # Keep the GPU subsequence in the same order as the global block LRU. + self._gpu_lru: OrderedDict[int, None] = OrderedDict() + self.counts = {"gpu": 0, "cpu": 0} + self.metrics = dict(demotions=0, gpu_hits=0, cpu_hits=0, skipped=0) + self.block_bytes = block_bytes + # Avoid rescanning every previous layer when validating the next one. + self._block_sizes: Dict[int, int] = {} + + def touch(self, key: int) -> None: + self.blocks.move_to_end(key) + if key in self._gpu_lru: + self._gpu_lru.move_to_end(key) + + def _drop(self, key: int) -> None: + self._gpu_lru.pop(key, None) + self.counts[self.tiers.pop(key)] -= 1 + del self.blocks[key] + del self.logical_dtypes[key] + del self._block_sizes[key] + self.ready.discard(key) + self.evicted.add(key) + + def _cpu_room(self, pinned: set) -> bool: + if self.counts["cpu"] < self.cpu_capacity: + return True + victim = next( + (k for k in self.blocks if self.tiers[k] == "cpu" and k not in pinned), + None, + ) + if victim is None: + return False + self._drop(victim) + return True + + @staticmethod + def _copy_to_cpu(layers: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + # A synchronous D2H copy per layer stalls the producer for every layer + # of every evicted block. Pack compatible layers and synchronize once + # per group. Each host allocation belongs to one block, so evicting a + # block releases its storage without retaining unrelated snapshots. + groups: Dict[ + Tuple[torch.device, torch.dtype], List[Tuple[str, torch.Tensor]] + ] = {} + host = {} + for layer, value in layers.items(): + if value.device.type != "cuda": + host[layer] = value.to("cpu", copy=True) + else: + groups.setdefault((value.device, value.dtype), []).append( + (layer, value) + ) + for group in groups.values(): + packed = torch.cat([value.reshape(-1) for _, value in group]).to("cpu") + values = packed.split([value.numel() for _, value in group]) + for (layer, original), value in zip(group, values): + host[layer] = value.reshape(original.shape) + return host + + def _admit(self, key: int, pinned: set) -> bool: + if self.gpu_capacity and self.counts["gpu"] >= self.gpu_capacity: + victim = next( + (k for k in self._gpu_lru if k not in pinned), + None, + ) + if victim is not None: + cpu_available = self.counts["cpu"] < self.cpu_capacity or any( + self.tiers[k] == "cpu" and k not in pinned for k in self.blocks + ) + if cpu_available: + # Copy before evicting CPU content or publishing the new tier. + host = self._copy_to_cpu(self.blocks[victim]) + self._cpu_room(pinned) + self.blocks[victim] = host + self.tiers[victim] = "cpu" + del self._gpu_lru[victim] + self.counts["gpu"] -= 1 + self.counts["cpu"] += 1 + self.metrics["demotions"] += 1 + else: + # CPU leases must not prevent reuse of an unleased GPU slot. + self._drop(victim) + tier = "gpu" if self.counts["gpu"] < self.gpu_capacity else "cpu" + if tier == "cpu" and not self._cpu_room(pinned): + self.metrics["skipped"] += 1 + return False + self.blocks[key] = {} + self.logical_dtypes[key] = {} + self._block_sizes[key] = 0 + self.tiers[key] = tier + if tier == "gpu": + self._gpu_lru[key] = None + self.counts[tier] += 1 + return True + + def stage( + self, + layer: str, + keys: List[int], + tensors: torch.Tensor, + logical_dtype: torch.dtype | None = None, + ): + # Reject the whole batch before changing content, placement, or LRU order. + for key, tensor in zip(keys, tensors): + existing = self.blocks.get(key, {}) + if layer in existing: + continue + used = self._block_sizes.get(key, 0) + if used + tensor.numel() * tensor.element_size() > self.block_bytes: + raise ValueError("Snapshot exceeds configured full-block size") + pinned = set().union(*self.leases.values()) if self.leases else set() + for key, tensor in zip(keys, tensors): + if key not in self.blocks and not self._admit(key, pinned): + continue + self.touch(key) + if layer not in self.blocks[key]: + device = ( + self.gpu_device if self.tiers[key] == "gpu" else torch.device("cpu") + ) + try: + self.blocks[key][layer] = ( + tensor.detach().to(device, copy=True).contiguous() + ) + self.logical_dtypes[key][layer] = logical_dtype or tensors.dtype + self._block_sizes[key] += tensor.numel() * tensor.element_size() + except Exception: + # A failed first layer must not consume an empty slot. + if not self.blocks[key]: + self._drop(key) + raise + + def stage_blocks(self, keys: List[int], layers: Dict[str, torch.Tensor]): + """Copy complete new blocks once per dtype, with independent ownership.""" + if not keys: + return + if len(keys) != len(set(keys)): + raise ValueError("Packed staging requires unique snapshot keys") + if not layers or any(len(value) != len(keys) for value in layers.values()): + raise ValueError("Snapshot keys and tensors differ in length") + groups: Dict[ + Tuple[torch.device, torch.dtype], List[Tuple[str, torch.Tensor]] + ] = {} + nbytes = sum( + value[0].numel() * value.element_size() for value in layers.values() + ) + if nbytes > self.block_bytes: + raise ValueError("Snapshot exceeds configured full-block size") + if any(key in self.blocks for key in keys): + raise ValueError("Packed staging requires new snapshot keys") + for name, value in layers.items(): + groups.setdefault((value.device, value.dtype), []).append((name, value)) + packed = [] + for group in groups.values(): + values = torch.cat( + [value.reshape(len(keys), -1) for _, value in group], dim=1 + ) + names = [name for name, _ in group] + shapes = [value.shape[1:] for _, value in group] + sizes = [value[0].numel() for _, value in group] + packed.append((names, shapes, sizes, values)) + dtypes = {name: value.dtype for name, value in layers.items()} + pinned = set().union(*self.leases.values()) if self.leases else set() + for index, key in enumerate(keys): + if not self._admit(key, pinned): + continue + try: + target = ( + self.gpu_device if self.tiers[key] == "gpu" else torch.device("cpu") + ) + content: Dict[str, torch.Tensor] = {} + for names, shapes, sizes, values in packed: + # Layers share only this block's allocation. Evicting one + # block cannot retain the rest of the staging batch. + block = values[index].detach().to(target, copy=True) + parts = ( + value.reshape(shape) + for value, shape in zip(block.split(sizes), shapes) + ) + content.update(zip(names, parts)) + self.blocks[key] = content + self.logical_dtypes[key] = dtypes.copy() + self._block_sizes[key] = nbytes + except Exception: + self._drop(key) + raise + + def reserve(self, lease: str, keys: List[int]) -> bool: + if not super().reserve(lease, keys): + return False + for key in keys: + if key in self._gpu_lru: + self._gpu_lru.move_to_end(key) + self.metrics[self.tiers[key] + "_hits"] += 1 + return True + + def read( + self, layer: str, keys: List[int], device: torch.device | None = None + ) -> torch.Tensor: + target = torch.device("cpu") if device is None else device + return torch.stack( + [self.blocks[key][layer].to(target) for key in keys] + ).contiguous() + + def stats(self) -> dict: + return dict( + self.metrics, + gpu_blocks=self.counts["gpu"], + cpu_blocks=self.counts["cpu"], + gpu_capacity=self.gpu_capacity, + cpu_capacity=self.cpu_capacity, + gpu_reserved_bytes=self.counts["gpu"] * self.block_bytes, + ) diff --git a/xinference/model/llm/xavier/block_tracker.py b/xinference/model/llm/xavier/block_tracker.py new file mode 100644 index 0000000000..582d025daa --- /dev/null +++ b/xinference/model/llm/xavier/block_tracker.py @@ -0,0 +1,153 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# +# 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 random +from typing import Dict, List, Optional, Set, Tuple + +import xoscar as xo + + +class BlockTracker(xo.StatelessActor): + @classmethod + def default_uid(cls): + return f"vllm-block-tracker-actor" + + def __init__(self): + super().__init__() + # engine -> hash -> (rank, block_id) + self._hash_to_rank_and_block_id: Dict[int, Dict[int, Set[Tuple[int, int]]]] = {} # type: ignore + # engine -> rank -> (hash, block_id) + self._rank_to_hash_and_block_id: Dict[int, Dict[int, Set[Tuple[int, int]]]] = {} # type: ignore + self._unavailable_ranks: Set[int] = set() # type: ignore + + def register_blocks( + self, virtual_engine: int, block_infos: List[Tuple[int, int]], rank: int + ): + if virtual_engine not in self._hash_to_rank_and_block_id: + self._hash_to_rank_and_block_id[virtual_engine] = {} + hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine] + + if virtual_engine not in self._rank_to_hash_and_block_id: + self._rank_to_hash_and_block_id[virtual_engine] = {} + rank_to_hash_and_block_id = self._rank_to_hash_and_block_id[virtual_engine] + if rank not in rank_to_hash_and_block_id: + rank_to_hash_and_block_id[rank] = set() + hash_and_block_id = rank_to_hash_and_block_id[rank] + + # Replace-on-register: an engine block id can be reused for new content + # (its KV cache slot is reallocated), so before registering the new + # (hash, block_id) mappings, drop any stale (old_hash, block_id) this + # rank still holds for the same block ids. Without this a stale hash + # keeps pointing at a block that now holds different KV, and a later + # query on that hash would load the wrong KV. This also bounds the + # tracker to one entry per (rank, block_id). + new_block_ids = {block_id for _, block_id in block_infos} + new_pairs = set(block_infos) + stale = [ + (stale_hash, stale_block_id) + for stale_hash, stale_block_id in hash_and_block_id + if stale_block_id in new_block_ids + and (stale_hash, stale_block_id) not in new_pairs + ] + for stale_hash, stale_block_id in stale: + hash_and_block_id.discard((stale_hash, stale_block_id)) + stale_entries = hash_to_rank_and_block_id.get(stale_hash) + if stale_entries is not None: + stale_entries.discard((rank, stale_block_id)) + if not stale_entries: + hash_to_rank_and_block_id.pop(stale_hash, None) + + for hash_content, block_id in block_infos: + hash_to_rank_and_block_id.setdefault(hash_content, set()).add( + (rank, block_id) + ) + hash_and_block_id.add((hash_content, block_id)) + + def query_blocks( + self, + virtual_engine: int, + hash_contents: List[Tuple[int, int]], + exclude_rank: Optional[int] = None, + ) -> Dict[int, Set[Tuple[int, int, int]]]: + if virtual_engine not in self._hash_to_rank_and_block_id: + return {} + hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine] + remote: Dict[int, Set[Tuple[int, int, int]]] = {} + for hash_content, _id in hash_contents: + if ( + hash_content in hash_to_rank_and_block_id + ) and hash_to_rank_and_block_id[hash_content]: + # exclude ranks that are in the recovery process or the caller's own rank + rank_and_block_id = [ + (r, b) + for r, b in hash_to_rank_and_block_id[hash_content] + if r not in self._unavailable_ranks and r != exclude_rank + ] + if rank_and_block_id: + # TODO: Randomly select here, and try to distribute requests as evenly as possible. + # There may be better methods in the future. + rank, block_id = random.choice(rank_and_block_id) + if rank not in remote: + remote[rank] = { + (hash_content, block_id, _id), + } + else: + remote[rank].add((hash_content, block_id, _id)) + return remote + + def unregister_block(self, virtual_engine: int, rank: int, block_id: int): + self.unregister_blocks(virtual_engine, rank, [block_id]) + + def unregister_blocks(self, virtual_engine: int, rank: int, block_ids: List[int]): + entries = self._rank_to_hash_and_block_id.get(virtual_engine, {}).get(rank) + hashes = self._hash_to_rank_and_block_id.get(virtual_engine, {}) + if not entries: + return + removed_ids = set(block_ids) + removed = { + (content_hash, block_id) + for content_hash, block_id in entries + if block_id in removed_ids + } + entries.difference_update(removed) + for content_hash, block_id in removed: + locations = hashes.get(content_hash) + if locations is not None: + locations.discard((rank, block_id)) + if not locations: + hashes.pop(content_hash, None) + + def unregister_rank(self, rank: int): + """ + This rank is in the recovery process, and its query results will be excluded. + """ + self._unavailable_ranks.add(rank) + + def register_rank(self, rank: int): + """ + After recovery is successful, clear all stale data of the rank and mark the rank as available. + """ + for _, rank_to_hash_and_block_id in self._rank_to_hash_and_block_id.items(): + rank_to_hash_and_block_id.pop(rank, None) + + for _, hash_to_rank_and_block_id in self._hash_to_rank_and_block_id.items(): + for _, rank_and_block_id in hash_to_rank_and_block_id.items(): + to_delete = [(r, b) for r, b in rank_and_block_id if r == rank] + if to_delete: + rank_and_block_id.difference_update(to_delete) + + self._unavailable_ranks.discard(rank) + + +# Keep the old class import and actor UID for existing deployments. +VLLMBlockTracker = BlockTracker diff --git a/xinference/model/llm/xavier/collective.py b/xinference/model/llm/xavier/collective.py new file mode 100644 index 0000000000..3098f5305f --- /dev/null +++ b/xinference/model/llm/xavier/collective.py @@ -0,0 +1,93 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# +# 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 List, Optional + +from .transport import get_transport_host + +logger = logging.getLogger(__name__) + + +class CollectiveRank: + def __init__( + self, + rank: int, + world_size: int, + rank_address: str, + store_address: str, + store_port: int, + world_addresses: List[str], + ): + self._rank = rank + self._world_size = world_size + self._rank_address = rank_address + self._world_addresses = world_addresses + self._store_address = store_address + self._store_port = store_port + self._device = None + self._tcp_store = None + self._context = None + + def init_rank(self): + from xoscar.collective import xoscar_pygloo as xp + + logger.debug( + f"init_rank called: rank={self._rank}, world_size={self._world_size}, " + f"rank_address={self._rank_address}, store_address={self._store_address}, " + f"store_port={self._store_port}, world_addresses={self._world_addresses}" + ) + + self._context = xp.rendezvous.Context(rank=self._rank, size=self._world_size) + + attr = xp.transport.tcp.attr(get_transport_host(self._rank_address)) + self._device = xp.transport.tcp.CreateDevice(attr) + + opt = xp.rendezvous.TCPStoreOptions() + opt.port = self._store_port + opt.numWorkers = self._world_size + opt.isServer = self._rank == 0 + opt.waitWorkers = False + + if self._store_address is None: + raise ValueError( + f"store_address is None for rank {self._rank}. " + f"This usually means the rank 0 model was not properly initialized before other ranks." + ) + if self._store_port is None: + raise ValueError( + f"store_port is None for rank {self._rank}. " + f"This usually means the rank 0 model was not properly initialized before other ranks." + ) + + self._tcp_store = xp.rendezvous.TCPStore(self._store_address, opt) + if self._world_addresses: + self.connect_full_mesh() + + def connect_full_mesh( + self, prefix: Optional[str] = None, world_addresses: Optional[List[str]] = None + ): + from xoscar.collective import xoscar_pygloo as xp + + assert self._device is not None + assert self._tcp_store is not None + assert self._context is not None + if world_addresses is not None: + self._world_addresses = world_addresses + prefix_store = xp.rendezvous.PrefixStore( + prefix or str(self._world_size), self._tcp_store + ) + self._context.connectFullMesh(prefix_store, self._device) + logger.debug( + f"Rank {self._rank} arrives successfully, world addresses: {self._world_addresses}" + ) diff --git a/xinference/model/llm/xavier/collective_manager.py b/xinference/model/llm/xavier/collective_manager.py new file mode 100644 index 0000000000..bb235618f8 --- /dev/null +++ b/xinference/model/llm/xavier/collective_manager.py @@ -0,0 +1,99 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. + +import asyncio +import logging +import traceback +from typing import Any, Dict, Optional + +import xoscar as xo + +from .block_tracker import BlockTracker + +logger = logging.getLogger(__name__) + + +def with_lock(method): + async def wrapper(self, *args, **kwargs): + async with self._lock: + return await method(self, *args, **kwargs) + + return wrapper + + +class CollectiveManager(xo.StatelessActor): + @classmethod + def default_uid(cls): + return f"xavier-collective-manager" + + def __init__(self, model_uid: str, transfer_actor_uid: str = "vllm-transfer-actor"): + super().__init__() + self._model_uid = model_uid + self._transfer_actor_uid = transfer_actor_uid + self._tracker_ref: Optional[xo.ActorRefType["BlockTracker"]] = None + self._rank_to_ref: Dict[int, xo.ActorRefType[Any]] = {} + self._lock = asyncio.Lock() + + async def __post_create__(self): + self._tracker_ref = await xo.actor_ref( + address=self.address, + uid=f"{BlockTracker.default_uid()}-{self._model_uid}", + ) + + async def unregister_rank(self, rank: int): + self._rank_to_ref.pop(rank, None) + await self._tracker_ref.unregister_rank(rank) # type: ignore + await asyncio.gather( + *( + ref.release_consumer_leases_v1(rank) + for other, ref in self._rank_to_ref.items() + if other != 0 + ), + return_exceptions=True, + ) + logger.debug(f"Unregister rank: {rank}") + + async def register_rank(self, rank: int, address: str, update: bool = False): + rank_ref = await xo.actor_ref( + address=address, uid=f"{self._transfer_actor_uid}-{rank}" + ) + self._rank_to_ref[rank] = rank_ref + logger.debug(f"Register rank: {rank}, address: {address}") + if update: + await self._update_world() + await self._tracker_ref.register_rank(rank) # type: ignore + + @with_lock + async def _update_world(self): + """ + Locking is used to prevent chaos when multiple replicas trigger recovery simultaneously. + """ + from ....core.utils import gen_random_string + + prefix = gen_random_string(6) + tasks = [] + rank_to_ref = self._rank_to_ref.copy() + world_addresses = [ref.address for _, ref in sorted(rank_to_ref.items())] + for rank, ref in rank_to_ref.items(): + tasks.append(ref.connect_full_mesh(prefix, world_addresses)) + try: + logger.debug( + f"Rebuild collective communication with world_addresses: {world_addresses}, prefix: {prefix}" + ) + await asyncio.gather(*tasks) + logger.debug( + f"Rebuild collective communication with world_addresses: {world_addresses}, prefix: {prefix} done." + ) + except Exception as e: + """ + The exception here is most likely due to another replica triggering recovery during the recovery process, + causing `connect_full_mesh` to time out. + Simply log the exception and + let the subsequent update process handle the reconstruction of the collective communication world. + """ + logger.error( + f"Rebuild collective communication with world_addresses: {world_addresses} failed. " + f"Exception: {e}" + ) + # Print the complete error stack + traceback.print_exception(type(e), e, e.__traceback__) diff --git a/xinference/model/llm/xavier/profiling.py b/xinference/model/llm/xavier/profiling.py new file mode 100644 index 0000000000..2e70043119 --- /dev/null +++ b/xinference/model/llm/xavier/profiling.py @@ -0,0 +1,57 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Opt-in, synchronizing diagnostics; never enable for throughput measurements.""" + +import json +import logging +import os +import time +from contextlib import contextmanager +from typing import Any, Iterator + +logger = logging.getLogger(__name__) +_ENABLED = os.environ.get("XINFERENCE_XAVIER_PROFILE") == "1" +if _ENABLED: + # Keep configured actor/file logging; EngineCore may have no handlers. + if not logger.hasHandlers(): + logger.addHandler(logging.StreamHandler()) + logger.setLevel(logging.INFO) + logger.propagate = True + + +@contextmanager +def profile_stage(stage: str, *, device: Any = None, **fields: Any) -> Iterator[None]: + if not _ENABLED: + yield + return + + def synchronize() -> None: + if device is not None and device.type == "cuda": + import torch + + torch.cuda.synchronize(device) + + # Drain preceding GPU work outside the timed region. This intentionally + # serializes execution; profile results are separate from benchmark runs. + synchronize() + started = time.perf_counter() + succeeded = False + try: + yield + synchronize() + succeeded = True + finally: + elapsed = time.perf_counter() - started + logger.info( + "Xavier profile: %s", + json.dumps( + dict( + stage=stage, + elapsed_s=elapsed, + pid=os.getpid(), + succeeded=succeeded, + **fields, + ), + sort_keys=True, + ), + ) diff --git a/xinference/model/llm/xavier/tests/__init__.py b/xinference/model/llm/xavier/tests/__init__.py new file mode 100644 index 0000000000..9766634c65 --- /dev/null +++ b/xinference/model/llm/xavier/tests/__init__.py @@ -0,0 +1,2 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. diff --git a/xinference/model/llm/xavier/tests/test_shared_core.py b/xinference/model/llm/xavier/tests/test_shared_core.py new file mode 100644 index 0000000000..fa3cdc0a51 --- /dev/null +++ b/xinference/model/llm/xavier/tests/test_shared_core.py @@ -0,0 +1,178 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import importlib +import pickle +import subprocess +import sys +import textwrap +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import xoscar as xo +from xoscar.backends.allocate_strategy import ProcessIndex + +from ..block_tracker import BlockTracker +from ..collective_manager import CollectiveManager + + +class AdapterPeerActor(xo.StatelessActor): + def __init__(self): + super().__init__() + self.world_addresses = [] + + def connect_full_mesh(self, prefix, world_addresses): + self.world_addresses = world_addresses + + def get_world_addresses(self): + return self.world_addresses + + +@pytest.mark.parametrize( + "name,backend", + [ + ("actor_loop", False), + ("block_tracker", False), + ("collective", False), + ("profiling", False), + ("utils", False), + ("snapshot", True), + ("tiered_snapshot", True), + ("request_transfer", True), + ("gpu_transfer", True), + ("direct_handoff", True), + ("direct_history", True), + ], +) +def test_legacy_modules_share_the_implementation_and_state(name, backend): + if backend: + pytest.importorskip("torch") + prefix = "xinference.model.llm.xavier" + if backend: + prefix += ".backends.torch" + shared = importlib.import_module(f"{prefix}.{name}") + legacy = importlib.import_module(f"xinference.model.llm.vllm.xavier.{name}") + assert legacy is shared + + +def test_legacy_actor_and_transport_exports_remain_compatible(): + from ...vllm.xavier.block_tracker import VLLMBlockTracker + from ...vllm.xavier.collective_manager import CollectiveManager as LegacyManager + from ...vllm.xavier.transport import ( + XAVIER_CONNECTOR_MODULE, + get_transport_host, + validate_gpu_cache_budget, + ) + from .. import transport + + assert VLLMBlockTracker is BlockTracker + assert BlockTracker.default_uid() == "vllm-block-tracker-actor" + assert LegacyManager is CollectiveManager + assert get_transport_host is transport.get_transport_host + assert validate_gpu_cache_budget is transport.validate_gpu_cache_budget + assert XAVIER_CONNECTOR_MODULE == "xinference.model.llm.vllm.xavier.v1_connector" + # Existing serialized actor class references use the old module and name. + legacy_reference = ( + b"cxinference.model.llm.vllm.xavier.block_tracker\nVLLMBlockTracker\n." + ) + assert pickle.loads(legacy_reference) is BlockTracker + + +def test_shared_coordination_imports_without_engine_or_device_dependencies(): + # A fresh interpreter exercises imports rather than reusing modules loaded + # by other tests. Xinference's normal parent-package registration runs first. + script = textwrap.dedent( + """ + import builtins + import importlib + import importlib.util + import sys + import xinference.model.llm + + prefix = "xinference.model.llm.xavier" + for name in list(sys.modules): + if name == prefix or name.startswith(prefix + "."): + del sys.modules[name] + + original_import = builtins.__import__ + def guarded_import(name, globals=None, locals=None, fromlist=(), level=0): + resolved = name + if level: + resolved = importlib.util.resolve_name( + "." * level + name, globals["__package__"] + ) + if resolved.split(".")[0] in {"torch", "vllm", "sglang", "mlx"}: + raise AssertionError("Unexpected dependency: " + resolved) + if resolved.startswith("xinference.model.llm.vllm"): + raise AssertionError("Unexpected engine adapter: " + resolved) + return original_import(name, globals, locals, fromlist, level) + + builtins.__import__ = guarded_import + for name in ( + "actor_loop", "block_tracker", "collective", "collective_manager", + "profiling", "transport", "utils", + ): + importlib.import_module(prefix + "." + name) + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=60 + ) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.asyncio +@pytest.mark.parametrize("uid", ["vllm-transfer-actor", "sglang-transfer-actor"]) +async def test_coordinator_registers_and_recovers_adapter_peers(monkeypatch, uid): + manager = CollectiveManager("model", transfer_actor_uid=uid) + peer = SimpleNamespace( + address="peer", + connect_full_mesh=AsyncMock(), + release_consumer_leases_v1=AsyncMock(), + ) + tracker = SimpleNamespace(register_rank=AsyncMock(), unregister_rank=AsyncMock()) + manager._tracker_ref = tracker + actor_ref = AsyncMock(return_value=peer) + monkeypatch.setattr(xo, "actor_ref", actor_ref) + + await manager.register_rank(1, "peer", update=True) + actor_ref.assert_awaited_once_with(address="peer", uid=f"{uid}-1") + peer.connect_full_mesh.assert_awaited_once() + assert peer.connect_full_mesh.await_args.args[1] == ["peer"] + tracker.register_rank.assert_awaited_once_with(1) + + await manager.unregister_rank(2) + tracker.unregister_rank.assert_awaited_once_with(2) + peer.release_consumer_leases_v1.assert_awaited_once_with(2) + + +@pytest.mark.asyncio +async def test_shared_actors_register_and_recover_in_a_subprocess(): + pool = await xo.create_actor_pool("127.0.0.1", n_process=1) + async with pool: + placement = { + "address": pool.external_address, + "allocate_strategy": ProcessIndex(1), + } + tracker = await xo.create_actor( + BlockTracker, uid=f"{BlockTracker.default_uid()}-model", **placement + ) + manager = await xo.create_actor( + CollectiveManager, + model_uid="model", + transfer_actor_uid="adapter-transfer", + **placement, + ) + peer = await xo.create_actor( + AdapterPeerActor, uid="adapter-transfer-1", **placement + ) + await manager.register_rank(1, peer.address, update=True) + assert await peer.get_world_addresses() == [peer.address] + await tracker.register_blocks(0, [(123, 7)], 1) + assert await tracker.query_blocks(0, [(123, 9)]) == {1: {(123, 7, 9)}} + await manager.unregister_rank(1) + assert await tracker.query_blocks(0, [(123, 9)]) == {} + await manager.register_rank(1, peer.address, update=True) + assert await tracker.query_blocks(0, [(123, 9)]) == {} + await tracker.register_blocks(0, [(456, 7)], 1) + assert await tracker.query_blocks(0, [(456, 9)]) == {1: {(456, 7, 9)}} diff --git a/xinference/model/llm/xavier/transport.py b/xinference/model/llm/xavier/transport.py new file mode 100644 index 0000000000..cbc9fd0adf --- /dev/null +++ b/xinference/model/llm/xavier/transport.py @@ -0,0 +1,21 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. + +from typing import Optional +from urllib.parse import urlsplit + + +def validate_gpu_cache_budget(value, enabled: bool, replicas: int): + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError("xavier_gpu_cache_bytes must be a non-negative integer") + if not enabled or replicas <= 1: + raise ValueError( + "xavier_gpu_cache_bytes requires Xavier with multiple replicas" + ) + return value + + +def get_transport_host(address: str) -> Optional[str]: + return urlsplit(address if "://" in address else "tcp://" + address).hostname diff --git a/xinference/model/llm/xavier/utils.py b/xinference/model/llm/xavier/utils.py new file mode 100644 index 0000000000..adc963164a --- /dev/null +++ b/xinference/model/llm/xavier/utils.py @@ -0,0 +1,60 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# +# 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 struct +from typing import List, Optional + +import xxhash + + +def hash_block_tokens( + is_first_block: bool, + prev_block_hash: Optional[int], + cur_block_token_ids: List[int], + extra_hash: Optional[int] = None, + none_hash: int = -1, +) -> int: + """ + Computes a stable hash value corresponding to the contents of a block + and the contents of the preceding block(s), using xxhash instead of + Python hash() for cross-process stability. + """ + + if prev_block_hash is None: + prev_block_hash = none_hash + if extra_hash is None: + extra_hash = none_hash + + buf = bytearray() + + # 0. hash version + buf += b"v1" + + # 1. is_first_block: bool -> 1 byte + buf += struct.pack(" Date: Mon, 5 Oct 2026 00:08:13 +0800 Subject: [PATCH 2/2] fix: share Xavier transfer actor UID across coordinator and adapter --- xinference/model/llm/vllm/xavier/transfer.py | 5 ++-- .../model/llm/xavier/collective_manager.py | 5 +++- xinference/model/llm/xavier/constants.py | 5 ++++ .../llm/xavier/tests/test_shared_core.py | 29 ++++++++++++++++++- 4 files changed, 40 insertions(+), 4 deletions(-) create mode 100644 xinference/model/llm/xavier/constants.py diff --git a/xinference/model/llm/vllm/xavier/transfer.py b/xinference/model/llm/vllm/xavier/transfer.py index 65b5e9d053..46999d21aa 100644 --- a/xinference/model/llm/vllm/xavier/transfer.py +++ b/xinference/model/llm/vllm/xavier/transfer.py @@ -23,6 +23,7 @@ from ...xavier.backends.torch.snapshot import KVSnapshotStore from ...xavier.collective import CollectiveRank +from ...xavier.constants import DEFAULT_TRANSFER_ACTOR_UID from ...xavier.profiling import profile_stage try: @@ -125,7 +126,7 @@ class TransferActor( ): @classmethod def default_uid(cls): - return f"vllm-transfer-actor" + return DEFAULT_TRANSFER_ACTOR_UID def __init__( self, @@ -722,7 +723,7 @@ class Rank0TransferActor(xo.StatelessActor, CollectiveRank): @classmethod def default_uid(cls): - return f"vllm-transfer-actor" + return DEFAULT_TRANSFER_ACTOR_UID def __init__( self, diff --git a/xinference/model/llm/xavier/collective_manager.py b/xinference/model/llm/xavier/collective_manager.py index bb235618f8..5c973c3d20 100644 --- a/xinference/model/llm/xavier/collective_manager.py +++ b/xinference/model/llm/xavier/collective_manager.py @@ -9,6 +9,7 @@ import xoscar as xo from .block_tracker import BlockTracker +from .constants import DEFAULT_TRANSFER_ACTOR_UID logger = logging.getLogger(__name__) @@ -26,7 +27,9 @@ class CollectiveManager(xo.StatelessActor): def default_uid(cls): return f"xavier-collective-manager" - def __init__(self, model_uid: str, transfer_actor_uid: str = "vllm-transfer-actor"): + def __init__( + self, model_uid: str, transfer_actor_uid: str = DEFAULT_TRANSFER_ACTOR_UID + ): super().__init__() self._model_uid = model_uid self._transfer_actor_uid = transfer_actor_uid diff --git a/xinference/model/llm/xavier/constants.py b/xinference/model/llm/xavier/constants.py new file mode 100644 index 0000000000..dcc721247f --- /dev/null +++ b/xinference/model/llm/xavier/constants.py @@ -0,0 +1,5 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. + +# Preserve the existing actor UID for deployments using the default adapter. +DEFAULT_TRANSFER_ACTOR_UID = "vllm-transfer-actor" diff --git a/xinference/model/llm/xavier/tests/test_shared_core.py b/xinference/model/llm/xavier/tests/test_shared_core.py index fa3cdc0a51..4228e4801d 100644 --- a/xinference/model/llm/xavier/tests/test_shared_core.py +++ b/xinference/model/llm/xavier/tests/test_shared_core.py @@ -14,6 +14,7 @@ from ..block_tracker import BlockTracker from ..collective_manager import CollectiveManager +from ..constants import DEFAULT_TRANSFER_ACTOR_UID class AdapterPeerActor(xo.StatelessActor): @@ -109,7 +110,7 @@ def guarded_import(name, globals=None, locals=None, fromlist=(), level=0): builtins.__import__ = guarded_import for name in ( - "actor_loop", "block_tracker", "collective", "collective_manager", + "actor_loop", "block_tracker", "collective", "collective_manager", "constants", "profiling", "transport", "utils", ): importlib.import_module(prefix + "." + name) @@ -121,6 +122,32 @@ def guarded_import(name, globals=None, locals=None, fromlist=(), level=0): assert result.returncode == 0, result.stdout + result.stderr +@pytest.mark.asyncio +@pytest.mark.parametrize("rank", [0, 1]) +async def test_default_coordinator_uid_matches_vllm_transfer_actors(monkeypatch, rank): + pytest.importorskip("torch") + from ...vllm.xavier.transfer import Rank0TransferActor, TransferActor + + manager = CollectiveManager("model") + assert DEFAULT_TRANSFER_ACTOR_UID == "vllm-transfer-actor" + assert ( + manager._transfer_actor_uid + == TransferActor.default_uid() + == Rank0TransferActor.default_uid() + == DEFAULT_TRANSFER_ACTOR_UID + ) + peer = SimpleNamespace(address="peer") + actor_ref = AsyncMock(return_value=peer) + monkeypatch.setattr(xo, "actor_ref", actor_ref) + + await manager.register_rank(rank, "peer") + actor_class = Rank0TransferActor if rank == 0 else TransferActor + actor_ref.assert_awaited_once_with( + address="peer", uid=f"{actor_class.default_uid()}-{rank}" + ) + assert manager._rank_to_ref[rank] is peer + + @pytest.mark.asyncio @pytest.mark.parametrize("uid", ["vllm-transfer-actor", "sglang-transfer-actor"]) async def test_coordinator_registers_and_recovers_adapter_peers(monkeypatch, uid):