Vulkan: keep a reshape that changes the last dim before a quantized linear

backends/vulkan/patterns/quantized_linear.py

@@ -27,6 +27,13 @@
 from torch.export.graph_signature import InputKind
 
 
+def _keeps_last_dim(view_node: torch.fx.Node) -> bool:
+    """True if `view_node` leaves its input's last dim unchanged."""
+    src = view_node.args[0]
+    assert isinstance(src, torch.fx.Node)
+    return str(src.meta["val"].shape[-1]) == str(view_node.meta["val"].shape[-1])
+
+
 class QuantizedLinearMatch(PatternMatch):
     def __init__(self, mm_node: torch.fx.Node) -> None:  # noqa: C901
         self.anchor_node = mm_node
@@ -98,8 +105,14 @@ def __init__(self, mm_node: torch.fx.Node) -> None:  # noqa: C901
         anchor_primary_input_node = self.anchor_node.args[input_arg_idx]
         assert isinstance(anchor_primary_input_node, torch.fx.Node)
 
-        # Skip potential view_copy between dq and linear
-        if utils.is_view_copy_node(anchor_primary_input_node):
+        # Skip potential view_copy between dq and linear. Only the view that
+        # flattens the leading dims of a decomposed N-D linear may be skipped: the
+        # fused op takes the N-D input and contracts its last dim, so a reshape that
+        # changes the last dim (e.g. merging heads, [S, H, D] -> [S, H * D]) must
+        # stay in the graph.
+        if utils.is_view_copy_node(anchor_primary_input_node) and _keeps_last_dim(
+            anchor_primary_input_node
+        ):
             self.all_nodes.append(anchor_primary_input_node)
             anchor_primary_input_node = anchor_primary_input_node.args[
                 0

backends/vulkan/test/test_vulkan_passes.py

@@ -712,6 +712,56 @@ def forward(self, x):
         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_fuse_q4gsw_linear_keeps_preceding_reshape(self):
+        """A reshape before a weight-only quantized linear that changes the last
+        dim (merging attention heads) must stay in the graph: the fused
+        et_vk.linear_q4gsw contracts its input's last dim, so skipping the view
+        hands it [S, H, D] instead of [S, H * D].
+        """
+        import executorch.backends.vulkan.custom_ops_lib  # noqa: registers et_vk ops
+        from torchao.quantization.granularity import PerGroup
+        from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
+        from torchao.utils import unwrap_tensor_subclass
+
+        class MergeHeadsLinear(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(64, 32, bias=False)
+
+            def forward(self, x):
+                # [4, 6, 16] -> [6, 4, 16] -> [6, 64] -> [6, 32]
+                return self.linear(x.transpose(0, 1).contiguous().view(-1, 64))
+
+        model = MergeHeadsLinear().eval()
+        quantize_(
+            model,
+            IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32)),
+        )
+        unwrap_tensor_subclass(model)
+
+        sample_inputs = (torch.randn(4, 6, 16),)
+        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_q4gsw.default"), 1)
+        # The head-merging reshape to [6, 64] is still there, before the linear.
+        self.assertEqual(op_node_count(gm, "view_copy.default"), 1)
+
+        fused_out = ep.module()(*sample_inputs)
+        self.assertEqual(tuple(fused_out.shape), (6, 32))
+        self.assertTrue(torch.allclose(fused_out, eager_ref, atol=1e-3, rtol=1e-3))
+
     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