From a503dfe89aa86a0af587699dbcb576af9da18d33 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Fri, 23 Oct 2026 16:27:11 +0200 Subject: [PATCH 1/2] Vulkan: keep a reshape that changes the input's last dim before a quantized linear QuantizedLinearMatch skipped any view_copy feeding the linear, assuming it is the leading-dim flatten from aten.linear's N-D decomposition. A reshape that merges attention heads ([S, H, D] -> [S, H * D]) was skipped too, so the fused et_vk.linear_q4gsw received the [S, H, D] tensor and lowering failed in the fake kernel. Skip the view only when it keeps the last dim. --- backends/vulkan/patterns/quantized_linear.py | 17 ++++++- backends/vulkan/test/test_vulkan_passes.py | 50 ++++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/patterns/quantized_linear.py b/backends/vulkan/patterns/quantized_linear.py index fcd58b78bf1..b89ddd5853e 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 From 1bf0cc8eab552e5d3434e293e7514f3fb4b425e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Mon, 26 Oct 2026 22:52:46 +0100 Subject: [PATCH 2/2] Fix lint --- backends/vulkan/patterns/quantized_linear.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/patterns/quantized_linear.py b/backends/vulkan/patterns/quantized_linear.py index b89ddd5853e..0003c995810 100644 --- a/backends/vulkan/patterns/quantized_linear.py +++ b/backends/vulkan/patterns/quantized_linear.py @@ -110,9 +110,9 @@ def __init__(self, mm_node: torch.fx.Node) -> None: # noqa: C901 # 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( + if utils.is_view_copy_node(anchor_primary_input_node) and _keeps_last_dim( 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