From ec8bf4662a6f81fb68607039f6f19b6b23a3c8fa Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Mon, 27 Jul 2026 12:41:52 -0700 Subject: [PATCH 1/4] [executorch][vulkan] Add payload-bounded constant sharding Pull Request resolved: https://github.com/pytorch/executorch/pull/21404 **Enable deterministic payload-bounded Vulkan constant sharding** **Problem** Large Vulkan exports can make inline PTE named data exceed downstream artifact limits, but `NamedDataStore` cannot externalize that data into bounded groups. **Solution** - **Before:** Vulkan constants remain inline with no payload-bounded externalization option. - **After:** an opt-in `uint64` raw-data cap externalizes constants after graph construction into deterministic shards. **Implementation** - `NamedDataStore.externalize_pte_data` groups aliases by backing buffer, preserves existing tags, and rejects oversize buffers or collisions atomically. - `parse_compile_options` and `parse_compile_spec` transport `external_constants_max_data_bytes` as a strict positive `uint64`. - `VulkanBackend.preprocess` applies sharding only when the option is present. **Constraints** The default path is unchanged; buffers are not copied or split. The cap bounds unique raw buffer bytes per external tag, not final `FlatTensor` file size including alignment and header overhead. Co-authored-with: Claude Code. ghstack-source-id: 407350478 @exported-using-ghexport Differential Revision: [D113608555](https://our.internmc.facebook.com/intern/diff/D113608555/) --- .../vulkan/partitioner/vulkan_partitioner.py | 15 ++ backends/vulkan/test/TARGETS | 14 ++ .../test/test_vulkan_compile_options.py | 96 ++++++++++- backends/vulkan/vulkan_preprocess.py | 24 +++ exir/_serialize/_named_data_store.py | 60 +++++++ exir/_serialize/test/test_named_data_store.py | 153 ++++++++++++++++++ 6 files changed, 361 insertions(+), 1 deletion(-) diff --git a/backends/vulkan/partitioner/vulkan_partitioner.py b/backends/vulkan/partitioner/vulkan_partitioner.py index fb51a0edfad..298581ebef7 100644 --- a/backends/vulkan/partitioner/vulkan_partitioner.py +++ b/backends/vulkan/partitioner/vulkan_partitioner.py @@ -287,6 +287,21 @@ def parse_compile_options(compile_options: Dict[str, Any]) -> List[CompileSpec]: compile_specs = [] for key, value in compile_options.items(): + if key == "external_constants_max_data_bytes": + # Validate at the user-facing option boundary. Preprocess and the + # data store repeat validation because they can be called directly. + if ( + isinstance(value, bool) + or not isinstance(value, int) + or value <= 0 + or value >= 1 << 64 + ): + raise ValueError( + "external_constants_max_data_bytes must be a positive uint64" + ) + compile_specs.append(CompileSpec(key, value.to_bytes(8, "little"))) + continue + if isinstance(value, (VkStorageType, VkMemoryLayout)): value_bytes = int(value).to_bytes(4, byteorder="little") compile_specs.append(CompileSpec(key, value_bytes)) diff --git a/backends/vulkan/test/TARGETS b/backends/vulkan/test/TARGETS index 7517f7d66f3..196c41d3ec3 100644 --- a/backends/vulkan/test/TARGETS +++ b/backends/vulkan/test/TARGETS @@ -50,6 +50,20 @@ python_unittest( ], ) +python_unittest( + name = "test_vulkan_compile_options", + srcs = [ + "test_vulkan_compile_options.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan:vulkan_preprocess", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/exir/_serialize:lib", + "//executorch/exir:lib", + ], +) + python_unittest( name = "test_serialization", srcs = [ diff --git a/backends/vulkan/test/test_vulkan_compile_options.py b/backends/vulkan/test/test_vulkan_compile_options.py index f44850d2915..f45cfdf12d9 100644 --- a/backends/vulkan/test/test_vulkan_compile_options.py +++ b/backends/vulkan/test/test_vulkan_compile_options.py @@ -6,11 +6,18 @@ import unittest from typing import Any, Dict +from unittest.mock import MagicMock, patch from executorch.backends.vulkan.partitioner.vulkan_partitioner import ( parse_compile_options, ) -from executorch.backends.vulkan.vulkan_preprocess import parse_compile_spec +from executorch.backends.vulkan.vulkan_preprocess import ( + parse_compile_spec, + VulkanBackend, +) +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir._serialize.data_serializer import DataEntry +from executorch.exir.backend.compile_spec_schema import CompileSpec class TestVulkanCompileOptions(unittest.TestCase): @@ -38,10 +45,97 @@ def test_force_fp16_round_trips(self) -> None: round_tripped = self._round_trip({"force_fp16": True}) self.assertTrue(round_tripped.get("force_fp16")) + def test_external_constants_max_data_bytes_round_trips_uint64_bounds( + self, + ) -> None: + for value in (1, (1 << 64) - 1): + with self.subTest(value=value): + self.assertEqual( + self._round_trip({"external_constants_max_data_bytes": value}).get( + "external_constants_max_data_bytes" + ), + value, + ) + + def test_external_constants_max_data_bytes_rejects_invalid_values(self) -> None: + invalid_values: list[Any] = [True, 0, -1, 1 << 64, 1.5, "10"] + for value in invalid_values: + with self.subTest(value=value), self.assertRaisesRegex( + ValueError, "positive uint64" + ): + parse_compile_options({"external_constants_max_data_bytes": value}) + + def test_external_constants_max_data_bytes_rejects_invalid_encoding( + self, + ) -> None: + for payload in (b"", b"\x01", b"\x01" * 7, b"\x01" * 9): + with self.subTest(payload=payload), self.assertRaisesRegex( + ValueError, "encoded as uint64" + ): + parse_compile_spec( + [CompileSpec("external_constants_max_data_bytes", payload)] + ) + with self.assertRaisesRegex(ValueError, "positive uint64"): + parse_compile_spec( + [CompileSpec("external_constants_max_data_bytes", b"\x00" * 8)] + ) + + def _preprocess_named_data(self, options: Dict[str, Any]): + store = NamedDataStore() + graph_builder = MagicMock() + graph_builder.named_data_store = store + + def build_graph(): + store.add_named_data("constant", b"constant", 16) + return MagicMock() + + graph_builder.build_graph.side_effect = build_graph + graph_builder.delegate_mapping_builder.get_delegate_mapping.return_value = {} + program = MagicMock() + + with patch.object( + store, "externalize_pte_data", wraps=store.externalize_pte_data + ) as externalize_pte_data, patch( + "executorch.backends.vulkan.vulkan_preprocess." + "unsafe_remove_auto_functionalized_pass", + side_effect=lambda value: value, + ), patch( + "executorch.backends.vulkan.vulkan_preprocess.apply_passes", + side_effect=lambda value, _passes: value, + ), patch( + "executorch.backends.vulkan.vulkan_preprocess.VkGraphBuilder", + return_value=graph_builder, + ), patch( + "executorch.backends.vulkan.vulkan_preprocess.serialize_vulkan_graph", + return_value=b"vk_graph", + ): + result = VulkanBackend.preprocess(program, parse_compile_options(options)) + return result.data_store_output, externalize_pte_data + + def test_external_constants_default_keeps_constants_inline(self) -> None: + output, externalize_pte_data = self._preprocess_named_data({}) + + self.assertEqual(output.buffers, [b"constant"]) + self.assertEqual(output.pte_data, {"constant": DataEntry(0, 16, None)}) + self.assertEqual(output.external_data, {}) + externalize_pte_data.assert_not_called() + + def test_external_constants_option_externalizes_constants(self) -> None: + output, externalize_pte_data = self._preprocess_named_data( + {"external_constants_max_data_bytes": 16} + ) + + self.assertEqual(output.buffers, [b"constant"]) + self.assertEqual(output.pte_data, {}) + self.assertEqual(len(output.external_data), 1) + self.assertEqual(list(next(iter(output.external_data.values()))), ["constant"]) + externalize_pte_data.assert_called_once_with(16, "vulkan_constants") + def test_unset_options_are_absent(self) -> None: round_tripped = self._round_trip({}) self.assertNotIn("small_texture_limits", round_tripped) self.assertNotIn("skip_memory_planning", round_tripped) + self.assertNotIn("external_constants_max_data_bytes", round_tripped) if __name__ == "__main__": diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index 34c68fe5828..d3954240880 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -90,6 +90,17 @@ def apply_passes(program: ExportedProgram, passes) -> ExportedProgram: return program +def _parse_external_constants_max_data_bytes(value_bytes: bytes) -> int: + # CompileSpec values can bypass parse_compile_options, so validate this + # serialized boundary independently. + if len(value_bytes) != 8: + raise ValueError("external_constants_max_data_bytes must be encoded as uint64") + value = int.from_bytes(value_bytes, byteorder="little") + if value <= 0: + raise ValueError("external_constants_max_data_bytes must be a positive uint64") + return value + + def parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]: options = {} for spec in compile_specs: @@ -119,6 +130,9 @@ def parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]: if spec.key == "skip_memory_planning": options[spec.key] = bool.from_bytes(spec.value, byteorder="little") + if spec.key == "external_constants_max_data_bytes": + options[spec.key] = _parse_external_constants_max_data_bytes(spec.value) + # Unhandled options are ignored return options @@ -246,6 +260,16 @@ def preprocess( # noqa: C901 force_fp16=force_fp16, ) vk_graph = graph_builder.build_graph() + external_constants_max_data_bytes = compile_options.get( + "external_constants_max_data_bytes" + ) + if external_constants_max_data_bytes is not None: + # VkGraphBuilder populates pte_data only from constant tensors; + # already-tagged named data remains in external_data. + graph_builder.named_data_store.externalize_pte_data( + external_constants_max_data_bytes, + "vulkan_constants", + ) return PreprocessResult( processed_bytes=serialize_vulkan_graph( diff --git a/exir/_serialize/_named_data_store.py b/exir/_serialize/_named_data_store.py index c2c67b4bc28..d0b18f9d6c2 100644 --- a/exir/_serialize/_named_data_store.py +++ b/exir/_serialize/_named_data_store.py @@ -210,6 +210,66 @@ def add_named_data( tensor_layout, ) + def externalize_pte_data( + self, + max_data_bytes: int, + tag_prefix: str, + ) -> None: + # Keep this generic API defensive because callers can bypass backend + # option parsing and serialized compile-spec validation. + if ( + isinstance(max_data_bytes, bool) + or not isinstance(max_data_bytes, int) + or max_data_bytes <= 0 + ): + raise ValueError("external data shard cap must be a positive integer") + if not tag_prefix: + raise ValueError("external data tag prefix must be nonempty") + + entries_by_buffer: Dict[int, Dict[str, DataEntry]] = {} + for key, entry in self.pte_data.items(): + entries_by_buffer.setdefault(entry.buffer_index, {})[key] = entry + + shards: List[Dict[str, DataEntry]] = [] + current_shard: Dict[str, DataEntry] = {} + current_bytes = 0 + # Prefer stable key order over size-based packing so equivalent stores + # always produce the same shards. + ordered_groups = sorted( + entries_by_buffer.items(), key=lambda item: tuple(sorted(item[1])) + ) + for buffer_index, entries in ordered_groups: + buffer_size = len(self.buffers[buffer_index]) + if buffer_size > max_data_bytes: + raise ValueError( + f"buffer {buffer_index} has {buffer_size} bytes and exceeds " + f"external data shard cap {max_data_bytes}" + ) + if current_shard and current_bytes + buffer_size > max_data_bytes: + shards.append(current_shard) + current_shard = {} + current_bytes = 0 + current_shard.update(entries) + current_bytes += buffer_size + if current_shard: + shards.append(current_shard) + + external_data = { + tag: dict(entries) for tag, entries in self.external_data.items() + } + for entries in shards: + keys = sorted(entries) + digest = hashlib.sha256("\0".join(keys).encode("utf-8")).hexdigest() + tag = f"{tag_prefix}_{digest}" + canonical_entries = {key: entries[key] for key in keys} + existing = external_data.get(tag) + if existing is not None and existing != canonical_entries: + raise ValueError(f"external data tag collision for {tag}") + external_data[tag] = canonical_entries + + self.external_data = external_data + self.pte_data = {} + def get_named_data_store_output(self) -> NamedDataStoreOutput: # Clean up empty maps inside self.external_data self.external_data = {k: v for k, v in self.external_data.items() if len(v) > 0} diff --git a/exir/_serialize/test/test_named_data_store.py b/exir/_serialize/test/test_named_data_store.py index a9e32afd0f1..929725dda6d 100644 --- a/exir/_serialize/test/test_named_data_store.py +++ b/exir/_serialize/test/test_named_data_store.py @@ -6,7 +6,10 @@ # pyre-strict +import copy +import hashlib import unittest +from typing import Any, cast import torch @@ -16,7 +19,157 @@ from executorch.exir.tensor_layout import TensorLayout +class _SizedBuffer: + def __init__(self, size: int) -> None: + self.size = size + + def __len__(self) -> int: + return self.size + + class TestNamedDataStore(unittest.TestCase): + def test_externalize_pte_data_counts_aliased_buffer_once(self) -> None: + store = NamedDataStore() + layout = TensorLayout(ScalarType.FLOAT, [1], [0]) + store.add_named_data("key_a", b"aaaaaa", 16, None, layout) + store.add_named_data("key_a_alias", b"aaaaaa", 32, None, layout) + store.add_named_data("key_b", b"bbbb", 16, None, layout) + expected_buffers = list(store.buffers) + expected_entries = copy.deepcopy(store.pte_data) + + store.externalize_pte_data(10, "test_constants") + output = store.get_named_data_store_output() + + self.assertEqual(output.buffers, expected_buffers) + self.assertEqual(output.pte_data, {}) + self.assertEqual(len(output.external_data), 1) + self.assertEqual(next(iter(output.external_data.values())), expected_entries) + + def test_externalize_pte_data_rollover_is_insertion_order_independent( + self, + ) -> None: + def externalize( + order: list[str], + ) -> dict[str, list[tuple[str, bytes, int, TensorLayout | None]]]: + data = { + "key_a": (b"a" * 6, 16), + "key_a_alias": (b"a" * 6, 32), + "key_b": (b"b" * 4, 32), + "key_c": (b"c" * 5, 64), + } + store = NamedDataStore() + for key in order: + payload, alignment = data[key] + store.add_named_data(key, payload, alignment) + store.externalize_pte_data(10, "test_constants") + output = store.get_named_data_store_output() + return { + tag: [ + ( + key, + output.buffers[entry.buffer_index], + entry.alignment, + entry.tensor_layout, + ) + for key, entry in entries.items() + ] + for tag, entries in output.external_data.items() + } + + forward = externalize(["key_a", "key_a_alias", "key_b", "key_c"]) + reverse = externalize(["key_c", "key_b", "key_a_alias", "key_a"]) + + self.assertEqual(len(forward), 2) + self.assertEqual(forward, reverse) + + def test_externalize_pte_data_rejects_oversized_buffer_atomically(self) -> None: + store = NamedDataStore() + store.add_named_data("external", b"ext", 8, "existing") + store.add_named_data("key_a", b"aaaa", 16) + store.add_named_data("key_z", b"z" * 11, 32) + before = ( + list(store.buffers), + copy.deepcopy(store.pte_data), + copy.deepcopy(store.external_data), + dict(store.key_to_buffer_idx), + ) + + with self.assertRaisesRegex(ValueError, "exceeds external data shard cap"): + store.externalize_pte_data(10, "test_constants") + + after = ( + list(store.buffers), + store.pte_data, + store.external_data, + store.key_to_buffer_idx, + ) + self.assertEqual(after, before) + + def test_externalize_pte_data_accepts_production_max_buffer_metadata( + self, + ) -> None: + store = NamedDataStore() + store.buffers = [cast(bytes, _SizedBuffer(1_174_405_120))] + store.pte_data = {"largest": DataEntry(0, 16, None)} + store.key_to_buffer_idx = {"largest": 0} + + store.externalize_pte_data(1_500_000_000, "vulkan_constants") + + output = store.get_named_data_store_output() + self.assertEqual(output.pte_data, {}) + self.assertEqual( + [list(entries) for entries in output.external_data.values()], + [["largest"]], + ) + + def test_externalize_pte_data_rejects_invalid_arguments(self) -> None: + store = NamedDataStore() + invalid_caps: list[Any] = [True, 0, -1, 1.5, "10"] + for cap in invalid_caps: + with self.subTest(cap=cap), self.assertRaisesRegex( + ValueError, "positive integer" + ): + store.externalize_pte_data(cap, "test_constants") + + with self.assertRaisesRegex(ValueError, "prefix must be nonempty"): + store.externalize_pte_data(10, "") + + def test_externalize_pte_data_rejects_tag_collision_atomically(self) -> None: + store = NamedDataStore() + key = "inline" + digest = hashlib.sha256(key.encode("utf-8")).hexdigest() + tag = f"test_constants_{digest}" + store.add_named_data("external", b"ext", 8, tag) + store.add_named_data(key, b"inline", 16) + before = ( + list(store.buffers), + copy.deepcopy(store.pte_data), + copy.deepcopy(store.external_data), + ) + + with self.assertRaisesRegex(ValueError, "external data tag collision"): + store.externalize_pte_data(10, "test_constants") + + self.assertEqual( + (list(store.buffers), store.pte_data, store.external_data), before + ) + + def test_externalize_pte_data_preserves_existing_external_data_and_is_idempotent( + self, + ) -> None: + store = NamedDataStore() + store.add_named_data("external", b"ext", 8, "existing") + store.add_named_data("inline", b"inline", 16) + + store.externalize_pte_data(10, "test_constants") + first = copy.deepcopy(store.get_named_data_store_output()) + store.externalize_pte_data(10, "test_constants") + second = store.get_named_data_store_output() + + self.assertEqual(first, second) + self.assertEqual(list(second.external_data["existing"]), ["external"]) + self.assertEqual(len(second.external_data), 2) + def test_add(self) -> None: store = NamedDataStore() store.add_named_data("key1", b"data1", None, None) From 65470d4287e8f83b34ea3c19737785304463fa79 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Mon, 27 Jul 2026 12:41:53 -0700 Subject: [PATCH 2/4] [executorch][webgpu] Preserve embedding packing on resize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21405 **Dynamic Q4 embedding reuse now preserves the complete packing configuration.** A resize used to rebuild the uniform field by field, which dropped `is_linear_weight`. The resize hook now captures the fully initialized `EmbeddingParams` and rewrites only `num_indices` and `total_blocks`. Key changes: - `EmbeddingQ4gsw.cpp` — pass a 32-byte base parameter struct by value through the resize callback. - Resize recomputation updates only dynamic counts, so current and future invariant fields are preserved by construction. - Mirrors Vulkan `EmbeddingQ4gsw.cpp`, which retains `is_linear_weight` in push-constant state across resize. Static and nonlinear embedding behavior is unchanged. Co-authored-with: Claude Code. ghstack-source-id: 407350572 @exported-using-ghexport Differential Revision: [D113627870](https://our.internmc.facebook.com/intern/diff/D113627870/) --- .../ops/embedding_q4gsw/EmbeddingQ4gsw.cpp | 47 ++++--------------- 1 file changed, 9 insertions(+), 38 deletions(-) diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp index 9ed4946eab6..64956121b17 100644 --- a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp +++ b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp @@ -43,11 +43,7 @@ void resize_embedding_q4gsw( WebGPUGraph& g, int indices_id, int out_id, - uint32_t embed_dim, - uint32_t blocks_per_row, - uint32_t gs_u, - uint32_t groups_per_row, - uint32_t bytes_per_row, + EmbeddingParams params, uint32_t wg_size, size_t dispatch_idx, WGPUBuffer params_buf) { @@ -56,23 +52,17 @@ void resize_embedding_q4gsw( if (ni == 0) { throw std::runtime_error("WebGPU embedding_q4gsw: zero indices"); } - const uint64_t total_blocks = ni * blocks_per_row; + const uint64_t total_blocks = ni * params.blocks_per_row; if (total_blocks > UINT32_MAX) { throw std::runtime_error( "WebGPU embedding_q4gsw: total_blocks exceeds uint32"); } std::vector od = id; - od.push_back(static_cast(embed_dim)); + od.push_back(static_cast(params.embed_dim)); g.set_cur_dims(out_id, od); - EmbeddingParams p = {}; - p.embed_dim = embed_dim; - p.blocks_per_row = blocks_per_row; - p.num_indices = static_cast(ni); - p.group_size = gs_u; - p.groups_per_row = groups_per_row; - p.bytes_per_row = bytes_per_row; - p.total_blocks = static_cast(total_blocks); - wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + params.num_indices = static_cast(ni); + params.total_blocks = static_cast(total_blocks); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, ¶ms, sizeof(params)); g.dispatch_at(dispatch_idx).workgroup_count_x = utils::compute_1d_workgroup_count( g.device(), @@ -230,32 +220,13 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { {bundle.pipeline, bundle.bind_group, workgroup_count, "embedding_q4gsw"}); // Dynamic shapes: recompute counts/dispatch; out = indices + [embed_dim]. - const uint32_t gs_u = static_cast(group_size); WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( indices_id, - [indices_id, - out_id, - embed_dim, - blocks_per_row, - gs_u, - groups_per_row, - bytes_per_row, - wg_size, - dispatch_idx, - params_buf](WebGPUGraph& g) { + [indices_id, out_id, params, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { resize_embedding_q4gsw( - g, - indices_id, - out_id, - embed_dim, - blocks_per_row, - gs_u, - groups_per_row, - bytes_per_row, - wg_size, - dispatch_idx, - params_buf); + g, indices_id, out_id, params, wg_size, dispatch_idx, params_buf); }); // Graph owns it so the resize hook can rewrite it; freed in the dtor. From 73171dee6797f71483b64ba944d62c514ef9cb92 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Mon, 27 Jul 2026 12:41:54 -0700 Subject: [PATCH 3/4] [executorch][webgpu] Test linear embedding dynamic reuse MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21406 **Dynamic-shape coverage now locks linear-packed Q4 embedding reuse.** The fixture exports a discriminating linear-packed model and reuses one loaded graph across N=16,8,1,16, covering the resize transition that previously reset nibble order. Key changes: - `test_dynamic_shape_export.py` — export deterministic true-packing artifacts and torch goldens for N=16,8,1. - `test_dynamic_shape.cpp` — add one-graph linear-packed reuse coverage while preserving the nonlinear control. The explicit true-vs-false output guard prevents a non-discriminating fixture. Co-authored-with: Claude Code. ghstack-source-id: 407350695 @exported-using-ghexport Differential Revision: [D113627869](https://our.internmc.facebook.com/intern/diff/D113627869/) --- .../webgpu/test/native/test_dynamic_shape.cpp | 25 ++++-- .../test_dynamic_shape_export.py | 76 +++++++++++++++---- 2 files changed, 80 insertions(+), 21 deletions(-) diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index f41adbf09ec..3752e4ee53f 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -193,31 +193,31 @@ void check_sdpa(int s) { constexpr int kEmbDim = 64; // Run emb_dyn at N tokens on an already-loaded module (so it can be reused // across N), and compare to the golden. -void run_embedding(Module& m, int n) { - const std::string b = g_dir + "/emb_dyn.S" + std::to_string(n) + "."; +void run_embedding(Module& m, int n, const char* prefix = "emb_dyn") { + const std::string b = g_dir + "/" + prefix + ".S" + std::to_string(n) + "."; std::ifstream f(b + "idx.bin", std::ios::binary | std::ios::ate); - ASSERT_TRUE(f.good()) << "missing emb_dyn.S" << n; + ASSERT_TRUE(f.good()) << "missing " << prefix << ".S" << n; const std::streamsize nb = f.tellg(); - ASSERT_GE(nb, 0) << "missing emb_dyn.S" << n; + ASSERT_GE(nb, 0) << "missing " << prefix << ".S" << n; f.seekg(0); std::vector idx(static_cast(nb) / sizeof(int64_t)); f.read(reinterpret_cast(idx.data()), nb); ASSERT_EQ(idx.size(), static_cast(n)) - << "wrong emb_dyn idx size S" << n; + << "wrong " << prefix << " idx size S" << n; auto golden = read_bin(b + "golden.bin"); auto t = make_tensor_ptr({n}, std::move(idx)); // int64 (Long) host input auto r = m.forward({EValue(t)}); ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor()) - << "emb N=" << n + << prefix << " N=" << n << " forward failed (err=" << (r.ok() ? 0 : (int)r.error()) << ")"; const auto& out = r.get()[0].toTensor(); const size_t numel = static_cast(n) * kEmbDim; ASSERT_EQ(static_cast(out.numel()), numel) - << "emb N=" << n << " output numel mismatch"; + << prefix << " N=" << n << " output numel mismatch"; std::vector got( out.const_data_ptr(), out.const_data_ptr() + numel); const float e = max_err(got, golden); - EXPECT_LT(e, 5e-3f) << "emb_dyn N=" << n << " max_err=" << e; + EXPECT_LT(e, 5e-3f) << prefix << " N=" << n << " max_err=" << e; } void check_embedding(int n) { @@ -459,6 +459,15 @@ TEST(DynamicShape, EmbeddingReusedGraph) { } } +// K3: linear-packed reuse must preserve nibble order across resizes. +TEST(DynamicShape, LinearPackedEmbeddingReusedGraph) { + Module m(g_dir + "/emb_dyn_linear.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load emb_dyn_linear.pte"; + for (int n : {16, 8, 1, 16}) { + run_embedding(m, n, "emb_dyn_linear"); + } +} + // L: dynamic RoPE (two outputs) at several seq-len S. TEST(DynamicShape, Rope) { for (int s : {16, 8, 1}) { diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index 969c5e06343..19cb41451e0 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -322,8 +322,53 @@ def cfg(s: int) -> "SdpaConfig": EMB_MAXN = 16 +class _LinearPackedEmbedding(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + packed = torch.arange(EMB_VOCAB * (EMB_DIM // 2), dtype=torch.int64).reshape( + EMB_VOCAB, EMB_DIM // 2 + ) + self.register_buffer("weight", (packed % 256).to(torch.uint8)) + self.register_buffer("scales", torch.ones(EMB_VOCAB, EMB_DIM // EMB_GROUP)) + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + return torch.ops.et_vk.embedding_q4gsw.default( + self.weight, self.scales, EMB_GROUP, indices, True + ) + + +def _write_embedding_goldens( + out_dir: str, + prefix: str, + weight: torch.Tensor, + scales: torch.Tensor, + group_size: int, + is_linear: bool, +) -> None: + for n in [EMB_MAXN, 8, 1]: + idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB + golden = torch.ops.et_vk.embedding_q4gsw.default( + weight, scales, group_size, idx, is_linear + ) + if is_linear: + nonlinear_golden = torch.ops.et_vk.embedding_q4gsw.default( + weight, scales, group_size, idx, False + ) + if torch.equal(golden, nonlinear_golden): + raise RuntimeError( + "emb_dyn_linear fixture does not distinguish nibble packing" + ) + idx.detach().numpy().astype(" None: - from executorch.backends.webgpu.test.ops.embedding_q4gsw.test_embedding_q4gsw import ( + from executorch.backends.webgpu.test.ops.test_embedding_q4gsw import ( _make_quantized_model, _quant_params, Shape, @@ -346,18 +391,23 @@ def _export_dynamic_embedding(out_dir: str) -> None: f.write(et.buffer) print("Exported emb_dyn.pte") weight, scales, group_size = _quant_params(qm) - for n in [EMB_MAXN, 8, 1]: - idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB - g = torch.ops.et_vk.embedding_q4gsw.default( - weight, scales, group_size, idx, False - ) - idx.detach().numpy().astype(" Date: Wed, 29 Jul 2026 18:57:30 +0000 Subject: [PATCH 4/4] Fix broken PT2E links in Arm VGF docs --- .../arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in | 4 ++-- docs/source/backends/arm-vgf/arm-vgf-overview.md | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in b/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in index 53af3e2499d..0fdfc60fe20 100644 --- a/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in +++ b/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in @@ -44,8 +44,8 @@ See [Partitioner API](arm-vgf-partitioner.md) for more information of the Partit ## Quantization -The VGF quantizer supports [Post Training Quantization (PT2E)](https://docs.pytorch.org/ao/main/tutorials_source/pt2e_quant_ptq.html) -and [Quantization-Aware Training (QAT)](https://docs.pytorch.org/ao/main/tutorials_source/pt2e_quant_qat.html). +The VGF quantizer supports [Post Training Quantization (PT2E)](https://docs.pytorch.org/ao/stable/pt2e_quantization/pt2e_quant_ptq.html) +and [Quantization-Aware Training (QAT)](https://docs.pytorch.org/ao/stable/pt2e_quantization/pt2e_quant_qat.html). Partial quantization is supported, allowing users to quantize only specific parts of the model while leaving others in floating-point. diff --git a/docs/source/backends/arm-vgf/arm-vgf-overview.md b/docs/source/backends/arm-vgf/arm-vgf-overview.md index 110c50af55c..a1dab94c082 100644 --- a/docs/source/backends/arm-vgf/arm-vgf-overview.md +++ b/docs/source/backends/arm-vgf/arm-vgf-overview.md @@ -112,8 +112,8 @@ See [Partitioner API](arm-vgf-partitioner.md) for more information of the Partit ## Quantization -The VGF quantizer supports [Post Training Quantization (PT2E)](https://docs.pytorch.org/ao/main/tutorials_source/pt2e_quant_ptq.html) -and [Quantization-Aware Training (QAT)](https://docs.pytorch.org/ao/main/tutorials_source/pt2e_quant_qat.html). +The VGF quantizer supports [Post Training Quantization (PT2E)](https://docs.pytorch.org/ao/stable/pt2e_quantization/pt2e_quant_ptq.html) +and [Quantization-Aware Training (QAT)](https://docs.pytorch.org/ao/stable/pt2e_quantization/pt2e_quant_qat.html). Partial quantization is supported, allowing users to quantize only specific parts of the model while leaving others in floating-point.