Vulkan: keep a reshape that follows a dynamically quantized linear

backends/vulkan/patterns/quantized_linear.py

@@ -185,8 +185,16 @@ def __init__(self, mm_node: torch.fx.Node) -> None:  # noqa: C901
         self.output_view_copy_node = None
         if len(self.output_node.users) == 1:
             cur_node = list(self.output_node.users)[0]
-            # Skip potential view_copy between linear and output quantize
-            if utils.is_view_copy_node(cur_node) and len(cur_node.users) == 1:
+            # Skip potential view_copy between linear and output quantize. Only
+            # the view that restores the leading dims of a decomposed N-D linear
+            # belongs to the pattern: the fused op replaces `output_node` and
+            # produces input.shape[:-1] + [N], so any other reshape that follows
+            # the linear must stay in the graph.
+            if (
+                utils.is_view_copy_node(cur_node)
+                and len(cur_node.users) == 1
+                and self._is_restoring_view(cur_node)
+            ):
                 self.output_view_copy_node = cur_node
                 self.all_nodes.append(self.output_view_copy_node)
                 self.output_node = self.output_view_copy_node
@@ -204,6 +212,16 @@ def __init__(self, mm_node: torch.fx.Node) -> None:  # noqa: C901
 
         self.match_found = True
 
+    def _is_restoring_view(self, view_node: torch.fx.Node) -> bool:
+        """True if `view_node` has the fused op's output shape, input.shape[:-1] + [N]."""
+        in_shape = self.pattern_input_node.meta["val"].shape
+        out_shape = self.anchor_node.meta["val"].shape
+        expected = list(in_shape[:-1]) + [out_shape[-1]]
+        actual = list(view_node.meta["val"].shape)
+        return len(expected) == len(actual) and all(
+            str(a) == str(b) for a, b in zip(expected, actual)
+        )
+
     def is_weight_only_quantized(self) -> bool:
         return self.dequantize_input_node is None
 

backends/vulkan/test/test_vulkan_passes.py

@@ -283,6 +283,62 @@ def forward(self, x):
         )
         self.assertTrue(torch.allclose(fused_out, eager_ref, atol=1e-3, rtol=1e-3))
 
+    def test_fuse_8da4w_linear_keeps_following_reshape(self):
+        """A reshape after a dynamically quantized linear that is not the view
+        restoring a decomposed N-D linear's leading dims must stay in the graph:
+        the fused et_vk.linear_dq8ca_q4gsw replaces the pattern's output node and
+        produces input.shape[:-1] + [N], not the reshaped shape.
+        """
+        import executorch.backends.vulkan.custom_ops_lib  # noqa: registers et_vk ops
+        from torchao.quantization.granularity import PerGroup
+        from torchao.quantization.quant_api import (
+            Int8DynamicActivationIntxWeightConfig,
+            quantize_,
+        )
+        from torchao.utils import unwrap_tensor_subclass
+
+        class LinearReshapeAdd(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(64, 32, bias=False)
+                self.register_buffer("pos", torch.randn(3, 32))
+
+            def forward(self, x):
+                # [6, 64] -> [6, 32] -> [2, 3, 32], then broadcast over dim 0
+                return self.linear(x).view(2, 3, 32) + self.pos
+
+        model = LinearReshapeAdd().eval()
+        quantize_(
+            model,
+            Int8DynamicActivationIntxWeightConfig(
+                weight_dtype=torch.int4, weight_granularity=PerGroup(32)
+            ),
+        )
+        unwrap_tensor_subclass(model)
+
+        sample_inputs = (torch.randn(6, 64),)
+        eager_ref = model(*sample_inputs)
+
+        program = torch.export.export(model, sample_inputs, strict=True)
+        edge_program = to_edge(
+            program,
+            compile_config=EdgeCompileConfig(_check_ir_validity=False),
+        )
+        ep = edge_program._edge_programs["forward"]
+        fuse_pass = FusePatternsPass()
+        fuse_pass._exported_program = ep
+        result = fuse_pass.call(ep.graph_module)
+
+        self.assertTrue(result.modified)
+        gm = ep.graph_module
+        self.assertEqual(op_node_count(gm, "linear_dq8ca_q4gsw.default"), 1)
+        # The reshape to [2, 3, 32] is still there, after the fused linear.
+        self.assertEqual(op_node_count(gm, "view_copy.default"), 1)
+
+        fused_out = ep.module()(*sample_inputs)
+        self.assertEqual(tuple(fused_out.shape), (2, 3, 32))
+        self.assertTrue(torch.allclose(fused_out, eager_ref, atol=1e-2, rtol=1e-2))
+
     def test_torchao_quantized_embedding_rejects_bad_embed_dim(self):
         """A torchao 4-bit quantized embedding whose embed_dim is not a multiple
         of 32 must NOT fuse: the runtime shader asserts embed_dim % 32 == 0