| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from unittest.mock import MagicMock |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm.distributed.eplb.eplb_state import ( |
| 10 | _commit_eplb_maps, |
| 11 | _commit_eplb_maps_for_layer, |
| 12 | ) |
| 13 | |
| 14 | |
| 15 | def _make_model_state( |
| 16 | phy2log: torch.Tensor, |
| 17 | log2phy: torch.Tensor, |
| 18 | logcnt: torch.Tensor, |
| 19 | phy2log_storage: torch.Tensor | None = None, |
| 20 | ) -> MagicMock: |
| 21 | """Build a minimal EplbModelState mock with its map tensors.""" |
| 22 | state = MagicMock() |
| 23 | state.physical_to_logical_map = phy2log |
| 24 | state.physical_to_logical_map_buffer = ( |
| 25 | phy2log if phy2log_storage is None else phy2log_storage |
| 26 | ) |
| 27 | state.logical_to_physical_map = log2phy |
| 28 | state.logical_replica_count = logcnt |
| 29 | return state |
| 30 | |
| 31 | |
| 32 | def test_commit_eplb_maps_shape_change(): |
| 33 | """When the number of physical experts changes, resize the active map within its |
| 34 | preallocated storage. |
| 35 | """ |
| 36 | num_layers, num_logical, num_physical = 2, 4, 6 |
| 37 | max_replicas = 3 |
| 38 | |
| 39 | # Build current state tensors |
| 40 | storage = torch.full((num_layers, num_physical + 2), -1, dtype=torch.long) |
| 41 | model_state = _make_model_state( |
| 42 | phy2log=storage[:, :num_physical], |
| 43 | phy2log_storage=storage, |
| 44 | log2phy=torch.full( |
| 45 | (num_layers, num_logical, max_replicas), -1, dtype=torch.long |
| 46 | ), |
| 47 | logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long), |
| 48 | ) |
| 49 | |
| 50 | # The new map has two more physical experts. These new physical experts will |
| 51 | # automatically map to the first two logical experts |
| 52 | new_phy2log_larger = ( |
| 53 | (torch.arange(num_physical + 2, dtype=torch.long) % num_logical) |
| 54 | .unsqueeze(0) |
| 55 | .expand(num_layers, -1) |
| 56 | ) |
| 57 | _commit_eplb_maps(model_state, new_phy2log_larger) |
| 58 | |
| 59 | # Check that the number of physical experts has been updated and that the values |
| 60 | # match |
| 61 | assert model_state.physical_to_logical_map.shape[1] == num_physical + 2 |
| 62 | assert model_state.physical_to_logical_map.data_ptr() == storage.data_ptr() |
| 63 | assert torch.equal(model_state.physical_to_logical_map, new_phy2log_larger) |
| 64 | |
| 65 | |
| 66 | def test_commit_eplb_maps_for_layer_logical_padding(): |
| 67 | """Test that logical_to_physical_map is padded with -1 to fill the |
| 68 | pre-allocated slots when the new map has fewer replicas than the max. |
| 69 | """ |
| 70 | num_layers, num_logical, num_physical = 2, 4, 6 |
| 71 | max_replicas = 3 |
| 72 | |
| 73 | model_state = _make_model_state( |
| 74 | phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long), |
| 75 | log2phy=torch.full( |
| 76 | (num_layers, num_logical, max_replicas), -1, dtype=torch.long |
| 77 | ), |
| 78 | logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long), |
| 79 | ) |
| 80 | |
| 81 | new_phy2log = ( |
| 82 | (torch.arange(num_physical, dtype=torch.long) % num_logical) |
| 83 | .unsqueeze(0) |
| 84 | .expand(num_layers, -1) |
| 85 | .contiguous() |
| 86 | ) |
| 87 | layer = 0 |
| 88 | _commit_eplb_maps_for_layer(model_state, new_phy2log[layer], layer) |
| 89 | |
| 90 | assert torch.all(model_state.logical_to_physical_map[layer, :, 2] == -1) |
| 91 | |
| 92 | |
| 93 | def test_commit_eplb_maps_for_layer_shape_assert(): |
| 94 | """Test that a mismatched number of physical experts triggers an assertion error.""" |
| 95 | num_layers, num_logical, num_physical = 2, 4, 6 |
| 96 | |
| 97 | model_state = _make_model_state( |
| 98 | phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long), |
| 99 | log2phy=torch.full((num_layers, num_logical, 2), -1, dtype=torch.long), |
| 100 | logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long), |
| 101 | ) |
| 102 | bad_phy2log = torch.zeros(num_layers, num_physical + 1, dtype=torch.long) |
| 103 | with pytest.raises(AssertionError): |
| 104 | _commit_eplb_maps_for_layer(model_state, bad_phy2log, layer=0) |
| 105 | |
| 106 | |
| 107 | def test_commit_eplb_maps(): |
| 108 | """Test that all values are copied correctly into model_state.""" |
| 109 | num_layers, num_logical, num_physical, max_replicas = 2, 3, 4, 2 |
| 110 | |
| 111 | model_state = _make_model_state( |
| 112 | phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long), |
| 113 | log2phy=torch.full( |
| 114 | (num_layers, num_logical, max_replicas), -1, dtype=torch.long |
| 115 | ), |
| 116 | logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long), |
| 117 | ) |
| 118 | |
| 119 | new_phy2log = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]], dtype=torch.long) |
| 120 | new_log2phy = torch.tensor( |
| 121 | [[[0, 3], [1, -1], [2, -1]], [[2, -1], [0, 3], [1, -1]]], dtype=torch.long |
| 122 | ) |
| 123 | new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long) |
| 124 | |
| 125 | _commit_eplb_maps(model_state, new_phy2log) |
| 126 | |
| 127 | assert torch.equal(model_state.physical_to_logical_map, new_phy2log) |
| 128 | assert torch.equal(model_state.logical_to_physical_map, new_log2phy) |
| 129 | assert torch.equal(model_state.logical_replica_count, new_logcnt) |
| 130 | |
| 131 | |
| 132 | def test_commit_eplb_maps_for_layer(): |
| 133 | """Test that only the target layer is updated""" |
| 134 | num_layers, num_logical, max_replicas = 2, 3, 2 |
| 135 | |
| 136 | original_phy2log = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long) |
| 137 | model_state = _make_model_state( |
| 138 | phy2log=original_phy2log.clone(), |
| 139 | log2phy=torch.full( |
| 140 | (num_layers, num_logical, max_replicas), -1, dtype=torch.long |
| 141 | ), |
| 142 | logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long), |
| 143 | ) |
| 144 | |
| 145 | new_phy2log = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]], dtype=torch.long) |
| 146 | new_log2phy = torch.tensor( |
| 147 | [[[0, 3], [1, -1], [2, -1]], [[2, -1], [0, 3], [1, -1]]], dtype=torch.long |
| 148 | ) |
| 149 | new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long) |
| 150 | |
| 151 | _commit_eplb_maps_for_layer(model_state, new_phy2log[0], layer=0) |
| 152 | |
| 153 | # Layer 0 updated |
| 154 | assert torch.equal(model_state.physical_to_logical_map[0], new_phy2log[0]) |
| 155 | assert torch.equal(model_state.logical_to_physical_map[0], new_log2phy[0]) |
| 156 | assert torch.equal(model_state.logical_replica_count[0], new_logcnt[0]) |
| 157 | |
| 158 | # Layer 1 untouched |
| 159 | assert torch.equal(model_state.physical_to_logical_map[1], original_phy2log[1]) |