[executorch][vulkan] Add payload-bounded constant sharding

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))

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 = [

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__":

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(

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<int64_t> od = id;
-  od.push_back(static_cast<int64_t>(embed_dim));
+  od.push_back(static_cast<int64_t>(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<uint32_t>(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<uint32_t>(total_blocks);
-  wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
+  params.num_indices = static_cast<uint32_t>(ni);
+  params.total_blocks = static_cast<uint32_t>(total_blocks);
+  wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &params, 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<int>& 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<uint32_t>(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.

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<int64_t> idx(static_cast<size_t>(nb) / sizeof(int64_t));
   f.read(reinterpret_cast<char*>(idx.data()), nb);
   ASSERT_EQ(idx.size(), static_cast<size_t>(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<size_t>(n) * kEmbDim;
   ASSERT_EQ(static_cast<size_t>(out.numel()), numel)
-      << "emb N=" << n << " output numel mismatch";
+      << prefix << " N=" << n << " output numel mismatch";
   std::vector<float> got(
       out.const_data_ptr<float>(), out.const_data_ptr<float>() + 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}) {