[Bugfix][EC Connector] Fix ECExampleConnector load device under TP>1 (#40416)
Signed-off-by: Hang Gong <kidogong@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Codex <noreply@openai.com>
Hang Gong committed
Sep 8, 2026 at 20:44 UTC
bf7eee9d7b2788bafe69f0c21fde6af22a2d3319
2 files changed
+89
-3
tests/v1/ec_connector/unit/test_ec_example_connector.py
+83
@@ -477,6 +477,89 @@ class TestCacheLoading:
477
# Cache should remain empty
478
assert len(encoder_cache) == 0
479
480
+ def test_start_load_caches_pins_device_to_current_rank(
481
+ self, mock_vllm_config_producer, mock_vllm_config_consumer
482
+ ):
483
+ """Load must pin to the selected device under tensor parallelism."""
484
+ # Produce a real cache file so the load path has something to open.
485
+ producer = ECExampleConnector(
486
+ vllm_config=mock_vllm_config_producer,
487
+ role=ECConnectorRole.WORKER,
488
+ )
489
+ mm_hash = "rank_pinning_hash"
490
+ producer.save_caches({mm_hash: torch.zeros(4, 8)}, mm_hash)
491
+
492
+ consumer = ECExampleConnector(
493
+ vllm_config=mock_vllm_config_consumer,
494
+ role=ECConnectorRole.WORKER,
495
+ )
496
+ metadata = ECExampleConnectorMetadata()
497
+ metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100))
498
+ consumer.bind_connector_metadata(metadata)
499
+
500
+ mock_platform = Mock()
501
+ mock_platform.device_type = "cuda"
502
+ mock_platform.is_cuda_alike.return_value = True
503
+ mock_platform.current_device.return_value = 2
504
+
505
+ # Simulate running on TP rank 2 without requiring a multi-GPU host.
506
+ captured: dict = {}
507
+
508
+ def fake_load_file(filename, device):
509
+ captured["device"] = device
510
+ return {"ec_cache": torch.zeros(4, 8)}
511
+
512
+ with (
513
+ patch("vllm.platforms.current_platform", mock_platform),
514
+ patch("safetensors.torch.load_file", side_effect=fake_load_file),
515
+ ):
516
+ encoder_cache: dict[str, torch.Tensor] = {}
517
+ consumer.start_load_caches(encoder_cache=encoder_cache)
518
+
519
+ assert captured["device"] == "cuda:2", (
520
+ f"Expected device='cuda:2' (pinned to current rank), "
521
+ f"got {captured['device']!r}. A bare 'cuda' string makes "
522
+ "safetensors load onto cuda:0, which breaks non-rank-0 TP workers."
523
+ )
524
+
525
+ def test_start_load_caches_non_cuda_platform_uses_device_type(
526
+ self, mock_vllm_config_producer, mock_vllm_config_consumer
527
+ ):
528
+ """Non-CUDA platforms pass `device_type` through unchanged (no rank suffix)."""
529
+ producer = ECExampleConnector(
530
+ vllm_config=mock_vllm_config_producer,
531
+ role=ECConnectorRole.WORKER,
532
+ )
533
+ mm_hash = "cpu_device_hash"
534
+ producer.save_caches({mm_hash: torch.zeros(2, 4)}, mm_hash)
535
+
536
+ consumer = ECExampleConnector(
537
+ vllm_config=mock_vllm_config_consumer,
538
+ role=ECConnectorRole.WORKER,
539
+ )
540
+ metadata = ECExampleConnectorMetadata()
541
+ metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100))
542
+ consumer.bind_connector_metadata(metadata)
543
+
544
+ mock_platform = Mock()
545
+ mock_platform.device_type = "cpu"
546
+ mock_platform.is_cuda_alike.return_value = False
547
+
548
+ captured: dict = {}
549
+
550
+ def fake_load_file(filename, device):
551
+ captured["device"] = device
552
+ return {"ec_cache": torch.zeros(2, 4)}
553
+
554
+ with (
555
+ patch("vllm.platforms.current_platform", mock_platform),
556
+ patch("safetensors.torch.load_file", side_effect=fake_load_file),
557
+ ):
558
+ encoder_cache: dict[str, torch.Tensor] = {}
559
+ consumer.start_load_caches(encoder_cache=encoder_cache)
560
+
561
+ assert captured["device"] == "cpu"
562
+
563
564
class TestFilenameGeneration:
565
"""Test filename and path generation."""
vllm/distributed/ec_transfer/ec_connector/example_connector.py
+6
-3
@@ -89,14 +89,17 @@ class ECExampleConnector(ECConnectorBase):
89
"In connector.start_load_caches, but the connector metadata is None"
90
)
91
return
92
+ # A bare "cuda" makes safetensors load onto cuda:0. Pin the load to the
93
+ # selected device so tensor-parallel workers load the cache locally.
94
+ device = current_platform.device_type
95
+ if current_platform.is_cuda_alike():
96
+ device = f"{device}:{current_platform.current_device()}"
97
# Load the EC for each mm data
98
for mm_data in metadata.mm_datas:
99
if mm_data.mm_hash in encoder_cache:
100
continue
101
filename = self._generate_filename_debug(mm_data.mm_hash)
97
- ec_cache = safetensors.torch.load_file(
98
- filename, device=current_platform.device_type
99
- )["ec_cache"]
102
+ ec_cache = safetensors.torch.load_file(filename, device=device)["ec_cache"]
103
encoder_cache[mm_data.mm_hash] = ec_cache
104
logger.debug("Success load encoder cache for hash %s", mm_data.mm_hash)
105