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 4, 06:15 UTC
[Bugfix][Rocm]Aiter MoE re-uses existing tensor addresses after weight update. (#40390)
Signed-off-by: Yuankai Chen <yuankach@amd.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Yuankai Chen committed
May 6, 2026 at 03:32 UTC
2e777d21a82c232ee8bc315c13400cd9d9afd489
2 files changed
+59
-18
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+36
-16
@@ -154,32 +154,52 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
154
w2: torch.Tensor,
155
) -> None:
156
# Shuffle weights to runtime format.
157
- w13, w2 = convert_to_unquantized_kernel_format(
157
+ w13_new, w2_new = convert_to_unquantized_kernel_format(
158
self.unquantized_backend,
159
layer=layer,
160
w13_weight=w13,
161
w2_weight=w2,
162
)
163
- replace_parameter(layer, "w13_weight", w13)
164
- replace_parameter(layer, "w2_weight", w2)
163
+ # `moe_kernel` is initialized to None in FusedMoEMethodBase.__init__;
164
+ # On the first call we replace the parameter normally. On subsequent
165
+ # calls (e.g. RL weight updates that re-trigger
166
+ # process_weights_after_loading) the moe kernel has already been set
167
+ # up and CUDA graphs may have captured the parameter addresses, so
168
+ # we copy the shuffled data into the existing storage instead of
169
+ # re-registering a new Parameter.
170
+ is_weight_update = self.moe_kernel is not None # type: ignore[has-type]
171
+ replace_parameter(layer, "w13_weight", w13_new, prefer_copy=is_weight_update)
172
+ replace_parameter(layer, "w2_weight", w2_new, prefer_copy=is_weight_update)
173
166
- # Setup moe kernel.
167
- self.moe_quant_config = self.get_fused_moe_quant_config(layer)
168
- assert self.moe_quant_config is not None
169
- assert self.experts_cls is not None
170
- self.moe_kernel = make_unquantized_moe_kernel(
171
- quant_config=self.moe_quant_config,
172
- moe_config=self.moe,
173
- backend=self.unquantized_backend,
174
- experts_cls=self.experts_cls,
175
- routing_tables=layer._maybe_init_expert_routing_tables(),
176
- shared_experts=layer.shared_experts,
177
- )
174
+ if not is_weight_update:
175
+ # Setup moe kernel only on the first call. For the unquantized
176
+ # method, moe_quant_config is either the constant
177
+ # FUSED_MOE_UNQUANTIZED_CONFIG or biased_moe_quant_config(...)
178
+ # which references layer.w{13,2}_bias; since weight updates
179
+ # mutate those bias tensors in place, the kernel does not need
180
+ # to be re-built.
181
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
182
+ assert self.moe_quant_config is not None
183
+ assert self.experts_cls is not None
184
+ self.moe_kernel = make_unquantized_moe_kernel(
185
+ quant_config=self.moe_quant_config,
186
+ moe_config=self.moe,
187
+ backend=self.unquantized_backend,
188
+ experts_cls=self.experts_cls,
189
+ routing_tables=layer._maybe_init_expert_routing_tables(),
190
+ shared_experts=layer.shared_experts,
191
+ )
192
193
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
194
super().process_weights_after_loading(layer)
195
182
- # Padding the weight for better performance on ROCm
196
+ # Padding the weight for better performance on ROCm.
197
+ # _maybe_pad_weight is idempotent: on the first call it allocates a
198
+ # padded storage and returns a strided view; on subsequent calls
199
+ # (weight updates) the stride condition no longer matches so it
200
+ # returns the input unchanged. The reassignment to .data is therefore
201
+ # a no-op on updates and preserves the storage address (data_ptr)
202
+ # used by captured CUDA graphs.
203
layer.w13_weight.data = self._maybe_pad_weight(layer.w13_weight.data)
204
layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
205
vllm/model_executor/utils.py
+23
-2
@@ -45,7 +45,10 @@ def set_weight_attrs(
45
46
47
def replace_parameter(
48
- layer: torch.nn.Module, param_name: str, new_data: torch.Tensor | None
48
+ layer: torch.nn.Module,
49
+ param_name: str,
50
+ new_data: torch.Tensor | None,
51
+ prefer_copy: bool = False,
52
):
53
"""
54
Replace a parameter of a layer while maintaining the ability to reload the weight.
@@ -57,6 +60,12 @@ def replace_parameter(
60
layer: Layer containing parameter to replace
61
param_name: Name of parameter to replace
62
new_data: New data of the new parameter, or None to set the parameter to None
63
+ prefer_copy: If True and the existing parameter is compatible with
64
+ ``new_data`` (same shape, dtype, and device), copy ``new_data``
65
+ into the existing parameter in place rather than re-registering
66
+ a new parameter. This preserves the parameter's storage address
67
+ (``data_ptr``), which is required for captured CUDA graphs to
68
+ remain valid across weight updates (e.g. in RL training loops).
69
"""
70
# should not be used on a tied/shared param
71
@@ -67,9 +76,21 @@ def replace_parameter(
76
77
if isinstance(new_data, torch.nn.Parameter):
78
new_data = new_data.data
70
- new_param = torch.nn.Parameter(new_data, requires_grad=False)
79
80
old_param: torch.nn.Parameter | None = getattr(layer, param_name, None)
81
+
82
+ if (
83
+ prefer_copy
84
+ and old_param is not None
85
+ and old_param.shape == new_data.shape
86
+ and old_param.dtype == new_data.dtype
87
+ and old_param.device == new_data.device
88
+ ):
89
+ old_param.copy_(new_data)
90
+ return
91
+
92
+ new_param = torch.nn.Parameter(new_data, requires_grad=False)
93
+
94
if old_param is not None and hasattr(old_param, "weight_loader"):
95
weight_loader = old_param.weight_loader
96
set_weight_attrs(new_param, {"weight_loader": weight_loader})