| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | from unittest.mock import Mock |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.distributed.device_communicators import cuda_communicator |
| 11 | from vllm.distributed.device_communicators import ( |
| 12 | flashinfer_pcie_ipc_all_reduce as pcie_ipc, |
| 13 | ) |
| 14 | from vllm.distributed.parallel_state import GroupCoordinator |
| 15 | |
| 16 | |
| 17 | def _uninitialized_comm() -> pcie_ipc.FlashInferPcieIpcAllReduce: |
| 18 | comm = pcie_ipc.FlashInferPcieIpcAllReduce.__new__( |
| 19 | pcie_ipc.FlashInferPcieIpcAllReduce |
| 20 | ) |
| 21 | comm.disabled = False |
| 22 | comm.group = object() |
| 23 | comm.tune_group = object() |
| 24 | comm.device = torch.device("cuda:0") |
| 25 | comm.world_size = 4 |
| 26 | comm.rank = 0 |
| 27 | comm.workspace = None |
| 28 | comm.hidden_dim = 0 |
| 29 | comm.dtype = None |
| 30 | return comm |
| 31 | |
| 32 | |
| 33 | def test_setup_uses_exact_graph_capacity(monkeypatch, tmp_path): |
| 34 | workspace = Mock() |
| 35 | factory = Mock(return_value=workspace) |
| 36 | monkeypatch.setattr( |
| 37 | pcie_ipc, |
| 38 | "flashinfer_comm", |
| 39 | SimpleNamespace(PcieIpcAllReduceWorkspace=factory), |
| 40 | ) |
| 41 | monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) |
| 42 | |
| 43 | comm = _uninitialized_comm() |
| 44 | cache = tmp_path / "tp4.json" |
| 45 | comm.setup( |
| 46 | hidden_dim=4096, |
| 47 | dtype=torch.bfloat16, |
| 48 | capture_sizes=[32, 8, 16, 32], |
| 49 | tune_cache=cache, |
| 50 | ) |
| 51 | |
| 52 | factory.assert_called_once_with( |
| 53 | group=comm.group, |
| 54 | max_numel=32 * 4096, |
| 55 | dtype=torch.bfloat16, |
| 56 | tune_batches=(8, 16, 32), |
| 57 | tune_cache=str(cache), |
| 58 | ) |
| 59 | workspace.tune.assert_called_once_with( |
| 60 | [4096], dtype=torch.bfloat16, tune_group=comm.tune_group |
| 61 | ) |
| 62 | workspace.prepare.assert_called_once_with( |
| 63 | [(8, 4096), (16, 4096), (32, 4096)], dtype=torch.bfloat16 |
| 64 | ) |
| 65 | assert workspace.rebind_stream.call_count == 2 |
| 66 | |
| 67 | |
| 68 | @pytest.mark.parametrize( |
| 69 | ("shape", "dtype", "contiguous", "expected"), |
| 70 | [ |
| 71 | ((16, 4096), torch.bfloat16, True, True), |
| 72 | ((16, 2048), torch.bfloat16, True, False), |
| 73 | ((16, 4096), torch.float16, True, False), |
| 74 | ((16, 4096), torch.bfloat16, False, False), |
| 75 | ], |
| 76 | ) |
| 77 | def test_should_use_only_prepared_hidden_shape(shape, dtype, contiguous, expected): |
| 78 | comm = _uninitialized_comm() |
| 79 | comm.hidden_dim = 4096 |
| 80 | comm.dtype = torch.bfloat16 |
| 81 | comm.workspace = Mock() |
| 82 | comm.workspace.supports.return_value = True |
| 83 | inp = SimpleNamespace( |
| 84 | is_cuda=True, |
| 85 | is_contiguous=lambda: contiguous, |
| 86 | dim=lambda: len(shape), |
| 87 | shape=shape, |
| 88 | dtype=dtype, |
| 89 | ) |
| 90 | |
| 91 | assert comm.should_use(inp) is expected |
| 92 | |
| 93 | |
| 94 | def test_capture_rebinds_workspace_even_on_error(monkeypatch): |
| 95 | comm = _uninitialized_comm() |
| 96 | comm.workspace = Mock() |
| 97 | synchronize = Mock() |
| 98 | monkeypatch.setattr(torch.accelerator, "synchronize", synchronize) |
| 99 | |
| 100 | with pytest.raises(RuntimeError, match="capture failed"), comm.capture(): |
| 101 | raise RuntimeError("capture failed") |
| 102 | |
| 103 | assert synchronize.call_count == 2 |
| 104 | assert comm.workspace.rebind_stream.call_count == 2 |
| 105 | |
| 106 | |
| 107 | def test_cuda_dispatch_prefers_pcie_ipc_before_existing_flashinfer(monkeypatch): |
| 108 | communicator = cuda_communicator.CudaCommunicator.__new__( |
| 109 | cuda_communicator.CudaCommunicator |
| 110 | ) |
| 111 | communicator.pynccl_comm = Mock(world_size=4) |
| 112 | communicator.qr_comm = None |
| 113 | communicator.fi_pcie_ipc_ar_comm = Mock() |
| 114 | communicator.fi_pcie_ipc_ar_comm.should_use.return_value = True |
| 115 | expected = object() |
| 116 | communicator.fi_pcie_ipc_ar_comm.all_reduce.return_value = expected |
| 117 | communicator.fi_ar_comm = Mock() |
| 118 | communicator.aiter_ar_comm = None |
| 119 | communicator.ca_comm = None |
| 120 | communicator.symm_mem_comm = None |
| 121 | monkeypatch.setattr( |
| 122 | cuda_communicator, "should_nccl_symm_mem_allreduce", lambda *_: False |
| 123 | ) |
| 124 | |
| 125 | assert communicator.all_reduce(object()) is expected |
| 126 | communicator.fi_ar_comm.should_use_fi_ar.assert_not_called() |
| 127 | |
| 128 | |
| 129 | def test_group_destroy_releases_communicator_before_process_groups(monkeypatch): |
| 130 | events = [] |
| 131 | coordinator = GroupCoordinator.__new__(GroupCoordinator) |
| 132 | coordinator.device_communicator = SimpleNamespace( |
| 133 | destroy=lambda: events.append("communicator") |
| 134 | ) |
| 135 | coordinator.device_group = object() |
| 136 | coordinator.cpu_group = object() |
| 137 | coordinator.mq_broadcaster = None |
| 138 | monkeypatch.setattr( |
| 139 | torch.distributed, |
| 140 | "destroy_process_group", |
| 141 | lambda _group: events.append("process_group"), |
| 142 | ) |
| 143 | |
| 144 | coordinator.destroy() |
| 145 | |
| 146 | assert events == ["communicator", "process_group", "process_group"] |