From a5d9d3ce197ef7c68709caeea28f6250d68a12ed Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 01:02:03 +0800 Subject: [PATCH 1/9] perf: prototype request-scoped Xavier GPU handoff --- xinference/core/pd_model.py | 14 +- xinference/model/llm/vllm/core.py | 13 +- .../model/llm/vllm/xavier/direct_handoff.py | 155 +++++++++++++++ .../model/llm/vllm/xavier/gpu_transfer.py | 34 +++- .../vllm/xavier/test/test_direct_handoff.py | 187 ++++++++++++++++++ .../model/llm/vllm/xavier/v1_connector.py | 98 ++++++++- 6 files changed, 488 insertions(+), 13 deletions(-) create mode 100644 xinference/model/llm/vllm/xavier/direct_handoff.py create mode 100644 xinference/model/llm/vllm/xavier/test/test_direct_handoff.py diff --git a/xinference/core/pd_model.py b/xinference/core/pd_model.py index ae521d512d..3b1ad94a43 100644 --- a/xinference/core/pd_model.py +++ b/xinference/core/pd_model.py @@ -16,6 +16,7 @@ import copy import json import logging +import os import time import uuid from abc import ABC, abstractmethod @@ -91,6 +92,7 @@ def __init__( self._model_uid = model_uid self._transport_backend = transport_backend + self._direct_test = os.getenv("XINFERENCE_XAVIER_DIRECT_TEST") == "1" # 使用字典存储副本:{replica_uid: actor_ref} self._prefill_replicas: Dict[str, xo.ActorRefType["ModelActor"]] = {} @@ -246,7 +248,11 @@ async def free_prefill_model_cache(self, request_id: str): ) if request_id in self._request_set: self._request_set.remove(request_id) - if self._prefill_replicas and self._transport_backend == "xavier": + if ( + self._prefill_replicas + and self._transport_backend == "xavier" + and not self._direct_test + ): await asyncio.gather( *[ actor_call( @@ -287,7 +293,7 @@ async def _infer(self, method, inputs, *args, **kwargs): if args and args[0] is not None and not isinstance(args[0], dict): raise TypeError("Generation config must be a dict or None") if ( - self._transport_backend == "nixl" + (self._transport_backend == "nixl" or self._direct_test) and args and args[0] and args[0].get("n", 1) != 1 @@ -302,7 +308,7 @@ async def _infer(self, method, inputs, *args, **kwargs): prefill_args = [{}] + prefill_args[1:] prefill_args[0]["max_tokens"] = 1 prefill_args[0]["stream"] = False - if self._transport_backend == "nixl": + if self._transport_backend == "nixl" or self._direct_test: prefill_args[0]["n"] = 1 prefill_args[0]["_pd_kv_transfer_params"] = { "do_remote_decode": True, @@ -328,7 +334,7 @@ async def _infer(self, method, inputs, *args, **kwargs): self._transport_backend, time.perf_counter() - prefill_start, ) - if self._transport_backend == "nixl": + if self._transport_backend == "nixl" or self._direct_test: payload = ( json.loads(result) if isinstance(result, (bytes, str)) else result ) diff --git a/xinference/model/llm/vllm/core.py b/xinference/model/llm/vllm/core.py index a11b0d44bc..bf7fd2122f 100644 --- a/xinference/model/llm/vllm/core.py +++ b/xinference/model/llm/vllm/core.py @@ -1995,7 +1995,10 @@ async def async_generate( sampling_params = SamplingParams(**sanitized_generate_config) if generate_config and "_pd_kv_transfer_params" in generate_config: - if self._nixl_config is None: + if ( + self._nixl_config is None + and os.getenv("XINFERENCE_XAVIER_DIRECT_TEST") != "1" + ): raise ValueError("KV handoff requires the NIXL PD backend") sampling_params.extra_args = { **(sampling_params.extra_args or {}), @@ -2178,7 +2181,13 @@ async def stream_results() -> AsyncGenerator[CompletionChunk, None]: completion = self._convert_request_output_to_completion( request_id, model=self.model_uid, request_output=final_output ) - if self._nixl_config is not None and self._nixl_config["role"] == "prefill": + if ( + self._nixl_config is not None and self._nixl_config["role"] == "prefill" + ) or ( + os.getenv("XINFERENCE_XAVIER_DIRECT_TEST") == "1" + and self._xavier_config + and self._xavier_config.get("role") == "prefill" + ): completion["_pd_kv_transfer_params"] = getattr( final_output, "kv_transfer_params", None ) diff --git a/xinference/model/llm/vllm/xavier/direct_handoff.py b/xinference/model/llm/vllm/xavier/direct_handoff.py new file mode 100644 index 0000000000..2398f958d7 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/direct_handoff.py @@ -0,0 +1,155 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Experimental request-scoped handoff; no persistent snapshot retention.""" + +import asyncio +import time +from dataclasses import dataclass + +import torch +import xoscar as xo + +from .gpu_transfer import GPUTransfer +from .request_transfer import LayerRead, batch_reads, unpack_reads + + +@dataclass +class DirectRequest: + request_id: str + blocks: set[int] + deadline: float + reading: bool = False + released: bool = False + + +class DirectGPUTransfer(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) + + def register_direct(self, ticket, request_id, blocks): + 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") + self.direct_requests[ticket] = DirectRequest( + request_id, set(blocks), time.monotonic() + 120 + ) + self.metrics["direct_registered"] += 1 + + def release_direct(self, ticket): + state = self.direct_requests.get(ticket) + if state is None: + return + state.released = True + if not state.reading: + self.finished_sending.add(state.request_id) + del self.direct_requests[ticket] + self.metrics["direct_finished"] += 1 + + def poll_direct(self): + for ticket, state in list(self.direct_requests.items()): + if not state.released and time.monotonic() >= state.deadline: + self.metrics["direct_expired"] += 1 + self.release_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.direct_requests.get(ticket) + if state is None or state.released: + raise RuntimeError("Direct KV handoff expired or was released") + 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() + 120 + 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, + torch.tensor(read.keys, device=self.device), + 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 + if state.released: + self.release_direct(ticket) + + async def load_direct(self, requests, tickets): + from .transfer import TransferActor + + self.metrics["load_calls"] += 1 + self.metrics["load_requests"] += len(requests) + async with self.recv_lock: + for ranks, ticket in zip(requests, tickets): + if len(ranks) != 1: + raise ValueError("Direct prototype 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}", + ) + 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) + await sender.send_direct_gpu_v1( + ticket, batch, self.recv_refs[slab_bytes], slab_bytes + ) + for read, blocks in unpack_reads( + self.recv_buffer[:size], batch + ): + cache = self.caches[read.layer] + cache[ + torch.tensor(read.destinations, device=self.device) + ] = 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) + await sender.release_direct_gpu_v1(ticket) diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index 15fb169d83..3ee2dccbde 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -460,7 +460,14 @@ def map_gpu_caches_v1(self, descriptors, budget): caches = { name: rebuild_cuda_tensor(*desc) for name, desc in descriptors.items() } - runtime = GPUTransfer(self, caches, budget) + import os + + runtime_type = GPUTransfer + if os.getenv("XINFERENCE_XAVIER_DIRECT_TEST") == "1": + from .direct_handoff import DirectGPUTransfer + + runtime_type = DirectGPUTransfer + runtime = runtime_type(self, caches, budget) self._gpu_transfer = runtime self._snapshot_store = runtime.store @@ -487,3 +494,28 @@ async def close_gpu_caches_v1(self): # 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): + self._gpu_transfer.register_direct(ticket, request_id, blocks) + + 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) diff --git a/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py new file mode 100644 index 0000000000..ecb339a455 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py @@ -0,0 +1,187 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import torch +import xoscar as xo + +from ..direct_handoff import DirectGPUTransfer +from ..request_transfer import LayerRead +from .test_gpu_transfer import runtime + + +def direct_runtime(monkeypatch): + r = runtime(monkeypatch, gpu_slots=0) + r.__class__ = DirectGPUTransfer + r.direct_requests, r.finished_sending = {}, set() + r.metrics.update(direct_registered=0, direct_finished=0, direct_expired=0) + return r + + +def peer_for(source): + async def send(*args): + return await source.run(source.send_direct, *args) + + return SimpleNamespace( + send_direct_gpu_v1=AsyncMock(side_effect=send), + release_direct_gpu_v1=AsyncMock(side_effect=source.release_direct), + ) + + +@pytest.mark.asyncio +async def test_direct_handoff_preserves_bits_and_uses_no_snapshots(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + source.caches["K"].view(torch.int16).copy_( + torch.tensor( + [ + [0, -32768], + [32767, -1], + [32640, -128], + [1, -32767], + [8, 9], + [10, 11], + [12, 13], + [14, 15], + ], + dtype=torch.int16, + ) + ) + source.register_direct("ticket", "producer", [1, 2]) + peer = peer_for(source) + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + + 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 torch.equal( + dest.caches["K"].view(torch.int16)[[5, 3]], + source.caches["K"].view(torch.int16)[[1, 2]], + ) + assert not source.store.blocks and not dest.store.blocks + assert source.poll_direct() == {"producer"} + assert not source.poll_direct() and not source.direct_requests + + +@pytest.mark.asyncio +async def test_release_during_read_waits_for_copy_and_fence(monkeypatch): + source = direct_runtime(monkeypatch) + source.register_direct("ticket", "producer", [1]) + entered, release = asyncio.Event(), asyncio.Event() + calls = [] + + async def copy(buffers, refs): + entered.set() + await release.wait() + refs[0].copy_(buffers[0]) + calls.append("copy") + + monkeypatch.setattr(xo, "copy_to", copy) + monkeypatch.setattr(torch.cuda, "synchronize", lambda device: calls.append("fence")) + read = LayerRead("K", [1], [0], (2,), torch.bfloat16) + task = asyncio.create_task( + source.run( + source.send_direct, "ticket", [read], torch.zeros(16, dtype=torch.uint8), 16 + ) + ) + await asyncio.wait_for(entered.wait(), 2) + source.release_direct("ticket") + assert not source.poll_direct() + release.set() + await task + assert calls == ["fence", "copy", "fence"] + assert source.poll_direct() == {"producer"} + + +@pytest.mark.asyncio +async def test_expired_request_cannot_read_reused_engine_blocks(monkeypatch): + source = direct_runtime(monkeypatch) + source.register_direct("ticket", "producer", [1]) + source.direct_requests["ticket"].deadline = 0 + assert source.poll_direct() == {"producer"} + with pytest.raises(RuntimeError, match="expired"): + await source.send_direct( + "ticket", [LayerRead("K", [1], [0], (2,), torch.bfloat16)], None, 16 + ) + assert not source.metrics["wire_bytes"] + + +@pytest.mark.asyncio +async def test_cancelled_decoder_drains_before_releasing_source(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + source.caches["K"].fill_(7) + source.register_direct("ticket", "producer", [1]) + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer_for(source))) + entered, release = asyncio.Event(), asyncio.Event() + + async def copy(buffers, refs): + entered.set() + await release.wait() + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(xo, "copy_to", copy) + task = asyncio.create_task( + dest.run(dest.load_direct, [{0: {"K": {1: 5}}}], ["ticket"]) + ) + await asyncio.wait_for(entered.wait(), 2) + task.cancel() + await asyncio.sleep(0) + assert not task.done() and not source.poll_direct() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert dest.caches["K"][5].tolist() == [7, 7] + assert source.poll_direct() == {"producer"} + + +def test_direct_metadata_preserves_local_prefix_offset(connector): + connector._direct_test = True + request = SimpleNamespace( + request_id="decoder", + kv_transfer_params={ + "xavier_direct": { + "rank": 0, + "ticket": "ticket", + "tokens": 32, + "blocks": [2, 3], + } + }, + ) + assert connector.get_num_new_matched_tokens(request, 16) == (16, True) + blocks = SimpleNamespace(get_block_ids=lambda: [[4, 5]]) + connector.update_state_after_alloc(request, blocks, 16) + load = connector._requests_need_load["decoder"] + assert load.local_transfers_by_group == {0: {0: {3: 5}}} + + +@pytest.mark.parametrize("prompt_tokens", [33, 30]) +def test_direct_producer_retains_original_engine_blocks( + connector, monkeypatch, prompt_tokens +): + import sys + + monkeypatch.setitem( + sys.modules, + "vllm.v1.request", + SimpleNamespace(RequestStatus=SimpleNamespace(FINISHED_ABORTED="aborted")), + ) + connector._direct_test = True + actor = SimpleNamespace(register_direct_gpu_v1=AsyncMock()) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = SimpleNamespace( + request_id="producer", + prompt_token_ids=list(range(prompt_tokens)), + status="finished", + kv_transfer_params={"do_remote_decode": True}, + ) + retained, metadata = connector.request_finished(request, [4, 5, 6]) + assert retained + handoff = metadata["xavier_direct"] + assert handoff["blocks"] == [4, 5] and handoff["tokens"] == prompt_tokens - 1 + actor.register_direct_gpu_v1.assert_awaited_once_with( + handoff["ticket"], "producer", [4, 5] + ) diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 0069c94f68..8900e7224a 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -14,6 +14,7 @@ import asyncio import logging import math +import os import uuid from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Tuple @@ -123,7 +124,12 @@ def __init__( raise ValueError( "Xavier V1 currently supports text-only models without LoRA" ) + self._direct_test = os.getenv("XINFERENCE_XAVIER_DIRECT_TEST") == "1" self._gpu_budget = self._xavier_config.get("gpu_cache_bytes") + if self._direct_test and ( + self._gpu_budget is None or len(kv_cache_config.kv_cache_groups) != 1 + ): + raise ValueError("Direct prototype requires GPU transport and one KV group") self._gpu_cache_mapped = False self._gpu_mapping_lock = asyncio.Lock() self._kv_schema: Optional[XavierKVSchema] = None @@ -217,8 +223,12 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non raise release_error def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: + sent = set() + if self._direct_test and self._is_producer and self._gpu_cache_mapped: + assert self._transfer_ref is not None + sent = self._call(self._transfer_ref.poll_direct_gpu_v1()) if not self._gpu_load_jobs: - return set(), set() + return sent, set() # Pump responses on the shared actor loop without waiting for any RPC. # CUDA writes and lease cleanup progress in the independent actor process. self._call(asyncio.sleep(0)) @@ -237,7 +247,7 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: # If the last request aborts while a load is pending, an idle EngineCore # may not poll again until new work arrives. Destination reclamation then # waits for that step; actor-side writes and lease release still progress. - return set(), received + return sent, received def wait_for_layer_load(self, layer_name: str) -> None: return @@ -298,6 +308,10 @@ def save_kv_layer( self._pending_store_requests[request.request_id] = request def wait_for_save(self): + if self._direct_test and self._is_producer: + self._call(self._ensure_gpu_cache_mapping()) + torch.cuda.synchronize() + return if not self._is_producer or not self._pending_store_requests: return @@ -324,6 +338,32 @@ def get_num_new_matched_tokens( raise ValueError( "Xavier V1 does not support LoRA, multimodal, embedding or salted prompts" ) + if self._direct_test: + params = getattr(request, "kv_transfer_params", None) or {} + handoff = params.get("xavier_direct") + if not self._is_consumer: + return 0, False + if handoff is None: + raise ValueError("Missing direct Xavier handoff metadata") + tokens = handoff["tokens"] + rank, ticket = handoff["rank"], handoff["ticket"] + if tokens <= num_computed_tokens: + if ticket: + + async def release(): + ref = await self._get_transfer_ref() + await ref.release_remote_direct_gpu_v1(rank, ticket) + + self._call(release()) + return 0, False + start = num_computed_tokens // self._block_size + blocks = handoff["blocks"] + self._requests_need_load[request.request_id] = XavierLoadRequest( + request.request_id, + {rank: {block: i for i, block in enumerate(blocks) if i >= start}}, + lease=ticket, + ) + return tokens - num_computed_tokens, True self._requests_need_load.pop(request.request_id, None) previous = self._leased_requests.pop(request.request_id, None) if previous is not None: @@ -424,7 +464,7 @@ def build_connector_meta( scheduler_output: SchedulerOutput, ) -> KVConnectorMetadata: meta = XavierConnectorMetadata() - if self._is_producer: + if self._is_producer and not self._direct_test: self._build_store_meta(scheduler_output, meta) if self._is_consumer: self._build_load_meta(scheduler_output, meta) @@ -435,6 +475,50 @@ def request_finished( request: "Request", block_ids: List[int], ) -> tuple[bool, dict[str, Any] | None]: + if self._direct_test: + params = getattr(request, "kv_transfer_params", None) or {} + if self._is_producer and params.get("do_remote_decode"): + from vllm.v1.request import RequestStatus + + if request.status == RequestStatus.FINISHED_ABORTED: + return False, None + token_count = max(len(request.prompt_token_ids) - 1, 0) + count = min( + len(block_ids), + (token_count + self._block_size - 1) // self._block_size, + ) + blocks = block_ids[:count] + ticket = uuid.uuid4().hex if blocks else "" + if blocks: + + async def register(): + ref = await self._get_transfer_ref() + await ref.register_direct_gpu_v1( + ticket, request.request_id, blocks + ) + + self._call(register()) + return bool(blocks), { + "do_remote_prefill": True, + "xavier_direct": { + "rank": self._rank, + "ticket": ticket, + "blocks": blocks, + "tokens": min(token_count, count * self._block_size), + }, + } + pending = self._requests_need_load.get(request.request_id) + if pending is not None and not pending.local_transfers_by_group: + self._requests_need_load.pop(request.request_id) + + async def release(): + ref = await self._get_transfer_ref() + await ref.release_remote_direct_gpu_v1( + next(iter(pending.transfers)), pending.lease + ) + + self._call(release()) + return False, None self._chunked_prefill.pop(request.request_id, None) pending = self._requests_need_load.get(request.request_id) if ( @@ -843,11 +927,13 @@ async def _submit_gpu_requests(self, requests): ranks[rank] = layers entries.append(ranks) torch.cuda.synchronize() - task = asyncio.create_task( - transfer.load_gpu_requests_v1( + if self._direct_test: + operation = transfer.load_direct_gpu_v1(entries, [r.lease for r in requests]) + else: + operation = transfer.load_gpu_requests_v1( entries, [(r.lease, r.transfers) for r in requests] ) - ) + task = asyncio.create_task(operation) except BaseException: # Worker metadata already owns these leases. Until task creation # succeeds, the TransferActor has no operation that can release them. From ed3af47022310b221e93de8deab927dfc16894fc Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 01:45:19 +0800 Subject: [PATCH 2/9] perf: reuse request-scoped CUDA indices for Xavier handoff --- .../model/llm/vllm/xavier/direct_handoff.py | 36 ++++++-- .../vllm/xavier/test/test_direct_handoff.py | 89 ++++++++++++++++++- 2 files changed, 118 insertions(+), 7 deletions(-) diff --git a/xinference/model/llm/vllm/xavier/direct_handoff.py b/xinference/model/llm/vllm/xavier/direct_handoff.py index 2398f958d7..7928af9bd3 100644 --- a/xinference/model/llm/vllm/xavier/direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/direct_handoff.py @@ -4,7 +4,7 @@ import asyncio import time -from dataclasses import dataclass +from dataclasses import dataclass, field import torch import xoscar as xo @@ -20,6 +20,7 @@ class DirectRequest: deadline: float reading: bool = False released: bool = False + indices: dict[tuple[int, ...], torch.Tensor] = field(default_factory=dict) class DirectGPUTransfer(GPUTransfer): @@ -27,7 +28,27 @@ 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) + self.metrics.update( + direct_registered=0, direct_finished=0, direct_expired=0, index_uploads=0 + ) + + 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): + await super()._close() + 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): if self.closing or ticket in self.direct_requests: @@ -49,6 +70,7 @@ def release_direct(self, ticket): return state.released = True if not state.reading: + state.indices.clear() self.finished_sending.add(state.request_id) del self.direct_requests[ticket] self.metrics["direct_finished"] += 1 @@ -94,7 +116,7 @@ async def send_direct(self, ticket, reads, remote_ref, slab_bytes): torch.index_select( cache, 0, - torch.tensor(read.keys, device=self.device), + self._index_tensor(state.indices, read.keys), out=target, ) offset = end @@ -123,6 +145,7 @@ async def load_direct(self, requests, tickets): address=self.actor._world_addresses[rank], uid=f"{TransferActor.default_uid()}-{rank}", ) + indices: dict[tuple[int, ...], torch.Tensor] = {} try: reads = [ LayerRead( @@ -145,11 +168,12 @@ async def load_direct(self, requests, tickets): self.recv_buffer[:size], batch ): cache = self.caches[read.layer] - cache[ - torch.tensor(read.destinations, device=self.device) - ] = blocks + 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() await sender.release_direct_gpu_v1(ticket) 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 ecb339a455..8611a32240 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py @@ -17,7 +17,9 @@ def direct_runtime(monkeypatch): r = runtime(monkeypatch, gpu_slots=0) r.__class__ = DirectGPUTransfer r.direct_requests, r.finished_sending = {}, set() - r.metrics.update(direct_registered=0, direct_finished=0, direct_expired=0) + r.metrics.update( + direct_registered=0, direct_finished=0, direct_expired=0, index_uploads=0 + ) return r @@ -89,11 +91,14 @@ async def copy(buffers, refs): ) ) await asyncio.wait_for(entered.wait(), 2) + state = source.direct_requests["ticket"] source.release_direct("ticket") + assert state.indices assert not source.poll_direct() release.set() await task assert calls == ["fence", "copy", "fence"] + assert not state.indices assert source.poll_direct() == {"producer"} @@ -185,3 +190,85 @@ def test_direct_producer_retains_original_engine_blocks( actor.register_direct_gpu_v1.assert_awaited_once_with( handoff["ticket"], "producer", [4, 5] ) + + +@pytest.mark.asyncio +async def test_indices_reused_across_layers_and_slabs_but_not_requests(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + for layer in ("K", "V", "Q"): + source.caches[layer] = torch.arange(16, dtype=torch.bfloat16).reshape(8, 2) + dest.caches[layer] = torch.zeros_like(source.caches[layer]) + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer_for(source))) + + async def copy(buffers, refs): + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(xo, "copy_to", copy) + for ticket, destinations in (("first", [5, 3]), ("second", [3, 5])): + source.register_direct(ticket, ticket, [1, 2]) + state = source.direct_requests[ticket] + mapping = {name: dict(zip([1, 2], destinations)) for name in source.caches} + await dest.run(dest.load_direct, [{0: mapping}], [ticket]) + for layer in source.caches: + assert torch.equal( + dest.caches[layer][destinations], source.caches[layer][[1, 2]] + ) + assert not state.indices + assert source.metrics["index_uploads"] == dest.metrics["index_uploads"] == 2 + assert dest.metrics["gpu_batches"] == 4 + assert source.poll_direct() == {"first", "second"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("finish", ["release", "expire", "close"]) +async def test_request_indices_reclaimed_on_lifecycle_end(monkeypatch, finish): + source = direct_runtime(monkeypatch) + source.register_direct("ticket", "request", [1, 2]) + state = source.direct_requests["ticket"] + index = source._index_tensor(state.indices, [1, 2]) + assert source._index_tensor(state.indices, [1, 2]) is index + assert source._index_tensor(state.indices, [2, 1]).tolist() == [2, 1] + assert index.dtype == torch.long + if finish == "release": + source.release_direct("ticket") + elif finish == "expire": + state.deadline = 0 + source.poll_direct() + else: + await source.close() + assert not source.direct_requests + assert not state.indices + + +@pytest.mark.asyncio +async def test_failed_later_slab_releases_both_index_sets(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + for layer in ("V", "Q"): + source.caches[layer] = source.caches["K"].clone() + dest.caches[layer] = dest.caches["K"].clone() + source.register_direct("ticket", "producer", [1, 2]) + state = source.direct_requests["ticket"] + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer_for(source))) + copies = 0 + receiver_indices = [] + original_index = dest._index_tensor + + def index(indices, values): + receiver_indices.append(indices) + return original_index(indices, values) + + async def copy(buffers, refs): + nonlocal copies + copies += 1 + if copies == 2: + raise RuntimeError("transfer failed") + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(dest, "_index_tensor", index) + monkeypatch.setattr(xo, "copy_to", copy) + mapping = {name: {1: 5, 2: 3} for name in source.caches} + with pytest.raises(RuntimeError, match="transfer failed"): + await dest.run(dest.load_direct, [{0: mapping}], ["ticket"]) + assert receiver_indices and all(not indices for indices in receiver_indices) + assert not state.indices + assert source.poll_direct() == {"producer"} From 6bdcd1dc84dba0382d119aee107e0d8588ca87ec Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 02:34:39 +0800 Subject: [PATCH 3/9] perf: retain bounded GPU history after direct KV handoff --- .../model/llm/vllm/xavier/direct_handoff.py | 37 ++- .../model/llm/vllm/xavier/direct_history.py | 235 +++++++++++++++++ .../model/llm/vllm/xavier/gpu_transfer.py | 21 +- .../vllm/xavier/test/test_direct_handoff.py | 1 + .../vllm/xavier/test/test_direct_history.py | 243 ++++++++++++++++++ .../model/llm/vllm/xavier/v1_connector.py | 103 +++++++- 6 files changed, 624 insertions(+), 16 deletions(-) create mode 100644 xinference/model/llm/vllm/xavier/direct_history.py create mode 100644 xinference/model/llm/vllm/xavier/test/test_direct_history.py diff --git a/xinference/model/llm/vllm/xavier/direct_handoff.py b/xinference/model/llm/vllm/xavier/direct_handoff.py index 7928af9bd3..eec6d5f8db 100644 --- a/xinference/model/llm/vllm/xavier/direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/direct_handoff.py @@ -1,6 +1,6 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Experimental request-scoped handoff; no persistent snapshot retention.""" +"""Experimental direct handoff with optional bounded GPU history.""" import asyncio import time @@ -9,6 +9,7 @@ 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 @@ -21,9 +22,13 @@ class DirectRequest: 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 + held_blocks: int = 0 -class DirectGPUTransfer(GPUTransfer): +class DirectGPUTransfer(DirectHistoryMixin, GPUTransfer): def __init__(self, actor, caches, budget): super().__init__(actor, caches, 0) self.direct_requests: dict[str, DirectRequest] = {} @@ -32,6 +37,8 @@ def __init__(self, actor, caches, budget): 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: @@ -44,13 +51,26 @@ def _index_tensor( 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() 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): + def register_direct( + self, ticket, request_id, blocks, hashes=None, held_blocks=None + ): if self.closing or ticket in self.direct_requests: raise ValueError("Invalid direct handoff registration") if not blocks or any( @@ -59,8 +79,14 @@ def register_direct(self, ticket, request_id, blocks): 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 + request_id, + set(blocks), + time.monotonic() + 120, + hashes=dict(zip(blocks, hashes or [])), + held_blocks=len(blocks) if held_blocks is None else held_blocks, ) self.metrics["direct_registered"] += 1 @@ -70,6 +96,8 @@ def release_direct(self, ticket): 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] @@ -79,6 +107,7 @@ def poll_direct(self): for ticket, state in list(self.direct_requests.items()): if not state.released and time.monotonic() >= state.deadline: self.metrics["direct_expired"] += 1 + state.retention_attempted = True self.release_direct(ticket) result = self.finished_sending self.finished_sending = set() diff --git a/xinference/model/llm/vllm/xavier/direct_history.py b/xinference/model/llm/vllm/xavier/direct_history.py new file mode 100644 index 0000000000..7c6cc5dc13 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/direct_history.py @@ -0,0 +1,235 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Bounded GPU history retained after request-scoped handoff.""" + +import asyncio +import logging +import time + +import torch + +from .request_transfer import LayerRead, batch_reads +from .tiered_snapshot import TieredKVSnapshotStore + +logger = logging.getLogger(__name__) + + +class GPUHistoryStore(TieredKVSnapshotStore): + """Independent, immutable GPU blocks; eviction never performs CPU copies.""" + + def __init__(self, budget, block_bytes, device): + super().__init__(1, budget, block_bytes, device) + self.capacity = self.gpu_capacity + self.cpu_capacity = 0 + + def _admit(self, key, pinned): + if not self.gpu_capacity: + self.metrics["skipped"] += 1 + return False + if len(self.blocks) >= self.gpu_capacity: + victim = next((k for k in self.blocks if k not in pinned), None) + if victim is None: + self.metrics["skipped"] += 1 + return False + self._drop(victim) + # No external directory in this phase: lookup is local to the P. + self.evicted.clear() + self.blocks[key] = {} + self.logical_dtypes[key] = {} + self._block_sizes[key] = 0 + self.tiers[key] = "gpu" + self._gpu_lru[key] = None + self.counts["gpu"] += 1 + return True + + +class DirectHistoryMixin: + def _init_history(self, budget): + block_bytes = sum(c[0].numel() * c.element_size() for c in self.caches.values()) + self.history = ( + GPUHistoryStore(budget, block_bytes, self.device) if budget else None + ) + self._history_task = None + self._history_closing = False + self._history_leases = {} + self._history_reading = set() + self.history_pending_bytes = min(budget, 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_loaded_requests=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: + assert self.history.reserve(lease, matched) + 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 + if state.held_blocks * self.history.block_bytes > self.history_pending_bytes: + self.metrics["history_skipped_requests"] += 1 + return False + if self._history_task is not None and not self._history_task.done(): + self.metrics["history_skipped_requests"] += 1 + return False + limit = min( + self.history.gpu_capacity, + self.history_pending_bytes // self.history.block_bytes, + ) + candidates = [] + for block, key in state.hashes.items(): + if key in self.history.ready: + self.history.touch(key) + elif len(candidates) < limit: + candidates.append((key, block)) + if not candidates: + 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=(type(exc), exc, exc.__traceback__), + ) + + task.add_done_callback(done) + return True + + async def _retain_history(self, ticket, state, candidates, deadline): + pending = [] + try: + # At most one writer and one bounded chunk. Yield between chunks; + # never queue the writer ahead of an already active handoff or load. + count = max(1, min(32, self.slab_bytes // self.history.block_bytes)) + for start in range(0, len(candidates), count): + while self.send_lock.locked() or self.recv_lock.locked(): + if self._history_closing or time.monotonic() >= deadline: + return + await asyncio.sleep(0.001) + if self._history_closing or time.monotonic() >= deadline: + return + async with self.send_lock: + self._expire_history_leases() + pairs = [ + (k, b) + for k, b in candidates[start : start + count] + 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) + pending = [] + del layers, index + 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): + 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 = [] + 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") + reads.append( + LayerRead( + name, + list(mapping), + list(mapping.values()), + 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 + await asyncio.to_thread(torch.cuda.synchronize, self.device) + self.metrics["history_hit_blocks"] += len( + self.history.leases[lease] + ) + self.metrics["history_loaded_requests"] += 1 + finally: + await asyncio.to_thread(torch.cuda.synchronize, self.device) + indices.clear() + self._history_reading.discard(lease) + self.release_history(lease) diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index 3ee2dccbde..026dfd9350 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -495,8 +495,12 @@ async def close_gpu_caches_v1(self): # 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): - self._gpu_transfer.register_direct(ticket, request_id, blocks) + def register_direct_gpu_v1( + self, ticket, request_id, blocks, hashes=None, held_blocks=None + ): + self._gpu_transfer.register_direct( + ticket, request_id, blocks, hashes, held_blocks + ) def release_direct_gpu_v1(self, ticket): self._gpu_transfer.release_direct(ticket) @@ -519,3 +523,16 @@ async def release_remote_direct_gpu_v1(self, rank, ticket): 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/vllm/xavier/test/test_direct_handoff.py b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py index 8611a32240..17d1c8be73 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py @@ -20,6 +20,7 @@ def direct_runtime(monkeypatch): r.metrics.update( direct_registered=0, direct_finished=0, direct_expired=0, index_uploads=0 ) + r._init_history(0) return r diff --git a/xinference/model/llm/vllm/xavier/test/test_direct_history.py b/xinference/model/llm/vllm/xavier/test/test_direct_history.py new file mode 100644 index 0000000000..0007bd8b97 --- /dev/null +++ b/xinference/model/llm/vllm/xavier/test/test_direct_history.py @@ -0,0 +1,243 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import asyncio +import threading + +import pytest +import torch + +from ..direct_history import GPUHistoryStore +from .test_direct_handoff import direct_runtime + + +def history_runtime(monkeypatch): + r = direct_runtime(monkeypatch) + r._init_history(32) + r.history_retention_seconds = 10 + r.caches["K"].copy_(torch.arange(16).reshape(8, 2)) + return r + + +@pytest.mark.asyncio +async def test_snapshot_waits_for_fence_and_survives_engine_reuse(monkeypatch): + r = history_runtime(monkeypatch) + entered, release = threading.Event(), threading.Event() + + def fence(*args): + entered.set() + assert release.wait(5) + + monkeypatch.setattr(torch.cuda, "synchronize", fence) + original = r.caches["K"][[1, 2]].clone() + r.register_direct("t", "p", [1, 2], [101, 102]) + r.release_direct("t") + try: + assert await asyncio.to_thread(entered.wait, 2) + assert not r.history.ready + assert not r.poll_direct() + assert r.direct_requests["t"].retaining + finally: + release.set() + await r._history_task + assert r.poll_direct() == {"p"} + r.caches["K"].zero_() + assert r.reserve_history("h", [101, 102, 999]) == [101, 102] + await r.load_history([{0: {"K": {101: 5, 102: 3}}}], ["h"]) + assert torch.equal(r.caches["K"][[5, 3]], original) + assert not r.history.leases and not r._history_leases + assert r.metrics["history_loaded_requests"] == 1 + assert r.history.stats()["cpu_blocks"] == 0 + + +@pytest.mark.asyncio +async def test_busy_writer_skips_next_request_without_backlog(monkeypatch): + r = history_runtime(monkeypatch) + await r.send_lock.acquire() + r.register_direct("a", "a", [1], [101]) + r.release_direct("a") + writer = r._history_task + r.register_direct("b", "b", [2], [102]) + r.release_direct("b") + assert r.poll_direct() == {"b"} + assert r._history_task is writer + assert r.metrics["history_skipped_requests"] == 1 + r.send_lock.release() + await writer + assert r.poll_direct() == {"a"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reason", ["held_budget", "deadline", "expired"]) +async def test_optional_retention_releases_engine_when_skipped(monkeypatch, reason): + r = history_runtime(monkeypatch) + r.register_direct( + "t", "p", [1], [101], held_blocks=9 if reason == "held_budget" else 1 + ) + if reason == "deadline": + r.history_retention_seconds = 0 + if reason == "expired": + r.direct_requests["t"].deadline = 0 + assert r.poll_direct() == {"p"} + else: + r.release_direct("t") + if r._history_task: + await r._history_task + assert r.poll_direct() == {"p"} + assert not r.history.blocks and not r.direct_requests + + +def test_gpu_eviction_respects_lease_without_cpu_demotion(): + store = GPUHistoryStore(8, 4, torch.device("cpu")) + + def save(keys): + store.stage_blocks(keys, {"K": torch.ones(len(keys), 2, dtype=torch.bfloat16)}) + store.publish(keys, {"K"}) + + save([1, 2]) + assert store.reserve("h", [1]) + save([3]) + assert store.ready == {1, 3} + assert store.reserve("j", [3]) + save([4]) + assert store.ready == {1, 3} + assert store.stats()["demotions"] == store.stats()["cpu_blocks"] == 0 + assert store.stats()["gpu_reserved_bytes"] == 8 + + +@pytest.mark.asyncio +async def test_failed_staging_never_publishes_partial_snapshot(monkeypatch): + r = history_runtime(monkeypatch) + original = r.history.stage_blocks + + def fail(keys, layers): + original(keys, layers) + raise RuntimeError("staging failed") + + monkeypatch.setattr(r.history, "stage_blocks", fail) + r.register_direct("t", "p", [1, 2], [101, 102]) + r.release_direct("t") + await r._history_task + assert not r.history.ready and not r.history.blocks + assert r.poll_direct() == {"p"} + assert r.metrics["history_failures"] == 1 + + +@pytest.mark.asyncio +async def test_invalid_history_read_releases_lease(monkeypatch): + r = history_runtime(monkeypatch) + r.register_direct("t", "p", [1], [101]) + r.release_direct("t") + await r._history_task + r.reserve_history("h", [101]) + with pytest.raises(ValueError, match="outside lease"): + await r.load_history([{0: {"K": {102: 5}}}], ["h"]) + assert not r._history_leases and not r.history.leases + + +def test_producer_history_maps_prefix_after_local_hit(connector): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + connector._direct_test = connector._history_enabled = True + actor = SimpleNamespace( + reserve_direct_history_v1=AsyncMock(), release_direct_history_v1=AsyncMock() + ) + request = SimpleNamespace(request_id="p", prompt_token_ids=list(range(65))) + keys = [ + key for key, _ in connector._build_xavier_hashes(request.prompt_token_ids[:64]) + ] + actor.reserve_direct_history_v1.return_value = keys[1:3] + connector._get_transfer_ref = AsyncMock(return_value=actor) + assert connector.get_num_new_matched_tokens(request, 16) == (32, True) + load = connector._requests_need_load["p"] + actor.reserve_direct_history_v1.assert_awaited_once_with(load.lease, keys[1:]) + connector.update_state_after_alloc( + request, SimpleNamespace(get_block_ids=lambda: [[4, 5, 6, 7]]), 32 + ) + assert connector._requests_need_load["p"].local_transfers_by_group == { + 0: {connector._rank: {keys[1]: 5, keys[2]: 6}} + } + + +def test_unallocated_history_abort_releases_lease(connector, monkeypatch): + import sys + from types import SimpleNamespace + from unittest.mock import AsyncMock + + monkeypatch.setitem( + sys.modules, + "vllm.v1.request", + SimpleNamespace(RequestStatus=SimpleNamespace(FINISHED_ABORTED="aborted")), + ) + connector._direct_test = connector._history_enabled = True + actor = SimpleNamespace( + reserve_direct_history_v1=AsyncMock(return_value=[101]), + release_direct_history_v1=AsyncMock(), + ) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = SimpleNamespace( + request_id="p", + prompt_token_ids=list(range(33)), + status="aborted", + kv_transfer_params={"do_remote_decode": True}, + ) + assert connector.get_num_new_matched_tokens(request, 0) == (16, True) + lease = connector._requests_need_load["p"].lease + assert connector.request_finished(request, []) == (False, None) + actor.release_direct_history_v1.assert_awaited_once_with(lease) + assert not connector._requests_need_load + + +@pytest.mark.asyncio +async def test_close_stops_waiting_writer_and_reclaims_snapshots(monkeypatch): + r = history_runtime(monkeypatch) + await r.recv_lock.acquire() + r.register_direct("t", "p", [1], [101]) + r.release_direct("t") + await r.close() + assert not r.direct_requests and not r.history.blocks + assert not r._history_leases + r.recv_lock.release() + + +@pytest.mark.asyncio +async def test_cancelled_restore_holds_lease_until_cuda_drains(monkeypatch): + r = history_runtime(monkeypatch) + r.register_direct("t", "p", [1], [101]) + r.release_direct("t") + await r._history_task + r.reserve_history("h", [101]) + entered, release = threading.Event(), threading.Event() + + def fence(*args): + entered.set() + assert release.wait(5) + + monkeypatch.setattr(torch.cuda, "synchronize", fence) + task = asyncio.create_task(r.run(r.load_history, [{0: {"K": {101: 5}}}], ["h"])) + try: + assert await asyncio.to_thread(entered.wait, 2) + task.cancel() + await asyncio.sleep(0) + r._history_leases["h"] = 0 + r._expire_history_leases() + r.release_history("h") + assert "h" in r.history.leases and not task.done() + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert not r.history.leases and not r._history_leases + assert r.caches["K"][5].tolist() == [2, 3] + + +def test_local_prefix_partial_tail_does_not_query_history(connector): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + connector._direct_test = connector._history_enabled = True + connector._get_transfer_ref = AsyncMock() + request = SimpleNamespace(request_id="p", prompt_token_ids=list(range(41))) + assert connector.get_num_new_matched_tokens(request, 32) == (0, False) + connector._get_transfer_ref.assert_not_called() + assert not connector._requests_need_load diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 8900e7224a..6b1056de48 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -136,6 +136,10 @@ def __init__( self._rank = int(self._xavier_config.get("rank", 0)) self._is_producer = self._kv_transfer_config.is_kv_producer self._is_consumer = self._kv_transfer_config.is_kv_consumer + self._history_enabled = self._direct_test and bool(self._gpu_budget) + if self._history_enabled and self._is_producer: + # P may restore its own independent history before computing a suffix. + self._is_consumer = True self._requests_need_load: Dict[str, XavierLoadRequest] = {} self._pending_store_requests: Dict[str, XavierStoreRequest] = {} self._leased_requests: Dict[str, XavierLoadRequest] = {} @@ -341,6 +345,43 @@ def get_num_new_matched_tokens( if self._direct_test: params = getattr(request, "kv_transfer_params", None) or {} handoff = params.get("xavier_direct") + if self._is_producer and getattr(self, "_history_enabled", False): + previous = self._requests_need_load.get(request.request_id) + if previous is not None: + if previous.local_transfers_by_group: + raise RuntimeError("History load already allocated") + self._requests_need_load.pop(request.request_id) + self._call(self._release_history_request(previous.lease)) + tokens = max(len(request.prompt_token_ids) - 1, 0) + # A local prefix hit can leave only a partial final block. + # Restoring that tiny tail costs more actor/layer work than + # computing it. Keep the local-cache fast path free of RPCs. + if tokens - num_computed_tokens < self._block_size: + return 0, False + start = num_computed_tokens // self._block_size + hashes = self._build_xavier_hashes(request.prompt_token_ids[:tokens])[ + start: + ] + lease = "history:" + uuid.uuid4().hex + + async def reserve(): + ref = await self._get_transfer_ref() + return await ref.reserve_direct_history_v1( + lease, [key for key, _ in hashes] + ) + + matched = self._call(reserve()) + if not matched: + return 0, False + self._requests_need_load[request.request_id] = XavierLoadRequest( + request.request_id, + {self._rank: {key: start + i for i, key in enumerate(matched)}}, + lease=lease, + ) + return ( + min(len(matched) * self._block_size, tokens - num_computed_tokens), + True, + ) if not self._is_consumer: return 0, False if handoff is None: @@ -476,6 +517,14 @@ def request_finished( block_ids: List[int], ) -> tuple[bool, dict[str, Any] | None]: if self._direct_test: + pending = self._requests_need_load.get(request.request_id) + if pending is not None and pending.lease.startswith("history:"): + if not pending.local_transfers_by_group: + self._requests_need_load.pop(request.request_id) + self._call(self._release_history_request(pending.lease)) + else: + # Allocated async loads must still drain after an abort. + return False, None params = getattr(request, "kv_transfer_params", None) or {} if self._is_producer and params.get("do_remote_decode"): from vllm.v1.request import RequestStatus @@ -493,9 +542,21 @@ def request_finished( async def register(): ref = await self._get_transfer_ref() - await ref.register_direct_gpu_v1( - ticket, request.request_id, blocks - ) + if getattr(self, "_history_enabled", False): + hashes = self._build_xavier_hashes( + request.prompt_token_ids[:token_count] + ) + await ref.register_direct_gpu_v1( + ticket, + request.request_id, + blocks, + [key for key, _ in hashes[:count]], + len(block_ids), + ) + else: + await ref.register_direct_gpu_v1( + ticket, request.request_id, blocks + ) self._call(register()) return bool(blocks), { @@ -580,6 +641,10 @@ async def _get_transfer_ref(self) -> xo.ActorRefType["TransferActor"]: self._transfer_ref = ref return self._transfer_ref + async def _release_history_request(self, lease): + transfer = await self._get_transfer_ref() + await transfer.release_direct_history_v1(lease) + async def _reserve_load_request(self, request): transfer = await self._get_transfer_ref() return await transfer.reserve_remote_blocks_v1(request.lease, request.transfers) @@ -927,13 +992,31 @@ async def _submit_gpu_requests(self, requests): ranks[rank] = layers entries.append(ranks) torch.cuda.synchronize() - if self._direct_test: - operation = transfer.load_direct_gpu_v1(entries, [r.lease for r in requests]) - else: - operation = transfer.load_gpu_requests_v1( - entries, [(r.lease, r.transfers) for r in requests] - ) - task = asyncio.create_task(operation) + async def load(): + if self._direct_test: + historical = [ + i for i, r in enumerate(requests) if r.lease.startswith("history:") + ] + direct = [ + i + for i, r in enumerate(requests) + if not r.lease.startswith("history:") + ] + if historical: + await transfer.load_direct_history_v1( + [entries[i] for i in historical], + [requests[i].lease for i in historical], + ) + if direct: + await transfer.load_direct_gpu_v1( + [entries[i] for i in direct], + [requests[i].lease for i in direct], + ) + else: + await transfer.load_gpu_requests_v1( + entries, [(r.lease, r.transfers) for r in requests] + ) + task = asyncio.create_task(load()) except BaseException: # Worker metadata already owns these leases. Until task creation # succeeds, the TransferActor has no operation that can release them. From 51611a8bd99b2f058adcc9fbcc110cb19ce0f921 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 02:57:10 +0800 Subject: [PATCH 4/9] perf: limit GPU history replacement with reuse-aware admission --- .../model/llm/vllm/xavier/direct_handoff.py | 1 + .../model/llm/vllm/xavier/direct_history.py | 71 +++++++++++--- .../vllm/xavier/test/test_direct_history.py | 95 +++++++++++++++++++ 3 files changed, 155 insertions(+), 12 deletions(-) diff --git a/xinference/model/llm/vllm/xavier/direct_handoff.py b/xinference/model/llm/vllm/xavier/direct_handoff.py index eec6d5f8db..243b5e8460 100644 --- a/xinference/model/llm/vllm/xavier/direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/direct_handoff.py @@ -63,6 +63,7 @@ async def _close(self): 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() diff --git a/xinference/model/llm/vllm/xavier/direct_history.py b/xinference/model/llm/vllm/xavier/direct_history.py index 7c6cc5dc13..e7f8a6de01 100644 --- a/xinference/model/llm/vllm/xavier/direct_history.py +++ b/xinference/model/llm/vllm/xavier/direct_history.py @@ -5,6 +5,7 @@ import asyncio import logging import time +from collections import OrderedDict import torch @@ -53,6 +54,15 @@ def _init_history(self, budget): 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.gpu_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(budget, 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. @@ -63,6 +73,9 @@ def _init_history(self, budget): history_failures=0, history_hit_blocks=0, history_loaded_requests=0, + history_admission_rejected_blocks=0, + history_admission_reused_blocks=0, + history_save_chunks=0, ) def _expire_history_leases(self): @@ -104,21 +117,37 @@ def _schedule_history(self, ticket, state): if state.held_blocks * self.history.block_bytes > self.history_pending_bytes: self.metrics["history_skipped_requests"] += 1 return False - if self._history_task is not None and not self._history_task.done(): - self.metrics["history_skipped_requests"] += 1 - return False limit = min( self.history.gpu_capacity, self.history_pending_bytes // self.history.block_bytes, ) candidates = [] + free = max(0, self.history.gpu_capacity - len(self.history.blocks)) for block, key in state.hashes.items(): if key in self.history.ready: self.history.touch(key) - elif len(candidates) < limit: - candidates.append((key, block)) + self._history_probation.pop(key, None) + continue + repeated = key in self._history_probation + 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 free or repeated: + if len(candidates) < limit: + candidates.append((key, block)) + free = max(0, free - 1) + if repeated: + self.metrics["history_admission_reused_blocks"] += 1 + else: + self.metrics["history_admission_rejected_blocks"] += 1 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( @@ -143,8 +172,27 @@ async def _retain_history(self, ticket, state, candidates, deadline): try: # At most one writer and one bounded chunk. Yield between chunks; # never queue the writer ahead of an already active handoff or load. - count = max(1, min(32, self.slab_bytes // self.history.block_bytes)) - for start in range(0, len(candidates), count): + 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 - len(self.history.blocks) + 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] + start += count while self.send_lock.locked() or self.recv_lock.locked(): if self._history_closing or time.monotonic() >= deadline: return @@ -153,11 +201,7 @@ async def _retain_history(self, ticket, state, candidates, deadline): return async with self.send_lock: self._expire_history_leases() - pairs = [ - (k, b) - for k, b in candidates[start : start + count] - if k not in self.history.ready - ] + 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)) @@ -170,6 +214,9 @@ async def _retain_history(self, ticket, state, candidates, deadline): 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) pending = [] del layers, index await asyncio.sleep(0) diff --git a/xinference/model/llm/vllm/xavier/test/test_direct_history.py b/xinference/model/llm/vllm/xavier/test/test_direct_history.py index 0007bd8b97..b0d0534233 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_history.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_history.py @@ -241,3 +241,98 @@ def test_local_prefix_partial_tail_does_not_query_history(connector): assert connector.get_num_new_matched_tokens(request, 32) == (0, False) connector._get_transfer_ref.assert_not_called() assert not connector._requests_need_load + + +@pytest.mark.asyncio +async def test_full_history_preserves_hot_content_until_new_content_repeats( + monkeypatch, +): + r = history_runtime(monkeypatch) + r._init_history(8) + r.history_retention_seconds = 10 + + async def finish(ticket, blocks, hashes): + r.register_direct(ticket, ticket, blocks, hashes) + r.release_direct(ticket) + if r._history_task: + await r._history_task + assert r.poll_direct() == {ticket} + + await finish("warm", [1, 2], [101, 102]) + assert r.history.ready == {101, 102} + await finish("cold", [3, 4], [103, 104]) + assert r.history.ready == {101, 102} + assert r.metrics["history_admission_rejected_blocks"] == 2 + await finish("repeat", [3, 4], [103, 104]) + assert r.history.ready == {103, 104} + assert r.reserve_history("hit", [103, 104]) == [103, 104] + assert not r._history_probation + assert r.metrics["history_saved_blocks"] == 4 + assert r.history.stats()["cpu_blocks"] == 0 + + +@pytest.mark.asyncio +async def test_probation_is_bounded_and_forgets_old_cold_content(monkeypatch): + r = history_runtime(monkeypatch) + r._init_history(4) + r.history_retention_seconds = 10 + for key in (100, 101, 102, 103, 101): + r.register_direct(str(key), str(key), [1], [key]) + r.release_direct(str(key)) + if r._history_task: + await r._history_task + r.poll_direct() + assert r.history.ready == {100} + assert list(r._history_probation) == [103, 101] + await r.close() + assert not r._history_probation + + +@pytest.mark.asyncio +async def test_snapshot_chunks_respect_byte_budget(monkeypatch): + r = history_runtime(monkeypatch) + r.history_chunk_bytes = r.history_fill_chunk_bytes = 8 + r.register_direct("t", "p", [1, 2, 3, 4, 5], [101, 102, 103, 104, 105]) + original = r.history.stage_blocks + sizes = [] + + def stage(keys, layers): + sizes.append(sum(t.numel() * t.element_size() for t in layers.values())) + return original(keys, layers) + + monkeypatch.setattr(r.history, "stage_blocks", stage) + r.release_direct("t") + await r._history_task + assert sizes == [8, 8, 4] + assert r.metrics["history_save_chunks"] == 3 + assert r.history.ready == {101, 102, 103, 104, 105} + assert r.poll_direct() == {"p"} + + +@pytest.mark.asyncio +async def test_fill_batches_grow_but_replacement_batches_stay_small(monkeypatch): + r = history_runtime(monkeypatch) + r._init_history(20) + r.history_retention_seconds = 10 + r.history_fill_chunk_bytes = 16 + r.history_chunk_bytes = 8 + sizes = [] + original = r.history.stage_blocks + + def stage(keys, layers): + sizes.append(len(keys)) + return original(keys, layers) + + monkeypatch.setattr(r.history, "stage_blocks", stage) + for ticket, hashes in ( + ("fill", list(range(101, 106))), + ("observe", list(range(201, 206))), + ("replace", list(range(201, 206))), + ): + r.register_direct(ticket, ticket, [1, 2, 3, 4, 5], hashes) + r.release_direct(ticket) + if r._history_task: + await r._history_task + assert r.poll_direct() == {ticket} + assert sizes == [4, 1, 2, 2, 1] + assert r.history.ready == set(range(201, 206)) From 5b4c036eced3884ab12b3182fdaa0780456575db Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 03:28:50 +0800 Subject: [PATCH 5/9] perf: default Xavier PD to direct handoff with tiered history --- .github/workflows/pd-gpu.yaml | 4 +- .github/workflows/python.yaml | 2 + .../LC_MESSAGES/user_guide/pd_separation.po | 42 +++++ .../LC_MESSAGES/user_guide/pd_separation.po | 41 +++++ .../LC_MESSAGES/user_guide/pd_separation.po | 42 +++++ .../LC_MESSAGES/user_guide/pd_separation.po | 42 +++++ .../LC_MESSAGES/user_guide/pd_separation.po | 38 ++++ .../LC_MESSAGES/user_guide/pd_separation.po | 37 ++++ .../LC_MESSAGES/user_guide/pd_separation.po | 41 +++++ .../LC_MESSAGES/user_guide/pd_separation.po | 36 ++++ .../LC_MESSAGES/user_guide/pd_separation.po | 36 ++++ doc/source/user_guide/pd_separation.rst | 12 +- xinference/core/pd_model.py | 21 +-- xinference/core/supervisor.py | 3 + xinference/core/tests/test_pd_launch.py | 15 ++ xinference/core/tests/test_pd_model.py | 28 ++- xinference/model/llm/vllm/core.py | 14 +- xinference/model/llm/vllm/tests/test_pd.py | 7 +- .../model/llm/vllm/xavier/direct_handoff.py | 4 +- .../model/llm/vllm/xavier/direct_history.py | 92 +++++----- .../model/llm/vllm/xavier/gpu_transfer.py | 6 +- .../vllm/xavier/test/test_direct_handoff.py | 17 +- .../vllm/xavier/test/test_direct_history.py | 169 +++++++++++++++++- .../llm/vllm/xavier/test/test_gpu_transfer.py | 17 ++ .../model/llm/vllm/xavier/test/test_pd_gpu.py | 29 ++- .../xavier/test/test_review_regressions.py | 4 +- xinference/model/llm/vllm/xavier/transport.py | 11 ++ .../model/llm/vllm/xavier/v1_connector.py | 37 ++-- 28 files changed, 722 insertions(+), 125 deletions(-) diff --git a/.github/workflows/pd-gpu.yaml b/.github/workflows/pd-gpu.yaml index a84fa49569..de8223d7e0 100644 --- a/.github/workflows/pd-gpu.yaml +++ b/.github/workflows/pd-gpu.yaml @@ -41,11 +41,11 @@ jobs: run: | python -m venv .venv-pd .venv-pd/bin/python -m pip install -U pip 'setuptools<82' - .venv-pd/bin/python -m pip install -e '.[dev]' "vllm==$VLLM_VERSION" "nixl==$NIXL_VERSION" openai + .venv-pd/bin/python -m pip install -e '.[dev]' "vllm==$VLLM_VERSION" "nixl==$NIXL_VERSION" "xoscar[nixl]>=0.11.1" openai - name: Verify PD with actual KV transfer run: | source .venv-pd/bin/activate - python -m pytest --timeout=1800 -vv xinference/model/llm/vllm/xavier/test/test_pd_gpu.py xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py + python -m pytest --timeout=1800 -vv xinference/model/llm/vllm/xavier/test/test_pd_gpu.py xinference/model/llm/vllm/xavier/test/test_v1_cache_layout.py xinference/model/llm/vllm/xavier/test/test_direct_history.py - name: Upload runtime logs if: always() uses: actions/upload-artifact@v4 diff --git a/.github/workflows/python.yaml b/.github/workflows/python.yaml index a06f0c7d9c..c4df0aeea6 100644 --- a/.github/workflows/python.yaml +++ b/.github/workflows/python.yaml @@ -380,6 +380,8 @@ jobs: xinference/model/llm/vllm/xavier/test/test_v1_engine.py \ xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py \ xinference/model/llm/vllm/xavier/test/test_kv_schema.py \ + xinference/model/llm/vllm/xavier/test/test_direct_handoff.py \ + xinference/model/llm/vllm/xavier/test/test_direct_history.py \ xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py \ xinference/model/llm/vllm/xavier/test/test_request_transfer.py \ xinference/model/llm/vllm/xavier/test/test_transfer_blocks.py \ diff --git a/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po index de58fe1b56..4d00af5892 100644 --- a/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po @@ -299,3 +299,45 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "Übertragungsfehler und inkompatible KV-Cache-Layouts lösen Fehler aus, statt die Anfrage stillschweigend neu zu berechnen." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D verwendet standardmäßig direkte GPU-zu-GPU-Übertragung. " +"Erforderlich sind Linux, NVIDIA-GPUs, vLLM >= 0.21.0 und " +"``xoscar[nixl]>=0.11.1`` in der Worker- und Modellumgebung. Verwenden Sie " +"eine erreichbare Hostadresse (nicht ``0.0.0.0``). Ohne NIXL schlägt der " +"Start fehl; es gibt keinen CPU-Fallback." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"``xavier_gpu_cache_bytes`` beträgt für P/D standardmäßig ``268435456`` (256 " +"MiB pro Replikat). Der Verlauf nutzt zuerst die GPU und wird bei erschöpfter" +" GPU-Kapazität auf die CPU ausgelagert. Der CPU-Verlauf ist durch die KV-" +"Blockkapazität der Engine begrenzt und kann im Hauptspeicher deren Cache-" +"Größe erreichen. CPU-Treffer stellen KV auf der GPU wieder her; Misses " +"werden lokal berechnet. ``0`` deaktiviert den Verlauf, behält aber die " +"direkte Übertragung bei. Übertragungsfehler lösen Fehler aus." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "Direkte Übertragung und gestufter Verlauf" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "Die direkte Übertragung erfordert derzeit ``n=1``." + diff --git a/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po index 1640945a04..6a427c5399 100644 --- a/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po @@ -297,3 +297,44 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "Los fallos de transferencia y los formatos de caché KV incompatibles generan errores en lugar de volver a calcular la solicitud silenciosamente." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D utiliza por defecto transferencia directa entre GPU. Requiere " +"Linux, GPU NVIDIA, vLLM >= 0.21.0 y ``xoscar[nixl]>=0.11.1`` en los entornos" +" del worker y del modelo. Use una dirección de host accesible (no " +"``0.0.0.0``). Sin NIXL, el inicio falla; no hay alternativa por CPU." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"Para P/D, ``xavier_gpu_cache_bytes`` vale por defecto ``268435456`` (256 MiB" +" por réplica). El historial usa primero GPU y pasa a CPU cuando se agota su " +"capacidad. El historial CPU está limitado por la capacidad de bloques KV del" +" motor y puede ocupar un tamaño similar en memoria del host. Los aciertos " +"CPU restauran KV en GPU; los fallos se calculan localmente. ``0`` desactiva " +"el historial y mantiene la transferencia directa. Los fallos de " +"transferencia generan errores." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "Transferencia directa e historial por niveles" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "La transferencia directa requiere actualmente ``n=1``." + diff --git a/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po index 8cfac69b64..857c32a38d 100644 --- a/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po @@ -300,3 +300,45 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "Les échecs de transfert et les dispositions incompatibles du cache KV provoquent des erreurs au lieu de recalculer silencieusement la requête." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D utilise par défaut le transfert direct entre GPU. Il nécessite " +"Linux, des GPU NVIDIA, vLLM >= 0.21.0 et ``xoscar[nixl]>=0.11.1`` dans les " +"environnements du worker et du modèle. Utilisez une adresse hôte accessible " +"(pas ``0.0.0.0``). Sans NIXL, le démarrage échoue ; aucun repli CPU n’est " +"prévu." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"Pour P/D, ``xavier_gpu_cache_bytes`` vaut par défaut ``268435456`` (256 MiB " +"par réplique). L’historique utilise d’abord le GPU puis déborde vers le CPU " +"lorsque sa capacité est épuisée. L’historique CPU est limité par la capacité" +" en blocs KV du moteur et peut approcher cette taille de cache en mémoire " +"hôte. Les accès CPU restaurent les KV sur le GPU ; les défauts sont calculés" +" localement. ``0`` désactive l’historique tout en conservant le transfert " +"direct. Les échecs de transfert produisent des erreurs." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "Transfert direct et historique hiérarchisé" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "Le transfert direct nécessite actuellement ``n=1``." + diff --git a/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po index 03dba70583..fb4dd1cb54 100644 --- a/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po @@ -297,3 +297,45 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "Gli errori di trasferimento e i layout incompatibili della cache KV generano errori anziché ricalcolare silenziosamente la richiesta." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D usa per impostazione predefinita il trasferimento diretto tra " +"GPU. Richiede Linux, GPU NVIDIA, vLLM >= 0.21.0 e ``xoscar[nixl]>=0.11.1`` " +"negli ambienti del worker e del modello. Usare un indirizzo host " +"raggiungibile (non ``0.0.0.0``). Senza NIXL, l’avvio fallisce; non è " +"previsto il ripiego sulla CPU." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"Per P/D, ``xavier_gpu_cache_bytes`` vale per impostazione predefinita " +"``268435456`` (256 MiB per replica). La cronologia usa prima la GPU e passa " +"alla CPU quando la capacità GPU è esaurita. La cronologia CPU è limitata " +"dalla capacità dei blocchi KV del motore e può avvicinarsi a tale dimensione" +" della cache nella memoria host. Gli hit CPU ripristinano KV sulla GPU; i " +"miss vengono calcolati localmente. ``0`` disabilita la cronologia mantenendo" +" il trasferimento diretto. I guasti di trasferimento generano errori." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "Trasferimento diretto e cronologia a livelli" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "Il trasferimento diretto richiede attualmente ``n=1``." + diff --git a/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po index 5e6b4c633f..eb2137a05a 100644 --- a/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po @@ -272,3 +272,41 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "転送の失敗や互換性のない KV キャッシュレイアウトは、リクエストを暗黙的に再計算するのではなく、エラーになります。" + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D はデフォルトで GPU 間の直接転送を使用します。Linux、NVIDIA GPU、vLLM >= " +"0.21.0、およびワーカーとモデルの両環境に ``xoscar[nixl]>=0.11.1`` " +"が必要です。到達可能なホストアドレスを使用してください(``0.0.0.0`` は不可)。NIXL がなければ起動に失敗し、CPU " +"へフォールバックしません。" + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"P/D の ``xavier_gpu_cache_bytes`` はデフォルトで ``268435456`` (レプリカ当たり 256 " +"MiB)です。履歴は GPU を優先し、容量不足時に CPU に退避します。CPU 履歴のブロック数はエンジン KV " +"ブロック容量が上限で、ホストメモリ使用量はそのキャッシュサイズに近づく場合があります。CPU ヒットでは KV を GPU " +"に復元し、ミスではローカル計算します。 ``0`` は履歴を無効にし、直接転送は維持します。転送失敗はエラーになります。" + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "直接転送と階層型履歴" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "直接転送は現在 ``n=1`` を必要とします。" + diff --git a/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po index 0c62940bb8..310058381f 100644 --- a/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po @@ -272,3 +272,40 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "전송 실패와 호환되지 않는 KV 캐시 레이아웃은 요청을 조용히 다시 계산하는 대신 오류를 발생시킵니다." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D는 기본적으로 GPU 간 직접 전송을 사용합니다. Linux, NVIDIA GPU, vLLM >= 0.21.0이 " +"필요하며 워커와 모델 환경 모두에 ``xoscar[nixl]>=0.11.1`` 을 설치해야 합니다. 연결 가능한 호스트 주소를 " +"사용하세요(``0.0.0.0`` 제외). NIXL이 없으면 시작에 실패하며 CPU로 대체하지 않습니다." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"P/D의 ``xavier_gpu_cache_bytes`` 기본값은 ``268435456`` (복제본당 256 MiB)입니다. 이력은 " +"GPU를 우선 사용하고 GPU 용량이 부족하면 CPU로 이동합니다. CPU 이력 블록 수는 엔진 KV 블록 용량으로 제한되며 호스트 " +"메모리 사용량이 해당 캐시 크기에 근접할 수 있습니다. CPU 적중 시 KV를 GPU로 복원하고 미스 시 로컬에서 계산합니다. ``0``" +" 은 이력을 비활성화하지만 직접 전송은 유지합니다. 전송 실패는 오류를 발생시킵니다." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "직접 전송 및 계층형 이력" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "직접 전송은 현재 ``n=1`` 을 요구합니다." + diff --git a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po index 4f97da875f..14c437d731 100644 --- a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po @@ -296,3 +296,44 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "Falhas de transferência e layouts incompatíveis do cache KV geram erros em vez de recalcular a solicitação silenciosamente." + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D usa transferência direta entre GPUs por padrão. Requer Linux, " +"GPUs NVIDIA, vLLM >= 0.21.0 e ``xoscar[nixl]>=0.11.1`` nos ambientes do " +"worker e do modelo. Use um endereço de host acessível (não ``0.0.0.0``). Sem" +" NIXL, a inicialização falha; não há alternativa por CPU." + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"Para P/D, ``xavier_gpu_cache_bytes`` assume ``268435456`` por padrão (256 " +"MiB por réplica). O histórico usa primeiro a GPU e passa para a CPU quando a" +" capacidade GPU se esgota. O histórico CPU é limitado pela capacidade de " +"blocos KV do motor e pode ocupar um tamanho próximo ao desse cache na " +"memória do host. Acertos CPU restauram KV na GPU; ausências são calculadas " +"localmente. ``0`` desativa o histórico e mantém a transferência direta. " +"Falhas de transferência geram erros." + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "Transferência direta e histórico em camadas" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "A transferência direta atualmente exige ``n=1``." + diff --git a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po index 4f096f26b7..26598627c7 100644 --- a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po @@ -273,3 +273,39 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "传输失败或 KV 缓存布局不兼容时会报错,而不会静默地重新计算请求。" + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D 默认使用 GPU 到 GPU 直传。要求 Linux、NVIDIA GPU、vLLM >= 0.21.0,并在 worker " +"和模型环境中都安装 ``xoscar[nixl]>=0.11.1``。请使用可达的主机地址(不能是 ``0.0.0.0``)。缺少 NIXL " +"时启动失败,不回退 CPU。" + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"P/D 的 ``xavier_gpu_cache_bytes`` 默认为 ``268435456`` (每副本 256 MiB)。历史缓存优先使用 " +"GPU,GPU 容量不足时下沉 CPU。CPU 历史块数受引擎 KV 块容量限制,主机内存占用可能接近该缓存大小。CPU 命中时将 KV 恢复到 " +"GPU,未命中时本地计算。 ``0`` 禁用历史缓存,仍保留直传。传输失败时报错。" + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "直传与分层历史缓存" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "直传目前要求 ``n=1``。" + diff --git a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po index 75e8ee75b9..85d5470d81 100644 --- a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po @@ -268,3 +268,39 @@ msgstr "" #: ../../source/user_guide/pd_separation.rst:145 msgid "Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request." msgstr "傳輸失敗或 KV 快取佈局不相容時會報錯,而不會默默地重新計算請求。" + +#: user_guide/pd_separation.rst:1 +msgid "" +"Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA " +"GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model " +"environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL " +"fails the launch; there is no CPU fallback." +msgstr "" +"Xavier P/D 預設使用 GPU 到 GPU 直傳。需要 Linux、NVIDIA GPU、vLLM >= 0.21.0,並在 worker " +"和模型環境中都安裝 ``xoscar[nixl]>=0.11.1``。請使用可連線的主機位址(不能是 ``0.0.0.0``)。缺少 NIXL " +"時啟動失敗,不退回 CPU。" + + +#: user_guide/pd_separation.rst:1 +msgid "" +"``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) " +"for P/D. History uses GPU first and spills to CPU when GPU capacity is " +"exhausted. CPU history is bounded by the engine KV block capacity and can " +"approach that cache size in host memory. CPU hits restore KV to GPU; misses " +"compute locally. ``0`` disables history, retaining direct transfer. Transfer" +" failures raise errors." +msgstr "" +"P/D 的 ``xavier_gpu_cache_bytes`` 預設為 ``268435456`` (每個副本 256 MiB)。歷史快取優先使用 " +"GPU,GPU 容量不足時移至 CPU。CPU 歷史區塊數受引擎 KV 區塊容量限制,主機記憶體用量可能接近該快取大小。CPU 命中時將 KV 還原至 " +"GPU,未命中時在本機計算。 ``0`` 停用歷史快取,仍保留直傳。傳輸失敗時回報錯誤。" + + +#: user_guide/pd_separation.rst:1 +msgid "Direct handoff and tiered history" +msgstr "直傳與分層歷史快取" + + +#: user_guide/pd_separation.rst:194 +msgid "Direct handoff currently requires ``n=1``." +msgstr "直傳目前要求 ``n=1``。" + diff --git a/doc/source/user_guide/pd_separation.rst b/doc/source/user_guide/pd_separation.rst index 091b7c07fc..42ddf7b058 100644 --- a/doc/source/user_guide/pd_separation.rst +++ b/doc/source/user_guide/pd_separation.rst @@ -11,7 +11,7 @@ No enterprise package or License is required. Launch ------ -Xavier requires NVIDIA GPUs, a reachable host address (not ``0.0.0.0``), and vLLM 0.21.0 or newer. Older deployments can use the legacy Xavier integration with vLLM below 0.11.0; versions 0.11 through 0.20 are not supported. +Xavier P/D uses GPU-to-GPU handoff by default. It requires Linux, NVIDIA GPUs, vLLM >= 0.21.0 and ``xoscar[nixl]>=0.11.1`` in both worker and model environments. Use a reachable host address (not ``0.0.0.0``). Missing NIXL fails the launch; there is no CPU fallback. The example starts one prefill replica on GPU 0 and one decode replica on GPU 1. Replace the worker address with the full ``ip:port`` reported by @@ -129,14 +129,12 @@ Xavier V1 requires one GPU per replica (TP=1, PP=1) and text-only models without Xavier's CPU cache can consume as much memory as the engine's GPU KV cache. Allow additional host memory for transfers when sizing your deployment. -GPU-first Xavier cache (experimental) -------------------------------------- - -GPU-first caching requires the Xavier backend and more than one replica. +Direct handoff and tiered history +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -For the Xavier backend, set ``xavier_gpu_cache_bytes`` to a non-negative integer, such as ``268435456`` for 256 MiB per replica. Omit it to use the existing CPU cache, or set it to ``0`` to use GPU-first caching with no GPU storage budget. This option does not select the native vLLM NIXL backend. +``xavier_gpu_cache_bytes`` defaults to ``268435456`` (256 MiB per replica) for P/D. History uses GPU first and spills to CPU when GPU capacity is exhausted. CPU history is bounded by the engine KV block capacity and can approach that cache size in host memory. CPU hits restore KV to GPU; misses compute locally. ``0`` disables history, retaining direct transfer. Transfer failures raise errors. -Install ``xoscar[nixl]>=0.11.1`` in both the worker and model environments. GPU-first caching requires Linux, NVIDIA CUDA, vLLM 0.21.0 or newer, a full-attention model and one GPU per replica (TP=1, PP=1). Use the same model, cache dtype and engine configuration on all replicas. +Direct handoff currently requires ``n=1``. The budget applies only to Xavier's GPU cache. Leave additional GPU memory for the model, vLLM's own KV cache and transfers. When this cache fills, Xavier stores reusable data in CPU memory; if no cache space is available, new data may not be retained for reuse. diff --git a/xinference/core/pd_model.py b/xinference/core/pd_model.py index 3b1ad94a43..4c45b71aac 100644 --- a/xinference/core/pd_model.py +++ b/xinference/core/pd_model.py @@ -16,7 +16,6 @@ import copy import json import logging -import os import time import uuid from abc import ABC, abstractmethod @@ -85,6 +84,7 @@ def __init__( [List[xo.ActorRefType["ModelActor"]]], SchedulingPolicy ] = RoundRobinSchedulingPolicy, transport_backend: str = "xavier", + direct_handoff: bool = False, ): super().__init__() # Prefill request map, used to skip the timeout task for specific request id. @@ -92,7 +92,7 @@ def __init__( self._model_uid = model_uid self._transport_backend = transport_backend - self._direct_test = os.getenv("XINFERENCE_XAVIER_DIRECT_TEST") == "1" + self._direct_handoff = transport_backend == "xavier" and direct_handoff # 使用字典存储副本:{replica_uid: actor_ref} self._prefill_replicas: Dict[str, xo.ActorRefType["ModelActor"]] = {} @@ -251,7 +251,7 @@ async def free_prefill_model_cache(self, request_id: str): if ( self._prefill_replicas and self._transport_backend == "xavier" - and not self._direct_test + and not self._direct_handoff ): await asyncio.gather( *[ @@ -293,14 +293,13 @@ async def _infer(self, method, inputs, *args, **kwargs): if args and args[0] is not None and not isinstance(args[0], dict): raise TypeError("Generation config must be a dict or None") if ( - (self._transport_backend == "nixl" or self._direct_test) + (self._transport_backend == "nixl" or self._direct_handoff) and args and args[0] and args[0].get("n", 1) != 1 ): - # NIXL releases the producer lease after one decoder (per TP rank) - # acknowledges the read, not after all parallel sampling children. - raise ValueError("Native NIXL PD currently requires n=1") + # Handoff leases cover one decoder, not parallel sampling children. + raise ValueError("PD KV handoff currently requires n=1") prefill = self._prefill_policy.schedule() decode = self._decode_policy.schedule() prefill_args = list(copy.deepcopy(args)) @@ -308,7 +307,7 @@ async def _infer(self, method, inputs, *args, **kwargs): prefill_args = [{}] + prefill_args[1:] prefill_args[0]["max_tokens"] = 1 prefill_args[0]["stream"] = False - if self._transport_backend == "nixl" or self._direct_test: + if self._transport_backend == "nixl" or self._direct_handoff: prefill_args[0]["n"] = 1 prefill_args[0]["_pd_kv_transfer_params"] = { "do_remote_decode": True, @@ -334,7 +333,7 @@ async def _infer(self, method, inputs, *args, **kwargs): self._transport_backend, time.perf_counter() - prefill_start, ) - if self._transport_backend == "nixl" or self._direct_test: + if self._transport_backend == "nixl" or self._direct_handoff: payload = ( json.loads(result) if isinstance(result, (bytes, str)) else result ) @@ -346,9 +345,7 @@ async def _infer(self, method, inputs, *args, **kwargs): if not isinstance(transfer, dict) or not transfer.get( "do_remote_prefill" ): - raise RuntimeError( - "NIXL prefill did not return KV transfer metadata" - ) + raise RuntimeError("PD prefill did not return KV transfer metadata") decode_args = list(copy.deepcopy(args)) if not decode_args or decode_args[0] is None: decode_args = [{}] + decode_args[1:] diff --git a/xinference/core/supervisor.py b/xinference/core/supervisor.py index c58025380d..3972cfac1c 100644 --- a/xinference/core/supervisor.py +++ b/xinference/core/supervisor.py @@ -3104,6 +3104,8 @@ async def launch_builtin_model( gpu_cache_bytes = validate_gpu_cache_budget( kwargs.pop("xavier_gpu_cache_bytes", None), enable_xavier, replica ) + if pd_enabled and enable_xavier and gpu_cache_bytes is None: + gpu_cache_bytes = 256 * 1024 * 1024 store_address = None store_port = None world_size = None @@ -3524,6 +3526,7 @@ async def _launch_model(): PDModelActor, model_uid, transport_backend=transport_backend, + direct_handoff=enable_xavier and gpu_cache_bytes is not None, address=self.address, uid=f"{model_uid}-{PDModelActor.default_uid()}", ) diff --git a/xinference/core/tests/test_pd_launch.py b/xinference/core/tests/test_pd_launch.py index 0fc30e7d98..dde35c29e8 100644 --- a/xinference/core/tests/test_pd_launch.py +++ b/xinference/core/tests/test_pd_launch.py @@ -331,3 +331,18 @@ async def test_gpu_pool_options_reach_actual_subpool(monkeypatch): ) assert env["UCX_MEMTYPE_CACHE"] == "n" worker._ensure_subpool_monitor.assert_awaited_once_with() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("budget", [None, 0, 268435456]) +async def test_gpu_pd_defaults_to_direct_handoff(launch_runtime, budget): + supervisor, workers, actors, _ = launch_runtime + await supervisor.launch_builtin_model( + **launch_kwargs(), xavier_gpu_cache_bytes=budget + ) + assert actors["PDModelActor"].constructor_kwargs["direct_handoff"] is True + for worker in workers: + config = worker.launch_builtin_model.call_args.kwargs + assert config["xavier_config"]["gpu_cache_bytes"] == ( + 268435456 if budget is None else budget + ) diff --git a/xinference/core/tests/test_pd_model.py b/xinference/core/tests/test_pd_model.py index c97b2dc0d7..7a7b16e6ff 100644 --- a/xinference/core/tests/test_pd_model.py +++ b/xinference/core/tests/test_pd_model.py @@ -279,11 +279,15 @@ async def test_abort_failure_still_cleans_request(router): @pytest.mark.asyncio @pytest.mark.parametrize("method", ["chat", "generate"]) -async def test_nixl_handoff_preserves_metadata_and_decode_settings(router, method): +@pytest.mark.parametrize("backend", ["nixl", "xavier"]) +async def test_nixl_handoff_preserves_metadata_and_decode_settings( + router, backend, method +): import json actor, prefill, decode = router - actor._transport_backend = "nixl" + actor._transport_backend = backend + actor._direct_handoff = backend == "xavier" transfer = { "do_remote_prefill": True, "remote_engine_id": "engine-p", @@ -309,9 +313,11 @@ async def test_nixl_handoff_preserves_metadata_and_decode_settings(router, metho @pytest.mark.asyncio -async def test_nixl_parallel_sampling_rejected_before_prefill(router): +@pytest.mark.parametrize("backend", ["nixl", "xavier"]) +async def test_nixl_parallel_sampling_rejected_before_prefill(router, backend): actor, prefill, decode = router - actor._transport_backend = "nixl" + actor._transport_backend = backend + actor._direct_handoff = backend == "xavier" with pytest.raises(ValueError, match="n=1"): await actor._infer("chat", [], {"n": 2}) prefill.chat.assert_not_awaited() @@ -321,9 +327,13 @@ async def test_nixl_parallel_sampling_rejected_before_prefill(router): @pytest.mark.asyncio @pytest.mark.parametrize("payload", [{}, None, b"null", b"[]", b"{}", b"42"]) -async def test_nixl_missing_metadata_does_not_silently_recompute(router, payload): +@pytest.mark.parametrize("backend", ["nixl", "xavier"]) +async def test_nixl_missing_metadata_does_not_silently_recompute( + router, backend, payload +): actor, prefill, decode = router - actor._transport_backend = "nixl" + actor._transport_backend = backend + actor._direct_handoff = backend == "xavier" prefill.chat.return_value = payload with pytest.raises(RuntimeError, match="KV transfer metadata"): await actor._infer("chat", [], {}) @@ -365,9 +375,11 @@ async def test_nixl_multiple_prefillers_and_decoders(): @pytest.mark.asyncio -async def test_nixl_stream_close_releases_decode_slot(router): +@pytest.mark.parametrize("backend", ["nixl", "xavier"]) +async def test_nixl_stream_close_releases_decode_slot(router, backend): actor, prefill, decode = router - actor._transport_backend = "nixl" + actor._transport_backend = backend + actor._direct_handoff = backend == "xavier" prefill.chat.return_value = {"_pd_kv_transfer_params": {"do_remote_prefill": True}} closed = [] diff --git a/xinference/model/llm/vllm/core.py b/xinference/model/llm/vllm/core.py index bf7fd2122f..2af1782380 100644 --- a/xinference/model/llm/vllm/core.py +++ b/xinference/model/llm/vllm/core.py @@ -1995,11 +1995,12 @@ async def async_generate( sampling_params = SamplingParams(**sanitized_generate_config) if generate_config and "_pd_kv_transfer_params" in generate_config: - if ( - self._nixl_config is None - and os.getenv("XINFERENCE_XAVIER_DIRECT_TEST") != "1" + from .xavier.transport import uses_direct_handoff + + if self._nixl_config is None and not uses_direct_handoff( + self._xavier_config ): - raise ValueError("KV handoff requires the NIXL PD backend") + raise ValueError("KV handoff requires NIXL or Xavier direct PD") sampling_params.extra_args = { **(sampling_params.extra_args or {}), "kv_transfer_params": generate_config["_pd_kv_transfer_params"], @@ -2181,11 +2182,12 @@ async def stream_results() -> AsyncGenerator[CompletionChunk, None]: completion = self._convert_request_output_to_completion( request_id, model=self.model_uid, request_output=final_output ) + from .xavier.transport import uses_direct_handoff + if ( self._nixl_config is not None and self._nixl_config["role"] == "prefill" ) or ( - os.getenv("XINFERENCE_XAVIER_DIRECT_TEST") == "1" - and self._xavier_config + uses_direct_handoff(self._xavier_config) and self._xavier_config.get("role") == "prefill" ): completion["_pd_kv_transfer_params"] = getattr( diff --git a/xinference/model/llm/vllm/tests/test_pd.py b/xinference/model/llm/vllm/tests/test_pd.py index a98f87d52d..3a9b6aff7e 100644 --- a/xinference/model/llm/vllm/tests/test_pd.py +++ b/xinference/model/llm/vllm/tests/test_pd.py @@ -62,10 +62,12 @@ def test_backend_selection(): @pytest.mark.parametrize("role", ["prefill", "decode"]) +@pytest.mark.parametrize("backend", ["nixl", "xavier"]) @pytest.mark.asyncio async def test_generate_passes_native_handoff_and_returns_producer_metadata( monkeypatch, role, + backend, ): from .. import core @@ -91,7 +93,10 @@ async def generate(self, prompt, params, request_id, **kwargs): monkeypatch.setattr(core, "VLLM_VERSION", Version("0.21.0")) monkeypatch.setattr(core, "VLLM_INSTALLED", False) model = object.__new__(core.VLLMModel) - model._nixl_config = {"role": role} + model._nixl_config = {"role": role} if backend == "nixl" else None + model._xavier_config = ( + {"role": role, "gpu_cache_bytes": 0} if backend == "xavier" else None + ) model._engine = Engine() model._active_request_ids = set() model.lora_requests = [] diff --git a/xinference/model/llm/vllm/xavier/direct_handoff.py b/xinference/model/llm/vllm/xavier/direct_handoff.py index 243b5e8460..d8f844b45e 100644 --- a/xinference/model/llm/vllm/xavier/direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/direct_handoff.py @@ -1,6 +1,6 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Experimental direct handoff with optional bounded GPU history.""" +"""Request-scoped GPU handoff with independent tiered history.""" import asyncio import time @@ -169,7 +169,7 @@ async def load_direct(self, requests, tickets): async with self.recv_lock: for ranks, ticket in zip(requests, tickets): if len(ranks) != 1: - raise ValueError("Direct prototype requires one producer") + raise ValueError("Direct handoff requires one producer") rank, layers = next(iter(ranks.items())) sender = await xo.actor_ref( address=self.actor._world_addresses[rank], diff --git a/xinference/model/llm/vllm/xavier/direct_history.py b/xinference/model/llm/vllm/xavier/direct_history.py index e7f8a6de01..5126433cf4 100644 --- a/xinference/model/llm/vllm/xavier/direct_history.py +++ b/xinference/model/llm/vllm/xavier/direct_history.py @@ -1,6 +1,6 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Bounded GPU history retained after request-scoped handoff.""" +"""Bounded GPU/CPU history retained after request-scoped handoff.""" import asyncio import logging @@ -15,40 +15,33 @@ logger = logging.getLogger(__name__) -class GPUHistoryStore(TieredKVSnapshotStore): - """Independent, immutable GPU blocks; eviction never performs CPU copies.""" +class HistoryStore(TieredKVSnapshotStore): + """Independent GPU-first history with bounded CPU overflow.""" - def __init__(self, budget, block_bytes, device): - super().__init__(1, budget, block_bytes, device) - self.capacity = self.gpu_capacity - self.cpu_capacity = 0 + 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 _admit(self, key, pinned): - if not self.gpu_capacity: - self.metrics["skipped"] += 1 - return False - if len(self.blocks) >= self.gpu_capacity: - victim = next((k for k in self.blocks if k not in pinned), None) - if victim is None: - self.metrics["skipped"] += 1 - return False - self._drop(victim) - # No external directory in this phase: lookup is local to the P. - self.evicted.clear() - self.blocks[key] = {} - self.logical_dtypes[key] = {} - self._block_sizes[key] = 0 - self.tiers[key] = "gpu" - self._gpu_lru[key] = None - self.counts["gpu"] += 1 - 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): + def _init_history(self, budget, cpu_capacity=None): block_bytes = sum(c[0].numel() * c.element_size() for c in self.caches.values()) + if cpu_capacity is None: + cpu_capacity = self.store.cpu_capacity self.history = ( - GPUHistoryStore(budget, block_bytes, self.device) if budget else None + HistoryStore(budget, block_bytes, self.device, cpu_capacity) + if budget + else None ) self._history_task = None self._history_closing = False @@ -59,11 +52,13 @@ def _init_history(self, budget): # 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.gpu_capacity if self.history else 0 + 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(budget, 64 * 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 @@ -72,6 +67,8 @@ def _init_history(self, budget): 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, @@ -118,11 +115,11 @@ def _schedule_history(self, ticket, state): self.metrics["history_skipped_requests"] += 1 return False limit = min( - self.history.gpu_capacity, + self.history.capacity, self.history_pending_bytes // self.history.block_bytes, ) candidates = [] - free = max(0, self.history.gpu_capacity - len(self.history.blocks)) + free = max(0, self.history.capacity - len(self.history.blocks)) for block, key in state.hashes.items(): if key in self.history.ready: self.history.touch(key) @@ -176,7 +173,7 @@ async def _retain_history(self, ticket, state, candidates, deadline): 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 - len(self.history.blocks) + free = self.history.gpu_capacity - self.history.counts["gpu"] chunk_bytes = ( self.history_fill_chunk_bytes if free > 0 @@ -248,15 +245,23 @@ async def load_history(self, requests, leases): cache = self.caches[name] if not set(mapping).issubset(self.history.leases[lease]): raise ValueError("History read outside lease") - reads.append( - LayerRead( - name, - list(mapping), - list(mapping.values()), - tuple(cache.shape[1:]), - cache.dtype, - ) - ) + 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( @@ -271,6 +276,9 @@ async def load_history(self, requests, leases): self._index_tensor(indices, read.destinations) ] = blocks await asyncio.to_thread(torch.cuda.synchronize, self.device) + for key in self.history.leases[lease]: + tier = self.history.tiers[key] + self.metrics[f"history_{tier}_hit_blocks"] += 1 self.metrics["history_hit_blocks"] += len( self.history.leases[lease] ) diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index 026dfd9350..fae03fc661 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -452,7 +452,7 @@ async def load_requests(self, requests): class GPUTransferMixin: - def map_gpu_caches_v1(self, descriptors, budget): + 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: @@ -460,10 +460,8 @@ def map_gpu_caches_v1(self, descriptors, budget): caches = { name: rebuild_cuda_tensor(*desc) for name, desc in descriptors.items() } - import os - runtime_type = GPUTransfer - if os.getenv("XINFERENCE_XAVIER_DIRECT_TEST") == "1": + if direct_handoff: from .direct_handoff import DirectGPUTransfer runtime_type = DirectGPUTransfer 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 17d1c8be73..77ca8e117f 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py @@ -145,7 +145,7 @@ async def copy(buffers, refs): def test_direct_metadata_preserves_local_prefix_offset(connector): - connector._direct_test = True + connector._direct_handoff = True request = SimpleNamespace( request_id="decoder", kv_transfer_params={ @@ -175,7 +175,7 @@ def test_direct_producer_retains_original_engine_blocks( "vllm.v1.request", SimpleNamespace(RequestStatus=SimpleNamespace(FINISHED_ABORTED="aborted")), ) - connector._direct_test = True + connector._direct_handoff = True actor = SimpleNamespace(register_direct_gpu_v1=AsyncMock()) connector._get_transfer_ref = AsyncMock(return_value=actor) request = SimpleNamespace( @@ -273,3 +273,16 @@ async def copy(buffers, refs): assert receiver_indices and all(not indices for indices in receiver_indices) assert not state.indices assert source.poll_direct() == {"producer"} + + +@pytest.mark.parametrize("role", ["prefill", "decode", "hybrid"]) +@pytest.mark.parametrize("budget", [None, 0, 268435456]) +def test_direct_handoff_selection_is_per_model(role, budget, monkeypatch): + from ..transport import uses_direct_handoff + + # No environment setting is consulted, including the old prototype flag. + monkeypatch.setenv("XINFERENCE_XAVIER_DIRECT_TEST", "1") + assert uses_direct_handoff({"role": role, "gpu_cache_bytes": budget}) is ( + role != "hybrid" and budget is not None + ) + assert not uses_direct_handoff(None) diff --git a/xinference/model/llm/vllm/xavier/test/test_direct_history.py b/xinference/model/llm/vllm/xavier/test/test_direct_history.py index b0d0534233..d0d0aa1c81 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_history.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_history.py @@ -6,13 +6,13 @@ import pytest import torch -from ..direct_history import GPUHistoryStore +from ..direct_history import HistoryStore from .test_direct_handoff import direct_runtime def history_runtime(monkeypatch): r = direct_runtime(monkeypatch) - r._init_history(32) + r._init_history(32, cpu_capacity=0) r.history_retention_seconds = 10 r.caches["K"].copy_(torch.arange(16).reshape(8, 2)) return r @@ -87,7 +87,7 @@ async def test_optional_retention_releases_engine_when_skipped(monkeypatch, reas def test_gpu_eviction_respects_lease_without_cpu_demotion(): - store = GPUHistoryStore(8, 4, torch.device("cpu")) + store = HistoryStore(8, 4, torch.device("cpu")) def save(keys): store.stage_blocks(keys, {"K": torch.ones(len(keys), 2, dtype=torch.bfloat16)}) @@ -138,7 +138,7 @@ def test_producer_history_maps_prefix_after_local_hit(connector): from types import SimpleNamespace from unittest.mock import AsyncMock - connector._direct_test = connector._history_enabled = True + connector._direct_handoff = connector._history_enabled = True actor = SimpleNamespace( reserve_direct_history_v1=AsyncMock(), release_direct_history_v1=AsyncMock() ) @@ -169,7 +169,7 @@ def test_unallocated_history_abort_releases_lease(connector, monkeypatch): "vllm.v1.request", SimpleNamespace(RequestStatus=SimpleNamespace(FINISHED_ABORTED="aborted")), ) - connector._direct_test = connector._history_enabled = True + connector._direct_handoff = connector._history_enabled = True actor = SimpleNamespace( reserve_direct_history_v1=AsyncMock(return_value=[101]), release_direct_history_v1=AsyncMock(), @@ -235,7 +235,7 @@ def test_local_prefix_partial_tail_does_not_query_history(connector): from types import SimpleNamespace from unittest.mock import AsyncMock - connector._direct_test = connector._history_enabled = True + connector._direct_handoff = connector._history_enabled = True connector._get_transfer_ref = AsyncMock() request = SimpleNamespace(request_id="p", prompt_token_ids=list(range(41))) assert connector.get_num_new_matched_tokens(request, 32) == (0, False) @@ -248,7 +248,7 @@ async def test_full_history_preserves_hot_content_until_new_content_repeats( monkeypatch, ): r = history_runtime(monkeypatch) - r._init_history(8) + r._init_history(8, cpu_capacity=0) r.history_retention_seconds = 10 async def finish(ticket, blocks, hashes): @@ -274,7 +274,7 @@ async def finish(ticket, blocks, hashes): @pytest.mark.asyncio async def test_probation_is_bounded_and_forgets_old_cold_content(monkeypatch): r = history_runtime(monkeypatch) - r._init_history(4) + r._init_history(4, cpu_capacity=0) r.history_retention_seconds = 10 for key in (100, 101, 102, 103, 101): r.register_direct(str(key), str(key), [1], [key]) @@ -312,7 +312,7 @@ def stage(keys, layers): @pytest.mark.asyncio async def test_fill_batches_grow_but_replacement_batches_stay_small(monkeypatch): r = history_runtime(monkeypatch) - r._init_history(20) + r._init_history(20, cpu_capacity=0) r.history_retention_seconds = 10 r.history_fill_chunk_bytes = 16 r.history_chunk_bytes = 8 @@ -336,3 +336,154 @@ def stage(keys, layers): assert r.poll_direct() == {ticket} assert sizes == [4, 1, 2, 2, 1] assert r.history.ready == set(range(201, 206)) + + +@pytest.mark.parametrize("matched_blocks", [1, 161, 277, 278]) +def test_fixed_history_prefix_length_and_destination_mapping(connector, matched_blocks): + """Keep the 4441-token divergence reproduction's exact prefix boundaries.""" + from types import SimpleNamespace + from unittest.mock import AsyncMock + + connector._direct_handoff = connector._history_enabled = True + request = SimpleNamespace(request_id="p", prompt_token_ids=list(range(4441))) + keys = [ + key for key, _ in connector._build_xavier_hashes(request.prompt_token_ids[:-1]) + ] + actor = SimpleNamespace( + reserve_direct_history_v1=AsyncMock(return_value=keys[:matched_blocks]) + ) + connector._get_transfer_ref = AsyncMock(return_value=actor) + expected = min(matched_blocks * 16, 4440) + assert connector.get_num_new_matched_tokens(request, 0) == (expected, True) + # Non-contiguous physical allocation must preserve logical token order. + destinations = list(reversed(range(278))) + connector.update_state_after_alloc( + request, SimpleNamespace(get_block_ids=lambda: [destinations]), expected + ) + load = connector._requests_need_load["p"] + assert load.local_transfers_by_group == { + 0: {connector._rank: dict(zip(keys[:matched_blocks], destinations))} + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("prefix_blocks", [1, 161, 277, 278]) +@pytest.mark.parametrize("gpu_blocks", [5, 278]) +async def test_fixed_prefix_restore_preserves_raw_bits( + monkeypatch, device, prefix_blocks, gpu_blocks +): + """Real CUDA fences when available; CPU CI exercises the same history path.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("requires a CUDA GPU") + synchronize = torch.cuda.synchronize + r = direct_runtime(monkeypatch) + if device == "cuda": + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + r.device = torch.device(device) + # 24 layer tensors, each including both K and V, with arbitrary BF16 bits. + generator = torch.Generator().manual_seed(42) + r.caches = { + str(layer): torch.randint( + -32768, 32768, (280, 2, 16, 2), dtype=torch.int16, generator=generator + ) + .view(torch.bfloat16) + .to(device) + for layer in range(24) + } + originals = { + name: cache[:prefix_blocks].clone() for name, cache in r.caches.items() + } + block_bytes = sum( + cache[0].numel() * cache.element_size() for cache in r.caches.values() + ) + r._init_history(gpu_blocks * block_bytes, cpu_capacity=278) + r.history_pending_bytes = 278 * block_bytes + r.history_retention_seconds = 30 + r.slab_bytes = 8 * 1024 * 1024 + keys = list(range(1000, 1000 + prefix_blocks)) + r.register_direct("save", "request", list(range(prefix_blocks)), keys) + r.release_direct("save") + await r._history_task + assert r.poll_direct() == {"request"} + assert r.metrics["history_failures"] == 0 + # Engine reuses every original slot before restoring to a reversed mapping. + for cache in r.caches.values(): + cache.zero_() + assert r.reserve_history("restore", keys + [-1]) == keys + destinations = list(reversed(range(2, prefix_blocks + 2))) + await r.load_history( + [{0: {name: dict(zip(keys, destinations)) for name in r.caches}}], ["restore"] + ) + for name, cache in r.caches.items(): + assert torch.equal( + cache[destinations].view(torch.int16), originals[name].view(torch.int16) + ) + assert not cache[:2].count_nonzero() + assert not r.history.leases and not r._history_leases + assert r.history.stats()["cpu_blocks"] == max(0, prefix_blocks - gpu_blocks) + assert r.metrics["history_cpu_hit_blocks"] == max(0, prefix_blocks - gpu_blocks) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +async def test_history_spills_to_cpu_and_restores_mixed_prefix(monkeypatch, device): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("requires a CUDA GPU") + synchronize = torch.cuda.synchronize + r = direct_runtime(monkeypatch) + if device == "cuda": + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + r.device = torch.device(device) + r.caches = { + name: (torch.arange(64, dtype=torch.int16).reshape(8, 2, 4) + shift) + .view(torch.bfloat16) + .to(device) + for name, shift in [("K", 32620), ("V", -32768)] + } + block_bytes = 32 + r._init_history(2 * block_bytes, cpu_capacity=2) + r.history_retention_seconds = 10 + original = {name: cache[:4].clone() for name, cache in r.caches.items()} + r.register_direct("save", "p", [0, 1, 2, 3], [100, 101, 102, 103]) + r.release_direct("save") + await r._history_task + assert r.poll_direct() == {"p"} + assert r.history.tiers == {100: "cpu", 101: "cpu", 102: "gpu", 103: "gpu"} + assert r.history.stats()["demotions"] == 2 + assert r.reserve_history("read", [100, 101, 102, 103, 999]) == [100, 101, 102, 103] + for cache in r.caches.values(): + cache.zero_() + await r.load_history( + [{0: {name: {100: 7, 101: 4, 102: 6, 103: 5} for name in r.caches}}], ["read"] + ) + for name, cache in r.caches.items(): + assert torch.equal( + cache[[7, 4, 6, 5]].view(torch.int16), original[name].view(torch.int16) + ) + assert not cache[:4].count_nonzero() + assert r.metrics["history_cpu_hit_blocks"] == 2 + assert r.metrics["history_gpu_hit_blocks"] == 2 + assert not r.history.leases and not r._history_leases + + +def test_history_eviction_only_after_both_tiers_full_and_preserves_leases(): + store = HistoryStore(8, 4, torch.device("cpu"), cpu_capacity=2) + + def save(key): + store.stage_blocks( + [key], {"K": torch.tensor([[key, -key]], dtype=torch.bfloat16)} + ) + store.publish([key], {"K"}) + + for key in [1, 2, 3, 4]: + save(key) + assert store.ready == {1, 2, 3, 4} + assert store.reserve("reader", [1, 3]) + save(5) + assert store.ready == {1, 3, 4, 5} + assert store.tiers[1] == "cpu" and store.tiers[3] == "gpu" + assert store.counts == {"gpu": 2, "cpu": 2} + assert store.reserve("all", list(store.ready)) + save(6) + assert store.ready == {1, 3, 4, 5} 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 ab9c04d283..d427560e2a 100644 --- a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py @@ -1490,3 +1490,20 @@ def fail(keys, layers): assert set(r.store.blocks) == r.store.ready == {7} assert r.store._block_sizes == {7: 4} assert_gpu_lru_consistent(r.store) + + +@pytest.mark.parametrize("version,nixl", [("0.11.0", True), ("0.11.1", False)]) +def test_gpu_transfer_missing_dependency_fails_without_cpu_fallback( + monkeypatch, version, nixl +): + import importlib.metadata + import importlib.util + + from ..transport import gpu_pool_options + + monkeypatch.setattr(importlib.metadata, "version", lambda name: version) + monkeypatch.setattr( + importlib.util, "find_spec", lambda name: object() if nixl else None + ) + with pytest.raises(RuntimeError, match=r"xoscar\[nixl\]>=0.11.1"): + gpu_pool_options("10.0.0.1:1234", {}) diff --git a/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py b/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py index f9bd64f5a4..c818411b36 100644 --- a/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py +++ b/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py @@ -71,8 +71,8 @@ def pd_cluster(monkeypatch, tmp_path, backend): @pytest.mark.parametrize( "backend,gpu_cache_bytes", [ - pytest.param("xavier", None, id="xavier-cpu"), - pytest.param("xavier", 256 * 1024**2, id="xavier-gpu-async"), + pytest.param("xavier", None, id="xavier-default-history"), + pytest.param("xavier", 0, id="xavier-no-history"), pytest.param("nixl", None, id="nixl"), ], ) @@ -115,6 +115,7 @@ def test_pd_gpu(pd_cluster, backend, gpu_cache_bytes): max_model_len=2048, gpu_memory_utilization=0.5, dtype="float16", + enable_prefix_caching=False, replica_config=[ { "role": role, @@ -164,18 +165,18 @@ def log_since(offset): responses[stream] = chat(str(stream), stream) evidence = log_since(offset) # Require actual KV transfer, not merely a successful recomputation. - if gpu_cache_bytes is not None: - assert "Finished Xavier async KV load: request=" in evidence - elif backend == "xavier": - assert "Stage Xavier V1 blocks" in evidence - assert "Load Xavier V1 blocks" in evidence + if backend == "xavier": + assert "Register Xavier direct handoff" in evidence + assert "Finished Xavier async KV load" in evidence else: assert "calling _read_blocks" in evidence assert re.search(r"and [1-9]\d* requests done recving", evidence) - # Repeated prompts may hit decode's local prefix cache, but must still - # yield the same deterministic answer without stale or corrupted KV. + # Compare exact text only with fixed zero-prefix computation. Historical + # prefix reuse can change floating-point shapes and greedy output. for stream in (False, True): - assert chat(str(stream), stream) == responses[stream] + repeated = chat(str(stream), stream) + if backend == "nixl" or gpu_cache_bytes == 0: + assert repeated == responses[stream] offset = Path(log_path).stat().st_size with ThreadPoolExecutor(max_workers=4) as executor: @@ -187,12 +188,8 @@ def log_since(offset): future.result(timeout=180) pattern = ( r"Finished Xavier async KV load: request=(\S+)" - if gpu_cache_bytes is not None - else ( - r"Load Xavier V1 blocks: request=(\S+)" - if backend == "xavier" - else r"with remote block size \d+ for req (\S+)" - ) + if backend == "xavier" + else r"with remote block size \d+ for req (\S+)" ) loaded_requests = set(re.findall(pattern, log_since(offset))) assert len(loaded_requests) == 4 diff --git a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py index 6b79b84bac..50ab9fd3a9 100644 --- a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py +++ b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py @@ -385,7 +385,9 @@ async def test_gpu_connector_mapping_staging_and_load_fences( connector._gpu_cache_mapped = False connector._registered_kv_caches = {"layer": cache} transfer = SimpleNamespace( - map_gpu_caches_v1=AsyncMock(side_effect=lambda *args: calls.append("map")), + map_gpu_caches_v1=AsyncMock( + side_effect=lambda *args, **kwargs: calls.append("map") + ), stage_gpu_requests_v1=AsyncMock( side_effect=lambda *args: calls.append("stage") ), diff --git a/xinference/model/llm/vllm/xavier/transport.py b/xinference/model/llm/vllm/xavier/transport.py index a4c5eb76e1..ac743f22cf 100644 --- a/xinference/model/llm/vllm/xavier/transport.py +++ b/xinference/model/llm/vllm/xavier/transport.py @@ -88,8 +88,19 @@ def gpu_pool_options(address: str, env: Dict[str, str]) -> Dict[str, str]: if not host: raise ValueError("Cannot determine Xavier NIXL worker host") host = resolve_nixl_host(host, "a reachable worker address") + if not host: + raise ValueError("Cannot determine Xavier NIXL worker host") if ":" in host: host = f"[{host}]" env["UCX_MEMTYPE_CACHE"] = "n" env.setdefault("UCX_TLS", os.environ.get("UCX_TLS", "tcp,cuda_copy,cuda_ipc")) return {"external_address": f"nixl://{host}:0"} + + +def uses_direct_handoff(config: Optional[Dict[str, Any]]) -> bool: + """GPU P/D uses request-scoped handoff; CPU and hybrid caches stay unchanged.""" + return bool( + config + and config.get("gpu_cache_bytes") is not None + and config.get("role") in {"prefill", "decode"} + ) diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 6b1056de48..6f772c2963 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -14,7 +14,6 @@ import asyncio import logging import math -import os import uuid from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Tuple @@ -35,6 +34,7 @@ from .profiling import profile_stage from .snapshot import block_major_view from .transfer import XAVIER_BF16_TRANSPORT_DTYPE, TransferActor +from .transport import uses_direct_handoff from .utils import hash_block_tokens if TYPE_CHECKING: @@ -124,19 +124,23 @@ def __init__( raise ValueError( "Xavier V1 currently supports text-only models without LoRA" ) - self._direct_test = os.getenv("XINFERENCE_XAVIER_DIRECT_TEST") == "1" + self._direct_handoff = uses_direct_handoff(self._xavier_config) self._gpu_budget = self._xavier_config.get("gpu_cache_bytes") - if self._direct_test and ( - self._gpu_budget is None or len(kv_cache_config.kv_cache_groups) != 1 + if self._direct_handoff and ( + self._gpu_budget is None + or len(kv_cache_config.kv_cache_groups) != 1 + or self._xavier_config.get("role") not in ("prefill", "decode") ): - raise ValueError("Direct prototype requires GPU transport and one KV group") + raise ValueError( + "Direct handoff requires P/D roles, GPU transport and one KV group" + ) self._gpu_cache_mapped = False self._gpu_mapping_lock = asyncio.Lock() self._kv_schema: Optional[XavierKVSchema] = None self._rank = int(self._xavier_config.get("rank", 0)) self._is_producer = self._kv_transfer_config.is_kv_producer self._is_consumer = self._kv_transfer_config.is_kv_consumer - self._history_enabled = self._direct_test and bool(self._gpu_budget) + self._history_enabled = self._direct_handoff and bool(self._gpu_budget) if self._history_enabled and self._is_producer: # P may restore its own independent history before computing a suffix. self._is_consumer = True @@ -228,7 +232,7 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: sent = set() - if self._direct_test and self._is_producer and self._gpu_cache_mapped: + if self._direct_handoff and self._is_producer and self._gpu_cache_mapped: assert self._transfer_ref is not None sent = self._call(self._transfer_ref.poll_direct_gpu_v1()) if not self._gpu_load_jobs: @@ -312,7 +316,7 @@ def save_kv_layer( self._pending_store_requests[request.request_id] = request def wait_for_save(self): - if self._direct_test and self._is_producer: + if self._direct_handoff and self._is_producer: self._call(self._ensure_gpu_cache_mapping()) torch.cuda.synchronize() return @@ -342,7 +346,7 @@ def get_num_new_matched_tokens( raise ValueError( "Xavier V1 does not support LoRA, multimodal, embedding or salted prompts" ) - if self._direct_test: + if self._direct_handoff: params = getattr(request, "kv_transfer_params", None) or {} handoff = params.get("xavier_direct") if self._is_producer and getattr(self, "_history_enabled", False): @@ -505,7 +509,7 @@ def build_connector_meta( scheduler_output: SchedulerOutput, ) -> KVConnectorMetadata: meta = XavierConnectorMetadata() - if self._is_producer and not self._direct_test: + if self._is_producer and not self._direct_handoff: self._build_store_meta(scheduler_output, meta) if self._is_consumer: self._build_load_meta(scheduler_output, meta) @@ -516,7 +520,7 @@ def request_finished( request: "Request", block_ids: List[int], ) -> tuple[bool, dict[str, Any] | None]: - if self._direct_test: + if self._direct_handoff: pending = self._requests_need_load.get(request.request_id) if pending is not None and pending.lease.startswith("history:"): if not pending.local_transfers_by_group: @@ -559,6 +563,11 @@ async def register(): ) self._call(register()) + logger.debug( + "Register Xavier direct handoff: request=%s blocks=%s", + request.request_id, + len(blocks), + ) return bool(blocks), { "do_remote_prefill": True, "xavier_direct": { @@ -956,7 +965,9 @@ async def _ensure_gpu_cache_mapping(self): "Xavier GPU transfer requires registered KV caches" ) torch.cuda.synchronize() - await transfer.map_gpu_caches_v1(descriptors, self._gpu_budget) + await transfer.map_gpu_caches_v1( + descriptors, self._gpu_budget, direct_handoff=self._direct_handoff + ) self._gpu_cache_mapped = True return transfer @@ -993,7 +1004,7 @@ async def _submit_gpu_requests(self, requests): entries.append(ranks) torch.cuda.synchronize() async def load(): - if self._direct_test: + if self._direct_handoff: historical = [ i for i, r in enumerate(requests) if r.lease.startswith("history:") ] From bf71d1cb226cc7559b6aa52821daf89eb87778a0 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 20:29:43 +0800 Subject: [PATCH 6/9] fix: preserve direct handoff cleanup after async load rebase --- xinference/model/llm/vllm/xavier/README.md | 12 +++++--- .../vllm/xavier/test/test_direct_handoff.py | 28 +++++++++++++++++++ .../xavier/test/test_review_regressions.py | 3 +- .../model/llm/vllm/xavier/v1_connector.py | 16 +++++++++-- 4 files changed, 52 insertions(+), 7 deletions(-) diff --git a/xinference/model/llm/vllm/xavier/README.md b/xinference/model/llm/vllm/xavier/README.md index d74065f013..a87fb69e54 100644 --- a/xinference/model/llm/vllm/xavier/README.md +++ b/xinference/model/llm/vllm/xavier/README.md @@ -6,13 +6,17 @@ This document covers implementation, diagnostics and contributor validation. For Xavier and vLLM's native `NixlConnector` are separate backends. Selecting `vllm_transfer_backend_type="nixl"` uses the native connector; configuring `xavier_gpu_cache_bytes` does not select it. -The Xavier V1 connector supports vLLM 0.21.0 or newer. The legacy V0 adapter remains for vLLM below 0.11.0; versions 0.11–0.20 are unsupported. GPU integration CI pins vLLM 0.21.0 and NIXL 1.1.0. Match vLLM, Torch and NIXL versions: importing vLLM alone does not validate engine startup. +The Xavier V1 connector supports vLLM 0.21.0 or newer. P/D requires xoscar[nixl]>=0.11.1 in both worker and model environments. The legacy V0 adapter remains for non-P/D deployments with vLLM below 0.11.0; versions 0.11–0.20 are unsupported. GPU integration CI pins vLLM 0.21.0 and NIXL 1.1.0. Match vLLM, Torch and NIXL versions: importing vLLM alone does not validate engine startup. The native backend passes producer handoff metadata to the selected decoder and lets NIXL manage transfer completion and cache release. Failed handoffs fail the request rather than silently recomputing. `VLLM_NIXL_ABORT_REQUEST_TIMEOUT` bounds producer retention if a decoder never consumes a completed prefill. This path does not create Xavier collective actors or force eager execution. Parallel sampling is rejected because child requests cannot safely share one producer transfer lease. ## Cache ownership and transfer -Xavier CPU snapshots are keyed by prompt content, not reusable engine block IDs. Readers reserve complete snapshots during transfer; unavailable snapshots or cache pressure can cause a cache miss and local computation. This differs from a transfer failure, which propagates as an error. Only complete layers are published for remote reuse. CPU capacity is bounded by the engine's KV block capacity. +P/D uses request-scoped GPU-to-GPU handoff by default, with no opt-in flag. The producer keeps its engine blocks until decoder reads and optional history copies finish. Independent history uses 256 MiB of GPU storage per replica by default and spills to CPU, bounded by engine KV block capacity. `xavier_gpu_cache_bytes=0` disables history while keeping direct handoff. CPU-history hits restore to the local GPU in per-layer batches; P-to-D traffic remains xoscar NIXL. Missing NIXL fails launch. + +History uses one writer, bounded staging and a soft deadline. When both tiers are full, new content must be observed again before replacing retained blocks. Active leases prevent eviction. History is restored on P; routing still visits P before D. Direct handoff supports `n=1`. + +The following snapshot-path details apply to hybrid deployments. Xavier CPU snapshots are keyed by prompt content, not reusable engine block IDs. Readers reserve complete snapshots during transfer; unavailable snapshots or cache pressure can cause a cache miss and local computation. This differs from a transfer failure, which propagates as an error. Only complete layers are published for remote reuse. CPU capacity is bounded by the engine's KV block capacity. In the GPU-first path, engine caches are shared with the local TransferActor through CUDA IPC. Retained GPU snapshots use xoscar NIXL; CPU snapshots use batched Gloo transfers. The worker sets `UCX_MEMTYPE_CACHE=n` and defaults `UCX_TLS` to `tcp,cuda_copy,cuda_ipc`, preserving an explicit `UCX_TLS` setting. @@ -20,7 +24,7 @@ The snapshot budget excludes the model, engine KV cache, persistent transfer buf GPU overflow demotes unleased snapshots to CPU. If both tiers have no evictable space, new snapshots are skipped. Layout mismatches and transfer failures propagate instead of silently converting data or promising recomputation. If staging fails, outstanding CUDA work is fenced before unpublished snapshots are discarded; synchronization errors still propagate. -When `xavier_gpu_cache_bytes` is set (including zero), KV loads execute asynchronously in the TransferActor so ready requests can continue decoding. Omitting the setting retains synchronous CPU loading. Cancellation holds destination blocks and source leases until writes complete. Cleanup attempts every lease release and preserves the original load exception if cleanup also fails. Shutdown drains outstanding operations before releasing buffers and snapshot storage, and logs cache placement and transfer counters. Repeated close calls share one close task; the closed runtime remains a guard against remapping. +P/D always uses asynchronous loads. In the hybrid snapshot path, when `xavier_gpu_cache_bytes` is set (including zero), KV loads execute asynchronously in the TransferActor so ready requests can continue decoding. Omitting the setting for hybrid replicas retains synchronous CPU loading. Cancellation holds destination blocks and source leases until writes complete. Cleanup attempts every lease release and preserves the original load exception if cleanup also fails. Shutdown drains outstanding operations before releasing buffers and snapshot storage, and logs cache placement and transfer counters. Repeated close calls share one close task; the closed runtime remains a guard against remapping. If the last request is aborted while an async load is pending, an idle vLLM EngineCore may not poll completion again until the next request arrives. Actor-side writes and source lease release continue independently, but destination block reclamation and connector-side completion cleanup wait for that next engine step or shutdown. @@ -65,6 +69,6 @@ XINFERENCE_TEST_PD_GPU=1 python -m pytest -v \ xinference/model/llm/vllm/xavier/test/test_pd_gpu.py ``` -The test runs synchronous CPU Xavier, GPU-first asynchronous Xavier and native NIXL in real subprocesses, with one GPU per role. It checks streaming and non-streaming responses, repeated prompts and four concurrent requests, then terminates the deployment. CPU Xavier requires producer staging and decoder KV-load log evidence; GPU-first Xavier requires successful async completion logs from `get_finished` after writes and lease cleanup; native NIXL requires `calling _read_blocks` and completed receive logs, including the concurrent requests. +The test runs default Xavier direct handoff with tiered history, Xavier with history disabled, and native NIXL in real subprocesses, with one GPU per role. It checks streaming and non-streaming responses, repeated prompts and four concurrent requests, then terminates the deployment. Local prefix caching is disabled. Exact repeated-text checks are restricted to native NIXL and Xavier with history disabled; fixed-prefix and raw-bit history tests cover history reuse separately. Xavier requires producer registration and successful async completion logs from `get_finished`; native NIXL requires `calling _read_blocks` and completed receive logs, including the concurrent requests. For locally cached weights, set `XINFERENCE_TEST_PD_MODEL_PATH`, `XINFERENCE_TEST_PD_MODEL_NAME` and `XINFERENCE_TEST_PD_MODEL_SIZE` to the path and registered model name/size. The manually triggered **PD GPU integration** GitHub Actions workflow runs the same test on a selected two-GPU runner. 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 77ca8e117f..6c92fbd899 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py @@ -286,3 +286,31 @@ def test_direct_handoff_selection_is_per_model(role, budget, monkeypatch): role != "hybrid" and budget is not None ) assert not uses_direct_handoff(None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lease", ["ticket", "history:ticket"]) +async def test_direct_submission_failure_releases_correct_lease( + connector, connector_module, lease +): + connector._direct_handoff = True + connector._ensure_gpu_cache_mapping = AsyncMock( + side_effect=RuntimeError("map failed") + ) + transfer = SimpleNamespace( + release_direct_history_v1=AsyncMock(), + release_remote_direct_gpu_v1=AsyncMock(), + release_remote_blocks_v1=AsyncMock(), + ) + connector._get_transfer_ref = AsyncMock(return_value=transfer) + request = connector_module.XavierLoadRequest("r", {2: {1: 0}}, lease=lease) + with pytest.raises(RuntimeError, match="map failed"): + await connector._submit_gpu_requests([request]) + if lease.startswith("history:"): + transfer.release_direct_history_v1.assert_awaited_once_with(lease) + transfer.release_remote_direct_gpu_v1.assert_not_awaited() + else: + transfer.release_remote_direct_gpu_v1.assert_awaited_once_with(2, lease) + transfer.release_direct_history_v1.assert_not_awaited() + transfer.release_remote_blocks_v1.assert_not_awaited() + assert not connector._gpu_load_jobs diff --git a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py index 50ab9fd3a9..e4e928876d 100644 --- a/xinference/model/llm/vllm/xavier/test/test_review_regressions.py +++ b/xinference/model/llm/vllm/xavier/test/test_review_regressions.py @@ -439,7 +439,8 @@ async def test_gpu_mapping_serializes_callers_and_retries( release = asyncio.Event() calls = [] - async def register(descriptors, budget): + async def register(descriptors, budget, *, direct_handoff): + assert direct_handoff == connector._direct_handoff calls.append((descriptors, budget)) if len(calls) == 1: entered.set() diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 6f772c2963..4d7a00eb81 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -660,7 +660,15 @@ async def _reserve_load_request(self, request): async def _release_load_request(self, request): transfer = await self._get_transfer_ref() - await transfer.release_remote_blocks_v1(request.lease, request.transfers) + if self._direct_handoff: + if request.lease.startswith("history:"): + await transfer.release_direct_history_v1(request.lease) + else: + await transfer.release_remote_direct_gpu_v1( + next(iter(request.transfers)), request.lease + ) + else: + await transfer.release_remote_blocks_v1(request.lease, request.transfers) async def _query_remote_blocks( self, @@ -1003,10 +1011,13 @@ async def _submit_gpu_requests(self, requests): ranks[rank] = layers entries.append(ranks) torch.cuda.synchronize() + async def load(): if self._direct_handoff: historical = [ - i for i, r in enumerate(requests) if r.lease.startswith("history:") + i + for i, r in enumerate(requests) + if r.lease.startswith("history:") ] direct = [ i @@ -1027,6 +1038,7 @@ async def load(): await transfer.load_gpu_requests_v1( entries, [(r.lease, r.transfers) for r in requests] ) + task = asyncio.create_task(load()) except BaseException: # Worker metadata already owns these leases. Until task creation From fdaffdc3a500ad00c07b1dfc6047ff9871bac144 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 20:30:35 +0800 Subject: [PATCH 7/9] test: distinguish decoder handoffs from producer history loads --- xinference/model/llm/vllm/xavier/test/test_pd_gpu.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py b/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py index c818411b36..de516e6051 100644 --- a/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py +++ b/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py @@ -191,7 +191,16 @@ def log_since(offset): if backend == "xavier" else r"with remote block size \d+ for req (\S+)" ) - loaded_requests = set(re.findall(pattern, log_since(offset))) + evidence = log_since(offset) + loaded_requests = set(re.findall(pattern, evidence)) + if backend == "xavier": + # P-side history restoration also logs async completion. Exclude + # producer IDs so all four completions must come from D handoffs. + producer_requests = set( + re.findall(r"Register Xavier direct handoff: request=(\S+)", evidence) + ) + assert len(producer_requests) == 4 + loaded_requests -= producer_requests assert len(loaded_requests) == 4 print("PD concurrent requests: 4 completed with remote KV loads", flush=True) client.terminate_model(uid) From 121c3f9fa236b463df9bf71d47b8d20cf0ea0498 Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 21:09:22 +0800 Subject: [PATCH 8/9] fix: handle Xavier handoff expiry and decoder preemption --- .../LC_MESSAGES/user_guide/pd_separation.po | 28 +++ .../LC_MESSAGES/user_guide/pd_separation.po | 28 +++ .../LC_MESSAGES/user_guide/pd_separation.po | 28 +++ .../LC_MESSAGES/user_guide/pd_separation.po | 28 +++ .../LC_MESSAGES/user_guide/pd_separation.po | 22 +++ .../LC_MESSAGES/user_guide/pd_separation.po | 22 +++ .../LC_MESSAGES/user_guide/pd_separation.po | 27 +++ .../LC_MESSAGES/user_guide/pd_separation.po | 20 ++ .../LC_MESSAGES/user_guide/pd_separation.po | 20 ++ doc/source/user_guide/pd_separation.rst | 6 +- xinference/core/pd_model.py | 104 +++++------ xinference/core/supervisor.py | 1 - xinference/core/tests/test_pd_launch.py | 2 +- xinference/core/tests/test_pd_model.py | 107 ++++++----- xinference/model/llm/vllm/xavier/README.md | 4 +- .../model/llm/vllm/xavier/direct_handoff.py | 65 ++++++- .../model/llm/vllm/xavier/direct_history.py | 48 ++++- .../model/llm/vllm/xavier/gpu_transfer.py | 18 +- .../vllm/xavier/test/test_direct_handoff.py | 172 +++++++++++++++++- .../vllm/xavier/test/test_direct_history.py | 74 ++++++-- .../model/llm/vllm/xavier/v1_connector.py | 67 ++++--- 21 files changed, 719 insertions(+), 172 deletions(-) diff --git a/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po index 4d00af5892..485dc4d9c6 100644 --- a/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po @@ -341,3 +341,31 @@ msgstr "Direkte Übertragung und gestufter Verlauf" msgid "Direct handoff currently requires ``n=1``." msgstr "Die direkte Übertragung erfordert derzeit ``n=1``." + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"Der Verlauf wird vor der Übergabe an D auf P wiederhergestellt. Bei langen " +"Prompts wird möglicherweise nur ein Präfix behalten; die Aufbewahrung aller " +"Cache-Daten ist nicht garantiert." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"Nicht angenommene Übergaben verfallen nach 120 Sekunden. Verfällt eine " +"Übergabe, bevor D sie annimmt, berechnet D lokal neu. Nach der Annahme " +"bleibt der Quellcache reserviert, bis das Laden abgeschlossen ist oder die " +"Anfrage vor dem Laden abgebrochen wird." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "" +"Direkte Übertragungen und das Laden des Verlaufs können parallel zur " +"Antwortgenerierung anderer Anfragen erfolgen." diff --git a/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po index 6a427c5399..9f9f26bf9a 100644 --- a/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po @@ -338,3 +338,31 @@ msgstr "Transferencia directa e historial por niveles" msgid "Direct handoff currently requires ``n=1``." msgstr "La transferencia directa requiere actualmente ``n=1``." + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"El historial se restaura en P antes de transferirlo a D. Los prompts largos " +"pueden conservar solo un prefijo; no se garantiza la retención de toda la " +"caché." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"Las transferencias no aceptadas caducan tras 120 segundos. Si caducan antes " +"de que D las acepte, D vuelve a calcular localmente. Tras la aceptación, la " +"caché de origen permanece reservada hasta que termine la carga o se cancele " +"la solicitud antes de cargar." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "" +"Las transferencias directas y las cargas del historial pueden solaparse con " +"la generación de respuestas de otras solicitudes." diff --git a/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po index 857c32a38d..76b813c13f 100644 --- a/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po @@ -342,3 +342,31 @@ msgstr "Transfert direct et historique hiérarchisé" msgid "Direct handoff currently requires ``n=1``." msgstr "Le transfert direct nécessite actuellement ``n=1``." + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"L’historique est restauré sur P avant le transfert vers D. Pour les longs " +"prompts, seul un préfixe peut être conservé ; la conservation de tout le " +"cache n’est pas garantie." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"Les transferts non acceptés expirent après 120 secondes. Si un transfert " +"expire avant que D ne l’accepte, D recalcule localement. Après acceptation, " +"le cache source reste réservé jusqu’à la fin du chargement ou à l’annulation" +" de la requête avant le chargement." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "" +"Les transferts directs et le chargement de l’historique peuvent se dérouler " +"en parallèle de la génération des réponses à d’autres requêtes." diff --git a/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po index fb4dd1cb54..ee8829d77b 100644 --- a/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po @@ -339,3 +339,31 @@ msgstr "Trasferimento diretto e cronologia a livelli" msgid "Direct handoff currently requires ``n=1``." msgstr "Il trasferimento diretto richiede attualmente ``n=1``." + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"La cronologia viene ripristinata su P prima del trasferimento a D. Per i " +"prompt lunghi può essere conservato solo un prefisso; la conservazione di " +"tutta la cache non è garantita." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"I trasferimenti non accettati scadono dopo 120 secondi. Se un trasferimento " +"scade prima che D lo accetti, D ricalcola localmente. Dopo l’accettazione, " +"la cache di origine rimane riservata fino al completamento del caricamento o" +" all’annullamento della richiesta prima del caricamento." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "" +"I trasferimenti diretti e il caricamento della cronologia possono avvenire " +"in parallelo alla generazione delle risposte ad altre richieste." diff --git a/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po index eb2137a05a..05d92e0d83 100644 --- a/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po @@ -310,3 +310,25 @@ msgstr "直接転送と階層型履歴" msgid "Direct handoff currently requires ``n=1``." msgstr "直接転送は現在 ``n=1`` を必要とします。" + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"履歴は D への引き渡し前に P で復元されます。長いプロンプトではプレフィックスのみが保持される場合があり、すべてのキャッシュの保持は保証されません。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"未受理の引き渡しは 120 秒で期限切れになります。D が受理する前に期限切れになると、D " +"がローカルで再計算します。受理後は、読み込みが完了するか、読み込み前にリクエストがキャンセルされるまで、転送元のキャッシュが予約されます。" + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "直接転送と履歴の読み込みは、他のリクエストの応答生成と並行して実行できます。" diff --git a/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po index 310058381f..a75868b2cd 100644 --- a/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po @@ -309,3 +309,25 @@ msgstr "직접 전송 및 계층형 이력" msgid "Direct handoff currently requires ``n=1``." msgstr "직접 전송은 현재 ``n=1`` 을 요구합니다." + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"기록은 D로 전달하기 전에 P에서 복원됩니다. 긴 프롬프트는 접두사만 보관될 수 있으며, 모든 캐시의 보관이 보장되지는 않습니다." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"수락되지 않은 전달은 120초 후 만료됩니다. D가 수락하기 전에 만료되면 D가 로컬에서 다시 계산합니다. D가 수락한 후에는 로드가 " +"완료되거나 로드 전에 요청이 취소될 때까지 원본 캐시가 예약됩니다." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "직접 전송과 기록 로드는 다른 요청의 응답 생성과 동시에 실행될 수 있습니다." diff --git a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po index 14c437d731..5eb11685a2 100644 --- a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po @@ -337,3 +337,30 @@ msgstr "Transferência direta e histórico em camadas" msgid "Direct handoff currently requires ``n=1``." msgstr "A transferência direta atualmente exige ``n=1``." + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "" +"O histórico é restaurado em P antes da transferência para D. Prompts longos " +"podem reter apenas um prefixo; a retenção de todo o cache não é garantida." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"Transferências não aceitas expiram após 120 segundos. Se expirarem antes de " +"D aceitá-las, D recalcula localmente. Após a aceitação, o cache de origem " +"permanece reservado até o carregamento terminar ou a solicitação ser " +"cancelada antes do carregamento." + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "" +"As transferências diretas e o carregamento do histórico podem ocorrer em " +"paralelo à geração de respostas de outras solicitações." diff --git a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po index 26598627c7..4a40e26d7c 100644 --- a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po @@ -309,3 +309,23 @@ msgstr "直传与分层历史缓存" msgid "Direct handoff currently requires ``n=1``." msgstr "直传目前要求 ``n=1``。" + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "历史缓存在 P 上恢复,然后交接给 D。长提示词可能只保留前缀;不保证保留所有缓存。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"尚未被接收的交接在 120 秒后过期。如果交接在 D 接收前过期,D 会在本地重新计算。D 接收后,源缓存会保留到加载完成,或请求在加载前被取消。" + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "直传和历史缓存加载可以与其他请求的响应生成重叠执行。" diff --git a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po index 85d5470d81..c240110e39 100644 --- a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po @@ -304,3 +304,23 @@ msgstr "直傳與分層歷史快取" msgid "Direct handoff currently requires ``n=1``." msgstr "直傳目前要求 ``n=1``。" + +#: user_guide/pd_separation.rst +msgid "" +"History is restored on P before the handoff to D. Long prompts may retain " +"only a prefix; cache retention is best effort." +msgstr "歷史快取在 P 上還原,然後交接給 D。長提示詞可能只保留前綴;不保證保留所有快取。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. If a handoff expires before D " +"accepts it, D recomputes locally. Once D accepts it, the source cache stays " +"reserved until the load finishes or the request is cancelled before loading." +msgstr "" +"尚未被接收的交接在 120 秒後過期。如果交接在 D 接收前過期,D 會在本機重新計算。D 接收後,來源快取會保留到載入完成,或請求在載入前被取消。" + +#: user_guide/pd_separation.rst +msgid "" +"Direct transfers and history loads can overlap with response generation for " +"other requests." +msgstr "直接傳輸和歷史快取載入可以與其他請求的回應生成重疊執行。" diff --git a/doc/source/user_guide/pd_separation.rst b/doc/source/user_guide/pd_separation.rst index 42ddf7b058..d4833a626c 100644 --- a/doc/source/user_guide/pd_separation.rst +++ b/doc/source/user_guide/pd_separation.rst @@ -136,9 +136,13 @@ Direct handoff and tiered history Direct handoff currently requires ``n=1``. +History is restored on P before the handoff to D. Long prompts may retain only a prefix; cache retention is best effort. + +Unclaimed handoffs expire after 120 seconds. If a handoff expires before D accepts it, D recomputes locally. Once D accepts it, the source cache stays reserved until the load finishes or the request is cancelled before loading. + The budget applies only to Xavier's GPU cache. Leave additional GPU memory for the model, vLLM's own KV cache and transfers. When this cache fills, Xavier stores reusable data in CPU memory; if no cache space is available, new data may not be retained for reuse. -With GPU-first caching, cache transfers can overlap with response generation for other requests. +Direct transfers and history loads can overlap with response generation for other requests. Transfer failures and incompatible KV cache layouts raise errors instead of silently recomputing the request. diff --git a/xinference/core/pd_model.py b/xinference/core/pd_model.py index 4c45b71aac..00e3f46805 100644 --- a/xinference/core/pd_model.py +++ b/xinference/core/pd_model.py @@ -84,15 +84,15 @@ def __init__( [List[xo.ActorRefType["ModelActor"]]], SchedulingPolicy ] = RoundRobinSchedulingPolicy, transport_backend: str = "xavier", - direct_handoff: bool = False, ): super().__init__() # Prefill request map, used to skip the timeout task for specific request id. self._request_set: Set[str] = set() + self._direct_transfers: dict[str, dict] = {} self._model_uid = model_uid self._transport_backend = transport_backend - self._direct_handoff = transport_backend == "xavier" and direct_handoff + self._direct_handoff = transport_backend == "xavier" # 使用字典存储副本:{replica_uid: actor_ref} self._prefill_replicas: Dict[str, xo.ActorRefType["ModelActor"]] = {} @@ -246,25 +246,22 @@ async def free_prefill_model_cache(self, request_id: str): logger.debug( f"[PDModelActor] Free prefill model cache for request {request_id}" ) + handoff = self._direct_transfers.pop(request_id, None) + if handoff and handoff.get("ticket"): + try: + from ..model.llm.vllm.xavier.transfer import TransferActor + + ref = await xo.actor_ref( + address=handoff["address"], + uid=f"{TransferActor.default_uid()}-{handoff['rank']}", + ) + # D owns claimed tickets until its writes drain. Only reclaim + # a handoff that was never accepted by D (router failure/abort). + await ref.abandon_direct_gpu_v1(handoff["ticket"]) + except Exception: + logger.warning("Failed to abandon Xavier handoff", exc_info=True) if request_id in self._request_set: self._request_set.remove(request_id) - if ( - self._prefill_replicas - and self._transport_backend == "xavier" - and not self._direct_handoff - ): - await asyncio.gather( - *[ - actor_call( - model, - "free_model_cache", - request_id, - _rpc_operation_request_id=request_id, - ) - for model in self._prefill_replicas.values() - ], - return_exceptions=True, - ) else: logger.warning( f"[request {request_id}] Prefill model cache has been freed already" @@ -292,12 +289,7 @@ async def _infer(self, method, inputs, *args, **kwargs): args = (kwargs.pop("generate_config"),) if args and args[0] is not None and not isinstance(args[0], dict): raise TypeError("Generation config must be a dict or None") - if ( - (self._transport_backend == "nixl" or self._direct_handoff) - and args - and args[0] - and args[0].get("n", 1) != 1 - ): + if args and args[0] and args[0].get("n", 1) != 1: # Handoff leases cover one decoder, not parallel sampling children. raise ValueError("PD KV handoff currently requires n=1") prefill = self._prefill_policy.schedule() @@ -307,12 +299,11 @@ async def _infer(self, method, inputs, *args, **kwargs): prefill_args = [{}] + prefill_args[1:] prefill_args[0]["max_tokens"] = 1 prefill_args[0]["stream"] = False - if self._transport_backend == "nixl" or self._direct_handoff: - prefill_args[0]["n"] = 1 - prefill_args[0]["_pd_kv_transfer_params"] = { - "do_remote_decode": True, - "do_remote_prefill": False, - } + prefill_args[0]["n"] = 1 + prefill_args[0]["_pd_kv_transfer_params"] = { + "do_remote_decode": True, + "do_remote_prefill": False, + } prefill_kwargs = copy.deepcopy(kwargs) if isinstance(prefill_kwargs.get("raw_params"), dict): prefill_kwargs["raw_params"].update(max_tokens=1, stream=False) @@ -325,41 +316,31 @@ async def _infer(self, method, inputs, *args, **kwargs): if hasattr(result, "__aiter__"): async for _ in result: pass - if request_id not in self._request_set: - raise asyncio.CancelledError(f"PD request {request_id} was aborted") logger.debug( "PD prefill complete: request=%s backend=%s elapsed_s=%.6f", request_id, self._transport_backend, time.perf_counter() - prefill_start, ) - if self._transport_backend == "nixl" or self._direct_handoff: - payload = ( - json.loads(result) if isinstance(result, (bytes, str)) else result - ) - transfer = ( - payload.get("_pd_kv_transfer_params") - if isinstance(payload, dict) - else None - ) - if not isinstance(transfer, dict) or not transfer.get( - "do_remote_prefill" - ): - raise RuntimeError("PD prefill did not return KV transfer metadata") - decode_args = list(copy.deepcopy(args)) - if not decode_args or decode_args[0] is None: - decode_args = [{}] + decode_args[1:] - decode_args[0]["_pd_kv_transfer_params"] = transfer - args = tuple(decode_args) - else: - await actor_call( - decode, - "set_unpin_handler", - self._model_uid, - request_id, - self.address, - _rpc_operation_request_id=request_id, - ) + payload = json.loads(result) if isinstance(result, (bytes, str)) else result + transfer = ( + payload.get("_pd_kv_transfer_params") + if isinstance(payload, dict) + else None + ) + if self._direct_handoff and isinstance(transfer, dict): + handoff = transfer.get("xavier_direct") + if handoff: + self._direct_transfers[request_id] = handoff + if not isinstance(transfer, dict) or not transfer.get("do_remote_prefill"): + raise RuntimeError("PD prefill did not return KV transfer metadata") + decode_args = list(copy.deepcopy(args)) + if not decode_args or decode_args[0] is None: + decode_args = [{}] + decode_args[1:] + decode_args[0]["_pd_kv_transfer_params"] = transfer + args = tuple(decode_args) + if request_id not in self._request_set: + raise asyncio.CancelledError(f"PD request {request_id} was aborted") result = await actor_call( decode, method, @@ -372,6 +353,8 @@ async def _infer(self, method, inputs, *args, **kwargs): await self.free_prefill_model_cache(request_id) raise if not hasattr(result, "__aiter__"): + # A successful D response has already completed its handoff. + self._direct_transfers.pop(request_id, None) await self.free_prefill_model_cache(request_id) return result @@ -379,6 +362,7 @@ async def stream(): try: async for chunk in result: yield chunk + self._direct_transfers.pop(request_id, None) finally: try: if hasattr(result, "aclose"): diff --git a/xinference/core/supervisor.py b/xinference/core/supervisor.py index 3972cfac1c..1687d695f9 100644 --- a/xinference/core/supervisor.py +++ b/xinference/core/supervisor.py @@ -3526,7 +3526,6 @@ async def _launch_model(): PDModelActor, model_uid, transport_backend=transport_backend, - direct_handoff=enable_xavier and gpu_cache_bytes is not None, address=self.address, uid=f"{model_uid}-{PDModelActor.default_uid()}", ) diff --git a/xinference/core/tests/test_pd_launch.py b/xinference/core/tests/test_pd_launch.py index dde35c29e8..c913732b9b 100644 --- a/xinference/core/tests/test_pd_launch.py +++ b/xinference/core/tests/test_pd_launch.py @@ -340,7 +340,7 @@ async def test_gpu_pd_defaults_to_direct_handoff(launch_runtime, budget): await supervisor.launch_builtin_model( **launch_kwargs(), xavier_gpu_cache_bytes=budget ) - assert actors["PDModelActor"].constructor_kwargs["direct_handoff"] is True + assert actors["PDModelActor"].constructor_kwargs["transport_backend"] == "xavier" for worker in workers: config = worker.launch_builtin_model.call_args.kwargs assert config["xavier_config"]["gpu_cache_bytes"] == ( diff --git a/xinference/core/tests/test_pd_model.py b/xinference/core/tests/test_pd_model.py index 7a7b16e6ff..d6002687b0 100644 --- a/xinference/core/tests/test_pd_model.py +++ b/xinference/core/tests/test_pd_model.py @@ -8,7 +8,7 @@ from ..pd_model import PDModelActor, RoundRobinSchedulingPolicy from ..replica_config import ReplicaConfig, validate_pd_replica_configs -from ..rpc_context import RPC_METADATA_KEY, get_current_rpc_metadata, rpc_context +from ..rpc_context import rpc_context def test_round_robin_updates(): @@ -56,16 +56,12 @@ async def router(): model.abort_request = AsyncMock(return_value="DONE") await actor.add_prefill_actor("p", prefill) await actor.add_decode_actor("d", decode) + prefill.chat.return_value = prefill.generate.return_value = { + "_pd_kv_transfer_params": {"do_remote_prefill": True} + } return actor, prefill, decode -def _assert_free_model_cache_called(model, request_id): - model.free_model_cache.assert_awaited_once() - call = model.free_model_cache.await_args - assert call.args == (request_id,) - assert call.kwargs[RPC_METADATA_KEY]["operation_request_id"] == request_id - - @pytest.mark.asyncio @pytest.mark.parametrize("method", ["chat", "generate"]) async def test_infer_preserves_decode_config(router, method): @@ -83,7 +79,7 @@ async def test_infer_preserves_decode_config(router, method): assert d_call.args[1]["max_tokens"] == 32 assert p_call.kwargs["request_id"] == d_call.kwargs["request_id"] == "r" assert not actor._request_set - _assert_free_model_cache_called(prefill, "r") + prefill.free_model_cache.assert_not_awaited() @pytest.mark.asyncio @@ -105,54 +101,20 @@ async def chunks(): assert await anext(stream) == b"first" await stream.aclose() assert closed and not actor._request_set - _assert_free_model_cache_called(prefill, "r") + prefill.free_model_cache.assert_not_awaited() @pytest.mark.asyncio -@pytest.mark.parametrize("failure", ["prefill", "decode", "handler", "cancel"]) +@pytest.mark.parametrize("failure", ["prefill", "decode", "cancel"]) async def test_failure_releases_cache(router, failure): actor, prefill, decode = router error = asyncio.CancelledError() if failure == "cancel" else RuntimeError("failed") target = prefill.chat if failure == "prefill" else decode.chat - if failure == "handler": - target = decode.set_unpin_handler target.side_effect = error with pytest.raises(type(error)): await actor._infer("chat", [], {"max_tokens": 32}, request_id="r") assert not actor._request_set - _assert_free_model_cache_called(prefill, "r") - - -@pytest.mark.asyncio -async def test_free_prefill_model_cache_propagates_operation_metadata(): - calls = [] - - class Prefill: - @rpc_context - async def free_model_cache(self, request_id, **kwargs): - calls.append((request_id, get_current_rpc_metadata(), kwargs)) - - actor = PDModelActor("pd") - actor._request_set.add("operation-id") - actor._prefill_replicas = {"p": Prefill()} - - await actor.free_prefill_model_cache( - "operation-id", - **{ - RPC_METADATA_KEY: { - "version": 1, - "correlation_id": "http-id", - "actor_call_id": "router-call", - } - }, - ) - - request_id, metadata, kwargs = calls[0] - assert request_id == "operation-id" - assert kwargs == {} - assert metadata.correlation_id == "http-id" - assert metadata.operation_request_id == "operation-id" - assert metadata.parent_call_id == "router-call" + prefill.free_model_cache.assert_not_awaited() @pytest.mark.asyncio @@ -181,6 +143,8 @@ async def free_model_cache(self, request_id): @xo.generator @rpc_context async def chat(self, messages, config, **kwargs): + if config.get("_pd_kv_transfer_params", {}).get("do_remote_decode"): + return {"_pd_kv_transfer_params": {"do_remote_prefill": True}} if not config.get("stream"): return b"non-stream" @@ -215,7 +179,7 @@ async def test_aborted_prefill_does_not_start_decode(router): async def wait(*args, **kwargs): started.set() await finish.wait() - return b"ok" + return {"_pd_kv_transfer_params": {"do_remote_prefill": True}} prefill.chat.side_effect = wait task = asyncio.create_task(actor._infer("chat", [], {}, request_id="r")) @@ -274,7 +238,7 @@ async def test_abort_failure_still_cleans_request(router): with pytest.raises(RuntimeError, match="replica offline"): await actor.abort_request("r") assert not actor._request_set - _assert_free_model_cache_called(prefill, "r") + prefill.free_model_cache.assert_not_awaited() @pytest.mark.asyncio @@ -410,3 +374,50 @@ async def test_duplicate_generation_config_rejected_before_dispatch(router): prefill.chat.assert_not_awaited() decode.chat.assert_not_awaited() assert not actor._request_set + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["decode", "cancel", "abort_before_decode"]) +async def test_direct_router_releases_unclaimed_handoff(router, monkeypatch, failure): + actor, prefill, decode = router + actor._direct_handoff = True + transfer = { + "do_remote_prefill": True, + "xavier_direct": {"ticket": "t", "rank": 0, "address": "127.0.0.1:1234"}, + } + + async def finish_prefill(*args, **kwargs): + if failure == "abort_before_decode": + actor._request_set.discard("r") + return {"_pd_kv_transfer_params": transfer} + + prefill.chat.side_effect = finish_prefill + decode.chat.side_effect = ( + RuntimeError("decode failed") + if failure == "decode" + else asyncio.CancelledError() + ) + peer = MagicMock(abandon_direct_gpu_v1=AsyncMock()) + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + with pytest.raises(RuntimeError if failure == "decode" else asyncio.CancelledError): + await actor._infer("chat", [], {}, request_id="r") + peer.abandon_direct_gpu_v1.assert_awaited_once_with("t") + assert not actor._direct_transfers and not actor._request_set + if failure == "abort_before_decode": + decode.chat.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_completed_direct_request_does_not_send_abandon_rpc(router, monkeypatch): + actor, prefill, decode = router + prefill.chat.return_value = { + "_pd_kv_transfer_params": { + "do_remote_prefill": True, + "xavier_direct": {"ticket": "t", "rank": 0, "address": "127.0.0.1:1234"}, + } + } + lookup = AsyncMock() + monkeypatch.setattr(xo, "actor_ref", lookup) + await actor._infer("chat", [], {}, request_id="r") + lookup.assert_not_awaited() + assert not actor._direct_transfers diff --git a/xinference/model/llm/vllm/xavier/README.md b/xinference/model/llm/vllm/xavier/README.md index a87fb69e54..fd9488e762 100644 --- a/xinference/model/llm/vllm/xavier/README.md +++ b/xinference/model/llm/vllm/xavier/README.md @@ -14,7 +14,9 @@ The native backend passes producer handoff metadata to the selected decoder and P/D uses request-scoped GPU-to-GPU handoff by default, with no opt-in flag. The producer keeps its engine blocks until decoder reads and optional history copies finish. Independent history uses 256 MiB of GPU storage per replica by default and spills to CPU, bounded by engine KV block capacity. `xavier_gpu_cache_bytes=0` disables history while keeping direct handoff. CPU-history hits restore to the local GPU in per-layer batches; P-to-D traffic remains xoscar NIXL. Missing NIXL fails launch. -History uses one writer, bounded staging and a soft deadline. When both tiers are full, new content must be observed again before replacing retained blocks. Active leases prevent eviction. History is restored on P; routing still visits P before D. Direct handoff supports `n=1`. +Unclaimed producer tickets expire after 120 seconds. D atomically claims a live ticket before reporting a remote hit; expiry is a scheduler miss, not a worker transfer error. Claimed tickets stay pinned until D releases them. Allocation marks the handoff consumed, so preempted decode requests recompute locally rather than rereading the ticket. Router cleanup abandons only unclaimed tickets, while allocated loads drain before release. + +History uses one writer, a 64 MiB candidate-prefix cap (also limited by total history capacity), and a 10 ms soft deadline. Long prompts retain the leading admissible prefix instead of being skipped entirely; all engine blocks remain owned until the active chunk drains. `history_capacity_limited_blocks` counts candidates beyond the cap, `history_deadline_dropped_blocks` and `history_closing_dropped_blocks` count remaining candidates abandoned at the deadline or shutdown, and `history_skipped_requests` counts a busy writer. When both tiers are full, new content must be observed again before replacing retained blocks. Active leases prevent eviction. History is restored on P; routing still visits P before D. Direct handoff supports `n=1`. The following snapshot-path details apply to hybrid deployments. Xavier CPU snapshots are keyed by prompt content, not reusable engine block IDs. Readers reserve complete snapshots during transfer; unavailable snapshots or cache pressure can cause a cache miss and local computation. This differs from a transfer failure, which propagates as an error. Only complete layers are published for remote reuse. CPU capacity is bounded by the engine's KV block capacity. diff --git a/xinference/model/llm/vllm/xavier/direct_handoff.py b/xinference/model/llm/vllm/xavier/direct_handoff.py index d8f844b45e..f7b23a205f 100644 --- a/xinference/model/llm/vllm/xavier/direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/direct_handoff.py @@ -3,6 +3,7 @@ """Request-scoped GPU handoff with independent tiered history.""" import asyncio +import logging import time from dataclasses import dataclass, field @@ -13,19 +14,21 @@ 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 - held_blocks: int = 0 class DirectGPUTransfer(DirectHistoryMixin, GPUTransfer): @@ -69,9 +72,7 @@ async def _close(self): self.direct_requests.clear() self.finished_sending.clear() - def register_direct( - self, ticket, request_id, blocks, hashes=None, held_blocks=None - ): + 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( @@ -87,10 +88,28 @@ def register_direct( set(blocks), time.monotonic() + 120, hashes=dict(zip(blocks, hashes or [])), - held_blocks=len(blocks) if held_blocks is None else held_blocks, ) self.metrics["direct_registered"] += 1 + def claim_direct(self, ticket): + state = self.direct_requests.get(ticket) + if state is None or state.released: + return False + if not state.claimed and time.monotonic() >= state.deadline: + self.metrics["direct_expired"] += 1 + state.retention_attempted = True + self.release_direct(ticket) + return False + # Ownership passes to D. Scheduler abort or load completion releases it; + # producer polling must not race queued allocation or a multi-slab read. + state.claimed = True + 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: @@ -106,7 +125,11 @@ def release_direct(self, ticket): def poll_direct(self): for ticket, state in list(self.direct_requests.items()): - if not state.released and time.monotonic() >= state.deadline: + if ( + not state.released + and not state.claimed + and time.monotonic() >= state.deadline + ): self.metrics["direct_expired"] += 1 state.retention_attempted = True self.release_direct(ticket) @@ -162,6 +185,35 @@ async def send_direct(self, ticket, reads, remote_ref, slab_bytes): self.release_direct(ticket) async def load_direct(self, requests, tickets): + failed = False + try: + await self._load_direct(requests, tickets) + except BaseException: + failed = True + raise + finally: + + async def release(ranks, ticket): + for rank in ranks: + from .transfer import TransferActor + + sender = await xo.actor_ref( + address=self.actor._world_addresses[rank], + uid=f"{TransferActor.default_uid()}-{rank}", + ) + await sender.release_direct_gpu_v1(ticket) + + results = await asyncio.gather( + *(release(ranks, ticket) for ranks, ticket in zip(requests, tickets)), + 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 @@ -206,4 +258,3 @@ async def load_direct(self, requests, tickets): finally: await asyncio.to_thread(torch.cuda.synchronize, self.device) indices.clear() - await sender.release_direct_gpu_v1(ticket) diff --git a/xinference/model/llm/vllm/xavier/direct_history.py b/xinference/model/llm/vllm/xavier/direct_history.py index 5126433cf4..44aced3407 100644 --- a/xinference/model/llm/vllm/xavier/direct_history.py +++ b/xinference/model/llm/vllm/xavier/direct_history.py @@ -34,10 +34,9 @@ def read(self, layer, keys, device=None): class DirectHistoryMixin: - def _init_history(self, budget, cpu_capacity=None): + def _init_history(self, budget): block_bytes = sum(c[0].numel() * c.element_size() for c in self.caches.values()) - if cpu_capacity is None: - cpu_capacity = self.store.cpu_capacity + cpu_capacity = self.store.cpu_capacity self.history = ( HistoryStore(budget, block_bytes, self.device, cpu_capacity) if budget @@ -73,6 +72,9 @@ def _init_history(self, budget, cpu_capacity=None): 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): @@ -93,7 +95,8 @@ def reserve_history(self, lease, keys): break matched.append(key) if matched: - assert self.history.reserve(lease, matched) + if not self.history.reserve(lease, matched): + return [] self._history_leases[lease] = time.monotonic() + 120 return matched @@ -111,9 +114,6 @@ def _schedule_history(self, ticket, state): state.retention_attempted = True if self._history_closing or self.history is None or not state.hashes: return False - if state.held_blocks * self.history.block_bytes > self.history_pending_bytes: - self.metrics["history_skipped_requests"] += 1 - return False limit = min( self.history.capacity, self.history_pending_bytes // self.history.block_bytes, @@ -136,6 +136,8 @@ def _schedule_history(self, ticket, state): free = max(0, free - 1) if repeated: self.metrics["history_admission_reused_blocks"] += 1 + else: + self.metrics["history_capacity_limited_blocks"] += 1 else: self.metrics["history_admission_rejected_blocks"] += 1 if not candidates: @@ -158,7 +160,7 @@ def done(future): exc = future.exception() logger.error( "GPU history retention failed", - exc_info=(type(exc), exc, exc.__traceback__), + exc_info=exc, ) task.add_done_callback(done) @@ -189,13 +191,21 @@ async def _retain_history(self, ticket, state, candidates, deadline): if free > 0: count = min(count, free) chunk = candidates[start : start + count] - start += count while self.send_lock.locked() or self.recv_lock.locked(): 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 await asyncio.sleep(0.001) 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 + 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] @@ -232,6 +242,25 @@ async def _retain_history(self, ticket, state, candidates, deadline): 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 = {} @@ -287,4 +316,3 @@ async def load_history(self, requests, leases): await asyncio.to_thread(torch.cuda.synchronize, self.device) indices.clear() self._history_reading.discard(lease) - self.release_history(lease) diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index fae03fc661..a4712d8f7d 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -493,12 +493,20 @@ async def close_gpu_caches_v1(self): # 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, held_blocks=None - ): - self._gpu_transfer.register_direct( - ticket, request_id, blocks, hashes, held_blocks + 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) 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 6c92fbd899..8bff920b9a 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py @@ -2,7 +2,7 @@ # Licensed under the Apache License, Version 2.0. import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest import torch @@ -157,6 +157,11 @@ def test_direct_metadata_preserves_local_prefix_offset(connector): } }, ) + connector._get_transfer_ref = AsyncMock( + return_value=SimpleNamespace( + claim_remote_direct_gpu_v1=AsyncMock(return_value=True) + ) + ) assert connector.get_num_new_matched_tokens(request, 16) == (16, True) blocks = SimpleNamespace(get_block_ids=lambda: [[4, 5]]) connector.update_state_after_alloc(request, blocks, 16) @@ -164,9 +169,10 @@ def test_direct_metadata_preserves_local_prefix_offset(connector): assert load.local_transfers_by_group == {0: {0: {3: 5}}} +@pytest.mark.parametrize("history", [False, True]) @pytest.mark.parametrize("prompt_tokens", [33, 30]) def test_direct_producer_retains_original_engine_blocks( - connector, monkeypatch, prompt_tokens + connector, monkeypatch, prompt_tokens, history ): import sys @@ -176,6 +182,8 @@ def test_direct_producer_retains_original_engine_blocks( SimpleNamespace(RequestStatus=SimpleNamespace(FINISHED_ABORTED="aborted")), ) connector._direct_handoff = True + connector._history_enabled = history + connector._gpu_budget = 32 if history else 0 actor = SimpleNamespace(register_direct_gpu_v1=AsyncMock()) connector._get_transfer_ref = AsyncMock(return_value=actor) request = SimpleNamespace( @@ -188,9 +196,16 @@ def test_direct_producer_retains_original_engine_blocks( assert retained handoff = metadata["xavier_direct"] assert handoff["blocks"] == [4, 5] and handoff["tokens"] == prompt_tokens - 1 + history_args = () + if history: + hashes = connector._build_xavier_hashes(request.prompt_token_ids[:-1]) + history_args = ([key for key, _ in hashes[:2]],) actor.register_direct_gpu_v1.assert_awaited_once_with( - handoff["ticket"], "producer", [4, 5] + handoff["ticket"], "producer", [4, 5], *history_args ) + metadata = connector.build_connector_meta(SimpleNamespace(num_scheduled_tokens={})) + assert metadata.direct_sends == {"producer"} + assert not metadata.direct_store and not connector._direct_sends @pytest.mark.asyncio @@ -277,11 +292,9 @@ async def copy(buffers, refs): @pytest.mark.parametrize("role", ["prefill", "decode", "hybrid"]) @pytest.mark.parametrize("budget", [None, 0, 268435456]) -def test_direct_handoff_selection_is_per_model(role, budget, monkeypatch): +def test_direct_handoff_selection_is_per_model(role, budget): from ..transport import uses_direct_handoff - # No environment setting is consulted, including the old prototype flag. - monkeypatch.setenv("XINFERENCE_XAVIER_DIRECT_TEST", "1") assert uses_direct_handoff({"role": role, "gpu_cache_bytes": budget}) is ( role != "hybrid" and budget is not None ) @@ -314,3 +327,150 @@ async def test_direct_submission_failure_releases_correct_lease( transfer.release_direct_history_v1.assert_not_awaited() transfer.release_remote_blocks_v1.assert_not_awaited() assert not connector._gpu_load_jobs + + +def decoder_request(): + return SimpleNamespace( + request_id="decoder", + kv_transfer_params={ + "do_remote_prefill": True, + "xavier_direct": { + "rank": 0, + "ticket": "ticket", + "tokens": 32, + "blocks": [2, 3], + }, + }, + ) + + +def test_preempted_decoder_does_not_reload_consumed_ticket(connector): + connector._direct_handoff = True + connector._is_producer = False + actor = SimpleNamespace(claim_remote_direct_gpu_v1=AsyncMock(return_value=True)) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = decoder_request() + assert connector.get_num_new_matched_tokens(request, 0) == (32, True) + connector.update_state_after_alloc( + request, SimpleNamespace(get_block_ids=lambda: [[4, 5]]), 32 + ) + connector._requests_need_load.clear() # worker consumed the scheduled metadata + assert connector.get_num_new_matched_tokens(request, 0) == (0, False) + actor.claim_remote_direct_gpu_v1.assert_awaited_once_with(0, "ticket") + assert not connector._requests_need_load + + +def test_expired_handoff_is_scheduler_miss(connector, monkeypatch): + source = direct_runtime(monkeypatch) + source.register_direct("ticket", "producer", [2, 3]) + source.direct_requests["ticket"].deadline = 0 + connector._direct_handoff = True + connector._is_producer = False + actor = SimpleNamespace( + claim_remote_direct_gpu_v1=AsyncMock( + side_effect=lambda rank, ticket: source.claim_direct(ticket) + ) + ) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = decoder_request() + assert connector.get_num_new_matched_tokens(request, 0) == (0, False) + assert not connector._requests_need_load + assert request.kv_transfer_params["do_remote_prefill"] is False + assert source.poll_direct() == {"producer"} + + +@pytest.mark.asyncio +async def test_claim_pins_queued_handoff_until_decoder_release(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + source.caches["K"].fill_(7) + source.register_direct("ticket", "producer", [2]) + assert source.claim_direct("ticket") + source.direct_requests["ticket"].deadline = 0 + source.abandon_direct("ticket") # router cancellation cannot release D's read + assert not source.poll_direct() + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer_for(source))) + + async def copy(buffers, refs): + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(xo, "copy_to", copy) + await dest.load_direct([{0: {"K": {2: 5}}}], ["ticket"]) + assert dest.caches["K"][5].tolist() == [7, 7] + assert source.poll_direct() == {"producer"} + + +@pytest.mark.parametrize("stage", ["unqueried", "claimed", "allocated"]) +def test_abort_before_load_releases_only_unallocated_ticket(connector, stage): + connector._direct_handoff = True + connector._is_producer = False + actor = SimpleNamespace( + claim_remote_direct_gpu_v1=AsyncMock(return_value=True), + release_remote_direct_gpu_v1=AsyncMock(), + ) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = decoder_request() + if stage != "unqueried": + assert connector.get_num_new_matched_tokens(request, 0) == (32, True) + if stage == "allocated": + connector.update_state_after_alloc( + request, SimpleNamespace(get_block_ids=lambda: [[4, 5]]), 32 + ) + assert connector.request_finished(request, []) == (False, None) + if stage == "allocated": + actor.release_remote_direct_gpu_v1.assert_not_awaited() + assert "decoder" in connector._requests_need_load + else: + actor.release_remote_direct_gpu_v1.assert_awaited_once_with(0, "ticket") + + +@pytest.mark.asyncio +async def test_direct_batch_failure_releases_later_tickets_and_preserves_error( + monkeypatch, +): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + for ticket in ("a", "b"): + source.register_direct(ticket, ticket, [1]) + peer = peer_for(source) + + async def release(ticket): + source.release_direct(ticket) + if ticket == "a": + raise RuntimeError("cleanup error") + + peer.release_direct_gpu_v1.side_effect = release + peer.send_direct_gpu_v1.side_effect = ValueError("original read error") + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + with pytest.raises(ValueError, match="original read error"): + await dest.load_direct([{0: {"K": {1: 5}}}] * 2, ["a", "b"]) + assert source.poll_direct() == {"a", "b"} + assert peer.release_direct_gpu_v1.await_count == 2 + + +def test_producer_metadata_gates_idle_rpc_and_fence( + connector, connector_module, monkeypatch +): + connector._direct_handoff = True + connector._is_consumer = False + metadata = connector_module.XavierConnectorMetadata() + monkeypatch.setattr( + connector, "_get_connector_metadata", lambda: metadata, raising=False + ) + sync = Mock() + monkeypatch.setattr(torch.cuda, "synchronize", sync) + connector._ensure_gpu_cache_mapping = AsyncMock() + connector._transfer_ref = SimpleNamespace( + poll_direct_gpu_v1=AsyncMock(return_value={"p"}) + ) + connector.start_load_kv(None) + connector.wait_for_save() + assert connector.get_finished(set()) == (set(), set()) + sync.assert_not_called() + connector._transfer_ref.poll_direct_gpu_v1.assert_not_awaited() + metadata.direct_sends = {"p"} + metadata.direct_store = True + connector.start_load_kv(None) + connector.wait_for_save() + assert connector.get_finished(set()) == ({"p"}, set()) + assert connector.get_finished(set()) == (set(), set()) + sync.assert_called_once() + connector._transfer_ref.poll_direct_gpu_v1.assert_awaited_once() diff --git a/xinference/model/llm/vllm/xavier/test/test_direct_history.py b/xinference/model/llm/vllm/xavier/test/test_direct_history.py index d0d0aa1c81..bdfb22fddf 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_history.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_history.py @@ -12,7 +12,8 @@ def history_runtime(monkeypatch): r = direct_runtime(monkeypatch) - r._init_history(32, cpu_capacity=0) + r.store.cpu_capacity = 0 + r._init_history(32) r.history_retention_seconds = 10 r.caches["K"].copy_(torch.arange(16).reshape(8, 2)) return r @@ -67,12 +68,10 @@ async def test_busy_writer_skips_next_request_without_backlog(monkeypatch): @pytest.mark.asyncio -@pytest.mark.parametrize("reason", ["held_budget", "deadline", "expired"]) +@pytest.mark.parametrize("reason", ["deadline", "expired"]) async def test_optional_retention_releases_engine_when_skipped(monkeypatch, reason): r = history_runtime(monkeypatch) - r.register_direct( - "t", "p", [1], [101], held_blocks=9 if reason == "held_budget" else 1 - ) + r.register_direct("t", "p", [1], [101]) if reason == "deadline": r.history_retention_seconds = 0 if reason == "expired": @@ -248,7 +247,8 @@ async def test_full_history_preserves_hot_content_until_new_content_repeats( monkeypatch, ): r = history_runtime(monkeypatch) - r._init_history(8, cpu_capacity=0) + r.store.cpu_capacity = 0 + r._init_history(8) r.history_retention_seconds = 10 async def finish(ticket, blocks, hashes): @@ -274,7 +274,8 @@ async def finish(ticket, blocks, hashes): @pytest.mark.asyncio async def test_probation_is_bounded_and_forgets_old_cold_content(monkeypatch): r = history_runtime(monkeypatch) - r._init_history(4, cpu_capacity=0) + r.store.cpu_capacity = 0 + r._init_history(4) r.history_retention_seconds = 10 for key in (100, 101, 102, 103, 101): r.register_direct(str(key), str(key), [1], [key]) @@ -312,7 +313,8 @@ def stage(keys, layers): @pytest.mark.asyncio async def test_fill_batches_grow_but_replacement_batches_stay_small(monkeypatch): r = history_runtime(monkeypatch) - r._init_history(20, cpu_capacity=0) + r.store.cpu_capacity = 0 + r._init_history(20) r.history_retention_seconds = 10 r.history_fill_chunk_bytes = 16 r.history_chunk_bytes = 8 @@ -397,8 +399,8 @@ async def test_fixed_prefix_restore_preserves_raw_bits( block_bytes = sum( cache[0].numel() * cache.element_size() for cache in r.caches.values() ) - r._init_history(gpu_blocks * block_bytes, cpu_capacity=278) - r.history_pending_bytes = 278 * block_bytes + r.store.cpu_capacity = 278 + r._init_history(gpu_blocks * block_bytes) r.history_retention_seconds = 30 r.slab_bytes = 8 * 1024 * 1024 keys = list(range(1000, 1000 + prefix_blocks)) @@ -442,7 +444,8 @@ async def test_history_spills_to_cpu_and_restores_mixed_prefix(monkeypatch, devi for name, shift in [("K", 32620), ("V", -32768)] } block_bytes = 32 - r._init_history(2 * block_bytes, cpu_capacity=2) + r.store.cpu_capacity = 2 + r._init_history(2 * block_bytes) r.history_retention_seconds = 10 original = {name: cache[:4].clone() for name, cache in r.caches.items()} r.register_direct("save", "p", [0, 1, 2, 3], [100, 101, 102, 103]) @@ -487,3 +490,52 @@ def save(key): assert store.reserve("all", list(store.ready)) save(6) assert store.ready == {1, 3, 4, 5} + + +@pytest.mark.asyncio +async def test_long_request_retains_bounded_leading_prefix(monkeypatch): + r = history_runtime(monkeypatch) + # Simulate a request larger than the retention cap, without allocating 8B KV. + r.history_pending_bytes = 2 * r.history.block_bytes + r.register_direct("t", "p", [0, 1, 2, 3], [100, 101, 102, 103]) + r.release_direct("t") + await r._history_task + assert r.history.ready == {100, 101} + assert r.metrics["history_capacity_limited_blocks"] == 2 + assert not r.metrics["history_skipped_requests"] + assert r.poll_direct() == {"p"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("closing", [False, True]) +async def test_abandoned_retention_counts_remaining_candidates(monkeypatch, closing): + r = history_runtime(monkeypatch) + r.register_direct("t", "p", [1, 2], [101, 102]) + r.history_retention_seconds = 0 + r.release_direct("t") + r._history_closing = closing + await r._history_task + reason = "closing" if closing else "deadline" + assert r.metrics[f"history_{reason}_dropped_blocks"] == 2 + assert r.poll_direct() == {"p"} + + +@pytest.mark.asyncio +async def test_history_batch_failure_releases_all_leases_preserving_error(monkeypatch): + r = history_runtime(monkeypatch) + r.register_direct("t", "p", [1, 2], [101, 102]) + r.release_direct("t") + await r._history_task + assert r.reserve_history("a", [101]) == [101] + assert r.reserve_history("b", [102]) == [102] + original = r.release_history + + def release(lease): + original(lease) + if lease == "a": + raise RuntimeError("cleanup error") + + monkeypatch.setattr(r, "release_history", release) + with pytest.raises(ValueError, match="outside lease"): + await r.load_history([{0: {"K": {999: 4}}}, {0: {"K": {102: 5}}}], ["a", "b"]) + assert not r.history.leases and not r._history_leases diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 4d7a00eb81..6885f67a10 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -80,6 +80,8 @@ class XavierLoadRequest: @dataclass class XavierConnectorMetadata(KVConnectorMetadata): + direct_sends: set[str] = field(default_factory=set) + direct_store: bool = False store_requests: List[XavierStoreRequest] = field(default_factory=list) load_requests: List[XavierLoadRequest] = field(default_factory=list) @@ -148,6 +150,7 @@ def __init__( self._pending_store_requests: Dict[str, XavierStoreRequest] = {} self._leased_requests: Dict[str, XavierLoadRequest] = {} self._gpu_load_jobs: Dict[asyncio.Task, List[str]] = {} + self._direct_sends: set[str] = set() self._num_cache_blocks = kv_cache_config.num_blocks self._request_staged_layers: Dict[str, set[str]] = {} self._registered_kv_caches: Dict[str, torch.Tensor | Sequence[torch.Tensor]] = ( @@ -189,12 +192,10 @@ def shutdown(self): self._loop = None def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: - if not self._is_consumer: - return - metadata = self._get_connector_metadata() assert isinstance(metadata, XavierConnectorMetadata) - if not metadata.load_requests: + self._direct_sends.update(metadata.direct_sends) + if not self._is_consumer or not metadata.load_requests: return if getattr(self, "_gpu_budget", None) is not None: @@ -232,9 +233,10 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> Non def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: sent = set() - if self._direct_handoff and self._is_producer and self._gpu_cache_mapped: + if self._direct_handoff and self._is_producer and self._direct_sends: assert self._transfer_ref is not None sent = self._call(self._transfer_ref.poll_direct_gpu_v1()) + self._direct_sends.difference_update(sent) if not self._gpu_load_jobs: return sent, set() # Pump responses on the shared actor loop without waiting for any RPC. @@ -317,8 +319,10 @@ def save_kv_layer( def wait_for_save(self): if self._direct_handoff and self._is_producer: - self._call(self._ensure_gpu_cache_mapping()) - torch.cuda.synchronize() + metadata = self._get_connector_metadata() + if metadata.direct_store: + self._call(self._ensure_gpu_cache_mapping()) + torch.cuda.synchronize() return if not self._is_producer or not self._pending_store_requests: return @@ -390,6 +394,10 @@ async def reserve(): return 0, False if handoff is None: raise ValueError("Missing direct Xavier handoff metadata") + if params.get("do_remote_prefill") is False: + # A preempted decode recomputes its prompt AND generated tokens. + # Its single-use producer ticket was consumed at allocation. + return 0, False tokens = handoff["tokens"] rank, ticket = handoff["rank"], handoff["ticket"] if tokens <= num_computed_tokens: @@ -400,6 +408,16 @@ async def release(): await ref.release_remote_direct_gpu_v1(rank, ticket) self._call(release()) + params["do_remote_prefill"] = False + return 0, False + + async def claim(): + ref = await self._get_transfer_ref() + return await ref.claim_remote_direct_gpu_v1(rank, ticket) + + if not self._call(claim()): + params["do_remote_prefill"] = False + self._requests_need_load.pop(request.request_id, None) return 0, False start = num_computed_tokens // self._block_size blocks = handoff["blocks"] @@ -498,6 +516,8 @@ def update_state_after_alloc( lease=load_request.lease, local_transfers_by_group=local_transfers_by_group, ) + if self._direct_handoff and not load_request.lease.startswith("history:"): + request.kv_transfer_params["do_remote_prefill"] = False logger.debug( "Xavier V1 allocated local blocks: request=%s, transfers=%s", request.request_id, @@ -509,6 +529,10 @@ def build_connector_meta( scheduler_output: SchedulerOutput, ) -> KVConnectorMetadata: meta = XavierConnectorMetadata() + if self._direct_handoff and self._is_producer: + meta.direct_sends = self._direct_sends + self._direct_sends = set() + meta.direct_store = bool(scheduler_output.num_scheduled_tokens) if self._is_producer and not self._direct_handoff: self._build_store_meta(scheduler_output, meta) if self._is_consumer: @@ -546,23 +570,18 @@ def request_finished( async def register(): ref = await self._get_transfer_ref() + history_args = () if getattr(self, "_history_enabled", False): hashes = self._build_xavier_hashes( request.prompt_token_ids[:token_count] ) - await ref.register_direct_gpu_v1( - ticket, - request.request_id, - blocks, - [key for key, _ in hashes[:count]], - len(block_ids), - ) - else: - await ref.register_direct_gpu_v1( - ticket, request.request_id, blocks - ) + history_args = ([key for key, _ in hashes[:count]],) + await ref.register_direct_gpu_v1( + ticket, request.request_id, blocks, *history_args + ) self._call(register()) + self._direct_sends.add(request.request_id) logger.debug( "Register Xavier direct handoff: request=%s blocks=%s", request.request_id, @@ -572,22 +591,28 @@ async def register(): "do_remote_prefill": True, "xavier_direct": { "rank": self._rank, + "address": self._xavier_config.get("rank_address"), "ticket": ticket, "blocks": blocks, "tokens": min(token_count, count * self._block_size), }, } pending = self._requests_need_load.get(request.request_id) - if pending is not None and not pending.local_transfers_by_group: - self._requests_need_load.pop(request.request_id) + if pending is not None and pending.local_transfers_by_group: + # The worker must drain all writes before freeing destinations. + return False, None + self._requests_need_load.pop(request.request_id, None) + handoff = params.get("xavier_direct") + if handoff and params.get("do_remote_prefill") is not False: async def release(): ref = await self._get_transfer_ref() await ref.release_remote_direct_gpu_v1( - next(iter(pending.transfers)), pending.lease + handoff["rank"], handoff["ticket"] ) self._call(release()) + params["do_remote_prefill"] = False return False, None self._chunked_prefill.pop(request.request_id, None) pending = self._requests_need_load.get(request.request_id) From e590a45ac166fc7a4216910ae02d54efdf3a323a Mon Sep 17 00:00:00 2001 From: qinxuye Date: Sun, 4 Oct 2026 21:49:53 +0800 Subject: [PATCH 9/9] fix: bound claimed Xavier leases and preserve reusable prefixes --- .../LC_MESSAGES/user_guide/pd_separation.po | 11 ++ .../LC_MESSAGES/user_guide/pd_separation.po | 11 ++ .../LC_MESSAGES/user_guide/pd_separation.po | 11 ++ .../LC_MESSAGES/user_guide/pd_separation.po | 11 ++ .../LC_MESSAGES/user_guide/pd_separation.po | 9 ++ .../LC_MESSAGES/user_guide/pd_separation.po | 9 ++ .../LC_MESSAGES/user_guide/pd_separation.po | 11 ++ .../LC_MESSAGES/user_guide/pd_separation.po | 9 ++ .../LC_MESSAGES/user_guide/pd_separation.po | 9 ++ doc/source/user_guide/pd_separation.rst | 2 +- xinference/core/tests/test_pd_model.py | 25 ++++ xinference/model/llm/vllm/xavier/README.md | 6 +- .../model/llm/vllm/xavier/direct_handoff.py | 69 ++++++----- .../model/llm/vllm/xavier/direct_history.py | 95 ++++++++++----- .../vllm/xavier/test/test_direct_handoff.py | 109 +++++++++++++++++- .../vllm/xavier/test/test_direct_history.py | 93 ++++++++++++++- .../model/llm/vllm/xavier/test/test_pd_gpu.py | 5 + .../model/llm/vllm/xavier/v1_connector.py | 20 ++-- 18 files changed, 439 insertions(+), 76 deletions(-) diff --git a/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po index 485dc4d9c6..05c7755774 100644 --- a/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po @@ -369,3 +369,14 @@ msgid "" msgstr "" "Direkte Übertragungen und das Laden des Verlaufs können parallel zur " "Antwortgenerierung anderer Anfragen erfolgen." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"Nicht angenommene Übergaben verfallen nach 120 Sekunden. Angenommene " +"Übergaben verfallen nach 10 Minuten ohne erneuten Planungsversuch oder " +"Übertragungsfortschritt. Laufende Kopien werden vor der Freigabe " +"abgeschlossen. Verfallene Übergaben werden lokal auf D neu berechnet." diff --git a/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po index 9f9f26bf9a..45c93832d5 100644 --- a/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po @@ -366,3 +366,14 @@ msgid "" msgstr "" "Las transferencias directas y las cargas del historial pueden solaparse con " "la generación de respuestas de otras solicitudes." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"Las transferencias no aceptadas caducan tras 120 segundos. Las aceptadas " +"caducan tras 10 minutos sin reintentos de planificación ni progreso de " +"transferencia. Las copias en curso terminan antes de liberar los recursos. D" +" vuelve a calcular localmente las transferencias caducadas." diff --git a/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po index 76b813c13f..369a021f39 100644 --- a/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po @@ -370,3 +370,14 @@ msgid "" msgstr "" "Les transferts directs et le chargement de l’historique peuvent se dérouler " "en parallèle de la génération des réponses à d’autres requêtes." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"Les transferts non acceptés expirent après 120 secondes. Les transferts " +"acceptés expirent après 10 minutes sans nouvelle tentative de planification " +"ni progression du transfert. Les copies en cours se terminent avant la " +"libération. Les transferts expirés sont recalculés localement sur D." diff --git a/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po index ee8829d77b..567aeab68c 100644 --- a/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po @@ -367,3 +367,14 @@ msgid "" msgstr "" "I trasferimenti diretti e il caricamento della cronologia possono avvenire " "in parallelo alla generazione delle risposte ad altre richieste." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"I trasferimenti non accettati scadono dopo 120 secondi. Quelli accettati " +"scadono dopo 10 minuti senza nuovi tentativi di pianificazione o avanzamento" +" del trasferimento. Le copie in corso terminano prima del rilascio. I " +"trasferimenti scaduti vengono ricalcolati localmente su D." diff --git a/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po index 05d92e0d83..0bf02ed39b 100644 --- a/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po @@ -332,3 +332,12 @@ msgid "" "Direct transfers and history loads can overlap with response generation for " "other requests." msgstr "直接転送と履歴の読み込みは、他のリクエストの応答生成と並行して実行できます。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"未受理の引き渡しは 120 秒で期限切れになります。受理済みの引き渡しも、スケジュールの再試行や転送の進行が 10 " +"分間なければ期限切れになります。実行中のコピーは解放前に完了します。期限切れの引き渡しは D がローカルで再計算します。" diff --git a/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po index a75868b2cd..ccad95fec5 100644 --- a/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po @@ -331,3 +331,12 @@ msgid "" "Direct transfers and history loads can overlap with response generation for " "other requests." msgstr "직접 전송과 기록 로드는 다른 요청의 응답 생성과 동시에 실행될 수 있습니다." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"수락되지 않은 전달은 120초 후 만료됩니다. 수락된 전달도 10분 동안 스케줄링 재시도나 전송 진행이 없으면 만료됩니다. 진행 중인 " +"복사는 해제 전에 완료됩니다. 만료된 전달은 D가 로컬에서 다시 계산합니다." diff --git a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po index 5eb11685a2..c04850c3bf 100644 --- a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po @@ -364,3 +364,14 @@ msgid "" msgstr "" "As transferências diretas e o carregamento do histórico podem ocorrer em " "paralelo à geração de respostas de outras solicitações." + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"Transferências não aceitas expiram após 120 segundos. As aceitas expiram " +"após 10 minutos sem nova tentativa de agendamento ou progresso na " +"transferência. Cópias em andamento terminam antes da liberação. " +"Transferências expiradas são recalculadas localmente em D." diff --git a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po index 4a40e26d7c..beb4a2832a 100644 --- a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po @@ -329,3 +329,12 @@ msgid "" "Direct transfers and history loads can overlap with response generation for " "other requests." msgstr "直传和历史缓存加载可以与其他请求的响应生成重叠执行。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"尚未被接收的交接在 120 秒后过期。已被接收的交接如果 10 分钟内没有调度重试或传输进展,也会过期。正在进行的拷贝会在释放前完成。过期的交接由 D " +"在本地重新计算。" diff --git a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po index c240110e39..3c36be2e2d 100644 --- a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po @@ -324,3 +324,12 @@ msgid "" "Direct transfers and history loads can overlap with response generation for " "other requests." msgstr "直接傳輸和歷史快取載入可以與其他請求的回應生成重疊執行。" + +#: user_guide/pd_separation.rst +msgid "" +"Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after " +"10 minutes without a scheduling retry or transfer progress. Active copies " +"finish before release. Expired handoffs are recomputed locally on D." +msgstr "" +"尚未被接收的交接在 120 秒後過期。已被接收的交接如果 10 分鐘內沒有排程重試或傳輸進展,也會過期。正在進行的複製會在釋放前完成。過期的交接由 D " +"在本機重新計算。" diff --git a/doc/source/user_guide/pd_separation.rst b/doc/source/user_guide/pd_separation.rst index d4833a626c..f335fffd6d 100644 --- a/doc/source/user_guide/pd_separation.rst +++ b/doc/source/user_guide/pd_separation.rst @@ -138,7 +138,7 @@ Direct handoff currently requires ``n=1``. History is restored on P before the handoff to D. Long prompts may retain only a prefix; cache retention is best effort. -Unclaimed handoffs expire after 120 seconds. If a handoff expires before D accepts it, D recomputes locally. Once D accepts it, the source cache stays reserved until the load finishes or the request is cancelled before loading. +Unclaimed handoffs expire after 120 seconds. Claimed handoffs expire after 10 minutes without a scheduling retry or transfer progress. Active copies finish before release. Expired handoffs are recomputed locally on D. The budget applies only to Xavier's GPU cache. Leave additional GPU memory for the model, vLLM's own KV cache and transfers. When this cache fills, Xavier stores reusable data in CPU memory; if no cache space is available, new data may not be retained for reuse. diff --git a/xinference/core/tests/test_pd_model.py b/xinference/core/tests/test_pd_model.py index d6002687b0..27c97e60d5 100644 --- a/xinference/core/tests/test_pd_model.py +++ b/xinference/core/tests/test_pd_model.py @@ -421,3 +421,28 @@ async def test_completed_direct_request_does_not_send_abandon_rpc(router, monkey await actor._infer("chat", [], {}, request_id="r") lookup.assert_not_awaited() assert not actor._direct_transfers + + +@pytest.mark.asyncio +async def test_completed_direct_stream_clears_handoff_without_abandon( + router, monkeypatch +): + actor, prefill, decode = router + prefill.chat.return_value = { + "_pd_kv_transfer_params": { + "do_remote_prefill": True, + "xavier_direct": {"ticket": "t", "rank": 0, "address": "127.0.0.1:1234"}, + } + } + + async def chunks(): + yield b"one" + yield b"two" + + decode.chat.return_value = chunks() + lookup = AsyncMock() + monkeypatch.setattr(xo, "actor_ref", lookup) + stream = await actor._infer("chat", [], {"stream": True}, request_id="r") + assert [chunk async for chunk in stream] == [b"one", b"two"] + assert not actor._direct_transfers and not actor._request_set + lookup.assert_not_awaited() diff --git a/xinference/model/llm/vllm/xavier/README.md b/xinference/model/llm/vllm/xavier/README.md index fd9488e762..26343a59d9 100644 --- a/xinference/model/llm/vllm/xavier/README.md +++ b/xinference/model/llm/vllm/xavier/README.md @@ -14,9 +14,9 @@ The native backend passes producer handoff metadata to the selected decoder and P/D uses request-scoped GPU-to-GPU handoff by default, with no opt-in flag. The producer keeps its engine blocks until decoder reads and optional history copies finish. Independent history uses 256 MiB of GPU storage per replica by default and spills to CPU, bounded by engine KV block capacity. `xavier_gpu_cache_bytes=0` disables history while keeping direct handoff. CPU-history hits restore to the local GPU in per-layer batches; P-to-D traffic remains xoscar NIXL. Missing NIXL fails launch. -Unclaimed producer tickets expire after 120 seconds. D atomically claims a live ticket before reporting a remote hit; expiry is a scheduler miss, not a worker transfer error. Claimed tickets stay pinned until D releases them. Allocation marks the handoff consumed, so preempted decode requests recompute locally rather than rereading the ticket. Router cleanup abandons only unclaimed tickets, while allocated loads drain before release. +Unclaimed producer tickets expire after 120 seconds. D atomically claims a live ticket before reporting a remote hit; expiry is a scheduler miss, not a worker transfer error. Claimed tickets have a 600-second idle lease refreshed by scheduling retries and slab transfers. Active reads cannot expire. If a claim expires before or between slabs, the worker reports every destination block of that request through vLLM's load-error callback together with receive completion; the default Xavier recompute policy retries locally without exposing partial KV or crashing EngineCore. Other transfer/layout failures still propagate. Allocation marks the handoff consumed, so preempted decode requests recompute locally rather than rereading the ticket. Router cleanup abandons only unclaimed tickets, while allocated loads drain before release. -History uses one writer, a 64 MiB candidate-prefix cap (also limited by total history capacity), and a 10 ms soft deadline. Long prompts retain the leading admissible prefix instead of being skipped entirely; all engine blocks remain owned until the active chunk drains. `history_capacity_limited_blocks` counts candidates beyond the cap, `history_deadline_dropped_blocks` and `history_closing_dropped_blocks` count remaining candidates abandoned at the deadline or shutdown, and `history_skipped_requests` counts a busy writer. When both tiers are full, new content must be observed again before replacing retained blocks. Active leases prevent eviction. History is restored on P; routing still visits P before D. Direct handoff supports `n=1`. +History uses one writer, a 64 MiB candidate-prefix cap (also limited by total history capacity), and a 10 ms soft deadline. Long prompts retain the leading admissible prefix instead of being skipped entirely; all engine blocks remain owned until the active chunk drains. `history_capacity_limited_blocks` counts candidates beyond the cap, `history_deadline_dropped_blocks` and `history_closing_dropped_blocks` count remaining candidates abandoned at the deadline or shutdown, and `history_skipped_requests` counts a busy writer. When both tiers are full, new content must be observed again before replacing retained blocks. Admission stops at the first rejected or capacity-limited prefix block. Cache and probation LRU order retain prefix heads longer than tails. History hit counters count unique blocks actually written by successful loads, not reservations or scheduling retries. Active leases prevent eviction. History is restored on P; routing still visits P before D. Direct handoff supports `n=1`. The following snapshot-path details apply to hybrid deployments. Xavier CPU snapshots are keyed by prompt content, not reusable engine block IDs. Readers reserve complete snapshots during transfer; unavailable snapshots or cache pressure can cause a cache miss and local computation. This differs from a transfer failure, which propagates as an error. Only complete layers are published for remote reuse. CPU capacity is bounded by the engine's KV block capacity. @@ -71,6 +71,6 @@ XINFERENCE_TEST_PD_GPU=1 python -m pytest -v \ xinference/model/llm/vllm/xavier/test/test_pd_gpu.py ``` -The test runs default Xavier direct handoff with tiered history, Xavier with history disabled, and native NIXL in real subprocesses, with one GPU per role. It checks streaming and non-streaming responses, repeated prompts and four concurrent requests, then terminates the deployment. Local prefix caching is disabled. Exact repeated-text checks are restricted to native NIXL and Xavier with history disabled; fixed-prefix and raw-bit history tests cover history reuse separately. Xavier requires producer registration and successful async completion logs from `get_finished`; native NIXL requires `calling _read_blocks` and completed receive logs, including the concurrent requests. +The test runs default Xavier direct handoff with tiered history, Xavier with history disabled, and native NIXL in real subprocesses, with one GPU per role. It checks streaming and non-streaming responses, repeated prompts and four concurrent requests, then terminates the deployment. Local prefix caching is disabled. Exact repeated-text checks are restricted to native NIXL and Xavier with history disabled; fixed-prefix and raw-bit history tests cover history reuse separately. Xavier requires producer registration and successful async completion logs from `get_finished`, and the default-history case requires a positive P-side history restore for each repeated prompt; native NIXL requires `calling _read_blocks` and completed receive logs, including the concurrent requests. For locally cached weights, set `XINFERENCE_TEST_PD_MODEL_PATH`, `XINFERENCE_TEST_PD_MODEL_NAME` and `XINFERENCE_TEST_PD_MODEL_SIZE` to the path and registered model name/size. The manually triggered **PD GPU integration** GitHub Actions workflow runs the same test on a selected two-GPU runner. diff --git a/xinference/model/llm/vllm/xavier/direct_handoff.py b/xinference/model/llm/vllm/xavier/direct_handoff.py index f7b23a205f..a365bac9c4 100644 --- a/xinference/model/llm/vllm/xavier/direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/direct_handoff.py @@ -91,18 +91,25 @@ def register_direct(self, ticket, request_id, blocks, hashes=None): ) self.metrics["direct_registered"] += 1 - def claim_direct(self, ticket): + def _live_direct(self, ticket): state = self.direct_requests.get(ticket) if state is None or state.released: - return False - if not state.claimed and time.monotonic() >= state.deadline: + 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 - # Ownership passes to D. Scheduler abort or load completion releases it; - # producer polling must not race queued allocation or a multi-slab read. + # 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): @@ -124,24 +131,17 @@ def release_direct(self, ticket): self.metrics["direct_finished"] += 1 def poll_direct(self): - for ticket, state in list(self.direct_requests.items()): - if ( - not state.released - and not state.claimed - and time.monotonic() >= state.deadline - ): - self.metrics["direct_expired"] += 1 - state.retention_attempted = True - self.release_direct(ticket) + 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.direct_requests.get(ticket) - if state is None or state.released: - raise RuntimeError("Direct KV handoff expired or was released") + 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: @@ -155,7 +155,7 @@ async def send_direct(self, ticket, reads, remote_ref, slab_bytes): ): raise ValueError("Direct KV source layout or block IDs differ") state.reading = True - state.deadline = time.monotonic() + 120 + state.deadline = time.monotonic() + 600 try: offset = 0 for read in reads: @@ -181,30 +181,26 @@ async def send_direct(self, ticket, reads, remote_ref, slab_bytes): # 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: - await self._load_direct(requests, tickets) + return await self._load_direct(requests, tickets) except BaseException: failed = True raise finally: - async def release(ranks, ticket): - for rank in ranks: - from .transfer import TransferActor - - sender = await xo.actor_ref( - address=self.actor._world_addresses[rank], - uid=f"{TransferActor.default_uid()}-{rank}", - ) - await sender.release_direct_gpu_v1(ticket) - results = await asyncio.gather( - *(release(ranks, ticket) for ranks, ticket in zip(requests, tickets)), + *( + 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: @@ -218,6 +214,7 @@ 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: @@ -243,9 +240,18 @@ async def _load_direct(self, requests, tickets): 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) - await sender.send_direct_gpu_v1( + 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 ): @@ -258,3 +264,4 @@ async def _load_direct(self, requests, tickets): finally: await asyncio.to_thread(torch.cuda.synchronize, self.device) indices.clear() + return invalid_blocks diff --git a/xinference/model/llm/vllm/xavier/direct_history.py b/xinference/model/llm/vllm/xavier/direct_history.py index 44aced3407..cf58f83545 100644 --- a/xinference/model/llm/vllm/xavier/direct_history.py +++ b/xinference/model/llm/vllm/xavier/direct_history.py @@ -23,6 +23,16 @@ def __init__(self, budget, block_bytes, device, cpu_capacity=0): 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. @@ -120,26 +130,41 @@ def _schedule_history(self, ticket, state): ) 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) - continue - repeated = key in self._history_probation - 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 free or repeated: - if len(candidates) < limit: - candidates.append((key, block)) - free = max(0, free - 1) - if repeated: - self.metrics["history_admission_reused_blocks"] += 1 - else: - self.metrics["history_capacity_limited_blocks"] += 1 else: - self.metrics["history_admission_rejected_blocks"] += 1 + 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 @@ -168,6 +193,7 @@ def done(future): 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. @@ -191,20 +217,16 @@ async def _retain_history(self, ticket, state, candidates, deadline): if free > 0: count = min(count, free) chunk = candidates[start : start + count] - while self.send_lock.locked() or self.recv_lock.locked(): + 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) - 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 start += len(chunk) async with self.send_lock: self._expire_history_leases() @@ -224,8 +246,20 @@ async def _retain_history(self, ticket, state, candidates, deadline): 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 @@ -269,6 +303,7 @@ async def _load_history(self, requests, 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] @@ -304,14 +339,22 @@ async def _load_history(self, requests, leases): self.caches[read.layer][ self._index_tensor(indices, read.destinations) ] = blocks + written.update(read.keys) await asyncio.to_thread(torch.cuda.synchronize, self.device) - for key in self.history.leases[lease]: + 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.metrics["history_hit_blocks"] += len( - self.history.leases[lease] - ) + 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() 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 8bff920b9a..c676d16505 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_handoff.py @@ -9,6 +9,7 @@ import xoscar as xo from ..direct_handoff import DirectGPUTransfer +from ..gpu_transfer import GPUTransferMixin from ..request_transfer import LayerRead from .test_gpu_transfer import runtime @@ -20,6 +21,12 @@ def direct_runtime(monkeypatch): r.metrics.update( direct_registered=0, direct_finished=0, direct_expired=0, index_uploads=0 ) + r.actor.default_uid = lambda: "transfer" + r.actor.release_remote_direct_gpu_v1 = ( + lambda rank, ticket: GPUTransferMixin.release_remote_direct_gpu_v1( + r.actor, rank, ticket + ) + ) r._init_history(0) return r @@ -109,10 +116,9 @@ async def test_expired_request_cannot_read_reused_engine_blocks(monkeypatch): source.register_direct("ticket", "producer", [1]) source.direct_requests["ticket"].deadline = 0 assert source.poll_direct() == {"producer"} - with pytest.raises(RuntimeError, match="expired"): - await source.send_direct( - "ticket", [LayerRead("K", [1], [0], (2,), torch.bfloat16)], None, 16 - ) + assert not await source.send_direct( + "ticket", [LayerRead("K", [1], [0], (2,), torch.bfloat16)], None, 16 + ) assert not source.metrics["wire_bytes"] @@ -385,7 +391,6 @@ async def test_claim_pins_queued_handoff_until_decoder_release(monkeypatch): source.caches["K"].fill_(7) source.register_direct("ticket", "producer", [2]) assert source.claim_direct("ticket") - source.direct_requests["ticket"].deadline = 0 source.abandon_direct("ticket") # router cancellation cannot release D's read assert not source.poll_direct() monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer_for(source))) @@ -474,3 +479,97 @@ def test_producer_metadata_gates_idle_rpc_and_fence( assert connector.get_finished(set()) == (set(), set()) sync.assert_called_once() connector._transfer_ref.poll_direct_gpu_v1.assert_awaited_once() + + +def test_claimed_but_never_read_ticket_expires(monkeypatch): + source = direct_runtime(monkeypatch) + source.register_direct("t", "p", [1]) + assert source.claim_direct("t") + source.direct_requests["t"].deadline = 0 + assert source.poll_direct() == {"p"} + assert not source.direct_requests + assert source.metrics["direct_expired"] == 1 + + +@pytest.mark.asyncio +async def test_active_read_cannot_expire_and_refreshes_idle_lease(monkeypatch): + source = direct_runtime(monkeypatch) + source.register_direct("t", "p", [1]) + assert source.claim_direct("t") + entered, finish = asyncio.Event(), asyncio.Event() + + async def copy(buffers, refs): + entered.set() + await finish.wait() + + monkeypatch.setattr(xo, "copy_to", copy) + task = asyncio.create_task( + source.send_direct( + "t", [LayerRead("K", [1], [0], (2,), torch.bfloat16)], None, 16 + ) + ) + await entered.wait() + source.direct_requests["t"].deadline = 0 + assert not source.poll_direct() + finish.set() + assert await task + assert source.direct_requests["t"].deadline > 0 + assert not source.poll_direct() + + +@pytest.mark.asyncio +async def test_expiry_between_slabs_reports_all_destinations_for_recompute(monkeypatch): + source, dest = direct_runtime(monkeypatch), direct_runtime(monkeypatch) + for name in ("V", "Q"): + source.caches[name] = source.caches["K"].clone() + dest.caches[name] = dest.caches["K"].clone() + source.register_direct("t", "p", [1, 2]) + assert source.claim_direct("t") + peer = peer_for(source) + original_send = peer.send_direct_gpu_v1.side_effect + + async def send(*args): + available = await original_send(*args) + if "t" in source.direct_requests: + source.direct_requests["t"].deadline = 0 + return available + + peer.send_direct_gpu_v1.side_effect = send + monkeypatch.setattr(xo, "actor_ref", AsyncMock(return_value=peer)) + + async def copy(buffers, refs): + refs[0].copy_(buffers[0]) + + monkeypatch.setattr(xo, "copy_to", copy) + invalid = await dest.load_direct( + [{0: {name: {1: 5, 2: 3} for name in source.caches}}], ["t"] + ) + assert invalid == {3, 5} + assert source.metrics["direct_expired"] == 1 + assert source.poll_direct() == {"p"} + + +def test_worker_reports_expired_destinations_with_finished_request(connector): + async def expired(): + return {3, 5} + + async def submit(): + task = asyncio.create_task(expired()) + connector._gpu_load_jobs[task] = ["decoder"] + await asyncio.sleep(0) + + connector._call(submit()) + assert connector.get_finished(set()) == (set(), {"decoder"}) + assert connector.get_block_ids_with_load_errors() == {3, 5} + assert connector.get_block_ids_with_load_errors() == set() + + +def test_full_local_hit_releases_unused_direct_ticket(connector): + connector._direct_handoff = True + actor = SimpleNamespace(release_remote_direct_gpu_v1=AsyncMock()) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = decoder_request() + assert connector.get_num_new_matched_tokens(request, 32) == (0, False) + actor.release_remote_direct_gpu_v1.assert_awaited_once_with(0, "ticket") + assert request.kv_transfer_params["do_remote_prefill"] is False + assert not connector._requests_need_load diff --git a/xinference/model/llm/vllm/xavier/test/test_direct_history.py b/xinference/model/llm/vllm/xavier/test/test_direct_history.py index bdfb22fddf..3a2d7db578 100644 --- a/xinference/model/llm/vllm/xavier/test/test_direct_history.py +++ b/xinference/model/llm/vllm/xavier/test/test_direct_history.py @@ -452,7 +452,8 @@ async def test_history_spills_to_cpu_and_restores_mixed_prefix(monkeypatch, devi r.release_direct("save") await r._history_task assert r.poll_direct() == {"p"} - assert r.history.tiers == {100: "cpu", 101: "cpu", 102: "gpu", 103: "gpu"} + assert r.history.tiers[100] == "gpu" + assert r.history.counts == {"gpu": 2, "cpu": 2} assert r.history.stats()["demotions"] == 2 assert r.reserve_history("read", [100, 101, 102, 103, 999]) == [100, 101, 102, 103] for cache in r.caches.values(): @@ -539,3 +540,93 @@ def release(lease): with pytest.raises(ValueError, match="outside lease"): await r.load_history([{0: {"K": {999: 4}}}, {0: {"K": {102: 5}}}], ["a", "b"]) assert not r.history.leases and not r._history_leases + + +@pytest.mark.asyncio +async def test_prefix_heads_survive_new_admission_and_reservation(monkeypatch): + r = history_runtime(monkeypatch) + r.store.cpu_capacity = 0 + r._init_history(12) # three blocks + r.history_retention_seconds = 10 + r.register_direct("a", "a", [0, 1, 2], [100, 101, 102]) + r.release_direct("a") + await r._history_task + assert list(r.history.blocks) == [102, 101, 100] + assert r.reserve_history("lease", [100, 101, 102]) == [100, 101, 102] + r.release_history("lease") + r._history_probation[200] = None + r.register_direct("b", "b", [3], [200]) + r.release_direct("b") + await r._history_task + assert r.reserve_history("read", [100, 101, 102]) == [100, 101] + + +@pytest.mark.asyncio +async def test_rejected_head_never_admits_repeated_tail(monkeypatch): + r = history_runtime(monkeypatch) + r._init_history(8) + r.history_retention_seconds = 10 + r.register_direct("a", "a", [0, 1], [100, 101]) + r.release_direct("a") + await r._history_task + r._history_probation[202] = None + r.register_direct("b", "b", [2, 3, 4], [200, 201, 202]) + r.release_direct("b") + assert r.history.ready == {100, 101} + assert r.metrics["history_admission_rejected_blocks"] == 3 + # A request longer than probation capacity retains its head, not its tail. + r._history_probation_limit = 2 + r.register_direct("c", "c", [2, 3, 4], [300, 301, 302]) + r.release_direct("c") + assert list(r._history_probation) == [301, 300] + + +@pytest.mark.asyncio +async def test_history_hit_counts_only_written_keys_once(monkeypatch): + r = history_runtime(monkeypatch) + r.register_direct("a", "a", [0, 1, 2], [100, 101, 102]) + r.release_direct("a") + await r._history_task + for lease in ("retry", "load"): + assert r.reserve_history(lease, [100, 101, 102]) == [100, 101, 102] + if lease == "retry": + r.release_history(lease) + assert r.history.metrics["gpu_hits"] == 0 + await r.load_history([{0: {"K": {101: 5}}}], ["load"]) + assert r.metrics["history_hit_blocks"] == 1 + assert r.metrics["history_gpu_hit_blocks"] == 1 + assert r.history.metrics["gpu_hits"] == 1 + + +def test_prefill_retry_releases_previous_history_reservation(connector): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + connector._direct_handoff = connector._history_enabled = True + actor = SimpleNamespace( + reserve_direct_history_v1=AsyncMock(return_value=[123]), + release_direct_history_v1=AsyncMock(), + ) + connector._get_transfer_ref = AsyncMock(return_value=actor) + request = SimpleNamespace( + request_id="p", prompt_token_ids=list(range(33)), kv_transfer_params={} + ) + assert connector.get_num_new_matched_tokens(request, 0) == (16, True) + previous = connector._requests_need_load["p"].lease + assert connector.get_num_new_matched_tokens(request, 0) == (16, True) + actor.release_direct_history_v1.assert_awaited_once_with(previous) + assert connector._requests_need_load["p"].lease != previous + + +@pytest.mark.asyncio +async def test_rejected_prefix_does_not_count_existing_tail_as_dropped(monkeypatch): + r = history_runtime(monkeypatch) + r._init_history(8) + r.history_retention_seconds = 10 + r.register_direct("a", "a", [0, 1], [100, 101]) + r.release_direct("a") + await r._history_task + r.register_direct("b", "b", [2, 3], [200, 101]) + r.release_direct("b") + assert r.metrics["history_admission_rejected_blocks"] == 1 + assert r.history.ready == {100, 101} diff --git a/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py b/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py index de516e6051..8b743f03ed 100644 --- a/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py +++ b/xinference/model/llm/vllm/xavier/test/test_pd_gpu.py @@ -174,9 +174,14 @@ def log_since(offset): # Compare exact text only with fixed zero-prefix computation. Historical # prefix reuse can change floating-point shapes and greedy output. for stream in (False, True): + offset = Path(log_path).stat().st_size repeated = chat(str(stream), stream) if backend == "nixl" or gpu_cache_bytes == 0: assert repeated == responses[stream] + else: + assert re.search( + r"Restored Xavier history: blocks=[1-9]\d*", log_since(offset) + ) offset = Path(log_path).stat().st_size with ThreadPoolExecutor(max_workers=4) as executor: diff --git a/xinference/model/llm/vllm/xavier/v1_connector.py b/xinference/model/llm/vllm/xavier/v1_connector.py index 6885f67a10..7078ef38a7 100644 --- a/xinference/model/llm/vllm/xavier/v1_connector.py +++ b/xinference/model/llm/vllm/xavier/v1_connector.py @@ -128,11 +128,7 @@ def __init__( ) self._direct_handoff = uses_direct_handoff(self._xavier_config) self._gpu_budget = self._xavier_config.get("gpu_cache_bytes") - if self._direct_handoff and ( - self._gpu_budget is None - or len(kv_cache_config.kv_cache_groups) != 1 - or self._xavier_config.get("role") not in ("prefill", "decode") - ): + if self._direct_handoff and len(kv_cache_config.kv_cache_groups) != 1: raise ValueError( "Direct handoff requires P/D roles, GPU transport and one KV group" ) @@ -151,6 +147,7 @@ def __init__( self._leased_requests: Dict[str, XavierLoadRequest] = {} self._gpu_load_jobs: Dict[asyncio.Task, List[str]] = {} self._direct_sends: set[str] = set() + self._invalid_block_ids: set[int] = set() self._num_cache_blocks = kv_cache_config.num_blocks self._request_staged_layers: Dict[str, set[str]] = {} self._registered_kv_caches: Dict[str, torch.Tensor | Sequence[torch.Tensor]] = ( @@ -245,7 +242,7 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: received = set() for task in list(self._gpu_load_jobs): if task.done(): - task.result() # A failed/partial transfer must never become ready. + self._invalid_block_ids.update(task.result() or ()) request_ids = self._gpu_load_jobs.pop(task) received.update(request_ids) for request_id in request_ids: @@ -259,6 +256,11 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: # waits for that step; actor-side writes and lease release still progress. return sent, received + def get_block_ids_with_load_errors(self) -> set[int]: + invalid = self._invalid_block_ids + self._invalid_block_ids = set() + return invalid + def wait_for_layer_load(self, layer_name: str) -> None: return @@ -353,7 +355,7 @@ def get_num_new_matched_tokens( if self._direct_handoff: params = getattr(request, "kv_transfer_params", None) or {} handoff = params.get("xavier_direct") - if self._is_producer and getattr(self, "_history_enabled", False): + if self._is_producer and self._history_enabled: previous = self._requests_need_load.get(request.request_id) if previous is not None: if previous.local_transfers_by_group: @@ -571,7 +573,7 @@ def request_finished( async def register(): ref = await self._get_transfer_ref() history_args = () - if getattr(self, "_history_enabled", False): + if self._history_enabled: hashes = self._build_xavier_hashes( request.prompt_token_ids[:token_count] ) @@ -1055,7 +1057,7 @@ async def load(): [requests[i].lease for i in historical], ) if direct: - await transfer.load_direct_gpu_v1( + return await transfer.load_direct_gpu_v1( [entries[i] for i in direct], [requests[i].lease for i in direct], )