From ae098abe3fffedd7acecc04265a10dc5624efdca Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Mon, 6 Jul 2026 20:40:23 +0100 Subject: [PATCH] [CI] Fix some errors on `main` (#47726) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/distributed/test_weight_transfer.py | 50 +++++++++++++++------ vllm/model_executor/warmup/kernel_warmup.py | 11 +---- 2 files changed, 39 insertions(+), 22 deletions(-) diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 8a13b24dc52..f3423745ca5 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/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") diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e31a14db663..b7f3c265704 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/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: