diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index 5d29ad692f..edddb6ec84 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import random -from types import SimpleNamespace -from unittest.mock import Mock import pytest import ray @@ -11,8 +9,6 @@ import torch import torch.distributed as dist from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa -from vllm.distributed.device_communicators import custom_all_reduce -from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce from vllm.distributed.parallel_state import get_tp_group, graph_capture from ..utils import ( @@ -27,59 +23,6 @@ for i, v in enumerate(test_sizes): test_sizes[i] -= v % 8 -def test_sp16_dispatches_only_to_mnnvl_lamport( - monkeypatch: pytest.MonkeyPatch, -): - """SP16 uses the MNNVL Lamport kernels and rejects same-host dispatch.""" - comm = CustomAllreduce.__new__(CustomAllreduce) - comm.disabled = False - comm.world_size = 16 - comm.fully_connected = False - comm.mnnvl_only = True - comm._IS_CAPTURING = False - comm.max_mnnvl_all_gather_size = 2 * 1024 * 1024 - comm.max_mnnvl_reduce_scatter_size = 16 * 1024 * 1024 - comm.mnnvl_multicast_ptr = 1 - comm.mnnvl_lamport_ag_local_ptr = 1 - comm.mnnvl_lamport_ag_multicast_ptr = 1 - comm.mnnvl_lamport_ag_epoch_ptr = 1 - comm.mnnvl_lamport_rs_local_ptr = 1 - comm.mnnvl_lamport_rs_epoch_ptr = 1 - comm.mnnvl_buffer_size = 32 * 1024 * 1024 - comm._ptr = 0 - - lamport_all_gather = Mock() - lamport_reduce_scatter = Mock() - monkeypatch.setattr(custom_all_reduce.current_platform, "is_cuda", lambda: True) - monkeypatch.setattr( - custom_all_reduce, - "ops", - SimpleNamespace( - mnnvl_lamport_all_gather=lamport_all_gather, - mnnvl_lamport_reduce_scatter=lamport_reduce_scatter, - ), - ) - - gathered = comm.custom_all_gather(torch.empty((8, 8), dtype=torch.bfloat16)) - scattered = comm.custom_reduce_scatter(torch.empty((16, 8), dtype=torch.bfloat16)) - - assert gathered is not None - assert scattered is not None - lamport_all_gather.assert_called_once() - lamport_reduce_scatter.assert_called_once() - assert not comm.should_custom_ar(torch.empty(8, dtype=torch.bfloat16)) - assert not comm.should_custom_all_gather(torch.empty((8, 8), dtype=torch.int32)) - assert not comm.should_custom_all_gather( - torch.empty((131073, 8), dtype=torch.bfloat16) - ) - - comm.mnnvl_only = False - assert not comm.should_custom_all_gather(torch.empty((8, 8), dtype=torch.bfloat16)) - assert not comm.should_custom_reduce_scatter( - torch.empty((16, 8), dtype=torch.bfloat16) - ) - - @ray.remote(num_gpus=1, max_calls=1) def graph_allreduce( monkeypatch: pytest.MonkeyPatch, @@ -137,32 +80,6 @@ def graph_allreduce( torch.testing.assert_close(out1, inp1) torch.testing.assert_close(out2, inp2) - fa = get_tp_group().device_communicator.ca_comm - tp_rank = rank % tp_size - with graph_capture(device=device) as graph_capture_context: - local = torch.full( - (512, 4096), tp_rank + 1, dtype=torch.bfloat16, device=device - ) - reduce_input = torch.full( - (512 * tp_size, 4096), - tp_rank + 1, - dtype=torch.bfloat16, - device=device, - ) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=graph_capture_context.stream): - gathered = fa.custom_all_gather(local) - scattered = fa.custom_reduce_scatter(reduce_input) - graph.replay() - assert gathered is not None - assert scattered is not None - expected_gather = torch.cat( - [torch.full_like(local, peer_rank + 1) for peer_rank in range(tp_size)] - ) - expected_scatter = torch.full_like(local, tp_size * (tp_size + 1) // 2) - torch.testing.assert_close(gathered, expected_gather) - torch.testing.assert_close(scattered, expected_scatter) - @ray.remote(num_gpus=1, max_calls=1) def eager_allreduce( @@ -193,29 +110,6 @@ def eager_allreduce( out = fa.all_reduce(out, registered=False) torch.testing.assert_close(out, inp * (tp_size**num_communication)) - group = get_tp_group().device_group - tp_rank = rank % tp_size - for dtype in [torch.float32, torch.float16, torch.bfloat16]: - local = torch.full((64, 4096), tp_rank + 1, dtype=dtype, device=device) - expected_gather = torch.empty( - (64 * tp_size, 4096), dtype=dtype, device=device - ) - dist.all_gather_into_tensor(expected_gather, local, group=group) - gathered = fa.custom_all_gather(local) - assert gathered is not None - torch.testing.assert_close(gathered, expected_gather) - - reduce_input = torch.full( - (64 * tp_size, 4096), tp_rank + 1, dtype=dtype, device=device - ) - expected_scatter = torch.empty((64, 4096), dtype=dtype, device=device) - dist.reduce_scatter_tensor( - expected_scatter, reduce_input.clone(), group=group - ) - scattered = fa.custom_reduce_scatter(reduce_input) - assert scattered is not None - torch.testing.assert_close(scattered, expected_scatter) - inp = torch.ones(sz * 4, dtype=torch.bfloat16, device=device) out = inp for _ in range(num_communication): @@ -236,14 +130,3 @@ def test_custom_allreduce( if world_size > torch.accelerator.device_count(): pytest.skip("Not enough GPUs to run the test.") multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target) - - -@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce]) -def test_custom_collectives_world_size_four( - monkeypatch: pytest.MonkeyPatch, - test_target, -): - """Exercise the four-rank kernel specialization used by Kimi SP.""" - if torch.accelerator.device_count() < 4: - pytest.skip("Not enough GPUs to run the test.") - multi_process_parallel(monkeypatch, 4, 1, test_target) diff --git a/tests/v1/shutdown/test_delete.py b/tests/v1/shutdown/test_delete.py index 39386f3fd6..991ad68918 100644 --- a/tests/v1/shutdown/test_delete.py +++ b/tests/v1/shutdown/test_delete.py @@ -132,4 +132,13 @@ def test_llm_delete_inprocess( wait_for_gpu_memory_to_clear( devices=[0], threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES, + # Activate the helper's ROCm idle-runtime floor. VllmRunner already + # performed the stable-memory wait during context exit. + threshold_ratio=0.01 if current_platform.is_rocm() else None, + # Fail in the child before the outer pytest timeout expires. + timeout_s=( + SHUTDOWN_TEST_TIMEOUT_SEC // 2 + if current_platform.is_rocm() + else SHUTDOWN_TEST_TIMEOUT_SEC + ), )