Split the cache's model geometry from its runtime config

[kiymetakdemir]

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.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22763

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 3 Unrelated Failures

As of commit cbaf358 with merge base 41fbae3 ():

NEW FAILURE - The following job has failed:

• MLX / test-mlx / test-mlx (gh)
  /Users/runner/work/executorch/executorch/pytorch/executorch/backends/mlx/test/op_test_runner.cpp:258:7: error: implicit conversion changes signedness: 'value_type' (aka 'int') to 'size_type' (aka 'unsigned long') [-Werror,-Wsign-conversion]

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• MLX / test-mlx-muse-glimmer-build / test-mlx-muse-glimmer-build (gh) (trunk failure)
  /Users/runner/work/executorch/executorch/pytorch/executorch/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.cpp:1258:55: error: reference to 'kActivationDtype' is ambiguous
• Test QNN Windows MSVC build / build-qnn-windows-arm64 (gh) (trunk failure)
• Test QNN Windows MSVC build / build-qnn-windows-x64 / windows-job (gh) (trunk failure)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 4m 4s —— View job

---

Code Review: Split cache geometry from runtime config

I read the split (cache.h), the reader/writer (model_metadata.h/.py, constants.h), the cache implementations, the call sites, and the tests. This is a clean, well-scoped refactor — the CacheGeometry / CacheConfig boundary is coherent, the overloaded valid() set is tidy, and the test coverage (especially the six malformed-geometry fixtures round-tripped through a real .pte) is excellent. A few observations, none blocking.

Things done well

• Dropping the "layers list of size 1 broadcasts to every layer" special case removes the cfg.layers.size() == 1 ? front() : [l] branch that was duplicated in every byte layer (MLXCellCache, MLXSequenceCache, print_cache_summary). Real simplification.
• read_cache_geometry and the Python write_cache_geometry validate the same invariants independently (nonempty, equal lengths, positive heads/dims, nonnegative windows), so a malformed program is caught on both the producer and consumer side.
• op_test_runner.cpp:251 — the explicit fields[1] <= 0 guard before layers.assign(fields[1], ...) is a good catch; a negative count would otherwise wrap to a huge size_type before valid() ever ran.

backends/mlx/llm/cache.py — this quietly fixes a latent bug

shared = getattr(text_config, "num_kv_shared_layers", 0)
if shared:
    layer_types = layer_types[:-shared]

The old if hasattr(...): layer_types = layer_types[:-text_config.num_kv_shared_layers] would evaluate layer_types[:-0] → layer_types[:0] → empty list whenever num_kv_shared_layers was present and 0. The if shared: guard sidesteps that. Worth a line in the PR description since it's a behavior change beyond the mechanical refactor.

Minor: redundant valid() in module_executor.cpp:241

cache::CacheConfig cfg{max_sessions * max_session_tokens, kv_dtype};
if (initial_capacity >= 0) { cfg.initial_capacity = initial_capacity; }
if (!cache::valid(*geometry, cfg)) { ... return Error::InvalidProgram; }
...
auto built = cache::CacheFactory::global().build(backend_id, cache_kind, *geometry, cfg);

CacheFactory::build already calls valid(geometry, cfg) (cache_registry.cpp:94) and returns InvalidArgument. The check here is harmless and gives a distinct error code/message, but it is a duplicate validation — worth a comment noting the intent (clearer diagnostic) or dropping it.

Design note: policy lives inside "immutable" geometry but is mutated at runtime

cache.h:148 calls LayerGeometry "Immutable per-layer architecture facts + cache policy," yet apply_window_override in run_llm_hf.cpp takes CacheGeometry& and rewrites geometry.layers[l].policy before the cache is built. That's fine as designed — the exported window is a default a runner may override — but the "immutable" phrasing in the comment is slightly aspirational given a mutable runtime override path exists. Consider softening the comment to "model-published defaults" or similar so the override isn't surprising to a later reader.

Nits

• model_metadata.py:90 write_cache_geometry returns dict[str, object] (tensors + an int) while the sibling writers return dict[str, int]; export_model_metadata.py's all_methods was correctly widened to dict[str, object] to match. Consistent, just noting the type widened.
• read_cache_geometry requires get_windows to always be present. That's guaranteed by write_cache_geometry always emitting all three tensors, so no compatibility gap today — just flagging the coupling in case a producer ever emits geometry without windows.

Overall this looks correct and ready; the items above are polish, not correctness blockers.
• branch cache-geometry