Split the cache's model geometry from its runtime config (pytorch#22763)

CacheConfig split in two: CacheGeometry holds the immutable per-layer
facts a model was built with, CacheConfig keeps only what a runtime
chooses — capacity, dtype, initial capacity, max_write. The export
publishes the geometry through write_cache_geometry; read_cache_geometry
reads it back at load, so run_llm_hf no longer takes --kv-cache or
--kv-cache-windows. The broadcast form, a layers list of size 1 standing
for every layer, is gone.

Read extension/llm/cache/cache.h for the split and
extension/llm/runner/model_metadata.h for the reader; the MLX files are
call sites following the new constructor signature, except
run_llm_hf.cpp, which loses its geometry flags.

Tests
extension/llm/runner/test/test_model_metadata.cpp 
extension/llm/runner/test/export_model_metadata.py 
llama-3.2-1B, gemma-3-1B and gemma-4-E2B all export and generate.

Author: kiymetakdemir

backends/mlx/examples/llm/export_llm_hf.py

@@ -45,6 +45,7 @@
 from executorch.extension.llm.export.model_metadata import (
     model_vocab_size,
     write_activation_dtype,
+    write_cache_geometry,
     write_logits_to_keep_mode,
     write_max_context_len,
     write_max_seq_len,
@@ -522,12 +523,7 @@ def _export_with_offgraph_cache(
         max_seq_len=prefill_chunk_size,
     )
     kv_metadata.update(
-        {
-            "get_n_caches": len(layer_types),
-            "get_kv_heads": torch.tensor(cache_kv_heads, dtype=torch.int32),
-            "get_head_dims": torch.tensor(cache_head_dims, dtype=torch.int32),
-            "get_windows": torch.tensor(cache_windows, dtype=torch.int32),
-        }
+        write_cache_geometry(cache_kv_heads, cache_head_dims, cache_windows)
     )
     logger.info(
         f"KV cache layout: {len(layer_types)} caches, "

backends/mlx/examples/llm/run_llm_hf.cpp

@@ -112,6 +112,7 @@ using ::executorch::extension::Module;
 using ::executorch::extension::llm::check_vocab_size;
 using ::executorch::extension::llm::LogitsToKeepMode;
 using ::executorch::extension::llm::read_activation_dtype;
+using ::executorch::extension::llm::read_cache_geometry;
 using ::executorch::extension::llm::read_logits_to_keep_mode;
 using ::executorch::extension::llm::read_max_seq_len;
 using ::executorch::extension::llm::read_vocab_size;
@@ -154,17 +155,6 @@ bool parse_int_list(
   return true;
 }
 
-// Constant methods the export publishes (get_n_caches and friends). They carry
-// no delegate, so reading them only needs the program loaded -- which is what
-// lets the cache be built before forward's backend init consumes its key.
-std::optional<int64_t> const_int(Module& module, const char* name) {
-  const auto r = module.execute(name);
-  if (!r.ok() || r->empty() || !r->at(0).isInt()) {
-    return std::nullopt;
-  }
-  return r->at(0).toInt();
-}
-
 // The sampler (sample_from_logits) fatally aborts on any other dtype, so an
 // unsupported logits type must be rejected at startup rather than at inference.
 bool is_supported_logits_type(::executorch::aten::ScalarType type) {
@@ -212,54 +202,6 @@ bool validate_forward_abi(
   return true;
 }
 
-std::optional<std::vector<int>> const_ints(Module& module, const char* name) {
-  const auto r = module.execute(name);
-  if (!r.ok() || r->empty() || !r->at(0).isTensor()) {
-    return std::nullopt;
-  }
-  const auto t = r->at(0).toTensor();
-  if (t.scalar_type() != ::executorch::aten::ScalarType::Int) {
-    return std::nullopt;
-  }
-  const int32_t* p = t.const_data_ptr<int32_t>();
-  return std::vector<int>(p, p + t.numel());
-}
-
-// Fill in the cache geometry the export published: get_n_caches, then one
-// entry per cache in get_kv_heads / get_head_dims / get_windows (0 = flat).
-// Capacity and dtype stay with the flags. False means this is not an off-graph
-// model.
-bool read_kv_layout(
-    Module& module,
-    int prefill_chunk,
-    cache::CacheConfig& cfg) {
-  const auto n_caches = const_int(module, "get_n_caches");
-  const auto kv_heads = const_ints(module, "get_kv_heads");
-  const auto head_dims = const_ints(module, "get_head_dims");
-  const auto windows = const_ints(module, "get_windows");
-  if (!n_caches || !kv_heads || !head_dims || !windows) {
-    return false;
-  }
-  cfg.max_write = prefill_chunk;
-  const size_t n = static_cast<size_t>(*n_caches);
-  if (kv_heads->size() != n || head_dims->size() != n || windows->size() != n) {
-    return false;
-  }
-  cfg.n_layers = static_cast<int>(n);
-  cfg.layers.clear();
-  cfg.layers.reserve(n);
-  for (size_t l = 0; l < n; ++l) {
-    cache::LayerConfig lc{};
-    lc.n_kv_heads = (*kv_heads)[l];
-    lc.head_dim = (*head_dims)[l];
-    lc.policy = (*windows)[l] > 0
-        ? cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, (*windows)[l]}
-        : cache::LayerPolicy{cache::LayerPolicy::Kind::Flat, 0};
-    cfg.layers.push_back(lc);
-  }
-  return true;
-}
-
 // Replace the model's own attention pattern with `spec`, a comma-separated list
 // of windows repeating over the caches (0 = flat). One entry makes every layer
 // sliding. Only the policy changes; each cache keeps the geometry the .pte
@@ -268,14 +210,17 @@ bool read_kv_layout(
 // The export sizes the chunk to the model's own window; narrowing the window
 // here would leave the ring (window + max_write - 1) sized by the chunk
 // instead, so the chunk follows the window down.
-bool apply_window_override(const std::string& spec, cache::CacheConfig& cfg) {
+bool apply_window_override(
+    const std::string& spec,
+    cache::CacheGeometry& geometry,
+    cache::CacheConfig& cfg) {
   std::vector<int> pattern;
   if (!parse_int_list(spec, ',', pattern) || pattern.empty()) {
     return false;
   }
-  for (size_t l = 0; l < cfg.layers.size(); ++l) {
+  for (size_t l = 0; l < geometry.layers.size(); ++l) {
     const int w = pattern[l % pattern.size()];
-    cfg.layers[l].policy = w > 0
+    geometry.layers[l].policy = w > 0
         ? cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, w}
         : cache::LayerPolicy{cache::LayerPolicy::Kind::Flat, 0};
   }
@@ -288,7 +233,7 @@ bool apply_window_override(const std::string& spec, cache::CacheConfig& cfg) {
   if (cfg.max_write && narrowest > 0 && narrowest < *cfg.max_write) {
     cfg.max_write = narrowest;
   }
-  return cache::valid(cfg);
+  return cache::valid(geometry, cfg);
 }
 
 // Human-readable name for a kv_dtype (an ET ScalarType int). Only the
@@ -310,14 +255,14 @@ std::string dtype_name(int st) {
 // Announce the cache shape: the same .pte runs under whatever config this
 // invocation asks for -- capacity, storage dtype, flat/ring layers -- with no
 // re-export. The footprint lines printed later then show it growing at runtime.
-void print_cache_summary(const cache::CacheConfig& cfg) {
+void print_cache_summary(
+    const cache::CacheGeometry& geometry,
+    const cache::CacheConfig& cfg) {
   // Ring layers grouped by window: --kv-windows can give each layer its own,
   // and the pools are sized per layer, so a single number would misreport them.
   std::map<int, int> ring;
   int flat = 0;
-  for (int l = 0; l < cfg.n_layers; ++l) {
-    const cache::LayerConfig& lc =
-        cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l];
+  for (const cache::LayerGeometry& lc : geometry.layers) {
     if (lc.policy.kind == cache::LayerPolicy::Kind::Ring) {
       ++ring[lc.policy.window];
     } else {
@@ -330,7 +275,8 @@ void print_cache_summary(const cache::CacheConfig& cfg) {
   if (cfg.max_write) {
     std::cout << " max_write=" << *cfg.max_write;
   }
-  std::cout << "\n        " << cfg.n_layers << " layers: " << flat << " flat";
+  std::cout << "\n        " << geometry.layers.size() << " layers: " << flat
+            << " flat";
   for (const auto& [window, n] : ring) {
     std::cout << " + " << n << " ring(window " << window << ")";
   }
@@ -743,16 +689,19 @@ int main(int argc, char** argv) {
                 << " (bf16|fp16|fp32)" << std::endl;
       return 1;
     }
-    if (!read_kv_layout(module, prefill_chunk, cfg)) {
-      std::cerr << "No KV cache layout in " << pte
+    auto geometry = read_cache_geometry(module);
+    if (!geometry.ok()) {
+      std::cerr << "No valid KV cache geometry in " << pte
                 << "; re-export with --use-offgraph-cache" << std::endl;
       return 1;
     }
-    if (!kv_windows.empty() && !apply_window_override(kv_windows, cfg)) {
+    cfg.max_write = prefill_chunk;
+    if (!kv_windows.empty() &&
+        !apply_window_override(kv_windows, *geometry, cfg)) {
       std::cerr << "Invalid --kv-windows: " << kv_windows << std::endl;
       return 1;
     }
-    if (!cache::valid(cfg)) {
+    if (!cache::valid(*geometry, cfg)) {
       std::cerr << "Invalid cache config" << std::endl;
       return 1;
     }
@@ -762,7 +711,7 @@ int main(int argc, char** argv) {
 
     const char* const cache_kind = cache::kind::kSingle;
     auto built = cache::CacheFactory::global().build(
-        ::executorch::backends::mlx::kMLXBackendId, cache_kind, cfg);
+        ::executorch::backends::mlx::kMLXBackendId, cache_kind, *geometry, cfg);
     if (!built.ok()) {
       std::cerr << "Failed to build cache: " << static_cast<int>(built.error())
                 << std::endl;
@@ -775,7 +724,7 @@ int main(int argc, char** argv) {
     // load_method() inside it.
     const cache::InstallGuard guard{kv};
 
-    print_cache_summary(cfg);
+    print_cache_summary(*geometry, cfg);
     if (guard.set_option(mlx_opts) != Error::Ok ||
         options_map.set_options(
             ::executorch::backends::mlx::kMLXBackendId, mlx_opts.view()) !=

backends/mlx/llm/cache.py

@@ -53,10 +53,23 @@ def resolve_hf_cache_layout(config):
     else:
         layer_types = list(layer_types)
 
-    if hasattr(text_config, "num_kv_shared_layers"):
-        layer_types = layer_types[: -text_config.num_kv_shared_layers]
-
-    if hasattr(text_config, "global_head_dim"):
+    shared = getattr(text_config, "num_kv_shared_layers", 0)
+    if shared:
+        layer_types = layer_types[:-shared]
+
+    per_layer = getattr(text_config, "per_layer_config", None)
+    if per_layer is not None:
+        # Entries index from 0, as layer_types does after the truncation.
+        head_dims = [per_layer[i].head_dim for i in range(len(layer_types))]
+        num_heads = [
+            getattr(
+                per_layer[i],
+                "num_key_value_heads",
+                per_layer[i].num_attention_heads,
+            )
+            for i in range(len(layer_types))
+        ]
+    elif hasattr(text_config, "global_head_dim"):
         head_dims = [
             (
                 text_config.global_head_dim

backends/mlx/runtime/MLXBackend.cpp

@@ -608,9 +608,11 @@ static auto success_with_compiler = register_backend(backend);
 // Cache it installs under a cache_key. Adding a kind is a new builder here.
 const int cache_builders_registered = [] {
   const Error single = cache::CacheFactory::global().register_builder(
-      kMLXBackendId, cache::kind::kSingle, [](const cache::CacheConfig& cfg) {
+      kMLXBackendId,
+      cache::kind::kSingle,
+      [](const cache::CacheGeometry& geometry, const cache::CacheConfig& cfg) {
         return std::shared_ptr<cache::Cache>(
-            std::make_shared<MLXSequenceCache>(cfg));
+            std::make_shared<MLXSequenceCache>(geometry, cfg));
       });
   ET_CHECK_MSG(
       single == Error::Ok,
@@ -620,9 +622,9 @@ const int cache_builders_registered = [] {
   const Error batched_cell = cache::CacheFactory::global().register_builder(
       kMLXBackendId,
       cache::kind::kBatchedCell,
-      [](const cache::CacheConfig& cfg) {
+      [](const cache::CacheGeometry& geometry, const cache::CacheConfig& cfg) {
         return std::shared_ptr<cache::Cache>(
-            std::make_shared<MLXCellCache>(cfg));
+            std::make_shared<MLXCellCache>(geometry, cfg));
       });
   ET_CHECK_MSG(
       batched_cell == Error::Ok,
@@ -632,9 +634,9 @@ const int cache_builders_registered = [] {
   const Error batched_seq = cache::CacheFactory::global().register_builder(
       kMLXBackendId,
       cache::kind::kBatchedSequence,
-      [](const cache::CacheConfig& cfg) {
+      [](const cache::CacheGeometry& geometry, const cache::CacheConfig& cfg) {
         return std::shared_ptr<cache::Cache>(
-            std::make_shared<MLXBatchedSequenceCache>(cfg));
+            std::make_shared<MLXBatchedSequenceCache>(geometry, cfg));
       });
   ET_CHECK_MSG(
       batched_seq == Error::Ok,
@@ -643,9 +645,11 @@ const int cache_builders_registered = [] {
       cache::kind::kBatchedSequence);
   // The layout kBatched points at, also registered under its own name above.
   const Error batched = cache::CacheFactory::global().register_builder(
-      kMLXBackendId, cache::kind::kBatched, [](const cache::CacheConfig& cfg) {
+      kMLXBackendId,
+      cache::kind::kBatched,
+      [](const cache::CacheGeometry& geometry, const cache::CacheConfig& cfg) {
         return std::shared_ptr<cache::Cache>(
-            std::make_shared<MLXBatchedSequenceCache>(cfg));
+            std::make_shared<MLXBatchedSequenceCache>(geometry, cfg));
       });
   ET_CHECK_MSG(
       batched == Error::Ok,

backends/mlx/runtime/MLXBatchedSequenceCache.h

@@ -33,8 +33,10 @@ class MLXBatchedSequenceCache : public cache::Cache,
                                 public cache::BatchControl,
                                 public MLXCache {
  public:
-  explicit MLXBatchedSequenceCache(const cache::CacheConfig& cfg)
-      : cfg_(checked(cfg)) {}
+  MLXBatchedSequenceCache(
+      const cache::CacheGeometry& geometry,
+      const cache::CacheConfig& cfg)
+      : geometry_(checked(geometry, cfg)), cfg_(cfg) {}
 
   int capacity() const override {
     return cfg_.capacity;
@@ -74,7 +76,7 @@ class MLXBatchedSequenceCache : public cache::Cache,
 
   std::optional<int32_t> seq_new() override {
     const int32_t id = free_id();
-    rows_.try_emplace(id, cfg_);
+    rows_.try_emplace(id, geometry_, cfg_);
     return id;
   }
 
@@ -193,7 +195,7 @@ class MLXBatchedSequenceCache : public cache::Cache,
       const Tensor& q,
       const Tensor& k,
       const Tensor& v) const {
-    if (layer < 0 || layer >= cfg_.n_layers) {
+    if (layer < 0 || layer >= static_cast<int>(geometry_.layers.size())) {
       throw std::out_of_range("attend: layer out of range");
     }
     if (q.ndim() != 4 || k.ndim() != 4 || v.ndim() != 4) {
@@ -278,13 +280,17 @@ class MLXBatchedSequenceCache : public cache::Cache,
     placed_ = false;
   }
 
-  static const cache::CacheConfig& checked(const cache::CacheConfig& cfg) {
-    if (!cache::valid(cfg)) {
-      throw std::runtime_error("MLXBatchedSequenceCache: invalid CacheConfig");
+  static const cache::CacheGeometry& checked(
+      const cache::CacheGeometry& geometry,
+      const cache::CacheConfig& cfg) {
+    if (!cache::valid(geometry, cfg)) {
+      throw std::runtime_error(
+          "MLXBatchedSequenceCache: invalid geometry or config");
     }
-    return cfg;
+    return geometry;
   }
 
+  cache::CacheGeometry geometry_;
   cache::CacheConfig cfg_;
   std::map<int32_t, MLXSequenceCache> rows_;
   // Set at init by the delegate that resolved this cache.

backends/mlx/runtime/MLXCellCache.h

@@ -34,16 +34,15 @@ namespace cache = ::executorch::extension::llm::cache;
 // prefill does, and neither is what MLX's fused "causal" describes.
 class MLXCellCache : public cache::CellCache, public MLXCache {
  public:
-  explicit MLXCellCache(const cache::CacheConfig& cfg)
-      : cache::CellCache(checked(cfg)) {
+  MLXCellCache(
+      const cache::CacheGeometry& geometry,
+      const cache::CacheConfig& cfg)
+      : cache::CellCache(checked(geometry, cfg), cfg) {
     const ::mlx::core::Dtype dt =
         resolve_dtype(static_cast<int8_t>(cfg.kv_dtype));
-    kpool_.reserve(static_cast<size_t>(cfg.n_layers));
-    vpool_.reserve(static_cast<size_t>(cfg.n_layers));
-    for (int l = 0; l < cfg.n_layers; ++l) {
-      // layers size 1 = one config broadcast to every layer, else per-layer.
-      const cache::LayerConfig& lc =
-          cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l];
+    kpool_.reserve(geometry.layers.size());
+    vpool_.reserve(geometry.layers.size());
+    for (const cache::LayerGeometry& lc : geometry.layers) {
       // A window bounds what a query attends, not where its token lives, so
       // every layer spans the whole cell table whatever its policy.
       kpool_.emplace_back(
@@ -116,12 +115,14 @@ class MLXCellCache : public cache::CellCache, public MLXCache {
 
   // Enforce the neutral contract as an exception, the failure mode this layer
   // already uses. Runs as the base initializer's argument because CellCache's
-  // own ctor indexes `layers` before this class's body does.
-  static const cache::CacheConfig& checked(const cache::CacheConfig& cfg) {
-    if (!cache::valid(cfg)) {
-      throw std::runtime_error("MLXCellCache: invalid CacheConfig");
+  // own constructor reads the geometry before this body runs.
+  static const cache::CacheGeometry& checked(
+      const cache::CacheGeometry& geometry,
+      const cache::CacheConfig& cfg) {
+    if (!cache::valid(geometry, cfg)) {
+      throw std::runtime_error("MLXCellCache: invalid geometry or config");
     }
-    return cfg;
+    return geometry;
   }
 
   std::vector<Pool> kpool_;