From d98a333991ec4b8c3dbcb7565a59a758550dbf15 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 26 Nov 2025 15:08:16 +0800 Subject: [PATCH 1/3] Allow define EC Tensor device Signed-off-by: knlnguyen1802 --- .../disaggregated_encoder/disagg_1e1p1d_example.sh | 2 ++ .../disaggregated_encoder/disagg_1e1pd_example.sh | 2 ++ .../ec_transfer/ec_connector/shared_storage_connector.py | 4 +++- 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/examples/online_serving/disaggregated_encoder/disagg_1e1p1d_example.sh b/examples/online_serving/disaggregated_encoder/disagg_1e1p1d_example.sh index 57489df64f51..4a22be1da33c 100644 --- a/examples/online_serving/disaggregated_encoder/disagg_1e1p1d_example.sh +++ b/examples/online_serving/disaggregated_encoder/disagg_1e1p1d_example.sh @@ -104,6 +104,7 @@ CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \ --ec-transfer-config '{ "ec_connector": "ECSharedStorageConnector", "ec_role": "ec_producer", + "ec_buffer_device": "cuda", "ec_connector_extra_config": { "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'" } @@ -128,6 +129,7 @@ vllm serve "$MODEL" \ --ec-transfer-config '{ "ec_connector": "ECSharedStorageConnector", "ec_role": "ec_consumer", + "ec_buffer_device": "cuda", "ec_connector_extra_config": { "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'" } diff --git a/examples/online_serving/disaggregated_encoder/disagg_1e1pd_example.sh b/examples/online_serving/disaggregated_encoder/disagg_1e1pd_example.sh index 6073e0580b11..344f4a2ffe58 100644 --- a/examples/online_serving/disaggregated_encoder/disagg_1e1pd_example.sh +++ b/examples/online_serving/disaggregated_encoder/disagg_1e1pd_example.sh @@ -98,6 +98,7 @@ CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \ --ec-transfer-config '{ "ec_connector": "ECSharedStorageConnector", "ec_role": "ec_producer", + "ec_buffer_device": "cuda", "ec_connector_extra_config": { "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'" } @@ -119,6 +120,7 @@ CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \ --ec-transfer-config '{ "ec_connector": "ECSharedStorageConnector", "ec_role": "ec_consumer", + "ec_buffer_device": "cuda", "ec_connector_extra_config": { "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'" } diff --git a/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py b/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py index c8388141dcc9..7c9fa245a138 100644 --- a/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py +++ b/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py @@ -51,10 +51,12 @@ def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole): # req_id -> index self._mm_datas_need_loads: dict[str, int] = {} transfer_config = vllm_config.ec_transfer_config + self.device = "cuda" if transfer_config is not None: self._storage_path = transfer_config.get_from_extra_config( "shared_storage_path", "/tmp" ) + self.device = transfer_config.ec_buffer_device logger.debug(transfer_config) logger.debug("Shared storage path is %s", self._storage_path) else: @@ -91,7 +93,7 @@ def start_load_caches(self, encoder_cache, **kwargs) -> None: if mm_data.mm_hash in encoder_cache: continue filename = self._generate_filename_debug(mm_data.mm_hash) - ec_cache = safetensors.torch.load_file(filename)["ec_cache"].cuda() + ec_cache = safetensors.torch.load_file(filename)["ec_cache"].to(self.device) encoder_cache[mm_data.mm_hash] = ec_cache logger.debug("Success load encoder cache for hash %s", mm_data.mm_hash) From a076bf7814abdf19e34540258926c83a470c86ff Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 26 Nov 2025 15:17:02 +0800 Subject: [PATCH 2/3] Fix pre-commit Signed-off-by: knlnguyen1802 --- .../ec_transfer/ec_connector/shared_storage_connector.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py b/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py index 7c9fa245a138..f3c9aac9f160 100644 --- a/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py +++ b/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py @@ -56,7 +56,8 @@ def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole): self._storage_path = transfer_config.get_from_extra_config( "shared_storage_path", "/tmp" ) - self.device = transfer_config.ec_buffer_device + if transfer_config.ec_buffer_device is not None: + self.device = transfer_config.ec_buffer_device logger.debug(transfer_config) logger.debug("Shared storage path is %s", self._storage_path) else: From 1ec5e2beae7264058575f45e614d828e87df030b Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 26 Nov 2025 15:22:13 +0800 Subject: [PATCH 3/3] Fix pre-commit Signed-off-by: knlnguyen1802 --- .../ec_transfer/ec_connector/shared_storage_connector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py b/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py index f3c9aac9f160..77211ae3b344 100644 --- a/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py +++ b/vllm/distributed/ec_transfer/ec_connector/shared_storage_connector.py @@ -57,7 +57,7 @@ def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole): "shared_storage_path", "/tmp" ) if transfer_config.ec_buffer_device is not None: - self.device = transfer_config.ec_buffer_device + self.device = transfer_config.ec_buffer_device logger.debug(transfer_config) logger.debug("Shared storage path is %s", self._storage_path) else: