[CI] Fix some errors on `main` (#47726)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jul 6, 2026 at 20:40 UTC
ae098abe3fffedd7acecc04265a10dc5624efdca
2 files changed
+39
-22
tests/distributed/test_weight_transfer.py
+37
-13
@@ -191,7 +191,7 @@ class TestNCCLEngineParsing:
191
return NCCLWeightTransferEngine(
192
config,
193
create_mock_vllm_config(),
194
- "cuda",
194
+ torch.device("cuda"),
195
MagicMock(spec=torch.nn.Module),
196
)
197
@@ -240,21 +240,30 @@ class TestEngineRegistry:
240
def test_create_engine_nccl(self):
241
config = WeightTransferConfig(backend="nccl")
242
engine = WeightTransferEngineFactory.create_engine(
243
- config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
243
+ config,
244
+ create_mock_vllm_config(),
245
+ torch.device("cuda"),
246
+ MagicMock(spec=torch.nn.Module),
247
)
248
assert isinstance(engine, NCCLWeightTransferEngine)
249
250
def test_create_engine_ipc(self):
251
config = WeightTransferConfig(backend="ipc")
252
engine = WeightTransferEngineFactory.create_engine(
250
- config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
253
+ config,
254
+ create_mock_vllm_config(),
255
+ torch.device("cuda"),
256
+ MagicMock(spec=torch.nn.Module),
257
)
258
assert isinstance(engine, IPCWeightTransferEngine)
259
260
def test_create_engine_sparse_nccl(self):
261
config = WeightTransferConfig(backend="sparse_nccl")
262
engine = WeightTransferEngineFactory.create_engine(
257
- config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
263
+ config,
264
+ create_mock_vllm_config(),
265
+ torch.device("cuda"),
266
+ MagicMock(spec=torch.nn.Module),
267
)
268
assert isinstance(engine, SparseNCCLWeightTransferEngine)
269
@@ -264,7 +273,7 @@ class TestEngineRegistry:
273
WeightTransferEngineFactory.create_engine(
274
config,
275
create_mock_vllm_config(),
267
- "cuda",
276
+ torch.device("cuda"),
277
MagicMock(spec=torch.nn.Module),
278
)
279
@@ -284,7 +293,7 @@ class TestSparseNCCLPatchApplication:
293
def _make_engine(self, model):
294
config = WeightTransferConfig(backend="sparse_nccl")
295
return SparseNCCLWeightTransferEngine(
287
- config, create_mock_vllm_config(), "cpu", model
296
+ config, create_mock_vllm_config(), torch.device("cpu"), model
297
)
298
299
def _make_model(self, numel: int = 8):
@@ -382,7 +391,10 @@ def test_nccl_receive_weights_without_init_raises():
391
392
config = WeightTransferConfig(backend="nccl")
393
engine = NCCLWeightTransferEngine(
385
- config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
394
+ config,
395
+ create_mock_vllm_config(),
396
+ torch.device("cuda"),
397
+ MagicMock(spec=torch.nn.Module),
398
)
399
400
update_info = NCCLWeightTransferUpdateInfo(
@@ -400,7 +412,10 @@ def test_sparse_nccl_receive_weights_without_init_raises():
412
413
config = WeightTransferConfig(backend="sparse_nccl")
414
engine = SparseNCCLWeightTransferEngine(
403
- config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
415
+ config,
416
+ create_mock_vllm_config(),
417
+ torch.device("cuda"),
418
+ MagicMock(spec=torch.nn.Module),
419
)
420
421
update_info = SparseNCCLWeightTransferUpdateInfo(
@@ -495,7 +510,9 @@ def inference_receive_tensor(
510
vllm_config.model_config = MagicMock()
511
512
recorder = Recorder()
498
- engine = NCCLWeightTransferEngine(config, vllm_config, "cuda", recorder)
513
+ engine = NCCLWeightTransferEngine(
514
+ config, vllm_config, torch.device("cuda"), recorder
515
+ )
516
# Transport-only test: bypass the set_current_vllm_config context that
517
# receive_weights enters, since vllm_config here is a mock.
518
import vllm.config as _vllm_config_mod
@@ -664,7 +681,9 @@ def inference_receive_sparse_tensor(
681
num_updates_list=[3],
682
)
683
667
- engine = SparseNCCLWeightTransferEngine(config, vllm_config, "cuda", model)
684
+ engine = SparseNCCLWeightTransferEngine(
685
+ config, vllm_config, torch.device("cuda"), model
686
+ )
687
from vllm.distributed.weight_transfer.nccl_common import (
688
NCCLWeightTransferInitInfo,
689
)
@@ -879,7 +898,7 @@ class TestIPCEngineParsing:
898
return IPCWeightTransferEngine(
899
config,
900
create_mock_vllm_config(),
882
- "cuda",
901
+ torch.device("cuda"),
902
MagicMock(spec=torch.nn.Module),
903
)
904
@@ -1068,7 +1087,9 @@ def inference_receive_ipc_tensor(
1087
vllm_config.model_config = MagicMock()
1088
1089
recorder = Recorder()
1071
- engine = IPCWeightTransferEngine(config, vllm_config, "cuda", recorder)
1090
+ engine = IPCWeightTransferEngine(
1091
+ config, vllm_config, _get_ray_assigned_device(), recorder
1092
+ )
1093
# Transport-only test: bypass the set_current_vllm_config context that
1094
# receive_weights enters, since vllm_config here is a mock.
1095
import vllm.config as _vllm_config_mod
@@ -1173,7 +1194,10 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():
1194
1195
config = WeightTransferConfig(backend="ipc")
1196
engine = IPCWeightTransferEngine(
1176
- config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module)
1197
+ config,
1198
+ create_mock_vllm_config(),
1199
+ torch.device("cuda:0"),
1200
+ MagicMock(spec=torch.nn.Module),
1201
)
1202
1203
dummy_tensor = torch.ones(10, 10, device="cuda:0")
vllm/model_executor/warmup/kernel_warmup.py
+2
-9
@@ -148,15 +148,8 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
148
149
use_persistent_cache = True
150
151
- deepep_a2a_backends = {
152
- "deepep_high_throughput",
153
- "deepep_low_latency",
154
- "deepep_v2",
155
- }
156
- if runner.vllm_config.parallel_config.all2all_backend in deepep_a2a_backends:
157
- # DeepEP dispatch/combine can timeout when only rank 0
158
- # performs autotune and falls behind other ranks.
159
- # Thus we skip persistent cache in this case.
151
+ # When distributed, tune on every rank so the collectives stay synchronized.
152
+ if get_world_group().world_size > 1:
153
use_persistent_cache = False
154
155
if not use_persistent_cache: