@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 305 lines 9.25 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Replay MoE weights after sleep/wake and verify captured graph replay.
4
5 This test exercises the failure mode directly at the MoE layer boundary:
6
7 1. Build a small ROCm unquantized `RoutedExperts` layer.
8 2. Load deterministic random expert weights and run post-load processing.
9 3. Capture a CUDA graph for `RoutedExperts.forward_modular`.
10 4. Put the CuMem weights pool through level-2 sleep/wake semantics.
11 5. Replay the same checkpoint-style weights through the layer weight loader.
12 6. Verify graph replay is deterministic across sleep/wake and weight replay.
13
14 The important invariant is that post-load processing must not replace parameter
15 storage after CUDA graph capture, because captured kernels hold device
16 addresses for the weight tensors.
17 """
18
19 import contextlib
20 import gc
21 import importlib.util
22 from collections.abc import Iterable
23
24 import pytest
25 import torch
26
27 from vllm.device_allocator.cumem import CuMemAllocator, cumem_available
28 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
29 from vllm.model_executor.layers.fused_moe.config import (
30 FusedMoEConfig,
31 FusedMoEParallelConfig,
32 RoutingMethodType,
33 )
34 from vllm.model_executor.layers.fused_moe.expert_map_manager import (
35 ExpertMapManager,
36 )
37 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
38 from vllm.platforms import current_platform
39
40 aiter_available = importlib.util.find_spec("aiter") is not None
41
42 pytestmark = pytest.mark.skipif(
43 not (current_platform.is_rocm() and aiter_available),
44 reason="MoE replay test requires ROCm with AITER installed",
45 )
46
47 DEVICE = "cuda"
48 DTYPE = torch.bfloat16
49 NUM_EXPERTS = 4
50 TOPK = 2
51 NUM_TOKENS = 16
52 HIDDEN_SIZE = 256
53 INTERMEDIATE_SIZE = 256
54
55
56 def _moe_topk_softmax_available() -> bool:
57 try:
58 importlib.import_module("vllm._moe_C_stable_libtorch")
59 except ImportError:
60 return False
61 return hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "topk_softmax")
62
63
64 def _set_backend_env(monkeypatch: pytest.MonkeyPatch, backend: str) -> None:
65 monkeypatch.setenv("VLLM_ROCM_MOE_PADDING", "1")
66 if backend == "aiter":
67 monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
68 monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
69 elif backend == "triton":
70 monkeypatch.delenv("VLLM_ROCM_USE_AITER", raising=False)
71 monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
72 else:
73 raise AssertionError(f"unknown backend {backend}")
74
75 from vllm._aiter_ops import rocm_aiter_ops
76
77 rocm_aiter_ops.refresh_env_variables()
78
79
80 def _make_moe_config() -> FusedMoEConfig:
81 return FusedMoEConfig(
82 num_experts=NUM_EXPERTS,
83 experts_per_token=TOPK,
84 hidden_dim=HIDDEN_SIZE,
85 intermediate_size=INTERMEDIATE_SIZE,
86 num_local_experts=NUM_EXPERTS,
87 num_logical_experts=NUM_EXPERTS,
88 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
89 activation=MoEActivation.SILU,
90 in_dtype=DTYPE,
91 device=DEVICE,
92 routing_method=RoutingMethodType.TopK,
93 max_num_tokens=NUM_TOKENS,
94 )
95
96
97 def _make_expert_map_manager(moe_config: FusedMoEConfig) -> ExpertMapManager:
98 return ExpertMapManager(
99 max_num_batched_tokens=NUM_TOKENS,
100 top_k=TOPK,
101 global_num_experts=NUM_EXPERTS,
102 num_redundant_experts=0,
103 num_expert_group=None,
104 moe_parallel_config=moe_config.moe_parallel_config,
105 placement_strategy="linear",
106 enable_eplb=False,
107 )
108
109
110 def _make_routed_experts() -> RoutedExperts:
111 moe_config = _make_moe_config()
112 return RoutedExperts(
113 "experts",
114 DTYPE,
115 moe_config,
116 quant_config=None,
117 expert_map_manager=_make_expert_map_manager(moe_config),
118 )
119
120
121 def _expert_weight_iterator(seed: int) -> Iterable[tuple[str, torch.Tensor]]:
122 generator = torch.Generator(device=DEVICE)
123 generator.manual_seed(seed)
124 for expert_id in range(NUM_EXPERTS):
125 yield (
126 f"{expert_id}.gate_proj.weight",
127 torch.randn(
128 INTERMEDIATE_SIZE,
129 HIDDEN_SIZE,
130 generator=generator,
131 device=DEVICE,
132 dtype=DTYPE,
133 ),
134 )
135 yield (
136 f"{expert_id}.up_proj.weight",
137 torch.randn(
138 INTERMEDIATE_SIZE,
139 HIDDEN_SIZE,
140 generator=generator,
141 device=DEVICE,
142 dtype=DTYPE,
143 ),
144 )
145 yield (
146 f"{expert_id}.down_proj.weight",
147 torch.randn(
148 HIDDEN_SIZE,
149 INTERMEDIATE_SIZE,
150 generator=generator,
151 device=DEVICE,
152 dtype=DTYPE,
153 ),
154 )
155
156
157 def _load_expert_weights(layer: RoutedExperts, seed: int) -> set[str]:
158 loaded = set(layer.load_weights(_expert_weight_iterator(seed)))
159 assert {"w13_weight", "w2_weight"} <= loaded
160 return loaded
161
162
163 def _process_weights_after_loading(layer: RoutedExperts) -> None:
164 layer.quant_method.process_weights_after_loading(layer)
165 torch.accelerator.synchronize()
166
167
168 def _make_static_inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
169 generator = torch.Generator(device=DEVICE)
170 generator.manual_seed(2026)
171 x = torch.randn(
172 NUM_TOKENS,
173 HIDDEN_SIZE,
174 generator=generator,
175 device=DEVICE,
176 dtype=DTYPE,
177 )
178 topk_ids = torch.arange(NUM_TOKENS, device=DEVICE, dtype=torch.int64).unsqueeze(1)
179 topk_ids = torch.cat(
180 [topk_ids % NUM_EXPERTS, (topk_ids + 1) % NUM_EXPERTS],
181 dim=1,
182 )
183 topk_weights = torch.full(
184 (NUM_TOKENS, TOPK),
185 1.0 / TOPK,
186 device=DEVICE,
187 dtype=DTYPE,
188 )
189 return x, topk_weights, topk_ids
190
191
192 def _forward(
193 layer: RoutedExperts,
194 x: torch.Tensor,
195 topk_weights: torch.Tensor,
196 topk_ids: torch.Tensor,
197 ) -> torch.Tensor:
198 return layer.forward_modular(
199 x,
200 topk_weights,
201 topk_ids,
202 shared_experts=None,
203 shared_experts_input=None,
204 )
205
206
207 def _capture_forward_graph(
208 layer: RoutedExperts,
209 x: torch.Tensor,
210 topk_weights: torch.Tensor,
211 topk_ids: torch.Tensor,
212 ) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]:
213 # Warm up kernel dispatch/JIT before capture so the captured region is only
214 # the steady-state forward path.
215 for _ in range(3):
216 _forward(layer, x, topk_weights, topk_ids)
217 torch.accelerator.synchronize()
218
219 graph = torch.cuda.CUDAGraph()
220 stream = torch.cuda.Stream()
221 with torch.cuda.graph(graph, stream=stream):
222 graph_output = _forward(layer, x, topk_weights, topk_ids)
223 torch.accelerator.synchronize()
224 return graph, graph_output
225
226
227 def _weight_pointers(layer: RoutedExperts) -> dict[str, int]:
228 return {
229 "w13_weight": layer.w13_weight.data_ptr(),
230 "w2_weight": layer.w2_weight.data_ptr(),
231 }
232
233
234 def _assert_backend(layer: RoutedExperts, backend: str) -> None:
235 selected = str(
236 getattr(
237 layer.quant_method.unquantized_backend,
238 "value",
239 layer.quant_method.unquantized_backend,
240 )
241 )
242 if backend == "aiter":
243 assert selected == "ROCm AITER"
244 else:
245 assert selected == "TRITON"
246
247
248 @pytest.mark.parametrize("backend", ["triton", "aiter"])
249 @torch.inference_mode()
250 def test_moe_cuda_graph_replay_after_sleep_wake_weight_update(
251 monkeypatch: pytest.MonkeyPatch,
252 default_vllm_config,
253 workspace_init,
254 backend: str,
255 ) -> None:
256 assert default_vllm_config is not None
257 assert workspace_init is None
258 if not cumem_available:
259 pytest.skip("sleep mode replay requires the cumem allocator extension")
260 if not _moe_topk_softmax_available():
261 pytest.skip("MoE C extension with topk_softmax is required")
262
263 _set_backend_env(monkeypatch, backend)
264
265 allocator = CuMemAllocator.get_instance()
266 layer = graph = graph_output = None
267 try:
268 with allocator.use_memory_pool("weights"), torch.device(DEVICE):
269 layer = _make_routed_experts()
270 _load_expert_weights(layer, seed=1)
271 _process_weights_after_loading(layer)
272 _assert_backend(layer, backend)
273
274 x, topk_weights, topk_ids = _make_static_inputs()
275 graph, graph_output = _capture_forward_graph(layer, x, topk_weights, topk_ids)
276 captured_pointers = _weight_pointers(layer)
277
278 graph.replay()
279 torch.accelerator.synchronize()
280 first_graph_output = graph_output.clone()
281
282 # Match vLLM level-2 sleep semantics: weights are discarded, then the
283 # same virtual addresses are remapped on wake and refilled by replay.
284 allocator.sleep(offload_tags=tuple())
285 allocator.wake_up()
286 _load_expert_weights(layer, seed=1)
287 _process_weights_after_loading(layer)
288
289 assert _weight_pointers(layer) == captured_pointers
290
291 graph_output.zero_()
292 graph.replay()
293 torch.accelerator.synchronize()
294
295 torch.testing.assert_close(
296 graph_output, first_graph_output, rtol=2e-2, atol=2e-2
297 )
298 finally:
299 del graph_output
300 del graph
301 del layer
302 gc.collect()
303 with contextlib.suppress(Exception):
304 allocator.release_pools()
305 torch.accelerator.empty_cache()