diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 746377372..4e5cb0689 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -80,6 +80,7 @@ def _make_request(skip_cache_refresh: bool = True): num_inference_steps=4, ) return SimpleNamespace( + request_id="req-test", prompts=["a prompt"], sampling_params=sampling_params, skip_cache_refresh=skip_cache_refresh, @@ -104,8 +105,9 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar runner.kv_transfer_manager = SimpleNamespace( receive_kv_cache=lambda req, target_device=None: None, receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None, - receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None, + receive_multi_kv_cache_distributed=lambda *a, **k: None, ) + runner._kv_prefetch_enabled = False return runner diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index fa26ae276..ba22743fc 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -184,7 +184,7 @@ class TestWorkerWrapperBaseDelegation: mock_reqs = [mocker.Mock()] result = wrapper.execute_model(mock_reqs, mock_od_config) - wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config) + wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config, kv_prefetch_jobs=None) assert result == mock_output def test_load_weights_delegation(self, mocker: MockerFixture, mock_od_config): diff --git a/tests/distributed/omni_connectors/test_kv_async_prefetch.py b/tests/distributed/omni_connectors/test_kv_async_prefetch.py new file mode 100644 index 000000000..b6feba26d --- /dev/null +++ b/tests/distributed/omni_connectors/test_kv_async_prefetch.py @@ -0,0 +1,356 @@ +"""Background KV prefetch — CPU-testable mechanism. + +Covers opt-in gating, the start_prefetch / consume_prefetched_kv round trip over a +mock connector, dedup, per-call sender_info isolation, role classification, and +miss/abort paths. GPU H2D and D2D paths need CUDA (integration tests). +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVCacheConfig, + OmniKVTransferManager, + ReceiveRole, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# start_prefetch() skips stubs without a sender endpoint, so submissions carry it. +_SENDER_INFO = {"host": "127.0.0.1", "zmq_port": 50051} + + +class MockConnector: + def __init__(self): + self.store = {} + self.get_calls = [] + + def put(self, from_stage, to_stage, put_key, data): + self.store[f"{from_stage}->{to_stage}:{put_key}"] = data + return True, len(str(data)), None + + def get(self, from_stage, to_stage, get_key, metadata=None): + self.get_calls.append({"get_key": get_key, "metadata": metadata}) + key = f"{from_stage}->{to_stage}:{get_key}" + if key in self.store: + return self.store[key], len(str(self.store[key])) + return None + + +def _make_sender_receiver(*, async_prefetch=True): + """Build sender+receiver managers sharing one mock connector, seed one KV.""" + connector = MockConnector() + + sender = OmniKVTransferManager( + OmniKVCacheConfig( + connector_config={"type": "mock"}, from_stage="sender", to_stage="receiver", need_send_cache=True + ) + ) + sender._connector = connector + + receiver = OmniKVTransferManager( + OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage="sender", + stage_id="receiver", + need_recv_cache=True, + recv_timeout=1.0, + ), + async_prefetch=async_prefetch, + ) + receiver._connector = connector + return sender, receiver, connector + + +def _seed_payload(sender, req_id, *, num_layers=2, block_size=4, num_heads=4, head_dim=8): + kv_caches = [torch.randn(2, 5, block_size, num_heads, head_dim) for _ in range(num_layers)] + sender.handle_finished_requests_kv_transfer( + {req_id: {"block_ids": [0, 1], "seq_len": 10}}, kv_caches, block_size, "float32" + ) + + +def _req(rid): + """A minimal request carrying just ``request_id`` (consume_prefetched_kv only needs that).""" + return SimpleNamespace(request_id=rid) + + +# --------------------------------------------------------------------------- # +# Opt-in gating +# --------------------------------------------------------------------------- # + + +def test_async_prefetch_defaults_false(): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + assert mgr._async_prefetch is False + + +def test_from_od_config_enables_prefetch_flag(): + od = SimpleNamespace(omni_kv_config={"enable_kv_async_prefetch": True}) + assert OmniKVTransferManager.from_od_config(od)._async_prefetch is True + od2 = SimpleNamespace(omni_kv_config=None) + assert OmniKVTransferManager.from_od_config(od2)._async_prefetch is False + od3 = SimpleNamespace(omni_kv_config={"enable_kv_async_prefetch": False}) + assert OmniKVTransferManager.from_od_config(od3)._async_prefetch is False + + +def test_prefetch_preserved_with_connector_config(): + # The flag must survive merging with a populated connector config. + od = SimpleNamespace( + omni_kv_config={ + "connector_config": {"type": "mock"}, + "need_recv_cache": True, + "enable_kv_async_prefetch": True, + }, + ) + assert OmniKVTransferManager.from_od_config(od)._async_prefetch is True + + +def test_prefetch_disabled_with_cfg_companion_collector(): + # A CFG companion collector forces the synchronous path (see _resolve_async_prefetch). + od = SimpleNamespace( + omni_kv_config={"enable_kv_async_prefetch": True}, + cfg_kv_collect_func=lambda *a, **k: {}, + ) + assert OmniKVTransferManager.from_od_config(od)._async_prefetch is False + + +def test_from_vllm_config_ar_path_is_always_sync(): + # The AR entry point never enables prefetch. + model_config = SimpleNamespace(omni_kv_config=None) + vllm_config = SimpleNamespace(kv_transfer_config=None) + assert OmniKVTransferManager.from_vllm_config(vllm_config, model_config)._async_prefetch is False + + +def test_start_prefetch_noop_when_disabled(): + _, receiver, _ = _make_sender_receiver(async_prefetch=False) + receiver.start_prefetch({"request_id": "r1", "kv_sender_info": None}) + assert receiver._prefetch_futures == {} + + +# --------------------------------------------------------------------------- # +# Round trip +# --------------------------------------------------------------------------- # + + +def test_prefetch_round_trip_returns_data(): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-1") + + receiver.start_prefetch({"request_id": "rid-1", "kv_sender_info": _SENDER_INFO}) + assert "rid-1" in receiver._prefetch_futures + + data, size = receiver.consume_prefetched_kv(_req("rid-1")) + assert data is not None + assert "layer_blocks" in data + assert data["metadata"]["seq_len"] == 10 + assert size > 0 + # Future is consumed (popped) on retrieval. + assert "rid-1" not in receiver._prefetch_futures + + +def test_consume_prefetched_kv_miss_without_prefetch_returns_none(): + _, receiver, _ = _make_sender_receiver(async_prefetch=False) + assert receiver.consume_prefetched_kv(_req("never-prefetched")) == (None, 0) + + +def test_start_prefetch_dedups_same_request(): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-dup") + receiver.start_prefetch({"request_id": "rid-dup", "kv_sender_info": _SENDER_INFO}) + fut1 = receiver._prefetch_futures["rid-dup"] + receiver.start_prefetch({"request_id": "rid-dup", "kv_sender_info": _SENDER_INFO}) + assert receiver._prefetch_futures["rid-dup"] is fut1 + + +def test_start_prefetch_skips_without_sender_info(): + # No endpoint => no submission; the sync path handles it. + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-nosi") + receiver.start_prefetch({"request_id": "rid-nosi", "kv_sender_info": None}) + assert receiver._prefetch_futures == {} + + +def test_consume_prefetched_kv_sweeps_stale_entries(): + # Consuming a slot drops every other tracked prefetch. + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-stale") + receiver.start_prefetch({"request_id": "rid-stale", "kv_sender_info": _SENDER_INFO}) + assert "rid-stale" in receiver._prefetch_futures + # Miss on the current request still sweeps stale entries. + receiver._async_prefetch = False + assert receiver.consume_prefetched_kv(_req("rid-current")) == (None, 0) + assert receiver._prefetch_futures == {} + + +# --------------------------------------------------------------------------- # +# consume_prefetched_kv: miss returns None, stale sweep +# --------------------------------------------------------------------------- # + + +def test_start_prefetch_sweeps_orphan(): + # An orphan prefetch (e.g. aborted) is dropped when the next prefetch starts. + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-orphan") + _seed_payload(sender, "rid-next") + receiver.start_prefetch({"request_id": "rid-orphan", "kv_sender_info": _SENDER_INFO}) + receiver.consume_prefetched_kv(_req("rid-orphan")) # complete the future + receiver.start_prefetch({"request_id": "rid-orphan", "kv_sender_info": _SENDER_INFO}) # re-add as leftover + receiver.start_prefetch({"request_id": "rid-next", "kv_sender_info": _SENDER_INFO}) + assert "rid-orphan" not in receiver._prefetch_futures + assert "rid-next" in receiver._prefetch_futures + + +def test_shutdown_prefetch_clears_state(): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-sd") + receiver.start_prefetch({"request_id": "rid-sd", "kv_sender_info": _SENDER_INFO}) + receiver.shutdown_prefetch() + assert receiver._prefetch_futures == {} + assert receiver._prefetch_executor is None + + +# --------------------------------------------------------------------------- # +# Correctness: per-call sender_info must not touch instance state +# --------------------------------------------------------------------------- # + + +def test_receive_with_sender_info_does_not_mutate_instance_state(): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-si") + assert receiver._sender_base_host is None + assert receiver._sender_base_zmq_port is None + + receiver.receive_kv_cache_for_request( + "rid-si", + target_device=None, + sender_info={"host": "1.2.3.4", "zmq_port": 50051}, + ) + # The per-call endpoint must not leak into instance fields. + assert receiver._sender_base_host is None + assert receiver._sender_base_zmq_port is None + + +# --------------------------------------------------------------------------- # +# Receive-role classification (recv_role) and the follower prefetch guard +# --------------------------------------------------------------------------- # + +import vllm_omni.diffusion.distributed.parallel_state as ps # noqa: E402 + + +def _patch_topo(monkeypatch, *, world_size, world_rank=0, cfg_size=1, cfg_rank=0, sp_size=1, sp_rank=0): + monkeypatch.setattr(ps, "get_world_group", lambda: SimpleNamespace(world_size=world_size, rank_in_group=world_rank)) + monkeypatch.setattr(ps, "get_classifier_free_guidance_world_size", lambda: cfg_size) + monkeypatch.setattr(ps, "get_classifier_free_guidance_rank", lambda: cfg_rank) + monkeypatch.setattr(ps, "get_cfg_group", lambda: None) + monkeypatch.setattr(ps, "get_sequence_parallel_world_size", lambda: sp_size) + monkeypatch.setattr(ps, "get_sequence_parallel_rank", lambda: sp_rank) + monkeypatch.setattr(ps, "get_sp_group", lambda: None) + + +def _set_tp(mgr, src, tgt): + mgr._tp_topo = SimpleNamespace(source_tp_size=src, target_tp_size=tgt, local_rank=0) + + +def _role(mgr): + """Recompute the receive role from the (patched) parallel state, bypassing the cache.""" + mgr._topo_config = None + return mgr.topo_config.role + + +def test_recv_role_world1_is_local(monkeypatch): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 1, 1) + _patch_topo(monkeypatch, world_size=1) + assert _role(mgr) is ReceiveRole.LOCAL + + +def test_recv_role_pure_tp_is_local(monkeypatch): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 2, 2) + _patch_topo(monkeypatch, world_size=2) + assert _role(mgr) is ReceiveRole.LOCAL + + +def test_recv_role_cfg_leader_and_follower(monkeypatch): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 2, 2) + _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=0) + assert _role(mgr) is ReceiveRole.LEADER + _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=1) + assert _role(mgr) is ReceiveRole.FOLLOWER + + +def test_recv_role_sp_leader_and_follower(monkeypatch): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 2, 2) + _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=0) + assert _role(mgr) is ReceiveRole.LEADER + _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=1) + assert _role(mgr) is ReceiveRole.FOLLOWER + + +def test_recv_role_tp_and_sp_leader_per_tp_rank(monkeypatch): + # TP=2 + SP=2: sp_rank 0 is leader (pulls), sp_rank 1 is follower (gets broadcast). + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 2, 2) + _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=0) + assert _role(mgr) is ReceiveRole.LEADER + _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=1) + assert _role(mgr) is ReceiveRole.FOLLOWER + + +def test_recv_role_hsdp_world_broadcast(monkeypatch): + # TP inactive, world > 1 (HSDP): rank-0 leader, others follower. + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 1, 1) + _patch_topo(monkeypatch, world_size=2, world_rank=0) + assert _role(mgr) is ReceiveRole.LEADER + _patch_topo(monkeypatch, world_size=2, world_rank=1) + assert _role(mgr) is ReceiveRole.FOLLOWER + + +def test_recv_role_uninitialized_defaults_local(monkeypatch): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + + def _boom(): + raise AssertionError("world group is not initialized") + + monkeypatch.setattr(ps, "get_world_group", _boom) + assert _role(mgr) is ReceiveRole.LOCAL + + +def test_start_prefetch_skips_follower(monkeypatch): + # A follower never pulls, so start_prefetch must not submit. + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-fol") + _set_tp(receiver, 2, 2) + _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=1) + receiver.start_prefetch({"request_id": "rid-fol", "kv_sender_info": _SENDER_INFO}) + assert receiver._prefetch_futures == {} + + +def test_start_prefetch_submits_for_owner(monkeypatch): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-own") + _set_tp(receiver, 2, 2) + _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=0) + receiver.start_prefetch({"request_id": "rid-own", "kv_sender_info": _SENDER_INFO}) + assert "rid-own" in receiver._prefetch_futures + + +def test_consume_then_apply_attaches_payload(): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-apply") + receiver.start_prefetch({"request_id": "rid-apply", "kv_sender_info": _SENDER_INFO}) + data, _ = receiver.consume_prefetched_kv(_req("rid-apply")) + assert data is not None + + req = OmniDiffusionRequest(prompts=["p"], sampling_params=OmniDiffusionSamplingParams(), request_id="rid-apply") + # Mirror consume_and_distribute_kv_cache's LOCAL apply path (CPU: record_stream no-op). + receiver._record_stream_for_prefetched(data) + receiver.apply_kv_cache_to_request(req, data) + assert req.past_key_values is not None diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py index 2cc19f1b3..8bf2b5d71 100644 --- a/tests/distributed/omni_connectors/test_kv_flow.py +++ b/tests/distributed/omni_connectors/test_kv_flow.py @@ -1,7 +1,6 @@ import json import struct -import numpy as np import pytest import torch @@ -206,17 +205,6 @@ def test_update_sender_info_uses_configured_source_stage(): assert manager.config.connector_config["sender_zmq_port"] == 50152 -def test_clone_received_payload_tensors_breaks_buffer_alias(): - payload, key_tensor = _make_serialized_payload() - raw = np.frombuffer(bytearray(payload), dtype=np.uint8) - data = KVCacheTransferData.from_bytes(memoryview(raw)) - - OmniKVTransferManager._clone_received_payload_tensors(data) - raw[:] = 0 - - assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor) - - def test_receive_kv_cache_uses_exponential_backoff(monkeypatch): config = OmniKVCacheConfig( connector_config={"type": "mock"}, @@ -526,6 +514,57 @@ def test_integration_flow(common_constants): assert req.kv_metadata["seq_len"] == 10 +def test_deferred_release_isolates_data_from_pool_buffer(): + payload, key_tensor = _make_serialized_payload() + + # ManagedBuffer-like object whose release() zeros storage (simulates reuse). + raw_array = bytearray(payload) + buf_tensor = torch.frombuffer(raw_array, dtype=torch.uint8) + + class _CorruptingManagedBuffer: + def __init__(self, storage, tensor): + self._storage = storage + self._tensor = tensor + + @property + def tensor(self): + return self._tensor + + def release(self): + for i in range(len(self._storage)): + self._storage[i] = 0 + + managed_buf = _CorruptingManagedBuffer(raw_array, buf_tensor) + + class _ManagedBufConnector: + def __init__(self, buf, from_stage, to_stage, req_id): + self._buf = buf + self._key = f"{from_stage}->{to_stage}:omni_{from_stage}_to_{to_stage}_kv_cache_{req_id}" + + def get(self, from_stage, to_stage, get_key, metadata=None): + if f"{from_stage}->{to_stage}:{get_key}" == self._key: + return self._buf, len(self._buf._storage) + return None + + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage="sender", + stage_id="receiver", + need_recv_cache=True, + recv_timeout=1.0, + ) + manager = OmniKVTransferManager(config) + req_id = "req-payload" + manager._connector = _ManagedBufConnector(managed_buf, "sender", "receiver", req_id) + + data, size = manager.receive_kv_cache_for_request(req_id, target_device=torch.device("cpu")) + + assert data is not None + assert size > 0 + # Returned data must survive the pool buffer being zeroed by release(). + assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor) + + def test_manager_extraction_no_connector(kv_config, common_constants): """Test extraction when connector is unavailable (should still return IDs).""" block_size = common_constants["block_size"] diff --git a/tests/distributed/omni_connectors/test_tp_rank_aware.py b/tests/distributed/omni_connectors/test_tp_rank_aware.py index 495c10a65..8f8783db2 100644 --- a/tests/distributed/omni_connectors/test_tp_rank_aware.py +++ b/tests/distributed/omni_connectors/test_tp_rank_aware.py @@ -586,8 +586,17 @@ class TestDistributedReceive: return_value=0, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size", + return_value=1, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank", + return_value=0, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None), ): - assert mgr.receive_multi_kv_cache_distributed(req) is True + assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True mgr.receive_multi_kv_cache.assert_called_once() assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu") @@ -642,6 +651,15 @@ class TestDistributedReceive: return_value=1, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size", + return_value=1, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank", + return_value=0, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None), ): assert mgr.receive_multi_kv_cache_distributed(req) is True @@ -666,7 +684,27 @@ class TestDistributedReceive: world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1) mgr.receive_multi_kv_cache = MagicMock(return_value=True) - with patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group): + with ( + patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", + return_value=1, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", + return_value=0, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=None), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size", + return_value=1, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank", + return_value=0, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None), + ): assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu")) @@ -722,7 +760,7 @@ class TestDistributedReceive: ), patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group), ): - assert mgr.receive_multi_kv_cache_distributed(req) is True + assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True # Verify behavior based on role if is_owner: @@ -817,7 +855,7 @@ class TestDistributedReceive: ), patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group), ): - assert mgr.receive_multi_kv_cache_distributed(req) is True + assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True # Verify behavior based on role if role == "owner": diff --git a/vllm_omni/deploy/hunyuan_image_3_moe.yaml b/vllm_omni/deploy/hunyuan_image_3_moe.yaml index 67ce8b66a..b661674c7 100644 --- a/vllm_omni/deploy/hunyuan_image_3_moe.yaml +++ b/vllm_omni/deploy/hunyuan_image_3_moe.yaml @@ -62,7 +62,7 @@ stages: max_tokens: 8192 detokenize: true skip_special_tokens: false - include_stop_str_in_output: true, + include_stop_str_in_output: true - stage_id: 1 max_num_seqs: 1 @@ -71,6 +71,8 @@ stages: distributed_executor_backend: "mp" omni_kv_config: need_recv_cache: true + enable_kv_async_prefetch: true + kv_prefetch_min_free_mem_ratio: 0.2 parallel_config: tensor_parallel_size: 2 enable_expert_parallel: true diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index d78377e50..cd5e02643 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -321,7 +321,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): new_req = scheduler_output.scheduled_new_reqs[0] result = self.collective_rpc( "execute_model", - args=(new_req.req, self.od_config), + args=(new_req.req, self.od_config, getattr(scheduler_output, "kv_prefetch_jobs", None)), unique_reply_rank=0, exec_all_ranks=True, ) diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index 812873a1c..fe7e21229 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -52,6 +52,7 @@ class _BaseScheduler(SchedulerInterface): self._running_sampling_params_key: SamplingParamsKey | None = None self._finished_req_ids: set[str] = set() self.max_num_running_reqs: int = 1 + self._prefetch_enabled: bool = False def initialize(self, od_config: OmniDiffusionConfig) -> None: self.od_config = od_config @@ -66,6 +67,8 @@ class _BaseScheduler(SchedulerInterface): self.max_num_running_reqs = max(1, int(max_num_seqs)) except (TypeError, ValueError): self.max_num_running_reqs = 1 + omni_kv = getattr(od_config, "omni_kv_config", None) or {} + self._prefetch_enabled = bool(omni_kv.get("enable_kv_async_prefetch", False)) self._reset_scheduler_state() def add_request(self, request: OmniDiffusionRequest) -> str: @@ -111,6 +114,22 @@ class _BaseScheduler(SchedulerInterface): else: scheduled_cached_request_ids.append(request_id) + # Expose the next waiting request (serial mode) so the runner can + # prefetch its KV during this forward. Skip a request without + # kv_sender_info (would target the wrong sender under multi-replica) or + # one already finished/aborted (would consume its sender buffer for + # nothing). + kv_prefetch_jobs: dict | None = None + if self._prefetch_enabled and self._waiting: + nxt = self._request_states.get(self._waiting[0]) + if nxt is not None and not nxt.is_finished(): + sender_info = getattr(nxt.req, "kv_sender_info", None) + if sender_info: + kv_prefetch_jobs = { + "request_id": nxt.request_id, + "kv_sender_info": sender_info, + } + scheduler_output = DiffusionSchedulerOutput( step_id=self._step_id, scheduled_new_reqs=scheduled_new_reqs, @@ -118,6 +137,7 @@ class _BaseScheduler(SchedulerInterface): finished_req_ids=set(self._finished_req_ids), num_running_reqs=len(self._running), num_waiting_reqs=len(self._waiting), + kv_prefetch_jobs=kv_prefetch_jobs, ) # update after schedule diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index b50e3555a..11dc66a21 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/vllm_omni/diffusion/sched/interface.py @@ -113,6 +113,8 @@ class DiffusionSchedulerOutput: finished_req_ids: set[str] num_running_reqs: int num_waiting_reqs: int + # next request to background-prefetch KV + kv_prefetch_jobs: dict | None = None @cached_property def scheduled_request_ids(self) -> list[str]: diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index fffea4a40..96a301366 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -80,9 +80,23 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # Cache for per-request stepwise state. self.state_cache: dict[str, DiffusionRequestState] = {} - # Initialize KV cache manager for connector management + # Initialize KV cache manager for connector management. self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config) + # Prefetch covers TP / SP / CFG-Parallel / HSDP. Disabled when a CFG + # companion KV collector is set (that KV is not backgrounded). + has_cfg_companion_kv = getattr(od_config, "cfg_kv_collect_func", None) is not None + + self._kv_prefetch_enabled = ( + bool(self.kv_transfer_manager.config.enable_kv_async_prefetch) + and not has_cfg_companion_kv + and self.kv_transfer_manager.config.need_recv_cache + ) + + @property + def target_device(self) -> torch.device | None: + return getattr(self.pipeline, "device", None) + def _compile_transformer(self, attr_name: str) -> None: """Compile a transformer attribute on the pipeline with torch.compile.""" model = getattr(self.pipeline, attr_name, None) @@ -269,7 +283,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): pool_overhead_gb / peak_reserved_gb * 100 if peak_reserved_gb > 0 else 0.0, ) - def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput: + def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput: """ Execute a forward pass for the given requests. @@ -293,12 +307,26 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): use_hsdp = self.od_config.parallel_config.use_hsdp grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() with grad_context: - # The manager handles the check for need_recv_cache internally - self.kv_transfer_manager.receive_multi_kv_cache_distributed( - req, - cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None), - target_device=getattr(self.pipeline, "device", None), - ) + # Receive AR KV (fetch → distribute → apply inside the entry). prefetch on: + # consume prior-forward payload, sync-fallback on miss; else sync receive. + kv_recv_t0 = time.perf_counter() + if self._kv_prefetch_enabled: + self.kv_transfer_manager.consume_and_distribute_kv_cache( + req, + target_device=self.target_device, + ) + else: + self.kv_transfer_manager.receive_multi_kv_cache_distributed( + req, + cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None), + target_device=self.target_device, + ) + kv_recv_ms = (time.perf_counter() - kv_recv_t0) * 1000 + logger.debug("KV recv for %s %.1fms", req.request_id, kv_recv_ms) + + # Kick off the next request's prefetch (+ H2D) to overlap this forward. + if self._kv_prefetch_enabled and kv_prefetch_jobs is not None: + self.kv_transfer_manager.start_prefetch(kv_prefetch_jobs, self.target_device) if req.sampling_params.generator is None and req.sampling_params.seed is not None: if req.sampling_params.generator_device is not None: @@ -404,7 +432,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): self.kv_transfer_manager.receive_multi_kv_cache_distributed( state_req, cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None), - target_device=getattr(self.pipeline, "device", None), + target_device=self.target_device, ) self.state_cache[request_id] = new_state resolved.append(new_state) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index db1a9fded..6b7e8863e 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -369,7 +369,12 @@ class DiffusionWorker: else: profiler.stop() - def execute_model(self, req: OmniDiffusionRequest, od_config: OmniDiffusionConfig) -> DiffusionOutput: + def execute_model( + self, + req: OmniDiffusionRequest, + od_config: OmniDiffusionConfig, + kv_prefetch_jobs: dict | None = None, + ) -> DiffusionOutput: """Execute a forward pass by delegating to the model runner.""" assert self.model_runner is not None, "Model runner not initialized" if self.lora_manager is not None: @@ -382,7 +387,7 @@ class DiffusionWorker: profiler = self._get_profiler() ctx = profiler.annotate_context_manager("diffusion_forward") if profiler else nullcontext() with ctx: - output = self.model_runner.execute_model(req) + output = self.model_runner.execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs) if profiler: profiler.step() return output @@ -646,6 +651,10 @@ class DiffusionWorker: def shutdown(self) -> None: """Shutdown the worker and cleanup distributed environment.""" + if self.model_runner is not None: + mgr = getattr(self.model_runner, "kv_transfer_manager", None) + if mgr is not None: + mgr.shutdown_prefetch() destroy_distributed_env() @@ -1077,18 +1086,24 @@ class WorkerWrapperBase: """ return self.worker.generate(requests) - def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput: + def execute_model( + self, + reqs: list[OmniDiffusionRequest], + od_config: OmniDiffusionConfig, + kv_prefetch_jobs: dict | None = None, + ) -> DiffusionOutput: """ Execute a forward pass. Args: reqs: List of diffusion requests od_config: OmniDiffusionConfig configuration + kv_prefetch_jobs: Optional next-request KV prefetch descriptor. Returns: DiffusionOutput with generated results """ - return self.worker.execute_model(reqs, od_config) + return self.worker.execute_model(reqs, od_config, kv_prefetch_jobs=kv_prefetch_jobs) def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: """Execute one diffusion step.""" diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 3b234ad1c..304489bbb 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -2,10 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unified OmniConnector and KV cache transfer management.""" +import enum import json import struct import time from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from dataclasses import asdict, dataclass from typing import Any @@ -13,6 +15,7 @@ import torch from vllm.logger import init_logger from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.platforms import current_omni_platform from .factory import OmniConnectorFactory from .utils.config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec @@ -32,11 +35,64 @@ from .utils.kv_utils import ( slice_layer_blocks, slice_received_rank_shard, ) +from .utils.serialization import OmniSerializer logger = init_logger(__name__) LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor] + +class KVPrefetchConsumeError(RuntimeError): + """Payload consumed from connector but post-get processing failed — sync retry impossible.""" + + +class ReceiveRole(enum.Enum): + """How a rank obtains KV: LOCAL pulls its own shard, LEADER pulls then + distributes, FOLLOWER receives via collective.""" + + LOCAL = "local" + LEADER = "leader" + FOLLOWER = "follower" + + +@dataclass(frozen=True) +class _TransferTopoConfig: + """Snapshot of the current rank's parallel topology for KV receive routing.""" + + role: ReceiveRole + tp_active: bool + cfg_size: int + cfg_rank: int + cfg_group: Any + sp_size: int + sp_rank: int + sp_group: Any + world: Any + + @property + def is_leader(self) -> bool: + return self.role == ReceiveRole.LEADER + + @property + def is_follower(self) -> bool: + return self.role == ReceiveRole.FOLLOWER + + @property + def is_local(self) -> bool: + return self.role == ReceiveRole.LOCAL + + @property + def cfg_active(self) -> bool: + return self.cfg_size > 1 + + @property + def sp_active(self) -> bool: + return self.sp_size > 1 + + +# Placeholder for the heavy primary KV in the side-payload dict; receiver swaps in the rebuilt object from the blob. +_KV_PLACEHOLDER = "__kv_placeholder__" + _SAFE_TORCH_DTYPES = { name: dtype for name in ( @@ -75,6 +131,8 @@ class OmniKVCacheConfig: recv_timeout: float = 30.0 from_tp: int = 1 to_tp: int = 1 + enable_kv_async_prefetch: bool = False + kv_prefetch_min_free_mem_ratio: float = 0.0 @dataclass @@ -290,7 +348,7 @@ class OmniKVTransferManager: - KV cache receiving with timeout """ - def __init__(self, config: OmniKVCacheConfig): + def __init__(self, config: OmniKVCacheConfig, *, async_prefetch: bool = False): self.config = config self._connector = None @@ -336,6 +394,19 @@ class OmniKVTransferManager: self._sender_base_host: str | None = None self._sender_base_zmq_port: int | None = None + # Prefetch + self._async_prefetch = async_prefetch + self._prefetch_min_free_mem_ratio: float = max(0.0, config.kv_prefetch_min_free_mem_ratio) + # Single-worker: serial mode ensures at most one outstanding prefetch + # and avoids stream-creation races on _bg_copy_stream. + self._prefetch_executor: ThreadPoolExecutor | None = ( + ThreadPoolExecutor(max_workers=1, thread_name_prefix="kv-prefetch") if async_prefetch else None + ) + self._prefetch_futures: dict[str, Any] = {} + self._bg_copy_stream: current_omni_platform.Stream | None = None + + self._topo_config: _TransferTopoConfig | None = None + if config.need_send_cache and config.connector_config: try: _ = self.connector @@ -343,15 +414,22 @@ class OmniKVTransferManager: except Exception as e: logger.warning("Failed to eagerly initialize sender connector: %s", e) + @property + def topo_config(self) -> _TransferTopoConfig: + """Lazily build the parallel topology snapshot on first access.""" + if self._topo_config is None: + self._topo_config = self._build_topo_config() + return self._topo_config + # ------------------------------------------------------------------ # # Factory helpers # ------------------------------------------------------------------ # @classmethod - def _create(cls, cfg: dict | None) -> "OmniKVTransferManager": + def _create(cls, cfg: dict | None, *, async_prefetch: bool = False) -> "OmniKVTransferManager": """Create manager from raw config dict.""" if not cfg or not isinstance(cfg, dict): - return cls(OmniKVCacheConfig()) + return cls(OmniKVCacheConfig(), async_prefetch=async_prefetch) rank_mapping = cfg.get("rank_mapping", {}) if not isinstance(rank_mapping, dict): @@ -369,38 +447,66 @@ class OmniKVTransferManager: recv_timeout=cfg.get("recv_timeout", 30.0), from_tp=int(rank_mapping.get("from_tp", 1)), to_tp=int(rank_mapping.get("to_tp", 1)), - ) + enable_kv_async_prefetch=async_prefetch, + kv_prefetch_min_free_mem_ratio=cfg.get("kv_prefetch_min_free_mem_ratio", 0.0), + ), + async_prefetch=async_prefetch, ) + @staticmethod + def _receiver_pool_on_device(omni_kv: Any) -> bool: + """True when the receiver connector pool is on a device (non-CPU).""" + if not isinstance(omni_kv, dict): + return False + cc = omni_kv.get("connector_config") + if not isinstance(cc, dict): + return False + return str(cc.get("memory_pool_device", "cpu")).lower() != "cpu" + + @classmethod + def _resolve_async_prefetch(cls, omni_kv: Any, *, has_companion: bool) -> bool: + """Effective prefetch flag: requested by config and not force-disabled.""" + if not isinstance(omni_kv, dict) or not bool(omni_kv.get("enable_kv_async_prefetch", False)): + return False + return not (has_companion or cls._receiver_pool_on_device(omni_kv)) + @classmethod def from_od_config(cls, config: Any) -> "OmniKVTransferManager": """Create from model or OmniDiffusion config.""" - return cls._create(getattr(config, "omni_kv_config", None)) + omni_kv = getattr(config, "omni_kv_config", None) + has_companion = getattr(config, "cfg_kv_collect_func", None) is not None + async_prefetch = cls._resolve_async_prefetch(omni_kv, has_companion=has_companion) + return cls._create(omni_kv, async_prefetch=async_prefetch) from_model_config = from_od_config @classmethod def from_vllm_config(cls, vllm_config: Any, model_config: Any) -> "OmniKVTransferManager": """Create from vllm config with fallback to kv_transfer_config.""" - # Primary: omni_kv_config from model_config - omni_kv = getattr(model_config, "omni_kv_config", None) - if isinstance(omni_kv, dict): - return cls._create(omni_kv) - - # Fallback: check kv_transfer_config - kv_cfg = getattr(vllm_config, "kv_transfer_config", None) - if kv_cfg: - direct = getattr(kv_cfg, "omni_connector_config", None) - if isinstance(direct, dict) and direct: - return cls._create({"connector_config": direct}) - extra = getattr(kv_cfg, "kv_connector_extra_config", None) - if isinstance(extra, dict): - omni = extra.get("omni_connector_config") - if isinstance(omni, dict) and omni: - return cls._create({"connector_config": omni}) + # Primary: omni_kv_config on model_config — same shape as from_od_config. + if isinstance(getattr(model_config, "omni_kv_config", None), dict): + return cls.from_od_config(model_config) + connector_cfg = cls._connector_cfg_from_kv_transfer(getattr(vllm_config, "kv_transfer_config", None)) + if connector_cfg is not None: + return cls._create({"connector_config": connector_cfg}) return cls(OmniKVCacheConfig()) + @staticmethod + def _connector_cfg_from_kv_transfer(kv_cfg: Any) -> dict | None: + """Extract the omni connector config dict nested in a kv_transfer_config.""" + if not kv_cfg: + return None + direct = getattr(kv_cfg, "omni_connector_config", None) + if isinstance(direct, dict) and direct: + return direct + extra = getattr(kv_cfg, "kv_connector_extra_config", None) + if isinstance(extra, dict): + omni = extra.get("omni_connector_config") + if isinstance(omni, dict) and omni: + return omni + return None + @property def connector(self): """Lazy initialization of connector.""" @@ -459,6 +565,88 @@ class OmniKVTransferManager: get_connector = property(lambda self: self.connector) + def _build_topo_config(self) -> _TransferTopoConfig: + """Build the parallel topology snapshot (called once at init).""" + from vllm_omni.diffusion.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_sp_group, + get_world_group, + ) + + try: + world = get_world_group() + world_size = world.world_size + world_rank = world.rank_in_group + except Exception: + logger.exception("World group unavailable; defaulting to LOCAL") + return _TransferTopoConfig( + role=ReceiveRole.LOCAL, + tp_active=False, + cfg_size=1, + cfg_rank=0, + cfg_group=None, + sp_size=1, + sp_rank=0, + sp_group=None, + world=None, + ) + + if world_size <= 1: + return _TransferTopoConfig( + role=ReceiveRole.LOCAL, + tp_active=False, + cfg_size=1, + cfg_rank=0, + cfg_group=None, + sp_size=1, + sp_rank=0, + sp_group=None, + world=world, + ) + + tp_active = self._tp_topo.source_tp_size > 1 or self._tp_topo.target_tp_size > 1 + + # CFG/SP groups may legitimately be uninitialized on pure-TP or + # world-only distributed paths; treat that as "not active" (size 1) + # rather than raising, so those paths can still take the local/world + # receive route. + try: + cfg_size = get_classifier_free_guidance_world_size() + cfg_rank = get_classifier_free_guidance_rank() + cfg_group = get_cfg_group() + except Exception: + cfg_size, cfg_rank, cfg_group = 1, 0, None # CFG-parallel not enabled + + try: + sp_size = get_sequence_parallel_world_size() + sp_rank = get_sequence_parallel_rank() + sp_group = get_sp_group() + except Exception: + sp_size, sp_rank, sp_group = 1, 0, None # SP not enabled + + if tp_active and cfg_size <= 1 and sp_size <= 1: + role = ReceiveRole.LOCAL + elif tp_active and (cfg_size > 1 or sp_size > 1): + role = ReceiveRole.LEADER if (cfg_rank == 0 and sp_rank == 0) else ReceiveRole.FOLLOWER + else: + role = ReceiveRole.LEADER if world_rank == 0 else ReceiveRole.FOLLOWER + + return _TransferTopoConfig( + role=role, + tp_active=tp_active, + cfg_size=cfg_size, + cfg_rank=cfg_rank, + cfg_group=cfg_group, + sp_size=sp_size, + sp_rank=sp_rank, + sp_group=sp_group, + world=world, + ) + def _resolve_sender_info( self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None ) -> dict[str, Any] | None: @@ -502,19 +690,6 @@ class OmniKVTransferManager: ) return None - @staticmethod - def _clone_received_payload_tensors(data: dict[str, Any]) -> dict[str, Any]: - if not isinstance(data, dict) or "layer_blocks" not in data: - return data - - layer_blocks = data["layer_blocks"] - for cache_name in ("key_cache", "value_cache"): - cache_list = layer_blocks.get(cache_name, []) - for idx, tensor in enumerate(cache_list): - if isinstance(tensor, torch.Tensor): - cache_list[idx] = tensor.clone() - return data - def _slice_transfer_data_for_target(self, kv_data: KVCacheTransferData, target_rank: int) -> KVCacheTransferData: """Pre-slice sender payload for one target rank when sender TP < receiver TP.""" topo = self._tp_topo @@ -984,24 +1159,183 @@ class OmniKVTransferManager: return False, 0, None + def _resolve_sender_base(self, sender_info: dict[str, Any] | None) -> tuple[str | None, int | None]: + """Resolve (host, base_zmq_port) from sender_info.""" + actual = self._resolve_sender_info(sender_info, sender_stage_id=self.recv_stages[0]) + if not actual or "host" not in actual: + return None, None + port = actual.get("zmq_port") + return actual.get("host"), (int(port) if port is not None else None) + + def _has_enough_prefetch_device_mem(self, device: torch.device | None) -> bool: + """False when free-device-memory fraction is below the prefetch threshold. + + CUDA/NPU only; other platforms report unconstrained (throttle skipped). + """ + ratio = self._prefetch_min_free_mem_ratio + if ratio <= 0.0 or device is None: + return True + try: + free, total = current_omni_platform.get_device_memory(device) + except Exception: + return True + return total <= 0 or (free / total) >= ratio + + @staticmethod + def _release_pool_buffers(buffers: list[Any]) -> None: + """Release all pool buffers in *buffers* and clear the list.""" + if not buffers: + return + for buf in buffers: + try: + buf.release() + except Exception: + logger.exception("Failed to release KV pool buffer") + buffers.clear() + + def start_prefetch( + self, kv_prefetch_jobs: dict[str, Any] | None, target_device: torch.device | None = None + ) -> None: + """Kick off a background KV load (non-blocking). No-op unless prefetch enabled.""" + if not (self._async_prefetch and self.config.need_recv_cache) or not kv_prefetch_jobs: + return + # Followers receive via collective distribute; bg pull would consume the owner's payload. + if self.topo_config.is_follower: + return + rid = kv_prefetch_jobs.get("request_id") + if not rid: + return + # Memory pressure → skip prefetch; sync receive handles it. + if not self._has_enough_prefetch_device_mem(target_device): + logger.warning( + "Skip KV prefetch for %s: device free mem below %.2f", rid, self._prefetch_min_free_mem_ratio + ) + return + # Serial mode: at most one outstanding prefetch; drop any leftover request's future. + for stale_rid in [k for k in self._prefetch_futures if k != rid]: + self._discard_future(stale_rid) + if rid in self._prefetch_futures: + return + sender_info = kv_prefetch_jobs.get("kv_sender_info") + if not sender_info: + # No explicit endpoint → bg receive would target wrong sender under multi-replica. + logger.debug("Skip KV prefetch for %s: stub has no kv_sender_info", rid) + return + try: + self._prefetch_futures[rid] = self._prefetch_executor.submit( + self._prefetch_payload, rid, sender_info, target_device + ) + except Exception: + logger.exception("Failed to submit KV prefetch for %s", rid) + + def _prefetch_payload( + self, + request_id: str, + sender_info: dict[str, Any] | None, + target_device: torch.device | None, + ) -> tuple[dict[str, Any] | None, int]: + """Bg-thread body: get + deserialize + H2D on the dedicated ``_bg_copy_stream``. + + Raises on failure (payload may be consumed → no sync retry). + """ + try: + on_device = target_device is not None and target_device.type != "cpu" + if on_device: + # bg thread doesn't inherit the main thread's current device. + torch.accelerator.set_device_index(target_device.index) + if self._bg_copy_stream is None: + assert self._prefetch_executor._max_workers == 1 + self._bg_copy_stream = current_omni_platform.Stream() + with current_omni_platform.stream(self._bg_copy_stream): + data, size = self.receive_kv_cache_for_request( + request_id, target_device=target_device, sender_info=sender_info + ) + else: + data, size = self.receive_kv_cache_for_request( + request_id, target_device=target_device, sender_info=sender_info + ) + except Exception: + logger.exception("KV prefetch payload failed for %s (payload may be lost)", request_id) + raise + + if data is None: + return None, 0 + return data, size + + def consume_prefetched_kv(self, req: Any) -> tuple[dict[str, Any] | None, int]: + """Consume a prefetched KV payload; (None, 0) on miss/recoverable failure. + + Raises ``KVPrefetchConsumeError`` when consumed but post-get failed (no fallback). + """ + request_id = self._resolve_request_id(req) + if not request_id: + return None, 0 + + fut = self._prefetch_futures.pop(request_id, None) + # Serial mode: any other request still in the table is an orphan — drop it. + for stale_rid in list(self._prefetch_futures): + self._discard_future(stale_rid) + + if fut is None: + return None, 0 + + try: + return fut.result() + except KVPrefetchConsumeError: + logger.exception("KV load failed for %s (payload consumed, cannot retry)", request_id) + raise + except Exception: + logger.exception("KV load failed for %s; falling back to sync receive", request_id) + return None, 0 + + def _discard_future(self, request_id: str) -> None: + """Cancel an unstarted prefetch or attach a callback to drop a running one.""" + fut = self._prefetch_futures.pop(request_id, None) + if fut is None: + return + if not fut.cancel(): + fut.add_done_callback(_drop_prefetch_result) + + def shutdown_prefetch(self) -> None: + """Cancel pending prefetches and stop the executor (call on teardown).""" + for rid in list(self._prefetch_futures): + self._discard_future(rid) + if self._prefetch_executor is not None: + try: + self._prefetch_executor.shutdown(wait=True, cancel_futures=True) + except Exception: + logger.exception("Failed to shut down KV prefetch executor") + self._prefetch_executor = None + + @staticmethod + def _record_stream_for_prefetched(data: dict[str, Any]) -> None: + """``record_stream(current_stream)`` on GPU tensors in *data*. + + Protects prefetch-thread tensors from allocator reuse on LEADER collective paths. + """ + if not isinstance(data, dict) or "layer_blocks" not in data: + return + if current_omni_platform.get_device_count() < 1: + return + current_stream = current_omni_platform.current_stream() + layer_blocks = data["layer_blocks"] + for cache_list in (layer_blocks.get("key_cache", []), layer_blocks.get("value_cache", [])): + for tensor in cache_list: + if isinstance(tensor, torch.Tensor) and tensor.device.type != "cpu": + tensor.record_stream(current_stream) + @torch.inference_mode() def receive_kv_cache_for_request( self, request_id: str, target_device: torch.device | None = None, + *, + sender_info: dict[str, Any] | None = None, ) -> tuple[dict[str, Any] | None, int]: - """Receive KV cache for a specific request. + """Receive KV cache for *request_id*; returns (data dict, size) or (None, 0). - This implements the receiving logic from gpu_diffusion_model_runner.py. - - Args: - request_id: The request ID to receive KV cache for - target_device: Optional device to move tensors to - - Returns: - Tuple of (data dict, size) if successful, (None, 0) otherwise + ``sender_info`` overrides instance ``_sender_base_*`` (per-call endpoint). """ - # Check if we should receive KV cache based on config if not self.config.need_recv_cache: logger.debug("Skip receiving KV cache for %s (need_recv_cache=False)", request_id) return None, 0 @@ -1031,6 +1365,13 @@ class OmniKVTransferManager: ) pending_pairs = list(recv_key_pairs) received_payloads: dict[str, tuple[dict[str, Any], int]] = {} + deferred_memory: list[Any] = [] + + # Per-call sender_info takes precedence over instance fields. + if sender_info is not None: + base_host, base_port = self._resolve_sender_base(sender_info) + else: + base_host, base_port = self._sender_base_host, self._sender_base_zmq_port logger.info( "Wait for KV cache for request %s from stage %s to %s via %s key(s)...", @@ -1044,16 +1385,21 @@ class OmniKVTransferManager: while True: link_start = time.perf_counter() for get_key, from_rank in list(pending_pairs): - # Construct per-rank metadata so the connector queries - # the correct sender endpoint (heterogeneous TP path). - # When from_rank is None (TP<=1), metadata stays None - # and the connector falls back to its default sender. rank_metadata: dict[str, Any] | None = None - if from_rank is not None and self._sender_base_host and self._sender_base_zmq_port is not None: + build_meta = from_rank is not None or sender_info is not None + if build_meta and base_host and base_port is not None: rank_metadata = { - "source_host": self._sender_base_host, - "source_port": self._sender_base_zmq_port + from_rank * KV_RANK_PORT_STRIDE, + "source_host": base_host, + "source_port": base_port + (from_rank or 0) * KV_RANK_PORT_STRIDE, } + elif sender_info is not None: + logger.error( + "KV receive for %s: unresolved per-call sender_info (host=%s, port=%s); aborting", + request_id, + base_host, + base_port, + ) + return None, 0 result = self.connector.get( from_stage=from_stage, @@ -1065,25 +1411,19 @@ class OmniKVTransferManager: continue raw_data, size = result - managed_buffer = None if hasattr(raw_data, "tensor") and hasattr(raw_data, "release"): - managed_buffer = raw_data try: buf_tensor = raw_data.tensor if getattr(buf_tensor.device, "type", "cpu") != "cpu": data = KVCacheTransferData.from_bytes_device(buf_tensor) raw_data.release() - managed_buffer = None else: data = KVCacheTransferData.from_bytes(memoryview(buf_tensor.numpy())) - data = self._clone_received_payload_tensors(data) - raw_data.release() - managed_buffer = None + deferred_memory.append(raw_data) except Exception as e: logger.error("Failed to deserialize KV cache from ManagedBuffer: %s", e) - if managed_buffer is not None: - raw_data.release() + raw_data.release() return None, 0 elif isinstance(raw_data, (bytes, bytearray)): data = KVCacheTransferData.from_bytes(raw_data) @@ -1110,20 +1450,27 @@ class OmniKVTransferManager: data = merge_received_rank_shards(ordered_payloads, merger=self.kv_payload_merger) data = slice_received_rank_shard(data, topo, slicer=self.kv_payload_slicer) + needs_clone = bool(deferred_memory) try: if isinstance(data, dict) and "layer_blocks" in data: layer_blocks = data["layer_blocks"] - for cache_list in [ + cache_lists = [ layer_blocks.get("key_cache", []), layer_blocks.get("value_cache", []), - ]: + ] + for cache_list in cache_lists: for i, tensor in enumerate(cache_list): if not isinstance(tensor, torch.Tensor): continue if target_device is not None and tensor.device != target_device: cache_list[i] = tensor.to(target_device).contiguous() - except Exception: - logger.exception("Failed to move KV cache tensors to target device") + elif needs_clone: + cache_list[i] = tensor.clone() + except Exception as exc: + logger.exception("Failed to detach/move KV cache tensors for %s", request_id) + raise KVPrefetchConsumeError( + f"Post-get processing failed for {request_id} (payload already consumed)" + ) from exc logger.info( "Successfully received KV cache for %s, %s bytes across %s key(s), wait=%.3fs, link=%.1fms", @@ -1142,9 +1489,13 @@ class OmniKVTransferManager: time.sleep(poll_interval) poll_interval = min(poll_interval * 2, max_poll_interval) + except KVPrefetchConsumeError: + raise except Exception: logger.exception("Error receiving KV cache for %s", request_id) return None, 0 + finally: + self._release_pool_buffers(deferred_memory) def apply_kv_cache_to_request(self, req: Any, data: dict[str, Any]) -> None: """Apply received KV cache data to a request object. @@ -1246,130 +1597,253 @@ class OmniKVTransferManager: return primary_ok + def _extract_primary_kv_obj(self, kv_payload: dict[str, Any]) -> Any | None: + """Return the primary KV object (``SimpleNamespace(key_cache, value_cache)``).""" + for key in ("past_key_values", "sp.past_key_values"): + obj = kv_payload.get(key) + if obj is not None and hasattr(obj, "key_cache"): + return obj + return None + + def _pack_kv_payload(self, kv_payload: dict[str, Any] | None, device: torch.device) -> torch.Tensor | None: + """Pack a KV payload into one uint8 tensor on device + Layout: ``[4B side_len][msgpack side_bytes][gpu blob]``. + Returns None on failure.""" + if not kv_payload: + return None + primary = self._extract_primary_kv_obj(kv_payload) + if primary is None: + return None + + key_cache = _move_to_device(getattr(primary, "key_cache", []) or [], device) + value_cache = _move_to_device(getattr(primary, "value_cache", []) or [], device) + if not any(isinstance(t, torch.Tensor) for t in key_cache): + return None + td = KVCacheTransferData( + request_id="", + layer_blocks={"key_cache": key_cache, "value_cache": value_cache}, + block_ids=[], + metadata={}, + ) + try: + blob = td.to_gpu_tensor() + except Exception: + logger.exception("D2D: failed to pack primary KV; falling back to object transport") + return None + + # Side-payload: replace the heavy primary with a sentinel, msgpack the rest. + side = dict(kv_payload) + for key in ("past_key_values", "sp.past_key_values"): + if side.get(key) is primary: + side[key] = _KV_PLACEHOLDER + + try: + side_bytes = OmniSerializer.serialize(side) + except Exception: + logger.exception("D2D: failed to serialize side-payload; falling back to object transport") + return None + side_len = len(side_bytes) + side_tensor = torch.frombuffer(bytearray(side_bytes), dtype=torch.uint8).to(device) + + # [4B side_len][side_bytes][blob] + header = torch.tensor([side_len], dtype=torch.int32, device=device).view(torch.uint8) + combined = torch.cat([header, side_tensor, blob]) + return combined + + def _unpack_kv_payload(self, combined: torch.Tensor) -> dict[str, Any] | None: + """Inverse of :meth:`_pack_kv_payload`.""" + if combined is None or combined.numel() == 0: + return None + + offset = 4 + side_len = int(combined[:offset].cpu().view(torch.int32).item()) + if side_len <= 0 and combined.numel() > offset: + # Corrupt/zeroed blob (sender-side failure) → treat as miss; raising + # would crash only the receiver and desync the group. + logger.warning( + "D2D: corrupt KV payload (side_len=%d, total_bytes=%d, likely sender-side failure); treating as miss", + side_len, + combined.numel(), + ) + return None + side_bytes = combined[offset : offset + side_len].cpu().numpy().tobytes() + offset += side_len + side: dict[str, Any] = OmniSerializer.deserialize(side_bytes) + + # Rebuild primary KV from the trailing blob and restore it in place. + primary: Any = None + blob = combined[offset:] + if blob.numel() > 0: + from types import SimpleNamespace + + data = KVCacheTransferData.from_bytes_device(blob) + lb = data["layer_blocks"] + primary = SimpleNamespace(key_cache=lb.get("key_cache"), value_cache=lb.get("value_cache")) + + for key in ("past_key_values", "sp.past_key_values"): + if side.get(key) == _KV_PLACEHOLDER: + side[key] = primary + return side + + def _broadcast_kv_payload( + self, group: Any, kv_payload: dict[str, Any] | None, device: torch.device, src: int = 0 + ) -> dict[str, Any] | None: + """Broadcast KV payload via packed tensor; falls back to ``broadcast_object`` on packing failure. + + The packed-tensor path uses a two-phase collective (size then data). + Both phases must always be called together — a rank entering only + one phase will deadlock the group. + """ + is_src = group.rank_in_group == src + if is_src: + combined = self._pack_kv_payload(kv_payload, device) + if combined is None: + group.broadcast_object(kv_payload, src=src) # packing failed -> object + return kv_payload + group.broadcast_object(int(combined.numel()), src=src) + group.broadcast(combined, src=src) + return kv_payload + + size_obj = group.broadcast_object(None, src=src) + if size_obj is None or not isinstance(size_obj, int): + return size_obj # src fell back to object + combined = torch.empty(size_obj, dtype=torch.uint8, device=device) + group.broadcast(combined, src=src) + return self._unpack_kv_payload(combined) + + def _send_kv_payload(self, group: Any, kv_payload: dict[str, Any] | None, dst: int, device: torch.device) -> None: + """Point-to-point send; falls back to ``send_object`` on packing failure.""" + combined = self._pack_kv_payload(kv_payload, device) + if combined is None: + group.send_object(kv_payload, dst) + return + size = int(combined.numel()) + group.send_object(size, dst) + try: + group.send(combined, dst) + except Exception: + logger.exception("D2D: send(combined) to rank %d failed; sending zero sentinel", dst) + combined.zero_() + group.send(combined, dst) + raise + + def _recv_kv_payload(self, group: Any, src: int, device: torch.device) -> dict[str, Any] | None: + """Inverse of :meth:`_send_kv_payload`.""" + size_or_obj = group.recv_object(src) + if not isinstance(size_or_obj, int): + return size_or_obj + combined = group.recv(torch.Size([size_or_obj]), torch.uint8, src) + combined = combined.to(device) + return self._unpack_kv_payload(combined) + def receive_multi_kv_cache_distributed( self, req: Any, cfg_kv_collect_func: Callable | None = None, target_device: torch.device | None = None, ) -> bool: - """Distributed wrapper around :meth:`receive_multi_kv_cache`. - - TP-aware path selection: - - world size 1: direct receive - - TP active, cfg size 1: each rank independently receives - - TP active, cfg size > 1: cfg-rank 0 receives, then broadcasts to - peers that share the same TP rank - - TP inactive: legacy rank-0 receive then world broadcast - """ - from vllm_omni.diffusion.distributed.parallel_state import ( - get_cfg_group, - get_classifier_free_guidance_rank, - get_classifier_free_guidance_world_size, - get_sequence_parallel_rank, - get_sequence_parallel_world_size, - get_sp_group, - get_world_group, - ) - - world = get_world_group() - - if world.world_size <= 1: - return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) - - topo = self._tp_topo - tp_active = topo.source_tp_size > 1 or topo.target_tp_size > 1 - cfg_size = 1 - cfg_rank = 0 - cfg_group = None - sp_size = 1 - sp_rank = 0 - sp_group = None - try: - cfg_size = get_classifier_free_guidance_world_size() - cfg_rank = get_classifier_free_guidance_rank() - cfg_group = get_cfg_group() - except Exception: - cfg_size = 1 - cfg_rank = 0 - cfg_group = None - - try: - sp_size = get_sequence_parallel_world_size() - sp_rank = get_sequence_parallel_rank() - sp_group = get_sp_group() - except Exception: - sp_size = 1 - sp_rank = 0 - sp_group = None - - if tp_active and (cfg_size <= 1 and sp_size <= 1): - logger.info( - "Rank-aware KV receive: rank %s independently receiving (from_tp=%s, to_tp=%s)", - topo.local_rank, - topo.source_tp_size, - topo.target_tp_size, - ) - return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) - - if tp_active and (cfg_size > 1 or sp_size > 1): - kv_payload: dict[str, object] | None = None - is_owner = cfg_rank == 0 and sp_rank == 0 - - # step1: only owner need to receive - if is_owner: - received = self.receive_multi_kv_cache( - req, - cfg_kv_collect_func, - torch.device("cpu"), - ) - if received: - if cfg_size > 1: # build cfg-local payloads - cfg_rank_payloads = self._build_cfg_rank_local_payloads( - req, - cfg_size, - ) - kv_payload = cfg_rank_payloads[0] - # send to other cfg - for dst_cfg_rank in range(1, cfg_size): - cfg_group.send_object( - cfg_rank_payloads[dst_cfg_rank], - dst_cfg_rank, - ) - elif sp_size > 1: - kv_payload = self._collect_request_kv_payload(req) - else: - # Owner receive failed: send None sentinel to CFG followers to avoid deadlock - if cfg_size > 1: - for dst_cfg_rank in range(1, cfg_size): - cfg_group.send_object(None, dst_cfg_rank) - elif sp_rank == 0 and cfg_size > 1: - kv_payload = cfg_group.recv_object(0) - # sp broadcast - if sp_size > 1 and sp_group is not None: - kv_payload = sp_group.broadcast_object( - kv_payload, - src=0, - ) - - if not kv_payload: - return False - + """Synchronous receive: fetch → distribute → apply (FOLLOWER skips fetch).""" + received = False + if not self.topo_config.is_follower: + received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) + kv_payload = self.distribute_kv_cache(req, target_device, received=received) + if kv_payload is not None: self._apply_request_kv_payload(req, kv_payload, target_device) - return True + return received or kv_payload is not None - kv_payload: dict[str, object] | None = None - if world.rank_in_group == 0: - received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu")) - if received: - kv_payload = self._collect_request_kv_payload(req) + def consume_and_distribute_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool: + """Consume prefetched KV → apply → distribute; sync-receive on miss (no retry on consume-error).""" + received = False + payload_consumed = False + if self._async_prefetch and not self.topo_config.is_follower: + try: + data, _ = self.consume_prefetched_kv(req) + if data is not None: + self._record_stream_for_prefetched(data) + self.apply_kv_cache_to_request(req, data) + received = True + except KVPrefetchConsumeError: + logger.error( + "KV prefetch consumed payload for %s but post-get failed; request cannot recover", + self._resolve_request_id(req), + ) + payload_consumed = True + # Only fall back to sync receive on miss, not when payload was + # already consumed from the connector (sync receive would block + # until timeout as the data is gone). + if not received and not self.topo_config.is_follower and not payload_consumed: + logger.debug("KV prefetch miss for %s; falling back to sync receive", self._resolve_request_id(req)) + received = self.receive_multi_kv_cache(req, None, target_device) + kv_payload = self.distribute_kv_cache(req, target_device, received=received) + if kv_payload is not None: + self._apply_request_kv_payload(req, kv_payload, target_device) + return received or kv_payload is not None - kv_payload = world.broadcast_object(kv_payload, src=0) + def _cfg_scatter( + self, + group: Any, + payloads: list[Any], + cfg_size: int, + device: torch.device | None, + request_id: str | None = None, + ) -> None: + """Scatter one payload per CFG follower.""" + for dst in range(1, cfg_size): + try: + self._send_kv_payload(group, payloads[dst], dst, device) + except Exception: + logger.warning("CFG scatter to rank %d failed for %s; that rank runs without KV", dst, request_id) + continue - if not kv_payload: - return False + def distribute_kv_cache( + self, req: Any, target_device: torch.device | None = None, *, received: bool = False + ) -> dict[str, Any] | None: + """Distribute KV LEADER→FOLLOWER. No-op for LOCAL/pure TP.""" + pt = self.topo_config - self._apply_request_kv_payload(req, kv_payload, target_device) - return True + # LOCAL: no distribution. + if pt.is_local: + return None + + device = target_device if target_device is not None else torch.device("cpu") + # TP + CFG/SP + if pt.tp_active and (pt.cfg_active or pt.sp_active): + kv_payload: dict[str, object] | None = None + if pt.is_leader: + if pt.cfg_active: + payloads = ( + self._build_cfg_rank_local_payloads(req, pt.cfg_size) if received else [None] * pt.cfg_size + ) + kv_payload = payloads[0] + self._cfg_scatter(pt.cfg_group, payloads, pt.cfg_size, device, self._resolve_request_id(req)) + elif pt.sp_active: + kv_payload = self._collect_request_kv_payload(req) if received else None + if pt.sp_active and pt.sp_group is not None: + kv_payload = self._broadcast_kv_payload(pt.sp_group, kv_payload, device, src=0) + elif pt.is_follower: + if pt.cfg_active and pt.sp_rank == 0: + kv_payload = self._recv_kv_payload(pt.cfg_group, 0, device) + if pt.sp_active and pt.sp_group is not None: + kv_payload = self._broadcast_kv_payload(pt.sp_group, kv_payload, device, src=0) + else: + logger.error("distribute_kv_cache: unexpected role %s in TP+CFG/SP path", pt.role) + return None + return kv_payload or None + + # TP inactive, world broadcast + kv_payload = None + if pt.is_leader and received: + kv_payload = self._collect_request_kv_payload(req) + kv_payload = self._broadcast_kv_payload(pt.world, kv_payload, device, src=0) + return kv_payload or None + + +def _drop_prefetch_result(fut: Any) -> None: + try: + if not fut.cancelled(): + fut.result() + except Exception: + pass def _move_to_device(obj: object, device: torch.device) -> object: diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 3902b11ba..97075c434 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -595,7 +595,6 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( "--boundary-ratio", diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index c96da69fc..38d74627c 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -213,6 +213,11 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): free, _ = torch.cuda.mem_get_info(device) return free + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + free, total = torch.cuda.mem_get_info(device) + return free, total + @classmethod def get_device_name(cls, device_id: int = 0) -> str: return torch.cuda.get_device_name(device_id) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 88f7eba7b..54725b553 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -161,6 +161,10 @@ class OmniPlatform(Platform): def get_free_memory(cls, device: torch.device | None = None) -> int: raise NotImplementedError + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + raise NotImplementedError + @classmethod def create_autocast_context( cls, diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 75f20a5a7..96e27edf3 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -146,17 +146,14 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): @classmethod def get_free_memory(cls, device: torch.device | None = None) -> int: - """Get the free memory on the MUSA device. - - Args: - device: Optional device to query. If None, uses current device. - - Returns: - Free memory in bytes. - """ free, _ = torch.musa.mem_get_info(device) return free + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + free, total = torch.musa.mem_get_info(device) + return free, total + @classmethod def get_device_name(cls, device_id: int = 0) -> str: return torch.musa.get_device_name(device_id) diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index cce70c7ec..f03fdb1bf 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -149,6 +149,11 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): free, _ = torch.npu.mem_get_info(device) return free + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + free, total = torch.npu.mem_get_info(device) + return free, total + @classmethod def get_device_total_memory(cls, device_id: int = 0) -> int: device_props = torch.npu.get_device_properties(device_id) diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index 1b50c328b..f13362eea 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/vllm_omni/platforms/rocm/platform.py @@ -142,6 +142,11 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform): free, _ = torch.cuda.mem_get_info(device) return free + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + free, total = torch.cuda.mem_get_info(device) + return free, total + @classmethod def set_device_control_env_var(cls, devices: str | int | None) -> None: import os diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 5500e013f..79d63b4b8 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -85,6 +85,11 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform): free, _ = torch.xpu.mem_get_info(device) return free + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + free, total = torch.xpu.mem_get_info(device) + return free, total + @classmethod def get_profiler_cls(cls) -> str: """Return XPU-specific profiler that handles XPU events."""