From 5610818183795abc12f88851a42c06b31df1334e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 00:05:24 -0700 Subject: [PATCH] Serialize only a view constant's own bytes in XNNPACK and Vulkan A constant tensor can be a view into a larger tensor, for example a buffer registered as `packed[8:16]`, or a slice that constant folding leaves as a view. The XNNPACK and Vulkan serializers copied the bytes of the whole underlying storage starting at its first byte. For a view that starts later in the storage, the delegate then got the wrong weights, and the .pte gave wrong results with no error: self.register_buffer("weight", torch.randn(3, 8, 8)[1]) F.linear(x, self.weight) # XNNPACK result was far from eager Both serializers now clone a weight whose storage is larger than the tensor before reading its bytes. This is the same rule the ExecuTorch emitter already uses for weights that are views. XNNPACK per-channel scales had the same problem, and also skipped the contiguous() step that weights get, so a transposed groupwise scale was written in the wrong order. Scales are small, so they are now always copied into contiguous storage of their own before serialization. The Vulkan path that serializes raw tensors, which only tests use today, now does the same. Tested on macOS arm64, each new test failing before this change and passing after: - an XNNPACK linear whose weight is a view (output mismatch before) - an XNNPACK groupwise int4 linear with a transposed scale (output mismatch before) - a Vulkan graph builder check of the serialized bytes of a view constant - the existing Vulkan serialization test, which now also serializes a view and a transposed view The default backends/xnnpack/test suite passes. Run on its own, test_linear.py passes except test_linear_qd8_{f16,f32}_per_channel_int4, which fail the same way without this change. A per-channel scale replaced by an equal view in an exported program was also checked by hand: the result differed by 4e4 before, and matches after. --- .../serialization/vulkan_graph_builder.py | 4 ++ .../serialization/vulkan_graph_serialize.py | 2 + backends/vulkan/test/test_serialization.py | 12 ++--- .../vulkan/test/test_vulkan_graph_builder.py | 21 ++++++++ backends/xnnpack/operators/node_visitor.py | 9 ++++ backends/xnnpack/test/ops/test_linear.py | 53 +++++++++++++++++++ 6 files changed, 92 insertions(+), 9 deletions(-) diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index e395e49ae78..bd425874169 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -162,6 +162,10 @@ def maybe_add_constant_tensor(self, node: Node) -> int: # Serialize tensor data to bytes tensor = tensor.contiguous() + # The bytes are read from the start of the storage, so a view of a + # larger tensor needs its own copy. + if tensor.nbytes != tensor.untyped_storage().nbytes(): + tensor = tensor.clone() size = tensor.untyped_storage().nbytes() if size > 0: diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index b20073493af..0871cb455a3 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -283,6 +283,8 @@ def serialize_constant_tensors( ): vk_graph.constants.append(VkBytes(current_offset, 0)) elif isinstance(tensor, torch.Tensor): + # The bytes are read from the start of the storage in memory order. + tensor = tensor.clone(memory_format=torch.contiguous_format) array_type = ctypes.c_char * tensor.untyped_storage().nbytes() array = ctypes.cast( tensor.untyped_storage().data_ptr(), diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index 8b95b1df2a1..2eb739a760e 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -6,7 +6,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import ctypes import random import unittest from types import SimpleNamespace @@ -267,6 +266,8 @@ def test_serialize_vulkan_binary(self) -> None: shaders=[], ) const_tensors = self._generate_random_const_tensors(5) + const_tensors.append(torch.randn(3, 4, 4)[1]) + const_tensors.append(torch.randn(3, 4, 4)[1].t()) serialized_binary = serialize_vulkan_graph(vk_graph, const_tensors, []) @@ -305,14 +306,7 @@ def test_serialize_vulkan_binary(self) -> None: constant_data_bytes = constant_data_payload[offset : offset + length] - array_type = ctypes.c_char * tensor.untyped_storage().nbytes() - array = ctypes.cast( - tensor.untyped_storage().data_ptr(), - ctypes.POINTER(array_type), - ).contents - - tensor_bytes = bytes(array) - self.assertEqual(constant_data_bytes, tensor_bytes) + self.assertEqual(constant_data_bytes, tensor.numpy().tobytes()) def test_serialize_deserialize_vkgraph(self) -> None: in_vk_graph = VkGraph( diff --git a/backends/vulkan/test/test_vulkan_graph_builder.py b/backends/vulkan/test/test_vulkan_graph_builder.py index aebb40f8ade..1332da0428d 100644 --- a/backends/vulkan/test/test_vulkan_graph_builder.py +++ b/backends/vulkan/test/test_vulkan_graph_builder.py @@ -108,6 +108,27 @@ def forward(self, first, second): self.assertEqual(len(builder.input_ids), 2) +class TestVkGraphBuilderConstants(unittest.TestCase): + def test_view_constant_serializes_only_its_own_bytes(self) -> None: + class AddViewBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("b", torch.randn(3, 4)[1]) + + def forward(self, x): + return x + self.b + + module = AddViewBuffer() + edge = to_edge(torch.export.export(module, (torch.randn(4),), strict=True)) + program = apply_passes(edge.exported_program(), [SpecPropPass()]) + builder = VkGraphBuilder( + program, DelegateMappingBuilder(generated_identifiers=True) + ) + builder.build_graph() + (data,) = builder.named_data_store.get_named_data_store_output().buffers + self.assertEqual(bytes(data), module.b.numpy().tobytes()) + + class TestVkGraphBuilderScalarLists(unittest.TestCase): """A scalar list serializes to exactly one Value. bool is a subclass of int, so a bool list also satisfies the all-int check, and an empty list diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py index 1963e4494fd..238bb992784 100644 --- a/backends/xnnpack/operators/node_visitor.py +++ b/backends/xnnpack/operators/node_visitor.py @@ -288,6 +288,10 @@ def get_quant_params( if quant_params.per_channel_group: scale = scale.to(torch.bfloat16) + # The bytes are read from the start of the storage in memory order. + # Scales are small, so always copy them into storage of their own. + scale = scale.clone(memory_format=torch.contiguous_format) + num_bytes = scale.untyped_storage().nbytes() scale_array = ctypes.cast( scale.untyped_storage().data_ptr(), @@ -619,6 +623,11 @@ def get_serialized_buffer_index( if quant_params is not None and quant_params.is_qc4w: const_val = self.convert_to_qc4w(const_val) + # The bytes are read from the start of the storage, so a view of a + # larger tensor needs its own copy. + if const_val.nbytes != const_val.untyped_storage().nbytes(): + const_val = const_val.clone() + size = const_val.untyped_storage().nbytes() array_type = ctypes.c_char * size array = ctypes.cast( diff --git a/backends/xnnpack/test/ops/test_linear.py b/backends/xnnpack/test/ops/test_linear.py index 154fa2fa559..e309bc514a7 100644 --- a/backends/xnnpack/test/ops/test_linear.py +++ b/backends/xnnpack/test/ops/test_linear.py @@ -781,6 +781,25 @@ def test_fp32_addmm(self): uses_bias=True, ) + def test_fp32_linear_view_weight(self): + class ViewWeightLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.randn(3, 8, 8)[1]) + + def forward(self, x): + return torch.nn.functional.linear(x, self.weight) + + ( + Tester(ViewWeightLinear(), (torch.randn(2, 8),)) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_fp32_linear_fused_relu(self): for use_bias in (True, False): for num_batch_dims in range(1, 3): @@ -898,6 +917,40 @@ def test_linear_qd8_per_token_groupwise_unsupported_groupsize(self): lin_mod, inputs, group_size=30, use_bias=False, atol=1e-2 ) + @unittest.skipIf(not torchao_installed, "Requires TorchAO") + def test_qd8_linear_transposed_scale(self): + mod = torch.nn.Linear(64, 4, bias=False) + with torch.no_grad(): + # Rows with very different scales make a misordered scale obvious. + mod.weight.mul_(torch.arange(1.0, 5.0).unsqueeze(1)) + quantize_( + mod, + Int8DynamicActivationIntxWeightConfig( + weight_dtype=torch.int4, weight_granularity=PerGroup(32) + ), + ) + unwrap_tensor_subclass(mod) + # XNNPACK stores groupwise scales as bf16, so round them for an exact comparison. + scale = mod.weight.scale.to(torch.bfloat16).float() + mod.weight.scale = scale.t().contiguous().t() + ( + Tester(mod, (torch.randn(2, 64),)) + .export() + .to_edge_transform_and_lower( + ToEdgeTransformAndLower( + [ + XnnpackPartitioner( + config_precisions=ConfigPrecisionType.DYNAMIC_QUANT + ) + ] + ) + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_qd8_per_channel_linear_parallel(self): in_size = 2 input_size = 4