Fix MDP partitioning for scoped ONNX aliases

QEfficient/compile/mdp_generator.py

@@ -561,6 +561,76 @@ def generate_disagg_mdp_partition_config(
     }
 
 
+def _repair_scoped_alias_partitioning(onnx_path: str, mdp_json: Dict[str, Any], compiler_dump: Dict[str, Any]) -> int:
+    """Align scoped aliases with the compiler dump partition."""
+    model = onnx.load(onnx_path, load_external_data=False)
+    names = {n.name for n in model.graph.node if n.name}
+    inlined_node_map, _ = _get_inlined_node_map(model)
+    for expanded_names in inlined_node_map.values():
+        names.update(expanded_names)
+    compiler_partition, compiler_position, compiler_nodes = {}, {}, set()
+    for partition_idx, partition in enumerate(compiler_dump.get("partitions", [])):
+        for node_pos, node_name in enumerate(partition.get("nodeList", [])):
+            if node_name:
+                compiler_nodes.add(node_name)
+                compiler_partition.setdefault(node_name, partition_idx)
+                compiler_position.setdefault(node_name, (partition_idx, node_pos))
+    current_partition = {}
+    for partition_idx, partition in enumerate(mdp_json.get("partitions", [])):
+        for node_name in partition.get("nodeList", []):
+            if node_name:
+                current_partition[node_name] = partition_idx
+    node_by_name = {node.name: node for node in model.graph.node if node.name}
+    consumers_by_value = {}
+    for node in model.graph.node:
+        if not node.name:
+            continue
+        for input_name in node.input:
+            if input_name:
+                consumers_by_value.setdefault(input_name, []).append(node.name)
+    repairs = []
+    for node_name, source_idx in current_partition.items():
+        if not node_name.startswith("/") or node_name.startswith("/language_model/"):
+            continue
+        peer_name = "/language_model" + node_name
+        target_idx = compiler_partition.get(node_name)
+        if peer_name not in names or peer_name in compiler_nodes or target_idx is None or target_idx == source_idx:
+            continue
+        node = node_by_name.get(node_name)
+        if node is not None:
+            consumer_partitions = {
+                current_partition[consumer_name]
+                for output_name in node.output
+                for consumer_name in consumers_by_value.get(output_name, [])
+                if consumer_name in current_partition
+            }
+            if any(consumer_idx < target_idx for consumer_idx in consumer_partitions):
+                logger.warning(
+                    "Skipping MDP alias repair for %r: compiler target Partition%d would follow "
+                    "an existing consumer partition",
+                    node_name,
+                    target_idx,
+                )
+                continue
+        repairs.append((node_name, source_idx, target_idx, compiler_position[node_name]))
+    for node_name, source_idx, target_idx, dump_pos in sorted(repairs, key=lambda item: item[3]):
+        mdp_json["partitions"][source_idx]["nodeList"].remove(node_name)
+        target_nodes = mdp_json["partitions"][target_idx]["nodeList"]
+        insert_at = len(target_nodes)
+        for idx, existing_name in enumerate(target_nodes):
+            if compiler_position.get(existing_name, (10**9, 10**9)) > dump_pos:
+                insert_at = idx
+                break
+        target_nodes.insert(insert_at, node_name)
+        logger.warning(
+            "MDP repaired compiler-scoped alias %r: moved from Partition%d to Partition%d",
+            node_name,
+            source_idx,
+            target_idx,
+        )
+    return len(repairs)
+
+
 def generate_disagg_mdp_intersection_config(
     onnx_path: str,
     compiler_dump_path: str,
@@ -576,7 +646,9 @@ def generate_disagg_mdp_intersection_config(
     QEff MDP, but the QEff MDP also contains nodes the compiler optimises away
     (constant-folded, etc.).  The intersection keeps only nodes that survive
     into the compiler IR, while preserving the **QEff MDP node order and
-    partition assignment** as the authoritative source of truth.
+    partition assignment** as the authoritative source of truth, except for
+    scoped aliases whose partition is aligned with the compiler dump when the
+    compiler materializes only the unscoped alias.
 
     Algorithm
     ---------
@@ -587,12 +659,16 @@ def generate_disagg_mdp_intersection_config(
        names the compiler will actually process.
     3. For each partition in the QEff MDP, retain only nodes whose names appear
        in the compiler dump set, in the original QEff order.
+    4. For scoped aliases with a compiler-only unscoped name, move the alias to
+       the compiler dump partition while preserving the compiler-relative order
+       needed by the compiler's partition validator.
 
     Why this is better than re-assigning compiler-dump nodes by heuristic:
     - Node order follows the ONNX topsort (compiler's SplitPlanMerge requires
       this); the compiler dump order is the compiler's internal IR order which
       may differ after constant-folding and other transforms.
-    - Partition boundaries are the same as the onnx strategy — no divergence.
+    - Partition boundaries are the same as the onnx strategy, except for the
+      narrow scoped-alias repair described above.
     - Every retained node is an exact-match hit in ``doManualPartitioning()``.
     - Output JSON is compact (~1-2 MB) because constant-folded nodes are pruned.
 
@@ -664,6 +740,17 @@ def generate_disagg_mdp_intersection_config(
             }
         )
 
+    mdp_json = {
+        "connections": qeff_mdp["connections"],
+        "partitions": partition_objs,
+    }
+
+    _repair_scoped_alias_partitioning(
+        onnx_path=onnx_path,
+        mdp_json=mdp_json,
+        compiler_dump=dump,
+    )
+
     logger.info(
         f"Intersection: {total_kept} nodes kept, {total_dropped} dropped "
         f"({total_dropped * 100 // max(qeff_total, 1)}% pruned by compiler)"
@@ -676,10 +763,7 @@ def generate_disagg_mdp_intersection_config(
             f"(compiler may have renamed them). Examples: {list(in_dump_not_qeff)[:5]}"
         )
 
-    return {
-        "connections": qeff_mdp["connections"],
-        "partitions": partition_objs,
-    }
+    return mdp_json
 
 
 def generate_disagg_mdp_config(

tests/unit_test/base/test_modeling_qeff_base.py

@@ -29,6 +29,8 @@
 from QEfficient.compile.mdp_generator import (
     _get_layer_num_from_inputs,
     _layer_partition_bounds,
+    generate_disagg_mdp_intersection_config,
+    precheck_and_autofix_mdp_partition_order,
 )
 from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM
 
@@ -789,6 +791,161 @@ def test_mdp_disagg_json_is_created(self, compile_workspace):
         disagg_files = list(compile_dir.glob("mdp_disagg_4d_2p.json"))
         assert len(disagg_files) == 1, "Expected exactly one mdp_disagg_4d_2p.json in compile_dir"
 
+    def test_intersection_repairs_scoped_alias_to_compiler_partition(self, tmp_path: Path):
+        """Intersection moves an unscoped alias to the compiler-assigned partition."""
+        onnx_path = tmp_path / "scoped_alias.onnx"
+        model = helper.make_model(
+            helper.make_graph(
+                [
+                    helper.make_node(
+                        "Identity",
+                        inputs=["input"],
+                        outputs=["hidden0"],
+                        name="/language_model/layers.0/hidden",
+                    ),
+                    helper.make_node(
+                        "Identity",
+                        inputs=["input"],
+                        outputs=["scoped_hidden"],
+                        name="/language_model/Gather_5",
+                    ),
+                    helper.make_node(
+                        "Identity",
+                        inputs=["input"],
+                        outputs=["hidden1"],
+                        name="/language_model/layers.1/hidden",
+                    ),
+                    helper.make_node(
+                        "Identity",
+                        inputs=["input"],
+                        outputs=["output"],
+                        name="/Gather_5",
+                    ),
+                ],
+                name="scoped_alias_graph",
+                inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1])],
+                outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [1])],
+            ),
+            opset_imports=[helper.make_opsetid("", 17)],
+        )
+        onnx.save(model, onnx_path)
+
+        compiler_dump_path = tmp_path / "compiler_dump.json"
+        compiler_dump_path.write_text(
+            json.dumps(
+                {
+                    "partitions": [
+                        {
+                            "nodeList": [
+                                "/language_model/layers.0/hidden",
+                                "/Gather_5",
+                            ]
+                        },
+                        {"nodeList": ["/language_model/layers.1/hidden"]},
+                    ]
+                }
+            )
+        )
+
+        mdp_json = generate_disagg_mdp_intersection_config(
+            onnx_path=str(onnx_path),
+            compiler_dump_path=str(compiler_dump_path),
+            num_devices=2,
+            num_partitions=2,
+            num_layers=2,
+        )
+
+        precheck_removed = precheck_and_autofix_mdp_partition_order(
+            onnx_path=str(onnx_path),
+            mdp_json=mdp_json,
+        )
+
+        assert precheck_removed == 0
+        assert "/Gather_5" in mdp_json["partitions"][0]["nodeList"]
+        assert mdp_json["partitions"][0]["nodeList"] == [
+            "/language_model/layers.0/hidden",
+            "/Gather_5",
+        ]
+        assert mdp_json["partitions"][1]["nodeList"] == ["/language_model/layers.1/hidden"]
+
+    def test_intersection_skips_alias_move_that_follows_consumer(self, tmp_path: Path):
+        """Intersection keeps a producer in place when relocation would reverse its dependency."""
+        onnx_path = tmp_path / "scoped_alias_dependency.onnx"
+        model = helper.make_model(
+            helper.make_graph(
+                [
+                    helper.make_node(
+                        "Identity",
+                        inputs=["input"],
+                        outputs=["hidden0"],
+                        name="/language_model/layers.0/hidden",
+                    ),
+                    helper.make_node(
+                        "Identity",
+                        inputs=["input"],
+                        outputs=["scoped_alias"],
+                        name="/language_model/Gather_9",
+                    ),
+                    helper.make_node(
+                        "Identity",
+                        inputs=["input"],
+                        outputs=["producer_output"],
+                        name="/Gather_9",
+                    ),
+                    helper.make_node(
+                        "Identity",
+                        inputs=["producer_output"],
+                        outputs=["consumer_output"],
+                        name="/language_model/Gather_11",
+                    ),
+                    helper.make_node(
+                        "Identity",
+                        inputs=["input"],
+                        outputs=["hidden1"],
+                        name="/language_model/layers.1/hidden",
+                    ),
+                ],
+                name="scoped_alias_dependency_graph",
+                inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1])],
+                outputs=[helper.make_tensor_value_info("consumer_output", TensorProto.FLOAT, [1])],
+            ),
+            opset_imports=[helper.make_opsetid("", 17)],
+        )
+        onnx.save(model, onnx_path)
+
+        compiler_dump_path = tmp_path / "compiler_dump.json"
+        compiler_dump_path.write_text(
+            json.dumps(
+                {
+                    "partitions": [
+                        {
+                            "nodeList": [
+                                "/language_model/layers.0/hidden",
+                                "/language_model/Gather_11",
+                            ]
+                        },
+                        {"nodeList": ["/Gather_9", "/language_model/layers.1/hidden"]},
+                    ]
+                }
+            )
+        )
+
+        mdp_json = generate_disagg_mdp_intersection_config(
+            onnx_path=str(onnx_path),
+            compiler_dump_path=str(compiler_dump_path),
+            num_devices=2,
+            num_partitions=2,
+            num_layers=2,
+        )
+
+        precheck_removed = precheck_and_autofix_mdp_partition_order(
+            onnx_path=str(onnx_path),
+            mdp_json=mdp_json,
+        )
+
+        assert precheck_removed == 0
+        assert "/Gather_9" in mdp_json["partitions"][0]["nodeList"]
+
     def test_mdp_disagg_json_has_two_partitions(self, compile_workspace):
         """mdp_disagg_4d_2p.json contains exactly two partitions."""
         compile_dir, _, _ = self._run_mdp_compile(compile_workspace)