diff --git a/backends/native/BUCK b/backends/native/BUCK index e0afa1058fc..204ff654beb 100644 --- a/backends/native/BUCK +++ b/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", ], ) diff --git a/backends/native/examples/hf_llm/BUCK b/backends/native/examples/hf_llm/BUCK index 4a3f6794f94..c8b68ddb7fd 100644 --- a/backends/native/examples/hf_llm/BUCK +++ b/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", diff --git a/backends/native/examples/hf_llm/export_hf_llm.py b/backends/native/examples/hf_llm/export_hf_llm.py index abb5c0925b8..e5a53d64f99 100644 --- a/backends/native/examples/hf_llm/export_hf_llm.py +++ b/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,8 +63,15 @@ ) +_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() @@ -62,7 +79,10 @@ def _load( 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() 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 d67e02a3362..349fab0144a 100644 --- a/backends/native/examples/hf_llm/test_export_hf_llm.py +++ b/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) diff --git a/backends/native/partitioner.py b/backends/native/partitioner.py index ca4580f8d40..cdedad795f2 100644 --- a/backends/native/partitioner.py +++ b/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