diff --git a/backends/vulkan/patterns/quantized_linear.py b/backends/vulkan/patterns/quantized_linear.py index fcd58b78bf1..0003c995810 100644 --- a/backends/vulkan/patterns/quantized_linear.py +++ b/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 diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index 95e3e9ebdd5..cf810a9e340 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/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