Consider offline static weight-layout folding separately

docs/execution_providers.md

@@ -142,11 +142,11 @@ absent or packed UINT8 zero points. Flat/row-major sidecars are supported;
 block size must be a power of two >= 16 and `accuracy_level` omitted or zero.
 Mixed bit widths, extra operands, and caller-owned function bodies are skipped.
 
-Parameter Concats retain checkpoint names until binding, then initializer
-Concats and matrix Transposes fold lazily without changing dtype. Missing leaves
-are checked even with `fold_constants=False`. Generic Reshape/Transpose layout
-folding is not part of this pipeline; those transformations remain explicit.
-Overridable inputs remain untouched.
+Post-weight folding is independent of experimental packing: parameter Concats
+retain checkpoint names until binding, then Concat/Reshape/Transpose chains fold
+lazily without changing dtype. Missing leaves are checked even with
+`fold_constants=False`. Overridable inputs and MoE/QMoE/QDQ storage-layout paths
+remain explicit to preserve loader names and byte layouts.
 
 **The INT4 numerical gate remains failed and unwaived:** large-dimension CUDA
 prefill/cached-decode logits exceeded `rtol=atol=1e-2`. Changing GEMM width can

src/mobius/_optimizations.py

@@ -21,8 +21,8 @@
 Post-weight passes
 ------------------
 :func:`fold_initializers_after_weights` should be called after weights are loaded.
-It lazily folds initializer matrix Transpose and Concat nodes for projection
-packing. Generic Reshape/Transpose layout chains remain explicit.
+It lazily folds initializer Transpose, Concat, and static Reshape/Transpose
+chains, independently of experimental projection packing.
 """
 
 from __future__ import annotations
@@ -47,6 +47,7 @@
 import contextlib
 import dataclasses
 import logging
+import math
 import warnings
 
 import onnx
@@ -62,6 +63,7 @@
 from mobius._passes import (
     FoldConcatInitializersPass,
     FoldTransposedInitializerPass,
+    FoldWeightLayoutPass,
     Fp8KvCachePass,
     RemoveDeadGraphInputsPass,
 )
@@ -795,12 +797,14 @@ def _should_inline(func: ir.Function) -> bool:
 
 
 def fold_initializers_after_weights(model: ir.Model) -> None:
-    """Lazily fold initializer matrix Transpose and Concat nodes.
+    """Lazily fold weight Transpose, Concat, and static layout chains.
 
-    Preserve the shipped root Transpose naming convention and fold projection
-    packing to a fixed point. Generic weight layout chains remain explicit.
+    Fold small integer shape expressions first without exposing weight payloads
+    to the general constant folder. Preserve the shipped root Transpose naming
+    convention, then fold layouts and packing to a fixed point.
     """
     common_passes.LiftConstantsToInitializersPass()(model)
+    hidden: dict[ir.Value, ir.TensorProtocol] = {}
     for graph in model.graphs():
         for value in graph.initializers.values():
             tensor = value.const_value
@@ -814,9 +818,29 @@ def fold_initializers_after_weights(model: ir.Model) -> None:
                     model.meta[LAZY_SOURCE_FILES_META_KEY] = frozenset(
                         (*model.meta.get(LAZY_SOURCE_FILES_META_KEY, ()), *source_files)
                     )
+            if tensor is not None and not (
+                tensor.dtype == ir.DataType.INT64
+                and all(isinstance(dim, int) for dim in tensor.shape)
+                and math.prod(tensor.shape.numpy()) <= 8192
+            ):
+                hidden[value] = tensor
+                value.const_value = None
+    try:
+        # Partial evaluators also simplify Identity without reading its input.
+        # In particular, an output Identity transfers its public output name
+        # to the producer value before lazy layout folding. Size limits and
+        # should_fold do not gate those upstream symbolic simplifications.
+        onnxscript.optimizer._constant_folding.FoldConstantsPass(
+            shape_inference=False, input_size_limit=8192, output_size_limit=8192
+        )(model)
+    finally:
+        for value, tensor in hidden.items():
+            value.const_value = tensor
+    common_passes.LiftConstantsToInitializersPass()(model)
     common_passes.RemoveUnusedNodesPass()(model)
     lazy_passes = (
         FoldTransposedInitializerPass(),
+        FoldWeightLayoutPass(),
         FoldConcatInitializersPass(),
     )
     while True:

src/mobius/_passes/__init__.py

@@ -36,18 +36,23 @@
     graph initializer) into a single pre-concatenated initializer.  Used to
     pre-pack QKV weight matrices after weight loading.
 
+:class:`FoldWeightLayoutPass`
+    Folds static Reshape/Transpose chains into lazy initializers, including
+    higher-rank layout conversions of quantized weights, scales, and zero points.
 """
 
 from __future__ import annotations
 
 __all__ = [
     "FoldTransposedInitializerPass",
     "FoldConcatInitializersPass",
+    "FoldWeightLayoutPass",
     "Fp8KvCachePass",
     "RemoveDeadGraphInputsPass",
 ]
 
 from mobius._passes._fold_concat import FoldConcatInitializersPass
 from mobius._passes._fold_transpose import FoldTransposedInitializerPass
+from mobius._passes._fold_weight_layout import FoldWeightLayoutPass
 from mobius._passes._fp8_kv_cache import Fp8KvCachePass
 from mobius._passes._remove_dead_inputs import RemoveDeadGraphInputsPass

src/mobius/_passes/_fold_transpose.py

@@ -74,8 +74,8 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
 
             # Register legacy ``weight_t`` names only for root-owned weights.
             # A Transpose inside If/Loop can capture a root weight and still use
-            # this pass. Skip child-local weights: registering them at the root
-            # would collide across sibling scopes.
+            # this pass. Child-local weights belong to FoldWeightLayoutPass:
+            # registering them at the root would collide across sibling scopes.
             if (
                 not inp.is_initializer()
                 or inp.is_graph_input()

src/mobius/_passes/_fold_transpose_test.py

@@ -660,58 +660,3 @@ def test_transpose_inside_loop_body_folds_against_root_initializer():
 
     w_t = model.graph.initializers["weight_t"]
     np.testing.assert_array_equal(w_t.const_value.numpy(), weight_data.T)
-
-
-def test_post_weight_pipeline_preserves_generic_layout_chain():
-    from mobius._optimizations import fold_initializers_after_weights
-
-    data = np.arange(24, dtype=np.uint8).reshape(2, 12)
-
-    def unexpected_read():
-        raise AssertionError("generic layout weights must remain untouched")
-
-    weight = ir.Value(
-        name="checkpoint.weight",
-        shape=ir.Shape(data.shape),
-        type=ir.TensorType(ir.DataType.UINT8),
-        const_value=ir.LazyTensor(
-            unexpected_read, dtype=ir.DataType.UINT8, shape=ir.Shape(data.shape)
-        ),
-    )
-    shapes = [
-        ir.Value(
-            name=name,
-            shape=ir.Shape([len(dimensions)]),
-            type=ir.TensorType(ir.DataType.INT64),
-            const_value=ir.tensor(np.array(dimensions, dtype=np.int64)),
-        )
-        for name, dimensions in (("split_shape", (2, 2, 6)), ("final_shape", (2, 12)))
-    ]
-    reshape = ir.Node("", "Reshape", [weight, shapes[0]], num_outputs=1)
-    transpose = ir.Node(
-        "",
-        "Transpose",
-        [reshape.outputs[0]],
-        attributes=[ir.AttrInt64s("perm", [0, 2, 1])],
-        num_outputs=1,
-    )
-    result = ir.Node("", "Reshape", [transpose.outputs[0], shapes[1]], num_outputs=1)
-    result.outputs[0].name = "reordered_weight"
-    nodes = [reshape, transpose, result]
-    model = ir.Model(
-        ir.Graph(
-            [],
-            [result.outputs[0]],
-            nodes=nodes,
-            initializers=[weight, *shapes],
-            opset_imports={"": 24},
-        ),
-        ir_version=10,
-    )
-
-    fold_initializers_after_weights(model)
-
-    assert list(model.graph) == nodes
-    assert model.graph.initializers["checkpoint.weight"] is weight
-    assert reshape.inputs[0] is weight
-    assert model.graph.outputs[0] is result.outputs[0]