From 1939db304159f617f9d0ca7ce5ca91c703810e0f Mon Sep 17 00:00:00 2001 From: Ethan Ng Date: Tue, 26 May 2026 13:38:56 -0700 Subject: [PATCH] Handle out_dtype in ReplacePT2DequantWithCadenceDequantPass (#19743) Summary: torchao's `convert_pt2e` adds `out_dtype` kwargs to dequant nodes for bf16 models. `cadence::dequantize_per_tensor` doesn't support this kwarg (it hardcodes float32 output), so `ReplacePT2DequantWithCadenceDequantPass` crashes when it forwards kwargs blindly to the cadence op. Strip `out_dtype` from kwargs before creating the cadence dequant node, and insert an `aten.to.dtype` cast after it to preserve the original output dtype semantics. Reviewed By: DrJessop Differential Revision: D105630451 --- backends/cadence/aot/replace_ops.py | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py index 4b60feb2121..50112a4eb66 100644 --- a/backends/cadence/aot/replace_ops.py +++ b/backends/cadence/aot/replace_ops.py @@ -162,14 +162,31 @@ def targets(self) -> list[EdgeOpOverload]: def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: ns = exir_ops.edge if isinstance(node.target, EdgeOpOverload) else torch.ops + out_dtype = node.kwargs.get("out_dtype") + kwargs = {k: v for k, v in node.kwargs.items() if k != "out_dtype"} with node.graph.inserting_before(node): new_node = node.graph.call_function( ns.cadence.dequantize_per_tensor.default, args=node.args, - kwargs=node.kwargs, + kwargs=kwargs, ) - new_node.meta = node.meta - node.replace_all_uses_with(new_node) + new_node.meta = node.meta.copy() + if ( + out_dtype is not None + and out_dtype != torch.float32 + and "val" in new_node.meta + ): + new_node.meta["val"] = new_node.meta["val"].to(torch.float32) + if out_dtype is not None and out_dtype != torch.float32: + with node.graph.inserting_after(new_node): + cast_node = node.graph.call_function( + ns.aten.to.dtype, + args=(new_node, out_dtype), + ) + cast_node.meta = node.meta.copy() + node.replace_all_uses_with(cast_node) + else: + node.replace_all_uses_with(new_node) return True