[Bugfix][MoE] Preserve unquantized weight storage on ROCm (#46009)
Signed-off-by: bingzhaodong <aaab8b@gmail.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: Zhaodong Bing <zhaobing@smc300x-ccs-aus-gpuf2c5.prov.aus.ccs.cpe.ice.amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Zhaodong Bing <zhaobing@smc300x-ccs-aus-gpuf2a0.prov.aus.ccs.cpe.ice.amd.com>
Zhaodong Bing committed
Sep 3, 2026 at 11:06 UTC
5d09eb2cf310486c03385dd8d97aac5ceee5abc1
4 files changed
+315
-9
.buildkite/test-amd.yaml
+4
@@ -2473,7 +2473,10 @@ steps:
2473
- vllm/model_executor/layers/quantization
2474
- vllm/_aiter_ops.py
2475
- vllm/platforms/rocm.py
2476
+ - vllm/model_executor/layers/fused_moe/oracle/unquantized.py
2477
+ - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
2478
- tests/quantization
2479
+ - tests/rocm/test_moe_weight_replay.py
2480
commands:
2481
2482
# temporary install here since we need nightly, will move to requirements/test.in
@@ -2486,6 +2489,7 @@ steps:
2489
- uv pip install --system torchao==0.17.0
2490
- uv pip install --system conch-triton-kernels
2491
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
2492
+ - pytest -v -s rocm/test_moe_weight_replay.py
2493
2494
#--------------------------------------------------------- mi300 · samplers ----------------------------------------------------------#
2495
tests/rocm/test_moe_weight_replay.py
new
+305
@@ -0,0 +1,305 @@
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()
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+1
@@ -336,6 +336,7 @@ def convert_to_unquantized_kernel_format(
336
w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(w13_weight, w2_weight)
337
w13_weight.is_shuffled = True
338
w2_weight.is_shuffled = True
339
+ return w13_weight, w2_weight
340
341
elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
342
if moe_config.is_act_and_mul:
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+5
-9
@@ -175,15 +175,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
175
def process_weights_after_loading(self, layer: "RoutedExperts") -> None:
176
super().process_weights_after_loading(layer)
177
178
- # Padding the weight for better performance on ROCm.
179
- # _maybe_pad_weight is idempotent: on the first call it allocates a
180
- # padded storage and returns a strided view; on subsequent calls
181
- # (weight updates) the stride condition no longer matches so it
182
- # returns the input unchanged. The reassignment to .data is therefore
183
- # a no-op on updates and preserves the storage address (data_ptr)
184
- # used by captured CUDA graphs.
185
- layer.w13_weight.data = self._maybe_pad_weight(layer.w13_weight.data)
186
- layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
178
+ # Padding may allocate a same-shaped strided view. Copy it back instead
179
+ # of rebinding `.data`, since CUDA graphs capture the parameter storage
180
+ # address.
181
+ layer.w13_weight.data.copy_(self._maybe_pad_weight(layer.w13_weight.data))
182
+ layer.w2_weight.data.copy_(self._maybe_pad_weight(layer.w2_weight.data))
183
184
if self.unquantized_backend in [
185
UnquantizedMoeBackend.TPU,