A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 22:15 UTC
[EPLB] Cleanup the transfer logic for the various eplb maps (#34520)
Signed-off-by: Sage Moore <sagmoore@redhat.com> Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Mar 27, 2026 at 02:18 UTC
497e234d38017d4adb83e072d1e5949c76b6d616
3 files changed
+247
-76
.buildkite/test_areas/expert_parallelism.yaml
+2
@@ -8,8 +8,10 @@ steps:
8
source_file_dependencies:
9
- vllm/distributed/eplb
10
- tests/distributed/test_eplb_algo.py
11
+ - tests/distributed/test_eplb_utils.py
12
commands:
13
- pytest -v -s distributed/test_eplb_algo.py
14
+ - pytest -v -s distributed/test_eplb_utils.py
15
16
- label: EPLB Execution
17
timeout_in_minutes: 20
tests/distributed/test_eplb_utils.py
new
+154
@@ -0,0 +1,154 @@
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
+) -> MagicMock:
20
+ """Build a minimal EplbModelState mock with only the three map tensors."""
21
+ state = MagicMock()
22
+ state.physical_to_logical_map = phy2log
23
+ state.logical_to_physical_map = log2phy
24
+ state.logical_replica_count = logcnt
25
+ return state
26
+
27
+
28
+def test_commit_eplb_maps_shape_change():
29
+ """
30
+ The normal path copies the physical_to_logical map in-place. When the number of
31
+ physical experts changes, the old map should be replaced entirely.
32
+ """
33
+ num_layers, num_logical, num_physical = 2, 4, 6
34
+ max_replicas = 3
35
+
36
+ # Build current state tensors
37
+ model_state = _make_model_state(
38
+ phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
39
+ log2phy=torch.full(
40
+ (num_layers, num_logical, max_replicas), -1, dtype=torch.long
41
+ ),
42
+ logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
43
+ )
44
+
45
+ # The new map has two more physical experts. These new physical experts will
46
+ # automatically map to the first two logical experts
47
+ new_phy2log_larger = (
48
+ (torch.arange(num_physical + 2, dtype=torch.long) % num_logical)
49
+ .unsqueeze(0)
50
+ .expand(num_layers, -1)
51
+ )
52
+ _commit_eplb_maps(model_state, new_phy2log_larger)
53
+
54
+ # Check that the number of physical experts has been updated and that the values
55
+ # match
56
+ assert model_state.physical_to_logical_map.shape[1] == num_physical + 2
57
+ assert torch.equal(model_state.physical_to_logical_map, new_phy2log_larger)
58
+
59
+
60
+def test_commit_eplb_maps_for_layer_logical_padding():
61
+ """
62
+ Test that logical_to_physical_map is padded with -1 to fill the
63
+ pre-allocated slots when the new map has fewer replicas than the max.
64
+ """
65
+ num_layers, num_logical, num_physical = 2, 4, 6
66
+ max_replicas = 3
67
+
68
+ model_state = _make_model_state(
69
+ phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
70
+ log2phy=torch.full(
71
+ (num_layers, num_logical, max_replicas), -1, dtype=torch.long
72
+ ),
73
+ logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
74
+ )
75
+
76
+ new_phy2log = (
77
+ (torch.arange(num_physical, dtype=torch.long) % num_logical)
78
+ .unsqueeze(0)
79
+ .expand(num_layers, -1)
80
+ .contiguous()
81
+ )
82
+ layer = 0
83
+ _commit_eplb_maps_for_layer(model_state, new_phy2log, layer)
84
+
85
+ assert torch.all(model_state.logical_to_physical_map[layer, :, 2] == -1)
86
+
87
+
88
+def test_commit_eplb_maps_for_layer_shape_assert():
89
+ """Test that a mismatched number of physical experts triggers an assertion error."""
90
+ num_layers, num_logical, num_physical = 2, 4, 6
91
+
92
+ model_state = _make_model_state(
93
+ phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
94
+ log2phy=torch.full((num_layers, num_logical, 2), -1, dtype=torch.long),
95
+ logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
96
+ )
97
+ bad_phy2log = torch.zeros(num_layers, num_physical + 1, dtype=torch.long)
98
+ with pytest.raises(AssertionError):
99
+ _commit_eplb_maps_for_layer(model_state, bad_phy2log, layer=0)
100
+
101
+
102
+def test_commit_eplb_maps():
103
+ """Test that all values are copied correctly into model_state."""
104
+ num_layers, num_logical, num_physical, max_replicas = 2, 3, 4, 2
105
+
106
+ model_state = _make_model_state(
107
+ phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
108
+ log2phy=torch.full(
109
+ (num_layers, num_logical, max_replicas), -1, dtype=torch.long
110
+ ),
111
+ logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
112
+ )
113
+
114
+ new_phy2log = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]], dtype=torch.long)
115
+ new_log2phy = torch.tensor(
116
+ [[[0, 3], [1, -1], [2, -1]], [[2, -1], [0, 3], [1, -1]]], dtype=torch.long
117
+ )
118
+ new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long)
119
+
120
+ _commit_eplb_maps(model_state, new_phy2log)
121
+
122
+ assert torch.equal(model_state.physical_to_logical_map, new_phy2log)
123
+ assert torch.equal(model_state.logical_to_physical_map, new_log2phy)
124
+ assert torch.equal(model_state.logical_replica_count, new_logcnt)
125
+
126
+
127
+def test_commit_eplb_maps_for_layer():
128
+ """Test that only the target layer is updated"""
129
+ num_layers, num_logical, max_replicas = 2, 3, 2
130
+
131
+ original_phy2log = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
132
+ model_state = _make_model_state(
133
+ phy2log=original_phy2log.clone(),
134
+ log2phy=torch.full(
135
+ (num_layers, num_logical, max_replicas), -1, dtype=torch.long
136
+ ),
137
+ logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
138
+ )
139
+
140
+ new_phy2log = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]], dtype=torch.long)
141
+ new_log2phy = torch.tensor(
142
+ [[[0, 3], [1, -1], [2, -1]], [[2, -1], [0, 3], [1, -1]]], dtype=torch.long
143
+ )
144
+ new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long)
145
+
146
+ _commit_eplb_maps_for_layer(model_state, new_phy2log, layer=0)
147
+
148
+ # Layer 0 updated
149
+ assert torch.equal(model_state.physical_to_logical_map[0], new_phy2log[0])
150
+ assert torch.equal(model_state.logical_to_physical_map[0], new_log2phy[0])
151
+ assert torch.equal(model_state.logical_replica_count[0], new_logcnt[0])
152
+
153
+ # Layer 1 untouched
154
+ assert torch.equal(model_state.physical_to_logical_map[1], original_phy2log[1])
vllm/distributed/eplb/eplb_state.py
+91
-76
@@ -729,13 +729,6 @@ class EplbState:
729
eplb_model_state.physical_to_logical_map.cpu(),
730
)
731
732
- num_logical_experts = global_expert_load_window.shape[-1]
733
- (new_logical_to_physical_map, new_logical_replica_count) = (
734
- compute_logical_maps(
735
- new_physical_to_logical_map, num_logical_experts
736
- )
737
- )
738
-
732
# Update expert weights
733
rearrange_expert_weights_inplace(
734
eplb_model_state.physical_to_logical_map,
@@ -747,39 +740,11 @@ class EplbState:
740
)
741
742
if not is_profile:
750
- if (
751
- eplb_model_state.physical_to_logical_map.shape[1]
752
- != new_physical_to_logical_map.shape[1]
753
- ):
754
- eplb_model_state.physical_to_logical_map = (
755
- new_physical_to_logical_map.to(
756
- eplb_model_state.physical_to_logical_map.device
757
- )
758
- )
759
- else:
760
- eplb_model_state.physical_to_logical_map.copy_(
761
- new_physical_to_logical_map
762
- )
763
- max_physical_slots = new_logical_to_physical_map.shape[-1]
764
- assert (
765
- max_physical_slots
766
- <= eplb_model_state.logical_to_physical_map.shape[-1]
767
- )
768
- new_logical_to_physical_map = torch.nn.functional.pad(
769
- new_logical_to_physical_map,
770
- (
771
- 0,
772
- eplb_model_state.logical_to_physical_map.shape[-1]
773
- - max_physical_slots,
774
- ),
775
- value=-1,
776
- )
777
- eplb_model_state.logical_to_physical_map.copy_(
778
- new_logical_to_physical_map
779
- )
780
- eplb_model_state.logical_replica_count.copy_(
781
- new_logical_replica_count
743
+ _commit_eplb_maps(
744
+ eplb_model_state,
745
+ new_physical_to_logical_map=new_physical_to_logical_map,
746
)
747
+
748
if is_main_rank:
749
assert start_event is not None
750
assert end_event is not None
@@ -829,42 +794,6 @@ class EplbState:
794
is_profile=is_profile,
795
)
796
832
- def _update_layer_mapping_from_new(
833
- self, model_state: EplbModelState, layer: int
834
- ) -> None:
835
- if model_state.new_physical_to_logical_map is None:
836
- return
837
-
838
- target_device = model_state.physical_to_logical_map.device
839
- new_physical = model_state.new_physical_to_logical_map
840
- # If the number of physical experts has changed, then the new map needs to
841
- # be copied synchronously to avoid a race condition with the async worker
842
- if model_state.physical_to_logical_map.shape[1] != new_physical.shape[1]:
843
- model_state.physical_to_logical_map = new_physical.to(target_device)
844
- else:
845
- model_state.physical_to_logical_map[layer].copy_(
846
- new_physical[layer].to(target_device, non_blocking=True)
847
- )
848
-
849
- num_logical_experts = model_state.logical_to_physical_map.shape[1]
850
- new_logical, new_replica_count = compute_logical_maps(
851
- new_physical[layer], num_logical_experts
852
- )
853
-
854
- logical_device = model_state.logical_to_physical_map.device
855
- max_slots = model_state.logical_to_physical_map.shape[-1]
856
- slot_delta = max_slots - new_logical.shape[-1]
857
- if slot_delta > 0:
858
- new_logical = torch.nn.functional.pad(
859
- new_logical, (0, slot_delta), value=-1
860
- )
861
- model_state.logical_to_physical_map[layer].copy_(new_logical.to(logical_device))
862
-
863
- replica_device = model_state.logical_replica_count.device
864
- model_state.logical_replica_count[layer].copy_(
865
- new_replica_count.to(replica_device)
866
- )
867
-
797
def _all_ranks_buffer_ready(self, model_state: EplbModelState) -> bool:
798
parallel_state = get_ep_group()
799
cpu_group = getattr(parallel_state, "cpu_group", None)
@@ -936,7 +865,12 @@ class EplbState:
865
model_state.buffer_consumed_event = consumed_event
866
867
transferred_layer = model_state.layer_to_transfer
939
- self._update_layer_mapping_from_new(model_state, transferred_layer)
868
+ assert model_state.new_physical_to_logical_map is not None
869
+ _commit_eplb_maps_for_layer(
870
+ model_state,
871
+ new_physical_to_logical_map=model_state.new_physical_to_logical_map,
872
+ layer=transferred_layer,
873
+ )
874
# After the main thread consumes, advance layer_to_transfer
875
model_state.layer_to_transfer += 1
876
model_state.ep_buffer_ready = 0
@@ -1175,3 +1109,84 @@ def compute_logical_maps(
1109
if per_layer:
1110
return logical_to_physical_map_out.squeeze(0), logical_replica_count.squeeze(0)
1111
return logical_to_physical_map_out, logical_replica_count
1112
+
1113
+
1114
+def _pad_out_tensor(src: torch.Tensor, dst: torch.Tensor) -> None:
1115
+ src_padding = dst.shape[-1] - src.shape[-1]
1116
+ assert src_padding >= 0
1117
+ new_src = torch.nn.functional.pad(src, (0, src_padding), value=-1)
1118
+ dst.copy_(new_src)
1119
+
1120
+
1121
+def _commit_eplb_maps_for_layer(
1122
+ model_state: EplbModelState,
1123
+ new_physical_to_logical_map: torch.Tensor,
1124
+ layer: int,
1125
+) -> None:
1126
+ """
1127
+ Per-layer version of _commit_eplb_maps that's used by the sync portion of EPLB
1128
+ when running async EPLB. Copies all of the new_* maps into model_state. After this
1129
+ function completes, the new mappings will become the current mappings and will be
1130
+ visible to the model.
1131
+ """
1132
+
1133
+ # Commit physical_to_logical_map
1134
+ src = new_physical_to_logical_map[layer]
1135
+ dst = model_state.physical_to_logical_map[layer]
1136
+ assert src.shape == dst.shape, (
1137
+ "The number of physical experts must stay the same while running Async EPLB. "
1138
+ f"Current number of physical experts: {dst.shape[0]}. New number of physical "
1139
+ f"experts {src.shape[0]}."
1140
+ )
1141
+ dst.copy_(src, non_blocking=True)
1142
+
1143
+ num_logical_experts = model_state.logical_to_physical_map.shape[1]
1144
+ new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts)
1145
+ # Commit logical_to_physical_map
1146
+ _pad_out_tensor(
1147
+ src=new_logical,
1148
+ dst=model_state.logical_to_physical_map[layer],
1149
+ )
1150
+
1151
+ # Commit logical_replica_count
1152
+ src = new_replica_count
1153
+ dst = model_state.logical_replica_count[layer]
1154
+ assert src.shape == dst.shape
1155
+ dst.copy_(src, non_blocking=True)
1156
+
1157
+
1158
+def _commit_eplb_maps(
1159
+ model_state: EplbModelState,
1160
+ new_physical_to_logical_map: torch.Tensor,
1161
+) -> None:
1162
+ """
1163
+ Copies all of the new_* maps into model_state. After this function completes,
1164
+ the new mappings will become the current mappings and will be visible to the
1165
+ model.
1166
+ """
1167
+
1168
+ # Commit physical_to_logical_map
1169
+ src = new_physical_to_logical_map
1170
+ dst = model_state.physical_to_logical_map
1171
+
1172
+ # Rare Case: When the number of physical experts has changed, discard the old
1173
+ # physical to logical expert map and use the new one. This only happens when the
1174
+ # number of GPUs available to vLLM changes while vLLM is running. Otherwise copy the
1175
+ # new map into the old one.
1176
+ if src.shape[1] != dst.shape[1]:
1177
+ model_state.physical_to_logical_map = src.to(dst.device)
1178
+ else:
1179
+ dst.copy_(src, non_blocking=True)
1180
+
1181
+ num_logical_experts = model_state.logical_to_physical_map.shape[1]
1182
+ new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts)
1183
+ # Commit logical_to_physical_map
1184
+ _pad_out_tensor(
1185
+ src=new_logical,
1186
+ dst=model_state.logical_to_physical_map,
1187
+ )
1188
+
1189
+ # Commit logical_replica_count
1190
+ src = new_replica_count
1191
+ dst = model_state.logical_replica_count
1192
+ dst.copy_(src, non_blocking=True)