[CI] Fix some errors on main (#47726)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

tests/distributed/test_weight_transfer.py

@@ -191,7 +191,7 @@ class TestNCCLEngineParsing:
         return NCCLWeightTransferEngine(
             config,
             create_mock_vllm_config(),
-            "cuda",
+            torch.device("cuda"),
             MagicMock(spec=torch.nn.Module),
         )
 
@@ -240,21 +240,30 @@ class TestEngineRegistry:
     def test_create_engine_nccl(self):
         config = WeightTransferConfig(backend="nccl")
         engine = WeightTransferEngineFactory.create_engine(
-            config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
+            config,
+            create_mock_vllm_config(),
+            torch.device("cuda"),
+            MagicMock(spec=torch.nn.Module),
         )
         assert isinstance(engine, NCCLWeightTransferEngine)
 
     def test_create_engine_ipc(self):
         config = WeightTransferConfig(backend="ipc")
         engine = WeightTransferEngineFactory.create_engine(
-            config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
+            config,
+            create_mock_vllm_config(),
+            torch.device("cuda"),
+            MagicMock(spec=torch.nn.Module),
         )
         assert isinstance(engine, IPCWeightTransferEngine)
 
     def test_create_engine_sparse_nccl(self):
         config = WeightTransferConfig(backend="sparse_nccl")
         engine = WeightTransferEngineFactory.create_engine(
-            config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
+            config,
+            create_mock_vllm_config(),
+            torch.device("cuda"),
+            MagicMock(spec=torch.nn.Module),
         )
         assert isinstance(engine, SparseNCCLWeightTransferEngine)
 
@@ -264,7 +273,7 @@ class TestEngineRegistry:
             WeightTransferEngineFactory.create_engine(
                 config,
                 create_mock_vllm_config(),
-                "cuda",
+                torch.device("cuda"),
                 MagicMock(spec=torch.nn.Module),
             )
 
@@ -284,7 +293,7 @@ class TestSparseNCCLPatchApplication:
     def _make_engine(self, model):
         config = WeightTransferConfig(backend="sparse_nccl")
         return SparseNCCLWeightTransferEngine(
-            config, create_mock_vllm_config(), "cpu", model
+            config, create_mock_vllm_config(), torch.device("cpu"), model
         )
 
     def _make_model(self, numel: int = 8):
@@ -382,7 +391,10 @@ def test_nccl_receive_weights_without_init_raises():
 
     config = WeightTransferConfig(backend="nccl")
     engine = NCCLWeightTransferEngine(
-        config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
+        config,
+        create_mock_vllm_config(),
+        torch.device("cuda"),
+        MagicMock(spec=torch.nn.Module),
     )
 
     update_info = NCCLWeightTransferUpdateInfo(
@@ -400,7 +412,10 @@ def test_sparse_nccl_receive_weights_without_init_raises():
 
     config = WeightTransferConfig(backend="sparse_nccl")
     engine = SparseNCCLWeightTransferEngine(
-        config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
+        config,
+        create_mock_vllm_config(),
+        torch.device("cuda"),
+        MagicMock(spec=torch.nn.Module),
     )
 
     update_info = SparseNCCLWeightTransferUpdateInfo(
@@ -495,7 +510,9 @@ def inference_receive_tensor(
     vllm_config.model_config = MagicMock()
 
     recorder = Recorder()
-    engine = NCCLWeightTransferEngine(config, vllm_config, "cuda", recorder)
+    engine = NCCLWeightTransferEngine(
+        config, vllm_config, torch.device("cuda"), recorder
+    )
     # Transport-only test: bypass the set_current_vllm_config context that
     # receive_weights enters, since vllm_config here is a mock.
     import vllm.config as _vllm_config_mod
@@ -664,7 +681,9 @@ def inference_receive_sparse_tensor(
         num_updates_list=[3],
     )
 
-    engine = SparseNCCLWeightTransferEngine(config, vllm_config, "cuda", model)
+    engine = SparseNCCLWeightTransferEngine(
+        config, vllm_config, torch.device("cuda"), model
+    )
     from vllm.distributed.weight_transfer.nccl_common import (
         NCCLWeightTransferInitInfo,
     )
@@ -879,7 +898,7 @@ class TestIPCEngineParsing:
         return IPCWeightTransferEngine(
             config,
             create_mock_vllm_config(),
-            "cuda",
+            torch.device("cuda"),
             MagicMock(spec=torch.nn.Module),
         )
 
@@ -1068,7 +1087,9 @@ def inference_receive_ipc_tensor(
     vllm_config.model_config = MagicMock()
 
     recorder = Recorder()
-    engine = IPCWeightTransferEngine(config, vllm_config, "cuda", recorder)
+    engine = IPCWeightTransferEngine(
+        config, vllm_config, _get_ray_assigned_device(), recorder
+    )
     # Transport-only test: bypass the set_current_vllm_config context that
     # receive_weights enters, since vllm_config here is a mock.
     import vllm.config as _vllm_config_mod
@@ -1173,7 +1194,10 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():
 
     config = WeightTransferConfig(backend="ipc")
     engine = IPCWeightTransferEngine(
-        config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
+        config,
+        create_mock_vllm_config(),
+        torch.device("cuda:0"),
+        MagicMock(spec=torch.nn.Module),
     )
 
     dummy_tensor = torch.ones(10, 10, device="cuda:0")

vllm/model_executor/warmup/kernel_warmup.py

@@ -148,15 +148,8 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
 
     use_persistent_cache = True
 
-    deepep_a2a_backends = {
-        "deepep_high_throughput",
-        "deepep_low_latency",
-        "deepep_v2",
-    }
-    if runner.vllm_config.parallel_config.all2all_backend in deepep_a2a_backends:
-        # DeepEP dispatch/combine can timeout when only rank 0
-        # performs autotune and falls behind other ranks.
-        # Thus we skip persistent cache in this case.
+    # When distributed, tune on every rank so the collectives stay synchronized.
+    if get_world_group().world_size > 1:
         use_persistent_cache = False
 
     if not use_persistent_cache: