From 793fad1bf4d69d232ec94b4f5a782df76f4cd48f Mon Sep 17 00:00:00 2001 From: Giulio Corradi Date: Sun, 13 Sep 2026 05:03:21 +0200 Subject: [PATCH] Pass the floating point activation to et_vk.linear_q8ta_q8csw The Vulkan implementation of linear_q8ta_q8csw quantizes the activation tensor itself: QuantizedLinear.cpp names the argument fp_input, allocates a temporary packed int8 tensor for it, and passes it to add_quantize_and_pack_4h4w_node. The AOT fusion was passing the quantized tensor instead. In QuantizedLinearMatch, pattern_input_node is set to the int8 output of the input quantize node, and only the dynamic quantization branch steps back past it to the floating point tensor. Statically quantized inputs therefore reached the op as int8, and the delegate failed at the first inference looking for a shader variant that takes an already quantized input: Could not find ShaderInfo with name clone_buffer_to_image_int8_int32 Add QuantizedLinearMatch.get_fp_input_node() to resolve the floating point activation, and use it when building the op. Where the activation is only available as a quantized tensor - it is a graph input, or it is produced by a preceding op that was already replaced with q8ta_linear - insert a dequantize node so the op still receives what it expects. q8ta_linear is unaffected: it takes packed_int8_input and continues to receive pattern_input_node. Note that no bundled quantizer produces this op today, which is why the mismatch went unnoticed. XNNPACKQuantizer also quantizes the linear's output, which selects q8ta_linear, and VulkanQuantizer has no static activation quantization mode. The added test therefore builds the quantize/dequantize pattern directly. Test Plan: python -m unittest backends.vulkan.test.test_vulkan_passes -v 12 tests pass. The new test, test_linear_q8ta_q8csw_takes_floating_point_input, fails without this change with "linear_q8ta_q8csw was given the quantized activation; it expects the floating point one". Also verified end to end on an AMD Radeon 8060S (RADV GFX1151), running the lowered .pte and comparing against the eager module: case before after single static-quant linear shader error max|d| 2.3e-05 2 chained quantized linears shader error max|d| 1.2e-04 3 chained quantized linears shader error max|d| 1.1e-04 int8 activation as a graph input shader error max|d| 1.8e-05 backends.vulkan.test.test_vulkan_delegate shows an identical set of results before and after the change. --- backends/vulkan/patterns/quantized_linear.py | 45 ++++++++- backends/vulkan/test/test_vulkan_passes.py | 96 ++++++++++++++++++++ 2 files changed, 140 insertions(+), 1 deletion(-) diff --git a/backends/vulkan/patterns/quantized_linear.py b/backends/vulkan/patterns/quantized_linear.py index 86a35298fa4..7050d060404 100644 --- a/backends/vulkan/patterns/quantized_linear.py +++ b/backends/vulkan/patterns/quantized_linear.py @@ -236,6 +236,27 @@ def is_weight_perchannel_quantized(self) -> bool: # scales should have same size as weight's output channels dim return scales_shape[0] == weight_shape[-2] + def get_fp_input_node(self) -> Optional[torch.fx.Node]: + """ + Returns the floating point tensor that the pattern's input quantization + consumes, or None if it is not available in the graph. + + Some custom ops (i.e. et_vk.linear_q8ta_q8csw) quantize the activation + tensor themselves, so they need the floating point input rather than the + quantized one. For dynamically quantized inputs pattern_input_node is + already the floating point tensor; for statically quantized inputs it is + the quantized output of a quantize node, so step back past it. + """ + if self.dequantize_input_node is None or self.quantize_input_node is not None: + return self.pattern_input_node + + if utils.is_quant_node(self.pattern_input_node): + return self.pattern_input_node.args[0] # pyre-ignore[7] + + # The quantized tensor is produced outside of the graph, so there is no + # floating point tensor to use. + return None + def is_input_static_per_tensor_quantized(self) -> bool: if self.dequantize_input_node is None: return False @@ -503,12 +524,34 @@ def make_linear_q8ta_q8csw_custom_op( data=sum_per_output_channel, ) + # This op quantizes the activation tensor itself, so it takes the floating + # point input rather than the quantized one. + fp_input_node = match.get_fp_input_node() + if fp_input_node is None: + # The input is only available as a quantized tensor, i.e. it is produced + # by a preceding quantized op. Dequantize it so that the op has the + # floating point tensor it expects. + with graph_module.graph.inserting_before(match.output_node): + fp_input_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=( + match.pattern_input_node, + match.input_scales_node, + match.input_zeros_node, + -128, + 127, + torch.int8, + ), + ) + fp_input_node.meta["val"] = match.dequantize_input_node.meta["val"] + with graph_module.graph.inserting_before(match.output_node): qlinear_node = graph_module.graph.create_node( "call_function", exir_ops.edge.et_vk.linear_q8ta_q8csw.default, args=( - match.pattern_input_node, + fp_input_node, match.input_scales_node, match.input_zeros_node, match.weight_node, diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index f030b9268a1..5d7a9a1290f 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -3,6 +3,8 @@ import torch +import executorch.backends.vulkan.utils as utils + from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass from executorch.exir import EdgeCompileConfig, EdgeProgramManager, to_edge @@ -637,6 +639,100 @@ def forward(self, x): f"q8ta_linear[{i}].output_zero_point should equal q8ta_linear[{i + 1}].input_zero_point", ) + def test_linear_q8ta_q8csw_takes_floating_point_input(self): + """et_vk.linear_q8ta_q8csw quantizes its activation itself, so it must be + given the floating point input rather than the quantized one. + + QuantizedLinear.cpp names that argument fp_input and passes it to + add_quantize_and_pack_4h4w_node. Before this was fixed the fusion passed + the int8 output of the input quantize node instead, and the delegate + failed at inference looking for a shader variant that takes an already + quantized input, e.g. clone_buffer_to_image_int8_int32. + """ + # The pattern is built directly rather than with a quantizer: + # XNNPACKQuantizer also quantizes the linear's output, which produces + # q8ta_linear instead, and VulkanQuantizer offers no static activation + # quantization mode. + qd = torch.ops.quantized_decomposed + in_features, out_features, batch = 256, 128, 4 + act_scale = 0.05 + + weight = torch.randn(out_features, in_features) + weight_scales = weight.abs().amax(dim=1).clamp(min=1e-8) / 127.0 + weight_zeros = torch.zeros(out_features, dtype=torch.int64) + weight_q = qd.quantize_per_channel.default( + weight, weight_scales, weight_zeros, 0, -127, 127, torch.int8 + ) + + class StaticQuantLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight_q", weight_q) + self.register_buffer("weight_scales", weight_scales) + self.register_buffer("weight_zeros", weight_zeros) + + def forward(self, x): + xq = qd.quantize_per_tensor.default( + x, act_scale, 0, -128, 127, torch.int8 + ) + xdq = qd.dequantize_per_tensor.default( + xq, act_scale, 0, -128, 127, torch.int8 + ) + wdq = qd.dequantize_per_channel.default( + self.weight_q, + self.weight_scales, + self.weight_zeros, + 0, + -127, + 127, + torch.int8, + ) + return torch.nn.functional.linear(xdq, wdq) + + model = StaticQuantLinear().eval() + sample_inputs = (torch.randn(batch, in_features),) + + edge_program = to_edge( + torch.export.export(model, sample_inputs, strict=True), + compile_config=EdgeCompileConfig( + _skip_dim_order=False, + _check_ir_validity=False, + ), + ) + + ep = edge_program._edge_programs["forward"] + fuse_pass = FusePatternsPass() + fuse_pass._exported_program = ep + self.assertTrue(fuse_pass.call(ep.graph_module).modified) + + gm = ep.graph_module + + # With no output quantization the linear becomes linear_q8ta_q8csw + # rather than q8ta_linear. + q8csw_nodes = [ + node + for node in gm.graph.nodes + if get_target_canonical_name(node) == "linear_q8ta_q8csw.default" + ] + self.assertEqual( + len(q8csw_nodes), + 1, + "Expected the output-unquantized linear to fuse to linear_q8ta_q8csw", + ) + + input_node = q8csw_nodes[0].args[0] + self.assertIsInstance(input_node, torch.fx.Node) + self.assertFalse( + utils.is_quant_node(input_node), + "linear_q8ta_q8csw was given the quantized activation; it expects the " + "floating point one", + ) + self.assertNotEqual( + input_node.meta["val"].dtype, + torch.int8, + "linear_q8ta_q8csw input must not be int8", + ) + def test_fuse_q8ta_linear_gemv_non_aligned_oc(self): """Test that quantized linear with non-aligned output channels (not multiple of 4) fuses correctly.""" from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (