[Native] Carry named metadata in native packages

backends/native/README.md

@@ -45,6 +45,15 @@ A `.ptn` is an uncompressed zip with fixed entry names:
 | `program.ptg` | Methods as graphs; FlatBuffer defined by `serialization/native_graph.fbs`. |
 | `program.safetensors` | Constants. Byte-identical immutable tensors are stored once. |
 | `aliases.json` | Deduplicated key to owner key; present only when needed. |
+| `metadata.json` | Producer-defined values by name; present only when given. |
+
+`to_native(ep, metadata={...})` writes `metadata.json`, the counterpart to a
+PTE's constant methods: a JSON object whose values are each a bool, number,
+string, or a flat list of those. Names mean nothing to the format; a runner
+reads the ones it knows with `Package::metadata_int` and `metadata_ints`, and
+`exir.native.read_metadata` reads them back in Python. LLM exports reuse the
+constant-method names in `extension/llm/export/model_metadata.py`, such as
+`get_windows` for the per-layer sliding windows of an off-graph KV cache.
 
 ## Samples (optional)
 

backends/native/runtime/deserialize/Package.cpp

@@ -9,6 +9,7 @@
 #include <algorithm>
 #include <array>
 #include <atomic>
+#include <limits>
 #include <stdexcept>
 #include <string_view>
 
@@ -72,6 +73,42 @@ std::unordered_map<std::string, std::string> parse_aliases(
   return aliases;
 }
 
+Json parse_metadata(std::string_view text) {
+  Json doc;
+  try {
+    doc = parse_json(text);
+  } catch (const Json::exception& error) {
+    throw std::runtime_error(
+        "package: invalid metadata.json: " + std::string(error.what()));
+  }
+  if (!doc.is_object()) {
+    throw std::runtime_error("package: metadata.json is not a JSON object");
+  }
+  for (auto entry = doc.begin(); entry != doc.end(); ++entry) {
+    const Json& value = entry.value();
+    const auto scalar = [](const Json& v) {
+      return v.is_boolean() || v.is_number() || v.is_string();
+    };
+    if (entry.key().empty() ||
+        !(scalar(value) ||
+          (value.is_array() &&
+           std::all_of(value.begin(), value.end(), scalar)))) {
+      throw std::runtime_error(
+          "package: metadata '" + entry.key() +
+          "' must be a JSON scalar or a list of scalars");
+    }
+  }
+  return doc;
+}
+
+// nlohmann parses an integer past int64_t's range as unsigned.
+bool is_metadata_int(const Json& value) {
+  return value.is_number_integer() &&
+      (!value.is_number_unsigned() ||
+       value.get<uint64_t>() <=
+           static_cast<uint64_t>(std::numeric_limits<int64_t>::max()));
+}
+
 } // namespace
 
 bool Package::looks_like_package(ByteSpan bytes) {
@@ -124,6 +161,7 @@ Package& Package::operator=(Package&& other) noexcept {
     tensors_ = std::move(other.tensors_);
     tensor_data_offset_ = other.tensor_data_offset_;
     aliases_ = std::move(other.aliases_);
+    metadata_ = std::move(other.metadata_);
   }
   return *this;
 }
@@ -179,6 +217,50 @@ void Package::load_metadata() {
     const std::vector<uint8_t> aliases = zip_->read(kAliasesEntry);
     aliases_ = parse_aliases(ByteSpan(aliases), *tensors_);
   }
+
+  const std::optional<size_t> metadata_size = zip_->member_size(kMetadataEntry);
+  if (metadata_size) {
+    if (*metadata_size > detail::kMaxJsonBytes) {
+      throw ResourceLimitError("package: metadata exceeds size limit");
+    }
+    const std::vector<uint8_t> metadata = zip_->read(kMetadataEntry);
+    metadata_.assign(metadata.begin(), metadata.end());
+    parse_metadata(metadata_);
+  }
+}
+
+std::optional<int64_t> Package::metadata_int(const std::string& name) const {
+  if (metadata_.empty()) {
+    return std::nullopt;
+  }
+  const Json doc = parse_metadata(metadata_);
+  const auto found = doc.find(name);
+  if (found == doc.end()) {
+    return std::nullopt;
+  }
+  if (!is_metadata_int(*found)) {
+    throw std::runtime_error(
+        "package: metadata '" + name + "' is not an integer");
+  }
+  return found->get<int64_t>();
+}
+
+std::optional<std::vector<int64_t>> Package::metadata_ints(
+    const std::string& name) const {
+  if (metadata_.empty()) {
+    return std::nullopt;
+  }
+  const Json doc = parse_metadata(metadata_);
+  const auto found = doc.find(name);
+  if (found == doc.end()) {
+    return std::nullopt;
+  }
+  if (!found->is_array() ||
+      !std::all_of(found->begin(), found->end(), is_metadata_int)) {
+    throw std::runtime_error(
+        "package: metadata '" + name + "' is not a list of integers");
+  }
+  return found->get<std::vector<int64_t>>();
 }
 
 std::optional<ConstantInfo> Package::constant_info(
@@ -254,8 +336,10 @@ void Package::verify() const {
   if (tensors_) {
     zip_->verify(kSafeTensorsEntry);
   }
-  if (zip_->member_size(kAliasesEntry)) {
-    zip_->verify(kAliasesEntry);
+  for (const char* member : {kAliasesEntry, kMetadataEntry}) {
+    if (zip_->member_size(member)) {
+      zip_->verify(member);
+    }
   }
 }
 

backends/native/runtime/deserialize/Package.h

@@ -27,6 +27,7 @@ namespace ptn {
 constexpr const char* kProgramEntry = "program.ptg";
 constexpr const char* kSafeTensorsEntry = "program.safetensors";
 constexpr const char* kAliasesEntry = "aliases.json";
+constexpr const char* kMetadataEntry = "metadata.json";
 
 // Metadata for one constant resolved out of a package. `sizes` is borrowed from
 // the Package and must not outlive it.
@@ -58,6 +59,8 @@ class Package {
   std::optional<SafeTensorsReader> tensors_;
   size_t tensor_data_offset_ = 0;
   std::unordered_map<std::string, std::string> aliases_;
+  // The metadata.json object, validated at load; empty when absent.
+  std::string metadata_;
 
   Package();
 
@@ -101,6 +104,12 @@ class Package {
   // Metadata for `key`, resolving an alias to its owner. nullopt when absent.
   std::optional<ConstantInfo> constant_info(const std::string& key) const;
 
+  // Producer-defined metadata (the metadata.json member) by name. nullopt when
+  // absent; throws std::runtime_error when the value has another type.
+  std::optional<int64_t> metadata_int(const std::string& name) const;
+  std::optional<std::vector<int64_t>> metadata_ints(
+      const std::string& name) const;
+
   // Load one constant into a new owning buffer. nullopt when absent.
   std::optional<OwnedBytes> acquire_constant(const std::string& key) const;
 

backends/native/test/runtime/deserialize/test_package.cpp

@@ -159,6 +159,46 @@ TEST(PackageTest, IgnoresSampleMembers) {
   EXPECT_NO_THROW(package.verify());
 }
 
+TEST(PackageTest, ReadsMetadataByName) {
+  const TempPackage file(
+      {},
+      {{kMetadataEntry,
+        R"({"get_n_caches":2,"get_windows":[0,4096],"arch":"llama","tied":true})"}});
+  const Package package = Package::load(file.path());
+
+  EXPECT_EQ(package.metadata_int("get_n_caches"), 2);
+  EXPECT_EQ(
+      package.metadata_ints("get_windows"), (std::vector<int64_t>{0, 4096}));
+  EXPECT_EQ(package.metadata_int("missing"), std::nullopt);
+  EXPECT_EQ(package.metadata_ints("missing"), std::nullopt);
+  EXPECT_THROW(package.metadata_int("get_windows"), std::runtime_error);
+  EXPECT_THROW(package.metadata_ints("get_n_caches"), std::runtime_error);
+  EXPECT_THROW(package.metadata_int("arch"), std::runtime_error);
+  EXPECT_THROW(package.metadata_int("tied"), std::runtime_error);
+  EXPECT_NO_THROW(package.verify());
+
+  const TempPackage none;
+  EXPECT_EQ(
+      Package::load(none.path()).metadata_int("get_n_caches"), std::nullopt);
+}
+
+TEST(PackageTest, RejectsMalformedMetadata) {
+  for (const std::string metadata :
+       {R"([1])",
+        R"({"nested":[[1]]})",
+        R"({"object":{"a":1}})",
+        R"({"":1})",
+        R"({"null":null})",
+        "{"}) {
+    const TempPackage file({}, {{kMetadataEntry, metadata}});
+    EXPECT_THROW(Package::load(file.path()), std::runtime_error) << metadata;
+  }
+  const TempPackage huge(
+      {}, {{kMetadataEntry, R"({"n":18446744073709551615})"}});
+  EXPECT_THROW(
+      Package::load(huge.path()).metadata_int("n"), std::runtime_error);
+}
+
 TEST(PackageTest, ReportsMissingConstantsAndWrongDestinations) {
   const TempPackage file;
   const Package package = Package::load(file.path());

exir/_serialize/_ptn.py

@@ -16,6 +16,7 @@
     program.ptg           native Program flatbuffer (file identifier "NPTG")
     program.safetensors   referenced constants, content-deduped (owners only)
     aliases.json          duplicate data_key -> owner key, only when duplicates exist
+    metadata.json         producer-defined name -> JSON scalar or list, when given
 
 Byte-identical *immutable* tensors (same dtype, shape, and content) are stored
 once: the first key owns the safetensors entry and the rest alias to it. An owner
@@ -50,6 +51,7 @@
 
 import hashlib
 import json
+import math
 import os
 import time
 import zipfile
@@ -59,6 +61,7 @@
 PTG_ENTRY = "program.ptg"
 SAFETENSORS_ENTRY = "program.safetensors"
 ALIASES_ENTRY = "aliases.json"
+METADATA_ENTRY = "metadata.json"
 _SAFETENSORS_METADATA_KEY = "__metadata__"
 
 # Tensor data starts on this boundary in the file, and each tensor is naturally
@@ -276,11 +279,39 @@ def _write_safetensors(
             stream.write(_raw(owners[key]))
 
 
+MetadataValue = bool | int | float | str | list[bool | int | float | str]
+
+
+def _check_metadata(metadata: dict[str, MetadataValue], where: str) -> None:
+    """Names are non-empty strings; values are JSON scalars or flat lists of them.
+
+    Floats must be finite: JSON has no infinity or NaN.
+    """
+    if not isinstance(metadata, dict):
+        raise TypeError(f"{where}: metadata must be a dict.")
+
+    def check_scalar(name: str, value: object) -> None:
+        if not isinstance(value, (bool, int, float, str)):
+            raise TypeError(
+                f"{where}: metadata {name!r} holds a {type(value).__name__}; "
+                "values are bool, int, float, str, or a list of those."
+            )
+        if isinstance(value, float) and not math.isfinite(value):
+            raise ValueError(f"{where}: metadata {name!r} holds a non-finite float.")
+
+    for name, value in metadata.items():
+        if not isinstance(name, str) or not name:
+            raise TypeError(f"{where}: metadata names must be non-empty strings.")
+        for item in value if isinstance(value, list) else [value]:
+            check_scalar(name, item)
+
+
 def write_ptn(
     path: str,
     ptg: bytes,
     constants: dict[str, torch.Tensor],
     mutable_keys: frozenset[str] = frozenset(),
+    metadata: dict[str, MetadataValue] | None = None,
 ) -> None:
     """Write a .ptn package, streaming the constants a tensor at a time.
 
@@ -293,15 +324,21 @@ def write_ptn(
             dedup, since equal bytes do not make two mutable buffers the same
             state. Passed in rather than derived here, so this stays independent of
             the graph format. Every mutable key must occur in ``constants``.
+        metadata: Producer-defined values a runtime reads by name, such as an
+            LLM's cache geometry; written as metadata.json when non-empty.
 
     Raises:
         KeyError: If a constant has a dtype safetensors cannot represent.
-        TypeError: If a data key is not a string or a value is not a tensor.
-        ValueError: If mutable-key policy is inconsistent or a data key is reserved.
+        TypeError: If a data key is not a string, a value is not a tensor, or
+            metadata is not a name -> scalar-or-list dict.
+        ValueError: If mutable-key policy is inconsistent, a data key is reserved,
+            or a metadata float is not finite.
     """
     # TODO: save atomically. This opens the destination directly, so a failure
     # part way through truncates an existing valid package. Write to a temporary
     # file beside the destination and os.replace it once complete.
+    metadata = metadata or {}
+    _check_metadata(metadata, "write_ptn")
     owners, aliases = _dedup(constants, mutable_keys)
     parent = os.path.dirname(path)
     if parent:
@@ -317,6 +354,26 @@ def write_ptn(
                 pkg.writestr(
                     ALIASES_ENTRY, json.dumps(aliases, indent=2, sort_keys=True)
                 )
+        if metadata:
+            pkg.writestr(
+                METADATA_ENTRY,
+                json.dumps(metadata, indent=2, sort_keys=True, allow_nan=False),
+            )
+
+
+def read_metadata(path: str) -> dict[str, MetadataValue]:
+    """The package's metadata.json, or an empty dict when it has none.
+
+    Raises:
+        TypeError: If metadata.json is not a name -> scalar-or-list object.
+        ValueError: If metadata.json is not valid JSON.
+    """
+    with zipfile.ZipFile(path) as pkg:
+        if METADATA_ENTRY not in pkg.namelist():
+            return {}
+        metadata = json.loads(pkg.read(METADATA_ENTRY))
+    _check_metadata(metadata, "read_metadata")
+    return metadata
 
 
 def read_ptn(path: str) -> tuple[bytes, dict[str, torch.Tensor]]:
@@ -371,4 +428,12 @@ def read_ptn(path: str) -> tuple[bytes, dict[str, torch.Tensor]]:
     return ptg, constants
 
 
-__all__ = ["ALIASES_ENTRY", "PTG_ENTRY", "SAFETENSORS_ENTRY", "read_ptn", "write_ptn"]
+__all__ = [
+    "ALIASES_ENTRY",
+    "METADATA_ENTRY",
+    "PTG_ENTRY",
+    "SAFETENSORS_ENTRY",
+    "read_metadata",
+    "read_ptn",
+    "write_ptn",
+]

exir/native/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -13,13 +13,15 @@
     SampleSet,
     write_samples,
 )
+from executorch.exir._serialize._ptn import read_metadata
 from executorch.exir.native.native import NativeProgramManager, to_native
 
 __all__ = [
     "Call",
     "NativeProgramManager",
     "Sample",
     "SampleSet",
+    "read_metadata",
     "read_samples",
     "to_native",
     "write_samples",