From 62644073964e2c27a5ddf43843f59ca8a8427d35 Mon Sep 17 00:00:00 2001 From: ssjia Date: Mon, 23 Mar 2026 07:22:55 -0700 Subject: [PATCH] [ET-VK] Support alias_copy in partitioner by removing it as a redundant op Models using QK normalization (e.g., Qwen3) generate aten.alias_copy.default nodes from RMSNorm. Without Vulkan support for this op, these nodes create partition boundaries that cause SpecViolationError during export. Since alias_copy is a no-op identity operation, the fix registers it in op_registry so the partitioner delegates it, and adds it to RemoveRedundantOpsTransform so preprocess removes it before serialization. The C++ runtime never needs to handle alias_copy. Differential Revision: [D97757921](https://our.internmc.facebook.com/intern/diff/D97757921/) [ghstack-poisoned] --- backends/vulkan/_passes/remove_redundant_ops.py | 2 ++ backends/vulkan/op_registry.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/backends/vulkan/_passes/remove_redundant_ops.py b/backends/vulkan/_passes/remove_redundant_ops.py index b95733021fc..545a25f69e9 100644 --- a/backends/vulkan/_passes/remove_redundant_ops.py +++ b/backends/vulkan/_passes/remove_redundant_ops.py @@ -28,6 +28,8 @@ class RemoveRedundantOpsTransform(ExportPass): exir_ops.edge.aten.clone.default, torch.ops.aten.alias.default, exir_ops.edge.aten.alias.default, + torch.ops.aten.alias_copy.default, + exir_ops.edge.aten.alias_copy.default, exir_ops.edge.aten.lift_fresh_copy.default, exir_ops.edge.dim_order_ops._to_dim_order_copy.default, exir_ops.edge.dim_order_ops._clone_dim_order.default, diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index f9881f637d3..ddb843e2335 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1125,6 +1125,20 @@ def register_clone_dim_order(): ) +# alias_copy is a no-op identity operation (same as clone/alias). It is removed +# by RemoveRedundantOpsTransform during preprocess, so the C++ runtime never sees +# it. Registering it here ensures the partitioner delegates it to Vulkan rather +# than creating partition boundaries that break the graph. +@update_features(exir_ops.edge.aten.alias_copy.default) +def register_alias_copy(): + return OpFeatures( + inputs_storage=utils.ANY_STORAGE, + inputs_dtypes=utils.FP_INT_BOOL_T, + supports_resize=True, + supports_highdim=True, + ) + + # ============================================================================= # Gather.cpp # =============================================================================