[Native] Export HF LLMs with an off-graph KV cache

backends/native/BUCK

@@ -47,6 +47,7 @@ fbcode_target(
         "//executorch/exir/backend:utils",
         "//executorch/exir/passes:cse_pass",
         "//executorch/exir/passes:lib",
+        "//executorch/extension/llm/cache:cache",
         "//executorch/extension/llm/export:gguf",
     ],
 )

backends/native/examples/hf_llm/BUCK

@@ -26,6 +26,8 @@ fbcode_target(
         "//executorch/backends/native:lib",
         "//executorch/exir:lib",
         "//executorch/exir/native:native",
+        "//executorch/extension/llm/cache:cache",
+        "//executorch/extension/llm/cache:hf_attention",
         "//executorch/extension/llm/export:load",
         "fbsource//third-party/pypi/safetensors:safetensors",
         "fbsource//third-party/pypi/transformers:transformers",

backends/native/examples/hf_llm/export_hf_llm.py

@@ -12,12 +12,16 @@
 ``(input_ids [1, T], cache_position [T])`` signature. Reports which operators
 the native partitioner leaves behind. ``--gguf`` takes the weights from a GGUF
 file of an architecture ``gguf_hf`` supports instead, keeping its quantized blocks
-packed in the PTN. ``--samples`` appends a ``golden`` sample set: eager logits
+packed in the PTN. ``--offgraph_cache`` exports attention as
+``kvcache::update_and_attend`` with no KV cache in the graph; the runner owns the
+cache. ``--samples`` appends a ``golden`` sample set: eager logits
 for a prefill and a few greedy decode steps per prompt.
 """
 
 import argparse
+import contextlib
 import logging
+from collections.abc import Iterator
 
 import torch
 import torch.nn as nn
@@ -27,6 +31,12 @@
 )
 from executorch.backends.native.examples.hf_llm.hf_utils import report_lowering
 from executorch.exir.native import Call, Sample, SampleSet, to_native, write_samples
+from executorch.extension.llm.cache.hf_attention import register_offgraph_attention
+from executorch.extension.llm.cache.reference_cache import (
+    CacheConfig,
+    SequenceReferenceCache,
+)
+from executorch.extension.llm.cache.update_and_attend import REGISTRY
 from executorch.extension.llm.export.load import assign_one, iter_checkpoint, RawIter
 from torch.export import Dim, export
 from torchao.quantization import (
@@ -53,16 +63,26 @@
 )
 
 
+_OFFGRAPH_ATTENTION = "offgraph"
+
+
 def _load(
-    path: str, layers: int, random_weights: bool, empty_weights: bool
+    path: str,
+    layers: int,
+    random_weights: bool,
+    empty_weights: bool,
+    offgraph_cache: bool = False,
 ) -> PreTrainedModel:
     config = AutoConfig.from_pretrained(path)
     text_config = config.get_text_config()
     if layers:
         text_config.num_hidden_layers = layers
         if getattr(text_config, "layer_types", None):
             text_config.layer_types = text_config.layer_types[:layers]
-    kwargs = {"torch_dtype": torch.float32, "attn_implementation": "sdpa"}
+    if offgraph_cache:
+        register_offgraph_attention(_OFFGRAPH_ATTENTION)
+    attn = _OFFGRAPH_ATTENTION if offgraph_cache else "sdpa"
+    kwargs = {"torch_dtype": torch.float32, "attn_implementation": attn}
     if empty_weights:
         with init_empty_weights():
             model = AutoModelForCausalLM.from_config(config, **kwargs)
@@ -185,6 +205,79 @@ def forward(
             logits_to_keep=1,
         ).logits
 
+    def reset_cache(self) -> None:
+        self.cache.reset()
+
+    def cache_scope(self) -> contextlib.AbstractContextManager:
+        return contextlib.nullcontext()
+
+
+class OffGraphLastTokenLM(nn.Module):
+    """Last-position logits with the KV cache outside the graph.
+
+    Each attention layer emits ``kvcache::update_and_attend`` over this step's
+    K/V. Eager calls (golden samples) run inside ``cache_scope()``, which makes
+    a reference cache of ``capacity`` positions the active one.
+    """
+
+    def __init__(self, model: PreTrainedModel, capacity: int) -> None:
+        super().__init__()
+        self.model = model
+        config = model.config.get_text_config()
+        # The cache op attends every position it holds, so a windowed or
+        # chunked layer would silently attend past its span. Without
+        # layer_types (e.g. Mistral), HF windows every layer when
+        # sliding_window is set.
+        layer_types = getattr(config, "layer_types", None) or (
+            ["sliding_attention"]
+            if getattr(config, "sliding_window", None) is not None
+            else []
+        )
+        unsupported = sorted(set(layer_types) - {"full_attention"})
+        if unsupported:
+            raise ValueError(
+                f"--offgraph_cache does not support {unsupported} layers "
+                f"({config.model_type})"
+            )
+        head_dim = (
+            getattr(config, "head_dim", None)
+            or config.hidden_size // config.num_attention_heads
+        )
+        self._reference = SequenceReferenceCache(
+            CacheConfig(
+                n_layers=config.num_hidden_layers,
+                n_kv_heads=config.num_key_value_heads,
+                head_dim=head_dim,
+                capacity=capacity,
+            )
+        )
+        self._key = f"export_hf_llm/{id(self)}"
+
+    def forward(
+        self, input_ids: torch.Tensor, cache_position: torch.Tensor
+    ) -> torch.Tensor:
+        return self.model(
+            input_ids=input_ids,
+            cache_position=cache_position,
+            # The model would derive positions from its (absent) cache's length.
+            position_ids=cache_position.unsqueeze(0),
+            use_cache=False,
+            past_key_values=None,
+            logits_to_keep=1,
+        ).logits
+
+    def reset_cache(self) -> None:
+        self._reference.reset()
+
+    @contextlib.contextmanager
+    def cache_scope(self) -> Iterator[None]:
+        REGISTRY.install(self._key, self._reference)
+        try:
+            with REGISTRY.active(self._key):
+                yield
+        finally:
+            REGISTRY.uninstall(self._key)
+
 
 # TEMPORARY: use torchao-based 8da4w quantization until PTN export has its own
 # quantization flow.
@@ -238,15 +331,17 @@ def _sample_prompts(
 
 
 def golden_samples(
-    wrapper: LastTokenLM, prompts: list[list[int]], decode_steps: int
+    wrapper: LastTokenLM | OffGraphLastTokenLM,
+    prompts: list[list[int]],
+    decode_steps: int,
 ) -> SampleSet:
     """Per tokenized prompt, from an empty cache: a prefill call, then
     `decode_steps` calls each fed the previous call's eager argmax."""
     quantized = any(isinstance(p, TorchAOBaseTensor) for p in wrapper.parameters())
     samples = []
-    with torch.no_grad():
+    with torch.no_grad(), wrapper.cache_scope():
         for ids in prompts:
-            wrapper.cache.reset()
+            wrapper.reset_cache()
             input_ids = torch.tensor([ids], dtype=torch.long)
             positions = torch.arange(len(ids), dtype=torch.long)
             calls = []
@@ -256,7 +351,7 @@ def golden_samples(
                 input_ids = logits[:, -1].argmax(dim=-1, keepdim=True)
                 positions = positions[-1:] + 1
             samples.append(Sample(calls))
-    wrapper.cache.reset()
+    wrapper.reset_cache()
     tol = _QUANTIZED_TOL if quantized else _FLOAT_TOL
     return SampleSet(samples, source="eager/cpu", tol=tol)
 
@@ -285,6 +380,12 @@ def _parse_args() -> argparse.Namespace:
         help="GGUF file with the model's weights; --model supplies the config.",
     )
     parser.add_argument("--kv_cache_dtype", choices=["fp32", "fp16"], default="fp32")
+    parser.add_argument(
+        "--offgraph_cache",
+        action="store_true",
+        help="Export attention as kvcache::update_and_attend with no KV cache in "
+        "the graph; the runner owns the cache and sets its capacity and dtype.",
+    )
     parser.add_argument(
         "--samples",
         type=int,
@@ -301,6 +402,13 @@ def _parse_args() -> argparse.Namespace:
             "--gguf loads every weight; it excludes --layers, --random_weights "
             "and --quantize"
         )
+    if args.offgraph_cache and (
+        args.max_context_length or args.kv_cache_dtype != "fp32"
+    ):
+        parser.error(
+            "--offgraph_cache leaves the KV cache to the runner; it excludes "
+            "--max_context_length and --kv_cache_dtype"
+        )
     if args.max_seq_length < 7:
         parser.error(
             "--max_seq_length must be at least 7: the example prompt has 6 tokens"
@@ -323,14 +431,23 @@ def main() -> None:
     torch.manual_seed(0)
 
     model = _load(
-        args.model, args.layers, args.random_weights, empty_weights=bool(args.gguf)
+        args.model,
+        args.layers,
+        args.random_weights,
+        empty_weights=bool(args.gguf),
+        offgraph_cache=args.offgraph_cache,
     )
     if args.gguf:
         load_gguf(model, args.gguf)
     if args.quantize == "8da4w":
         _quantize_8da4w(model, args.group_size, args.embedding_group_size)
-    kv_cache_dtype = torch.float16 if args.kv_cache_dtype == "fp16" else torch.float32
-    wrapper = LastTokenLM(model, cache_len, kv_cache_dtype).eval()
+    if args.offgraph_cache:
+        wrapper = OffGraphLastTokenLM(model, cache_len).eval()
+    else:
+        kv_cache_dtype = (
+            torch.float16 if args.kv_cache_dtype == "fp16" else torch.float32
+        )
+        wrapper = LastTokenLM(model, cache_len, kv_cache_dtype).eval()
 
     input_ids = torch.zeros((1, 6), dtype=torch.long)
     cache_position = torch.arange(input_ids.shape[1], dtype=torch.long)
@@ -351,3 +468,7 @@ def main() -> None:
         samples = golden_samples(wrapper, prompts, _DECODE_STEPS)
         write_samples(args.output, "golden", samples)
         print(f"wrote {len(prompts)} golden samples")
+
+
+if __name__ == "__main__":
+    main()

backends/native/examples/hf_llm/test_export_hf_llm.py

@@ -26,7 +26,10 @@
     golden_samples,
     LastTokenLM,
     load_gguf,
+    OffGraphLastTokenLM,
 )
+from executorch.extension.llm.cache.hf_attention import register_offgraph_attention
+from executorch.extension.llm.cache.update_and_attend import REGISTRY
 from transformers import (
     AutoModelForCausalLM,
     Gemma2Config,
@@ -269,6 +272,53 @@ def test_golden_samples_teacher_force_eager_argmax(self) -> None:
                 self.assertTrue(torch.equal(call.outputs[0], previous))
         self.assertFalse(any(layer.keys.any() for layer in wrapper.cache.layers))
 
+    def test_offgraph_cache_matches_the_static_cache(self) -> None:
+        torch.manual_seed(7)
+        model = _tiny()
+        register_offgraph_attention("offgraph")
+        offgraph_model = copy.deepcopy(model)
+        offgraph_model.config._attn_implementation = "offgraph"
+        ref = LastTokenLM(model, 32).eval()
+        offgraph = OffGraphLastTokenLM(offgraph_model, 32).eval()
+        prompts = [[3, 4, 5, 6, 7, 8], [9, 10]]
+
+        samples = golden_samples(offgraph, prompts, 2)
+        want = golden_samples(ref, prompts, 2)
+
+        for got, expected in zip(samples.samples, want.samples, strict=True):
+            for call, ref_call in zip(got.calls, expected.calls, strict=True):
+                self.assertTrue(torch.equal(call.inputs[0], ref_call.inputs[0]))
+                self.assertTrue(
+                    torch.allclose(
+                        call.outputs[0], ref_call.outputs[0], atol=1e-5, rtol=1e-5
+                    )
+                )
+        seq = torch.export.Dim("seq", min=1, max=31)
+        ids = torch.tensor([[3, 4, 5]])
+        with torch.no_grad():
+            ep = torch.export.export(
+                offgraph,
+                (ids, torch.arange(3)),
+                dynamic_shapes={"input_ids": {1: seq}, "cache_position": {0: seq}},
+                strict=True,
+            )
+        ops = _ops_except_casts(ep)
+        self.assertEqual(ops["kvcache.update_and_attend.default"], 2)
+        self.assertNotIn("aten.index_copy_.default", ops)
+        self.assertFalse(ep.graph_signature.buffers_to_mutate)
+
+    def test_offgraph_cache_excludes_in_graph_cache_flags(self) -> None:
+        for flags in (
+            ["--max_context_length", "64"],
+            ["--kv_cache_dtype", "fp16"],
+        ):
+            argv = ["export_hf_llm", "--model", "unused", "--offgraph_cache", *flags]
+            with self.subTest(flags=flags), mock.patch.object(sys, "argv", argv):
+                with mock.patch.object(export_hf_llm, "_load") as load:
+                    with self.assertRaises(SystemExit):
+                        export_hf_llm.main()
+                load.assert_not_called()
+
     def test_samples_flag_needs_a_count_in_range_and_an_output(self) -> None:
         for flags in (
             ["--samples", str(len(_PROMPTS) + 1), "--output", "m.ptn"],
@@ -319,3 +369,19 @@ def test_golden_samples_tolerance_follows_weight_quantization(self) -> None:
     def test_fp16_kv_cache_rejects_sliding_window_layers(self) -> None:
         with self.assertRaisesRegex(ValueError, "cannot change the KV dtype"):
             LastTokenLM(_tiny("gemma3"), 32, torch.float16)
+
+    def test_offgraph_cache_rejects_sliding_window_layers(self) -> None:
+        with self.assertRaisesRegex(ValueError, "sliding_attention"):
+            OffGraphLastTokenLM(_tiny("gemma3"), 32)
+
+    def test_offgraph_cache_rejects_other_attention_layer_types(self) -> None:
+        model = _tiny()
+        model.config.layer_types = ["chunked_attention", "full_attention"]
+        with self.assertRaisesRegex(ValueError, "chunked_attention"):
+            OffGraphLastTokenLM(model, 32)
+
+    def test_offgraph_cache_scope_uninstalls_its_reference_cache(self) -> None:
+        wrapper = OffGraphLastTokenLM(_tiny(), 32)
+        with wrapper.cache_scope():
+            self.assertIn(wrapper._key, REGISTRY._caches)
+        self.assertNotIn(wrapper._key, REGISTRY._caches)

backends/native/partitioner.py

@@ -16,7 +16,9 @@
 import operator
 from typing import Callable, final, List, Mapping, Optional, Tuple
 
-# Registers torch.ops.torchao.dequantize_gguf, referenced in _SUPPORTED_NON_CORE_OPS.
+# Register torch.ops.kvcache.update_and_attend and torch.ops.torchao.dequantize_gguf,
+# referenced in _SUPPORTED_NON_CORE_OPS.
+import executorch.extension.llm.cache.update_and_attend  # noqa: F401
 import executorch.extension.llm.export.gguf  # noqa: F401
 import torch
 from executorch.backends.native.custom_ops import rope_op
@@ -50,6 +52,8 @@
     torch.ops.torchao.dequantize_affine.default,
     rope_op,
     torch.ops.aten.rms_norm.default,
+    # Attention over a runtime-owned KV cache; engines bind the cache at load.
+    torch.ops.kvcache.update_and_attend.default,
 ]
 
 # A SymInt, and any op arg that references one, is an execution-time value. An