[Native] Publish the context length with off-graph packages

backends/native/examples/hf_llm/export_hf_llm.py

@@ -42,7 +42,10 @@
 from executorch.extension.llm.cache.update_and_attend import REGISTRY
 from executorch.extension.llm.export.gguf import iter_gguf
 from executorch.extension.llm.export.load import assign_one, iter_checkpoint, RawIter
-from executorch.extension.llm.export.model_metadata import write_cache_geometry
+from executorch.extension.llm.export.model_metadata import (
+    write_cache_geometry,
+    write_max_context_len,
+)
 from torch.export import Dim, export
 from torchao.quantization import (
     Int8DynamicActivationIntxWeightConfig,
@@ -319,7 +322,7 @@ class OffGraphLastTokenLM(nn.Module):
 
     Each attention layer emits ``kvcache::update_and_attend`` over this step's
     K/V; the cache, not the graph, applies a sliding layer's window, so the
-    runner learns the windows from ``cache_metadata()``. Eager calls (golden
+    runner learns the windows from ``metadata()``. Eager calls (golden
     samples) run inside ``cache_scope()``, which makes a reference cache of
     ``capacity`` positions the active one.
     """
@@ -334,6 +337,10 @@ def __init__(self, model: nn.Module, capacity: int) -> None:
             or config.hidden_size // config.num_attention_heads
         )
         self._kv_heads = config.num_key_value_heads
+        # The longest context the model supports, else the reference cache's.
+        self._max_context_len = (
+            getattr(config, "max_position_embeddings", None) or capacity
+        )
         self._reference = SequenceReferenceCache(
             CacheConfig(
                 n_layers=config.num_hidden_layers,
@@ -348,16 +355,17 @@ def __init__(self, model: nn.Module, capacity: int) -> None:
         )
         self._key = f"export_hf_llm/{id(self)}"
 
-    def cache_metadata(self) -> dict[str, int | list[int]]:
-        """The cache geometry, under the names LLM runners read it by."""
+    def metadata(self) -> dict[str, int | list[int]]:
+        """The cache geometry and the model's context length, under the names
+        LLM runners read them by."""
         n = len(self._windows)
         geometry = write_cache_geometry(
             [self._kv_heads] * n, [self._head_dim] * n, self._windows
         )
         return {
             name: value.tolist() if isinstance(value, torch.Tensor) else value
             for name, value in geometry.items()
-        }
+        } | write_max_context_len(self._max_context_len)
 
     def forward(
         self, input_ids: torch.Tensor, cache_position: torch.Tensor
@@ -572,7 +580,7 @@ def main() -> None:
     report_lowering(ep, "forward")
 
     if args.output:
-        metadata = wrapper.cache_metadata() if args.offgraph_cache else None
+        metadata = wrapper.metadata() if args.offgraph_cache else None
         to_native(ep, metadata=metadata).save(args.output)
         print(f"saved {args.output}")
     if prompts:

backends/native/examples/hf_llm/test_export_hf_llm.py

@@ -395,7 +395,7 @@ def test_muse_glimmer_exports_text_only_with_an_offgraph_cache(self) -> None:
                 with torch.no_grad():
                     want = reference(torch.tensor([sequence])).logits[:, -1:]
                 torch.testing.assert_close(call.outputs[0], want, atol=1e-5, rtol=1e-5)
-        self.assertEqual(offgraph.cache_metadata()["get_windows"], [3, 3, 3, 0])
+        self.assertEqual(offgraph.metadata()["get_windows"], [3, 3, 3, 0])
         seq = torch.export.Dim("seq", min=1, max=31)
         with torch.no_grad():
             ep = torch.export.export(
@@ -511,7 +511,7 @@ def test_offgraph_cache_windows_sliding_layers_like_hf(self) -> None:
                             call.outputs[0], want, atol=1e-5, rtol=1e-5
                         )
 
-    def test_offgraph_cache_metadata_gives_each_layer_its_window(self) -> None:
+    def test_offgraph_metadata_gives_each_layer_its_window(self) -> None:
         for model, windows in (
             (_tiny(), [0, 0]),
             (_tiny("mistral", sliding_window=8), [8, 8]),
@@ -529,12 +529,13 @@ def test_offgraph_cache_metadata_gives_each_layer_its_window(self) -> None:
             with self.subTest(arch=model.config.model_type):
                 wrapper = OffGraphLastTokenLM(model, 32)
                 self.assertEqual(
-                    wrapper.cache_metadata(),
+                    wrapper.metadata(),
                     {
                         "get_n_caches": 2,
                         "get_kv_heads": [2, 2],
                         "get_head_dims": [16, 16],
                         "get_windows": windows,
+                        "get_max_context_len": 64,
                     },
                 )
                 self.assertEqual(