[Weight processing] Copy over new_data attributes in replace_parameter (#49601)

Signed-off-by: Felix Marty <Felix.Marty@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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

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)
 

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,

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,

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:

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

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

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):

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

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