From 6bc1e38ddb93612ef36d6b0ee78ebbf9672ed424 Mon Sep 17 00:00:00 2001 From: Asmita Goswami Date: Sat, 10 Oct 2026 12:12:05 +0530 Subject: [PATCH 1/4] Fix MDP partitioning for scoped ONNX aliases Signed-off-by: Asmita Goswami --- QEfficient/compile/mdp_generator.py | 59 +++++++++++++++++++++++++++-- 1 file changed, 55 insertions(+), 4 deletions(-) diff --git a/QEfficient/compile/mdp_generator.py b/QEfficient/compile/mdp_generator.py index a37a62a9be..b129e6a80a 100644 --- a/QEfficient/compile/mdp_generator.py +++ b/QEfficient/compile/mdp_generator.py @@ -561,6 +561,49 @@ 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} + 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 + 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 + 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, @@ -664,6 +707,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 +730,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( From 30d0c30f625c5e67f16e5dc92e3f8cd33b0a0a61 Mon Sep 17 00:00:00 2001 From: Asmita Goswami Date: Tue, 13 Oct 2026 10:28:44 +0530 Subject: [PATCH 2/4] Added tests and updated doc strings Signed-off-by: Asmita Goswami --- QEfficient/compile/mdp_generator.py | 13 +++- .../unit_test/base/test_modeling_qeff_base.py | 71 +++++++++++++++++++ 2 files changed, 82 insertions(+), 2 deletions(-) diff --git a/QEfficient/compile/mdp_generator.py b/QEfficient/compile/mdp_generator.py index b129e6a80a..d9da942a01 100644 --- a/QEfficient/compile/mdp_generator.py +++ b/QEfficient/compile/mdp_generator.py @@ -565,6 +565,9 @@ def _repair_scoped_alias_partitioning(onnx_path: str, mdp_json: Dict[str, Any], """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", [])): @@ -619,7 +622,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 --------- @@ -630,12 +635,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. diff --git a/tests/unit_test/base/test_modeling_qeff_base.py b/tests/unit_test/base/test_modeling_qeff_base.py index 73f0d54168..dac80ed5ab 100644 --- a/tests/unit_test/base/test_modeling_qeff_base.py +++ b/tests/unit_test/base/test_modeling_qeff_base.py @@ -29,6 +29,7 @@ from QEfficient.compile.mdp_generator import ( _get_layer_num_from_inputs, _layer_partition_bounds, + generate_disagg_mdp_intersection_config, ) from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM @@ -789,6 +790,76 @@ 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=["hidden0"], + outputs=["hidden1"], + name="/language_model/Gather_5", + ), + helper.make_node( + "Identity", + inputs=["hidden1"], + outputs=["hidden2"], + name="/language_model/layers.1/hidden", + ), + helper.make_node( + "Identity", + inputs=["hidden2"], + 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, + ) + + 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_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) From 37b20b036bc4ce9c20d1c400297ff2b0a9f9ba0c Mon Sep 17 00:00:00 2001 From: Asmita Goswami Date: Tue, 13 Oct 2026 11:42:25 +0530 Subject: [PATCH 3/4] Addressed Comments Signed-off-by: Asmita Goswami --- .../unit_test/base/test_modeling_qeff_base.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/tests/unit_test/base/test_modeling_qeff_base.py b/tests/unit_test/base/test_modeling_qeff_base.py index dac80ed5ab..f95e90f973 100644 --- a/tests/unit_test/base/test_modeling_qeff_base.py +++ b/tests/unit_test/base/test_modeling_qeff_base.py @@ -30,6 +30,7 @@ _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 @@ -804,19 +805,19 @@ def test_intersection_repairs_scoped_alias_to_compiler_partition(self, tmp_path: ), helper.make_node( "Identity", - inputs=["hidden0"], - outputs=["hidden1"], + inputs=["input"], + outputs=["scoped_hidden"], name="/language_model/Gather_5", ), helper.make_node( "Identity", - inputs=["hidden1"], - outputs=["hidden2"], + inputs=["input"], + outputs=["hidden1"], name="/language_model/layers.1/hidden", ), helper.make_node( "Identity", - inputs=["hidden2"], + inputs=["input"], outputs=["output"], name="/Gather_5", ), @@ -854,6 +855,13 @@ def test_intersection_repairs_scoped_alias_to_compiler_partition(self, tmp_path: 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", From ff19cf7026833d1341c87105ba1b423864218722 Mon Sep 17 00:00:00 2001 From: Asmita Goswami Date: Tue, 13 Oct 2026 20:41:32 +0530 Subject: [PATCH 4/4] CI Failure fix Signed-off-by: Asmita Goswami --- QEfficient/compile/mdp_generator.py | 24 ++++++ .../unit_test/base/test_modeling_qeff_base.py | 78 +++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/QEfficient/compile/mdp_generator.py b/QEfficient/compile/mdp_generator.py index d9da942a01..1594e14702 100644 --- a/QEfficient/compile/mdp_generator.py +++ b/QEfficient/compile/mdp_generator.py @@ -580,6 +580,14 @@ def _repair_scoped_alias_partitioning(onnx_path: str, mdp_json: Dict[str, Any], 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/"): @@ -588,6 +596,22 @@ def _repair_scoped_alias_partitioning(onnx_path: str, mdp_json: Dict[str, Any], 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) diff --git a/tests/unit_test/base/test_modeling_qeff_base.py b/tests/unit_test/base/test_modeling_qeff_base.py index f95e90f973..f4c645c2aa 100644 --- a/tests/unit_test/base/test_modeling_qeff_base.py +++ b/tests/unit_test/base/test_modeling_qeff_base.py @@ -868,6 +868,84 @@ def test_intersection_repairs_scoped_alias_to_compiler_partition(self, tmp_path: ] 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)