From 81be2e09aebfd1c45b3ed9f73d2850da8a72984c Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Thu, 6 Aug 2026 19:03:37 +0200 Subject: [PATCH] [Weight processing] Copy over `new_data` attributes in `replace_parameter` (#49601) Signed-off-by: Felix Marty Co-authored-by: Claude Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/model_executor/test_utils.py | 283 ++++++++++++++++++ .../fused_moe/experts/aiter_mxfp8_moe.py | 5 - .../layers/fused_moe/oracle/fp8.py | 4 + .../layers/fused_moe/oracle/mxfp4.py | 6 + .../layers/fused_moe/oracle/unquantized.py | 2 + .../fused_moe/unquantized_fused_moe_method.py | 5 - .../model_executor/layers/quantization/fp8.py | 6 - .../layers/quantization/mxfp4.py | 17 +- .../layers/quantization/quark/quark_moe.py | 9 - vllm/model_executor/utils.py | 48 ++- 10 files changed, 334 insertions(+), 51 deletions(-) create mode 100644 tests/model_executor/test_utils.py diff --git a/tests/model_executor/test_utils.py b/tests/model_executor/test_utils.py new file mode 100644 index 0000000000..1ef6f03813 --- /dev/null +++ b/tests/model_executor/test_utils.py @@ -0,0 +1,283 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import pytest +import torch + +from vllm.model_executor.parameter import ModelWeightParameter, PackedvLLMParameter +from vllm.model_executor.utils import replace_parameter + + +@pytest.fixture +def single_rank_tp(monkeypatch: pytest.MonkeyPatch) -> None: + """`BasevLLMParameter.__init__` queries the TP group, which is not + initialized in a unit test. Pin it to a single rank. + """ + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0 + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1 + ) + + +@pytest.mark.parametrize("prefer_copy", [False, True]) +@pytest.mark.parametrize("wrap_in_parameter", [False, True]) +def test_replace_parameter_preserves_custom_attribute( + prefer_copy: bool, wrap_in_parameter: bool +) -> None: + """`replace_parameter` must carry over attributes attached to the + replacement tensor (e.g. the `is_shuffled` flag set by AITER weight + preprocessing in `Fp8MoEMethod.process_weights_after_loading`). + + The replacement is passed both as a plain `torch.Tensor` and as a + `torch.nn.Parameter`, since the latter is unwrapped through `.data` + internally, which does not carry attributes over either. + """ + layer = torch.nn.Module() + layer.register_parameter( + "weight", torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) + ) + original_data_ptr = layer.weight.data_ptr() + + new_data: torch.Tensor = torch.ones(4, 4).t() + if wrap_in_parameter: + new_data = torch.nn.Parameter(new_data, requires_grad=False) + new_data.is_shuffled = True + + # Sanity check on the assumption above: tensor internals are not + # reachable through `__dict__`, only user-set custom attributes are. + assert new_data.__dict__.keys() == {"is_shuffled"} + + replace_parameter(layer, "weight", new_data, prefer_copy=prefer_copy) + + assert isinstance(layer.weight, torch.nn.Parameter) + assert layer.weight.is_shuffled is True + assert torch.equal(layer.weight.data, new_data) + assert layer.weight.device == new_data.device + + if prefer_copy: + # The existing storage is reused, so addresses captured in CUDA graphs + # stay valid across the update. + assert layer.weight.data_ptr() == original_data_ptr + else: + assert layer.weight.stride() == new_data.stride() + assert layer.weight.data_ptr() == new_data.data_ptr() + + +@pytest.mark.parametrize("prefer_copy", [False, True]) +def test_replace_parameter_preserves_weight_loader(prefer_copy: bool) -> None: + """The reload path must survive replacement: the old parameter's + `weight_loader` is carried over as-is, so it is still invoked as + `weight_loader(param, loaded_weight)`. + + Real loaders are frequently bound methods of the layer (`RoutedExperts` + hands `self.weight_loader` to `create_weights`); what must not happen is + the carry-over re-binding the loader to the new parameter, which would + shift `loaded_weight` into the `param` slot. + """ + calls: list[tuple[torch.Tensor, torch.Tensor]] = [] + + def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + calls.append((param, loaded_weight)) + + layer = torch.nn.Module() + old_param = torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) + old_param.weight_loader = weight_loader + layer.register_parameter("weight", old_param) + + replace_parameter(layer, "weight", torch.ones(4, 4), prefer_copy=prefer_copy) + + # Re-binding to the new parameter would shift `loaded_weight` into the + # `param` slot, so the plain function must not have grown a `__self__`. + assert not hasattr(layer.weight.weight_loader, "__self__") + + loaded_weight = torch.full((4, 4), 2.0) + layer.weight.weight_loader(layer.weight, loaded_weight) + + assert len(calls) == 1 + assert calls[0][0] is layer.weight + assert calls[0][1] is loaded_weight + + +def test_replace_parameter_weight_loader_comes_from_old_parameter() -> None: + """`weight_loader` is deliberately excluded from the attribute carry-over: + the old parameter's loader is authoritative. A loader riding along on the + replacement tensor must neither shadow it nor trip the overwrite assertion + in `set_weight_attrs`, and the other attributes must still be carried over. + + `_weight_loader` is excluded for the same reason: it is the backing field + of `BasevLLMParameter.weight_loader`, so leaving it in would smuggle a + stale loader past the `weight_loader` exclusion. + """ + calls: list[str] = [] + + def old_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + calls.append("old") + + def stale_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + calls.append("stale") + + layer = torch.nn.Module() + old_param = torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) + old_param.weight_loader = old_weight_loader + layer.register_parameter("weight", old_param) + + new_data = torch.ones(4, 4) + new_data.weight_loader = stale_weight_loader + new_data._weight_loader = stale_weight_loader + new_data.is_shuffled = True + + replace_parameter(layer, "weight", new_data) + + layer.weight.weight_loader(layer.weight, torch.full((4, 4), 2.0)) + + assert calls == ["old"] + assert not hasattr(layer.weight, "_weight_loader") + assert layer.weight.is_shuffled is True + + +def test_replace_parameter_does_not_rebind_plain_function_attribute() -> None: + """A plain function stored on the replacement tensor must be carried over + verbatim; re-binding it as a method of the new parameter would silently + shift the arguments of every subsequent call. + """ + + def scale_for(group_size: int) -> int: + return group_size * 2 + + layer = torch.nn.Module() + layer.register_parameter( + "weight", torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) + ) + + new_data = torch.ones(4, 4) + new_data.scale_for = scale_for + + replace_parameter(layer, "weight", new_data) + + assert layer.weight.scale_for is scale_for + assert layer.weight.scale_for(3) == 6 + + +def test_replace_parameter_preserves_bound_method_attribute() -> None: + """A callable already bound to another object must keep pointing at that + object after replacement. + """ + + class KernelDispatcher: + def __init__(self) -> None: + self.calls: list[int] = [] + + def record(self, value: int) -> None: + self.calls.append(value) + + dispatcher = KernelDispatcher() + + layer = torch.nn.Module() + layer.register_parameter( + "weight", torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) + ) + + new_data = torch.ones(4, 4) + new_data.record = dispatcher.record + + replace_parameter(layer, "weight", new_data) + + assert layer.weight.record.__self__ is dispatcher + layer.weight.record(7) + assert dispatcher.calls == [7] + + +@pytest.mark.parametrize("param_kind", ["plain", "model_weight", "packed"]) +def test_replace_parameter_attributes_from_the_layers_own_parameter( + single_rank_tp: None, param_kind: str +) -> None: + """Several callers hand back the parameter they were given, after mutating + its `.data` in place: the HUMMING branch of + `convert_to_fp8_moe_kernel_format` (`w13 = layer.w13_weight`), the + `AITER_MXFP4_BF16` branch of + `convert_gpt_oss_weight_to_mxfp4_moe_kernel_format`, the no-transpose + branch of `XPUFP8ScaledMM` (`layer_weight = w`), and `auto_awq`/`auto_gptq`, + which pass whatever is currently registered -- possibly still a + `BasevLLMParameter` subclass rather than a plain `Parameter`. + + `new_data is old_param` there, so every attribute of the replacement is by + definition an attribute of the old parameter; pin exactly which survive. + + For the vLLM parameter classes only the private backing fields come along. + The public names weight loading branches on -- `output_dim`, `input_dim`, + `packed_dim` (`getattr(param, "output_dim", None)` in `linear.py`) -- are + class-level properties, so they do not survive onto the plain replacement + and cannot be read back stale after the layout has changed. + """ + + class Layer(torch.nn.Module): + def weight_loader( + self, param: torch.Tensor, loaded_weight: torch.Tensor + ) -> None: + pass + + layer = Layer() + data = torch.zeros(4, 4) + loader = layer.weight_loader + vllm_param_attrs = { + "_weight_loader": loader, + "_input_dim": 1, + "_output_dim": 0, + "tp_rank": 0, + "tp_size": 1, + } + + old_param: torch.nn.Parameter + if param_kind == "plain": + old_param = torch.nn.Parameter(data, requires_grad=False) + old_param.weight_loader = loader + # MoE scale marker, read back by `RoutedExperts.weight_loader`. + old_param.quant_method = "block" + source_attrs = {"weight_loader": loader, "quant_method": "block"} + elif param_kind == "model_weight": + old_param = ModelWeightParameter( + data=data, input_dim=1, output_dim=0, weight_loader=loader + ) + source_attrs = dict(vllm_param_attrs) + else: + old_param = PackedvLLMParameter( + data=data, + input_dim=1, + output_dim=0, + packed_dim=0, + packed_factor=8, + weight_loader=loader, + ) + source_attrs = dict(vllm_param_attrs) + source_attrs |= { + "_packed_dim": 0, + "_packed_factor": 8, + "_marlin_tile_size": None, + } + + layer.register_parameter("weight", old_param) + old_param.data = torch.ones(4, 4) + + # Pin what the caller hands back: on the vLLM classes `weight_loader` is a + # property, so the loader sits under `_weight_loader` instead. + assert dict(old_param.__dict__) == source_attrs + + replace_parameter(layer, "weight", old_param) + + # Everything carries over by value, except that the loader is re-read from + # the old parameter and lands under its public name. Comparing values (not + # just keys) is what rules out the loader being re-bound to the new + # parameter, which would shift `loaded_weight` into the `param` slot. + expected = dict(source_attrs) + expected.pop("_weight_loader", None) + expected["weight_loader"] = loader + + assert type(layer.weight) is torch.nn.Parameter + assert dict(layer.weight.__dict__) == expected + + if param_kind != "plain": + for public_name in ("output_dim", "input_dim", "packed_dim", "packed_factor"): + assert getattr(layer.weight, public_name, None) is None diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index 59115df13b..c1fbb11a81 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -145,11 +145,6 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase): from vllm._aiter_ops import rocm_aiter_ops - # Re-tag the preshuffled weights: replace_parameter drops the - # is_shuffled flag, without which aiter picks a broken CK kernel. - w1.is_shuffled = True - w2.is_shuffled = True - limit = self.quant_config.gemm1_clamp_limit swiglu_limit = 0.0 if limit is None else float(limit) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 28e36b8b0e..7cb57a44eb 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -483,10 +483,14 @@ def convert_to_fp8_moe_kernel_format( ) elif fp8_backend == Fp8MoeBackend.AITER: w13, w2 = rocm_aiter_ops.shuffle_weights(w13, w2) + w13.is_shuffled = True + w2.is_shuffled = True elif fp8_backend == Fp8MoeBackend.AITER_MXFP8: w13, w2, w13_scale, w2_scale = rocm_aiter_ops.shuffle_mxfp8_moe_weights( w13, w2, w13_scale, w2_scale ) + w13.is_shuffled = True + w2.is_shuffled = True elif fp8_backend == Fp8MoeBackend.HUMMING: from vllm.model_executor.layers.quantization.utils.humming_utils import ( convert_to_humming_moe_kernel_format, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 0d8dc570bd..359c4e55ee 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1081,6 +1081,9 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( .view(-1, n) ) + w13_weight.is_shuffled = True + w2_weight.is_shuffled = True + return ( w13_weight, w2_weight, @@ -1504,6 +1507,9 @@ def convert_weight_to_mxfp4_moe_kernel_format( False, ) + w13_weight.is_shuffled = True + w2_weight.is_shuffled = True + return ( w13_weight, w2_weight, diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 9e09f1dd38..b2041f29fb 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -332,6 +332,8 @@ def convert_to_unquantized_kernel_format( ) -> tuple[torch.Tensor, torch.Tensor]: if unquantized_backend == UnquantizedMoeBackend.AITER: w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(w13_weight, w2_weight) + w13_weight.is_shuffled = True + w2_weight.is_shuffled = True elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS: if moe_config.is_act_and_mul: diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 9f7e1b1f7f..6ed1730737 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -146,11 +146,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): replace_parameter(layer, "w13_weight", w13_new, prefer_copy=is_weight_update) replace_parameter(layer, "w2_weight", w2_new, prefer_copy=is_weight_update) - # AITER backend requires weights to be marked as shuffled. - if self.unquantized_backend == UnquantizedMoeBackend.AITER: - layer.w13_weight.is_shuffled = True - layer.w2_weight.is_shuffled = True - if not is_weight_update: # Setup moe kernel only on the first call. For the unquantized # method, moe_quant_config carries no quantized scales -- only diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 98d126ce44..b377a16eba 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -28,7 +28,6 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( - Fp8MoeBackend, convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, make_fp8_moe_quant_config, @@ -671,11 +670,6 @@ class Fp8MoEMethod(FusedMoEMethodBase): replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale) replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale) - # AITER backend requires weights to be marked as shuffled. - if self.fp8_backend == Fp8MoeBackend.AITER: - layer.w13_weight.is_shuffled = True - layer.w2_weight.is_shuffled = True - self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.moe_quant_config is not None assert self.experts_cls is not None diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 47d47614c4..476921b63e 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -366,11 +366,6 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase): self.w13_precision_config = w13_scale self.w2_precision_config = w2_scale - # AITER backend requires weights to be marked as shuffled. - if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: - layer.w13_weight.is_shuffled = True - layer.w2_weight.is_shuffled = True - if w13_bias is not None and w2_bias is not None: replace_parameter(layer, "w13_bias", w13_bias) replace_parameter(layer, "w2_bias", w2_bias) @@ -760,14 +755,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self.w13_precision_config = w13_scale self.w2_precision_config = w2_scale - # AITER backend requires weights to be marked as shuffled. - if ( - self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 - and not uses_triton_weight_format - ): - layer.w13_weight.is_shuffled = True - layer.w2_weight.is_shuffled = True - if w13_bias is not None and w2_bias is not None: replace_parameter(layer, "w13_bias", w13_bias) replace_parameter(layer, "w2_bias", w2_bias) @@ -816,6 +803,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), num_experts, guinterleave ) w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1])) + + w13.is_shuffled = True + w2.is_shuffled = True + return w13, w2, w13_scale, w2_scale def process_weights_after_loading(self, layer): diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index f1a22c8903..f9468cf5da 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -27,7 +27,6 @@ from vllm.model_executor.layers.fused_moe.config import ( ocp_mx_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( - Fp8MoeBackend, convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, make_fp8_moe_quant_config, @@ -447,10 +446,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): replace_parameter(layer, "w13_weight_scale", w13_scale) replace_parameter(layer, "w2_weight_scale", w2_scale) - if self.fp8_backend == Fp8MoeBackend.AITER: - layer.w13_weight.is_shuffled = True - layer.w2_weight.is_shuffled = True - self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.moe_quant_config is not None assert self.experts_cls is not None @@ -1328,10 +1323,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): replace_parameter(layer, "w13_bias", w13_bias) replace_parameter(layer, "w2_bias", w2_bias) - if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: - layer.w13_weight.is_shuffled = True - layer.w2_weight.is_shuffled = True - torch.accelerator.empty_cache() # Build quant config and kernel diff --git a/vllm/model_executor/utils.py b/vllm/model_executor/utils.py index 70acb33e36..180820d61b 100644 --- a/vllm/model_executor/utils.py +++ b/vllm/model_executor/utils.py @@ -47,21 +47,29 @@ def set_weight_attrs( def replace_parameter( layer: torch.nn.Module, param_name: str, - new_data: torch.Tensor | None, + new_tensor: torch.Tensor | None, prefer_copy: bool = False, ): """ Replace a parameter of a layer while maintaining the ability to reload the weight. Called within implementations of the `process_weights_after_loading` method. + Custom attributes set on ``new_tensor`` (e.g. kernel dispatch flags such as + ``is_shuffled``) are carried over to the replacement parameter, except + ``weight_loader``, which is always taken from the existing parameter. + + Attributes of the existing parameter are otherwise dropped when a new + parameter is registered, but kept when ``prefer_copy`` reuses it in place. + This function should not be called on weights which are tied/shared Args: layer: Layer containing parameter to replace param_name: Name of parameter to replace - new_data: New data of the new parameter, or None to set the parameter to None + new_tensor: New data of the new parameter, or None to set the parameter + to None prefer_copy: If True and the existing parameter is compatible with - ``new_data`` (same shape, dtype, and device), copy ``new_data`` + ``new_tensor`` (same shape, dtype, and device), copy ``new_tensor`` into the existing parameter in place rather than re-registering a new parameter. This preserves the parameter's storage address (``data_ptr``), which is required for captured CUDA graphs to @@ -69,27 +77,41 @@ def replace_parameter( """ # should not be used on a tied/shared param - # If new_data is None, set the parameter to None - if new_data is None: + # If new_tensor is None, set the parameter to None + if new_tensor is None: setattr(layer, param_name, None) return - if isinstance(new_data, torch.nn.Parameter): - new_data = new_data.data - old_param: torch.nn.Parameter | None = getattr(layer, param_name, None) + new_tensor_attrs = dict(new_tensor.__dict__) + + # `weight_loader` is the only attribute not ported over from new_tensor, + # old_param.weight_loader only is supported. `_weight_loader` is + # `BasevLLMParameter`'s backing field for its `weight_loader` property, so + # a loader riding along under that name has to be dropped as well. + new_tensor_attrs.pop("weight_loader", None) + new_tensor_attrs.pop("_weight_loader", None) + + if isinstance(new_tensor, torch.nn.Parameter): + new_tensor = new_tensor.data + if ( prefer_copy and old_param is not None - and old_param.shape == new_data.shape - and old_param.dtype == new_data.dtype - and old_param.device == new_data.device + and old_param.shape == new_tensor.shape + and old_param.dtype == new_tensor.dtype + and old_param.device == new_tensor.device ): - old_param.copy_(new_data) + old_param.copy_(new_tensor) + for attr_name, attr in new_tensor_attrs.items(): + setattr(old_param, attr_name, attr) return - new_param = torch.nn.Parameter(new_data, requires_grad=False) + new_param = torch.nn.Parameter(new_tensor, requires_grad=False) + + for attr_name, attr in new_tensor_attrs.items(): + setattr(new_param, attr_name, attr) if old_param is not None and hasattr(old_param, "weight_loader"): weight_loader = old_param.weight_loader