diff --git a/backends/native/examples/hf_llm/export_hf_llm.py b/backends/native/examples/hf_llm/export_hf_llm.py index 7335a1951cd..92d0dc3dfd2 100644 --- a/backends/native/examples/hf_llm/export_hf_llm.py +++ b/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,8 +355,9 @@ 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 @@ -357,7 +365,7 @@ def cache_metadata(self) -> dict[str, int | list[int]]: 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: diff --git a/backends/native/examples/hf_llm/test_export_hf_llm.py b/backends/native/examples/hf_llm/test_export_hf_llm.py index 0daeb0a9bff..489571c5e34 100644 --- a/backends/native/examples/hf_llm/test_export_hf_llm.py +++ b/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(