[BugFix][Model] Remove redundant SFA o_proj TP weights (vllm-project#11068)

### What this PR does / why we need it?

This PR fixes the SFA v1 DSA-CP mixed `o_proj` path so the original TP
shard remains the only persistent source of truth.

SFA DSA-CP mixed execution already couples DSA-CP with the `o_proj` TP
path. This PR keeps that behavior and documents the intended storage
invariant instead of adding a new user-facing switch:

- Decode-only batches exchange SFA outputs with TP all-to-all, then run
the original TP-sharded `o_proj`.
- Prefill/mixed batches temporarily all-gather the TP-sharded `o_proj`
weight and input-sharded quantization parameters into full-weight
buffers before the `o_proj` forward.
- `o_proj_tp_*` tensors alias the original parameter storage instead of
cloning it.
- `o_proj_full_*` tensors remain temporary gather destinations for the
prefill/mixed path only.
- Input-sharded `o_proj` quant params are discovered by `input_dim == 1`
instead of hard-coded parameter names.
- The DSA-CP mixed `o_proj` data path is documented in
`docs/source/developer_guide/Design_Documents/context_parallel.md`.
- Load-time `.contiguous()` is kept for W4A4 MXFP4 and W8A8 MXFP8 linear
transpose paths that feed NPU quant matmul. W8A8 restore now uses
`reshape` instead of `view` for the restored scale shape so the
contiguous transformed scale can still be restored for weight loading.

The previous clone-based setup made the mixed path look like two
persistent `o_proj` weights. This patch keeps the existing runtime
behavior but removes the redundant persistent TP copy and records the
intended storage invariant in the design doc and code comments.

### Does this PR introduce _any_ user-facing change?

No. There is no new config switch or public API change. The existing SFA
DSA-CP mixed behavior remains coupled with the `o_proj` TP path.

### How was this patch tested?

Static/unit checks run in an A5 environment with this PR checkout:

```bash
python -m pytest tests/ut/attention/test_sfa_o_proj_tp.py -q
# 2 passed

python -m pytest \
  tests/ut/quantization/methods/test_w4a4_mxfp4.py \
  tests/ut/quantization/methods/test_w8a8_mxfp8.py -q
# 18 passed

python -m ruff check \
  vllm_ascend/attention/sfa_v1.py \
  vllm_ascend/quantization/methods/w4a4_mxfp4.py \
  vllm_ascend/quantization/methods/w8a8_mxfp8.py \
  tests/ut/attention/test_sfa_o_proj_tp.py \
  tests/ut/quantization/methods/test_w4a4_mxfp4.py \
  tests/ut/quantization/methods/test_w8a8_mxfp8.py

python -m ruff format --check \
  vllm_ascend/attention/sfa_v1.py \
  vllm_ascend/quantization/methods/w4a4_mxfp4.py \
  vllm_ascend/quantization/methods/w8a8_mxfp8.py \
  tests/ut/attention/test_sfa_o_proj_tp.py \
  tests/ut/quantization/methods/test_w4a4_mxfp4.py \
  tests/ut/quantization/methods/test_w8a8_mxfp8.py

python -m py_compile vllm_ascend/quantization/methods/w8a8_mxfp8.py
```

NPU validation:

| Case | Checkpoint | TP | DSA-CP | Loading model weights | Available /
current KV cache |
| --- | --- | ---: | --- | ---: | ---: |
| W4A4 baseline | GLM-5.1-W4A4C8-mxfp4-rot | 8 | off | 52.3663 GB |
27.59 GiB / 27.59 GiB |
| W4A4 DSA-CP | GLM-5.1-W4A4C8-mxfp4-rot | 8 | on | 57.6788 GB | 22.26
GiB / 22.26 GiB |

The W4A4 before/after DSA-CP comparison keeps the same TP size, model
length, batch-token limit, quantization mode, and memory utilization.
The DSA-CP run has higher persistent weight accounting than the non-DSA
baseline because DSA-CP changes the attention-side layout beyond
`o_proj`; the PR removes the redundant persistent `o_proj` TP clone
within that mixed path.

- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@dc68bd8

Signed-off-by: maoxx241 <maomaoyu870@gmail.com>

Author: maoxx241

docs/source/developer_guide/Design_Documents/context_parallel.md

@@ -121,6 +121,25 @@ By predefining the maximum amount of KV cache processed per round, we sequential
 
 ![PCP-ChunkedPrefill](../../assets/cp/chunkedprefill.png)
 
+### SFA DSA-CP Mixed `o_proj` Path
+
+SFA DSA-CP mixed execution intentionally reuses the normal TP-sharded `o_proj`.
+This is part of the DSA-CP mixed data path, not a standalone user-facing `o_proj` TP switch.
+The mixed path is used when one instance may handle both decode-only and prefill/mixed batches, so `o_proj` must support two layouts at runtime:
+
+- **Decode-only batches** keep the decode TP path.
+  SFA outputs are exchanged with an all-to-all in the TP group, then the original TP-sharded `o_proj` runs normally.
+- **Prefill or mixed batches** produce SFA outputs that are not directly compatible with the TP-sharded `o_proj` input layout.
+  Before `o_proj` forward, each rank all-gathers the TP-sharded `o_proj` weight and all input-sharded quantization parameters into temporary full-weight buffers.
+  The full-weight `o_proj` forward runs once for that batch, and the module is then restored to the TP parameter aliases.
+
+The storage invariant is that the original TP-sharded `o_proj` parameter remains the only persistent source of truth.
+`o_proj_tp_*` tensors are aliases of the original parameter storage.
+`o_proj_full_*` tensors are reusable communication buffers for prefill/mixed full-gather execution only.
+They must not become a second persistent copy of the TP weight.
+
+This coupling preserves the existing decode TP behavior, supports prefill/mixed DSA-CP batches, and avoids adding an extra configuration path whose state can drift from DSA-CP mixed execution.
+
 ### Related Files
 
 - slot_mapping computation: `vllm_ascend/worker/block_table.py`

tests/ut/attention/test_sfa_o_proj_tp.py

@@ -0,0 +1,103 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+import sys
+from unittest.mock import MagicMock
+
+import torch
+
+from tests.ut.base import TestBase
+
+if "torch_npu._inductor" not in sys.modules:
+    sys.modules["torch_npu._inductor"] = MagicMock()
+
+from vllm_ascend.attention.sfa_v1 import AscendSFAImpl
+
+
+class TestAscendSFAOProjTPParams(TestBase):
+    class _OProj(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.randn(4, 3), requires_grad=False)
+            self.aclnn_input_scale = torch.nn.Parameter(torch.randn(3), requires_grad=False)
+            self.weight_scale_second = torch.nn.Parameter(torch.randn(4, 2), requires_grad=False)
+            self.weight_scale_second.input_dim = 1
+            self.weight_offset_second = torch.nn.Parameter(torch.randn(4, 2), requires_grad=False)
+            self.weight_offset_second.input_dim = 1
+            self.extra_input_scale = torch.nn.Parameter(torch.randn(4, 2), requires_grad=False)
+            self.extra_input_scale.input_dim = 1
+            self.weight_scale = torch.nn.Parameter(torch.randn(4), requires_grad=False)
+
+    def setUp(self):
+        AscendSFAImpl.o_proj_full_pools.clear()
+
+    def _make_impl(self):
+        impl = AscendSFAImpl.__new__(AscendSFAImpl)
+        impl.tp_size = 2
+        impl.o_proj = self._OProj()
+        impl._is_o_proj_unquantized = lambda: False
+        return impl
+
+    def test_o_proj_tp_params_alias_original_storage(self):
+        impl = self._make_impl()
+        o_proj = impl.o_proj
+
+        impl._init_o_proj_tp_full_params()
+
+        self.assertEqual(impl.o_proj_tp_weight.data_ptr(), o_proj.weight.data_ptr())
+        self.assertEqual(
+            impl.o_proj_tp_aclnn_input_params["aclnn_input_scale"].data_ptr(),
+            o_proj.aclnn_input_scale.data_ptr(),
+        )
+        self.assertEqual(
+            impl.o_proj_tp_input_sharded_quant_params["weight_scale_second"].data_ptr(),
+            o_proj.weight_scale_second.data_ptr(),
+        )
+        self.assertEqual(
+            impl.o_proj_tp_input_sharded_quant_params["weight_offset_second"].data_ptr(),
+            o_proj.weight_offset_second.data_ptr(),
+        )
+        self.assertEqual(
+            impl.o_proj_tp_input_sharded_quant_params["extra_input_scale"].data_ptr(),
+            o_proj.extra_input_scale.data_ptr(),
+        )
+        self.assertNotIn("weight_scale", impl.o_proj_tp_input_sharded_quant_params)
+
+    def test_o_proj_full_weight_forward_restores_tp_storage(self):
+        impl = self._make_impl()
+        impl._init_o_proj_tp_full_params()
+        original_weight_ptr = impl.o_proj.weight.data_ptr()
+        original_scale_ptr = impl.o_proj.weight_scale_second.data_ptr()
+        full_weight_ptr = impl.o_proj_full_pool.data_ptr()
+        full_scale_ptr = impl.o_proj_full_input_sharded_quant_params["weight_scale_second"].data_ptr()
+
+        def _apply_with_full_weight(_attn_output):
+            self.assertEqual(impl.o_proj.weight.data_ptr(), full_weight_ptr)
+            self.assertEqual(impl.o_proj.weight_scale_second.data_ptr(), full_scale_ptr)
+            return torch.ones(2, 4)
+
+        impl._apply_o_proj_full_weight = MagicMock(side_effect=_apply_with_full_weight)
+
+        output, require_o_proj_forward = impl._handle_o_proj_weight_switch_and_forward(
+            attn_output=torch.randn(2, 3),
+            output=torch.empty(2, 4),
+            o_proj_full_handle=None,
+            o_proj_full_param_handles=[],
+            should_shard_weight=True,
+        )
+
+        self.assertEqual(impl.o_proj.weight.data_ptr(), original_weight_ptr)
+        self.assertEqual(impl.o_proj.weight_scale_second.data_ptr(), original_scale_ptr)
+        self.assertFalse(require_o_proj_forward)
+        self.assertTrue(torch.equal(output, torch.ones(2, 4)))

tests/ut/quantization/methods/test_w4a4_mxfp4.py

@@ -40,6 +40,8 @@ def test_process_weights_after_loading_transposes(self):
         self.scheme.process_weights_after_loading(layer)
         self.assertEqual(layer.weight.shape, (128, 128))
         self.assertEqual(layer.weight_scale.shape[0], 4)
+        self.assertTrue(layer.weight.data.is_contiguous())
+        self.assertTrue(layer.weight_scale.data.is_contiguous())
 
     @patch("vllm_ascend.quantization.methods.w4a4_mxfp4.torch_npu")
     def test_apply_3d_input(self, mock_npu):
@@ -102,6 +104,10 @@ def test_process_weights_transposes_weights(self):
         self.scheme.process_weights_after_loading(layer)
         self.assertEqual(layer.w13_weight.shape, (8, 64, 256))
         self.assertEqual(layer.w13_weight_scale.shape, (8, 2, 256, 2))
+        self.assertFalse(layer.w13_weight.data.is_contiguous())
+        self.assertFalse(layer.w2_weight.data.is_contiguous())
+        self.assertFalse(layer.w13_weight_scale.data.is_contiguous())
+        self.assertFalse(layer.w2_weight_scale.data.is_contiguous())
 
     @patch("vllm_ascend.quantization.methods.w4a4_mxfp4.torch_npu")
     @patch("vllm_ascend.quantization.methods.w4a4_mxfp4._EXTRA_CTX")

tests/ut/quantization/methods/test_w8a8_mxfp8.py

@@ -47,6 +47,8 @@ def test_process_weights_stores_original_shapes(self):
         self.assertEqual(layer._mxfp8_original_shapes["weight"], (128, 256))
         self.assertTrue(layer._mxfp8_transformed)
         self.assertEqual(layer.weight_scale.shape, (4, 128, 2))
+        self.assertTrue(layer.weight.data.is_contiguous())
+        self.assertTrue(layer.weight_scale.data.is_contiguous())
 
     def test_restore_after_process_returns_original_shape(self):
         layer = nn.Module()
@@ -120,6 +122,10 @@ def test_process_weights_stores_original_shapes(self):
         self.assertTrue(hasattr(layer, "_mxfp8_original_shapes"))
         self.assertIn("w13_weight", layer._mxfp8_original_shapes)
         self.assertEqual(layer.w13_weight.shape, (original_shape[0], original_shape[2], original_shape[1]))
+        self.assertFalse(layer.w13_weight.data.is_contiguous())
+        self.assertFalse(layer.w2_weight.data.is_contiguous())
+        self.assertFalse(layer.w13_weight_scale.data.is_contiguous())
+        self.assertFalse(layer.w2_weight_scale.data.is_contiguous())
 
     def test_restore_weights_for_rl_loading(self):
         layer = create_mxfp_moe_layer(

vllm_ascend/attention/sfa_v1.py

@@ -72,7 +72,6 @@
     "aclnn_input_scale_reciprocal",
     "aclnn_input_offset",
 )
-O_PROJ_INPUT_SHARDED_QUANT_PARAMS = ("weight_scale_second", "weight_scale")
 
 
 def _get_indexer_types(configs: tuple[Any, ...]) -> Any | None:
@@ -581,8 +580,12 @@ def __init__(
         # Enable layer sharding via DSA-CP on the P node in the PD-disaggregated setup.
         self.enable_dsa_cp_with_layer_shard = enable_dsa_cp_with_layer_shard()
 
-        # use original TP o_proj weight in PD mix stage, and full gather
-        # for o_proj weight for prefill stage.
+        # SFA DSA-CP mixed deployments keep o_proj in the existing TP layout.
+        # Decode can use the TP-sharded o_proj directly after an activation
+        # all-to-all, while prefill/mixed batches temporarily gather the TP
+        # shards into a full-weight buffer because their SFA output is not
+        # TP-sharded. This is part of the DSA-CP mixed-mode data path rather
+        # than an independent user-facing feature switch.
         self.enable_dsa_cp_with_o_proj_tp = enable_dsa_cp_with_o_proj_tp()
 
         if self.enable_dsa_cp:
@@ -657,7 +660,7 @@ def process_weights_after_loading(self, act_dtype: torch.dtype):
                 for layer in self.layer_sharding_kwargs or []:
                     if is_hidden_layer(layer):
                         post_process_after_loading_for_shard_weight_series(layer)
-            else:
+            elif self.enable_dsa_cp_with_o_proj_tp:
                 self._init_o_proj_tp_full_params()
 
         if self.enable_mlapo:
@@ -861,12 +864,20 @@ def rope_single(
 
     def _init_o_proj_tp_full_params(self):
         """
-        Initialize TP-mode and Full-mode parameters for o_proj weight,
-        preparing for weight switching in PD mix stage.
-
-        For PD mix stage:
-        - Use original TP o_proj weight for decode phase
-        - Need full-gather o_proj weight from all TP ranks for prefill phase
+        Initialize TP-mode aliases and Full-mode buffers for DSA-CP o_proj.
+
+        In SFA DSA-CP mixed execution, the same model instance can run both
+        decode-only and prefill/mixed batches:
+        - Decode-only batches all-to-all the SFA output in the TP group, then
+          run the original TP-sharded o_proj.
+        - Prefill/mixed batches produce SFA output that is not directly
+          compatible with TP-sharded o_proj, so each rank all-gathers the TP
+          o_proj shards and input-sharded quant params before running o_proj.
+
+        The original TP parameter storage remains the persistent source of
+        truth. The o_proj_tp_* tensors below alias that storage, while the
+        o_proj_full_* tensors are temporary gather destinations reused across
+        forwards. They are not a second persistent copy of the TP weight.
         """
         sample = self.o_proj.weight
         self.o_proj_full_weight_gather_dim = 1 if self._is_o_proj_unquantized() else 0
@@ -895,36 +906,41 @@ def _init_o_proj_tp_full_params(self):
         else:
             self.o_proj_full_pool = self.o_proj_full_gather_pool.transpose(0, 1)
 
-        # Save TP-mode parameters (original sharded weights)
-        self.o_proj_tp_weight = self.o_proj.weight.clone().detach()
+        # TP tensors alias the original parameter storage. The TP shard remains
+        # the single source of truth; full-weight tensors below are temporary
+        # gather destinations only.
+        self.o_proj_tp_weight = self.o_proj.weight.detach()
         if self.o_proj_full_weight_gather_dim == 0:
             self.o_proj_tp_weight_gather_input = self.o_proj_tp_weight
         else:
+            # Communication scratch only: all_gather_into_tensor concatenates on
+            # dim0, while unquantized row-parallel o_proj is sharded on dim1.
             self.o_proj_tp_weight_gather_input = self.o_proj_tp_weight.transpose(0, 1).contiguous()
         self.o_proj_tp_aclnn_input_params = {}
         self.o_proj_full_aclnn_input_params = {}
         for param_name in O_PROJ_ACLNN_INPUT_PARAMS:
             param = getattr(self.o_proj, param_name, None)
             if param is None:
                 continue
-            self.o_proj_tp_aclnn_input_params[param_name] = param.clone().detach()
+            self.o_proj_tp_aclnn_input_params[param_name] = param.detach()
             self.o_proj_full_aclnn_input_params[param_name] = param.repeat(self.tp_size)
 
         self.o_proj_tp_input_sharded_quant_params = {}
         self.o_proj_full_input_sharded_quant_params = {}
-        for param_name in O_PROJ_INPUT_SHARDED_QUANT_PARAMS:
-            param = getattr(self.o_proj, param_name, None)
-            if param is None or getattr(param, "input_dim", None) != 1:
-                continue
-            self.o_proj_tp_input_sharded_quant_params[param_name] = param.clone().detach()
+        for param_name, param in self._iter_o_proj_input_sharded_quant_params():
+            self.o_proj_tp_input_sharded_quant_params[param_name] = param.detach()
             self.o_proj_full_input_sharded_quant_params[param_name] = torch.empty(
                 (param.shape[0] * self.tp_size, *param.shape[1:]), dtype=param.dtype, device=param.device
             )
 
-        # Initially switch to TP mode for graph capture
-        self.o_proj.weight.set_(self.o_proj_tp_weight)
-        self._switch_o_proj_params(self.o_proj_tp_aclnn_input_params)
-        self._switch_o_proj_params(self.o_proj_tp_input_sharded_quant_params)
+    def _iter_o_proj_input_sharded_quant_params(self):
+        if not isinstance(self.o_proj, nn.Module):
+            return
+        for param_name, param in self.o_proj.named_parameters(recurse=False):
+            if param_name == "weight" or param_name in O_PROJ_ACLNN_INPUT_PARAMS:
+                continue
+            if getattr(param, "input_dim", None) == 1:
+                yield param_name, param
 
     def _switch_o_proj_params(self, params: dict[str, torch.Tensor]):
         for param_name, param in params.items():
@@ -960,15 +976,13 @@ def _handle_o_proj_weight_switch_and_forward(
                 if handle is not None:
                     handle.wait()
 
-            # Switch o_proj to Full-mode (gathered weight from all TP ranks)
+            # Temporarily switch o_proj to the gathered full-weight view for
+            # prefill/mixed DSA-CP, whose attention output is not TP-sharded.
             self.o_proj.weight.set_(self.o_proj_full_pool)
             self._switch_o_proj_params(self.o_proj_full_aclnn_input_params)
             self._switch_o_proj_params(self.o_proj_full_input_sharded_quant_params)
-
-            # Apply quantization method and execute forward computation
             output[...] = self._apply_o_proj_full_weight(attn_output)
-
-            # Switch o_proj back to TP-mode for subsequent decode operations
+            # Restore TP aliases so later decode batches keep using TP storage.
             self.o_proj.weight.set_(self.o_proj_tp_weight)
             self._switch_o_proj_params(self.o_proj_tp_aclnn_input_params)
             self._switch_o_proj_params(self.o_proj_tp_input_sharded_quant_params)
@@ -1312,8 +1326,8 @@ def forward(
         # all-gather o_proj weight for prefill stage of PD mix node
         o_proj_full_handle = None
         o_proj_full_param_handles = None
-        # if is PD mix stage, using original TP o_proj weight, and also need to full gather for o_proj
-        # weight for prefill stage.
+        # Prefill/mixed DSA-CP computes o_proj with a temporary full weight.
+        # Decode keeps the original TP path and only exchanges activations.
         full_gather_o_proj_enabled = self.enable_dsa_cp_with_o_proj_tp and attn_metadata.attn_state not in {
             AscendAttentionState.DecodeOnly,
             AscendAttentionState.SpecDecoding,
@@ -1621,9 +1635,9 @@ def forward(
         )
 
         if self.enable_dsa_cp_with_o_proj_tp:
-            # When using SFA-CP with pd mixed, o_proj has two cases:
-            # 1. prefill: o_proj is a TP weight, we need to all-gather o_proj weight to switch TP=1.
-            # 2. decode: all-to-all the hidden_state before the o_proj forward.
+            # SFA DSA-CP mixed mode keeps o_proj weight sharded in the TP domain:
+            # 1. prefill/mixed: gather TP shards into a temporary full weight.
+            # 2. decode-only: all-to-all hidden states, then run TP o_proj.
             result, require_o_proj_forward = self._handle_o_proj_weight_switch_and_forward(
                 attn_output=attn_output,
                 output=output,

vllm_ascend/quantization/methods/w4a4_mxfp4.py

@@ -111,8 +111,8 @@ def process_weights_after_loading(self, layer):
 
         n_dim, k_dim = layer.weight_scale.data.shape
         layer.weight_scale.data = layer.weight_scale.data.reshape(n_dim, k_dim // 2, 2)
-        layer.weight.data = layer.weight.data.transpose(0, 1)
-        layer.weight_scale.data = layer.weight_scale.data.transpose(0, 1)
+        layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
+        layer.weight_scale.data = layer.weight_scale.data.transpose(0, 1).contiguous()
 
 
 @register_scheme("W4A4_MXFP4", "moe")
@@ -252,6 +252,8 @@ def process_weights_after_loading(self, layer):
         layer.w13_weight_scale.data = layer.w13_weight_scale.data.reshape(g_num, n_size, k_size // 2, 2)
         g_num, n_size, k_size = layer.w2_weight_scale.shape
         layer.w2_weight_scale.data = layer.w2_weight_scale.data.reshape(g_num, n_size, k_size // 2, 2)
+        # The A5 MXFP4 fused grouped-matmul-swiglu op relies on the
+        # transpose stride to interpret packed FP4 weights as logical K.
         layer.w13_weight.data = layer.w13_weight.data.transpose(1, 2)
         layer.w2_weight.data = layer.w2_weight.data.transpose(1, 2)
         layer.w13_weight_scale.data = layer.w13_weight_scale.data.transpose(1, 2)

vllm_ascend/quantization/methods/w8a8_mxfp8.py

@@ -138,8 +138,8 @@ def process_weights_after_loading(self, layer):
             layer.weight_scale.data = layer.weight_scale.data.reshape(n_dim, k_dim // 2 + 1, 2)
         else:
             layer.weight_scale.data = layer.weight_scale.data.reshape(n_dim, k_dim // 2, 2)
-        layer.weight.data = layer.weight.data.transpose(0, 1)
-        layer.weight_scale.data = layer.weight_scale.data.transpose(0, 1)
+        layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
+        layer.weight_scale.data = layer.weight_scale.data.transpose(0, 1).contiguous()
 
         # Mark as transformed
         layer._mxfp8_transformed = True
@@ -184,7 +184,7 @@ def restore_weights_for_rl_loading(self, layer):
         # Current shape: (k_dim//2, n_dim, 2)
         # Target shape: (n_dim, k_dim)
         target_scale = layer.weight_scale.data.transpose(0, 1).reshape(orig_scale_shape).contiguous()
-        layer.weight_scale.data = layer.weight_scale.data.transpose(0, 1).view(orig_scale_shape)
+        layer.weight_scale.data = layer.weight_scale.data.transpose(0, 1).reshape(orig_scale_shape)
         layer.weight_scale.data.copy_(target_scale)
 
         # Mark as not transformed (ready for weight loading)