Skip to content

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

Draft
metascroy wants to merge 1 commit into
offgraph-kv-hf-attentionfrom
offgraph-kv-native-export
Draft

[Native] Export HF LLMs with an off-graph KV cache#24635
metascroy wants to merge 1 commit into
offgraph-kv-hf-attentionfrom
offgraph-kv-native-export

Conversation

@metascroy

@metascroy metascroy commented Nov 5, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Stack

Off-graph KV cache for the native ggml engine, through serving Muse Glimmer on it. Review bottom-up; each PR builds and passes its tests on its own.

Summary

Adds --offgraph_cache to the native backend's HF LLM export (backends/native/examples/hf_llm/export_hf_llm.py). With it, attention runs through kvcache::update_and_attend, so the exported .ptn holds no KV cache buffers; the runtime supplies the cache.

  • The native partitioner accepts torch.ops.kvcache.update_and_attend.default.
  • OffGraphLastTokenLM wraps the model with use_cache=False and positions taken from cache_position. While computing golden samples, it installs the reference SequenceReferenceCache in the Python REGISTRY; export only traces the op and needs no cache.
  • --max_context_length and --kv_cache_dtype only describe an in-graph cache, so combining them with --offgraph_cache is rejected.
  • Models with any attention layer other than full attention are rejected: sliding-window layers (Gemma 3, Mistral, or Qwen 2/3 with use_sliding_window) and chunked layers (Llama 4). The cache op attends every position it holds, so they would silently attend past their span, in eager and at runtime alike.
  • The reference cache is installed in the global REGISTRY only inside cache_scope(), and uninstalled on exit.
  • Adds the missing if __name__ == "__main__" guard. Without it, python -m ...export_hf_llm exited without exporting anything.

Example (Llama 3.2 1B, Q4_K_M GGUF):

python -m executorch.backends.native.examples.hf_llm.export_hf_llm \
  --model unsloth/Llama-3.2-1B-Instruct \
  --gguf Llama-3.2-1B-Instruct-Q4_K_M.gguf \
  --offgraph_cache --max_seq_length 512 --samples 2 \
  --output llama1b_q4km_offgraph.ptn

Test plan

  • test_export_hf_llm.py adds five tests:
    • test_offgraph_cache_matches_the_static_cache: off-graph eager logits match the static cache within 1e-5. The exported graph has one update_and_attend per layer, no index_copy_, and no mutated buffers.
    • test_offgraph_cache_excludes_in_graph_cache_flags.
    • test_offgraph_cache_rejects_sliding_window_layers (tiny Gemma 3).
    • test_offgraph_cache_rejects_other_attention_layer_types (chunked_attention).
    • test_offgraph_cache_scope_uninstalls_its_reference_cache.
  • pytest backends/native/examples/hf_llm/test_export_hf_llm.py backends/native/test/test_partitioner.py backends/native/test/test_serialize.py: 132 passed.

@pytorch-bot

pytorch-bot Bot commented Nov 5, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 134 New Failures, 8 Unrelated Failures, 10 Unclassified Failures

As of commit c3d3ba5 with merge base c7f7931 (image):

NEW FAILURES - The following jobs have failed:

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 5, 2026
This was referenced Nov 5, 2026
Draft
Draft
Draft
Draft
Draft
@metascroy
metascroy force-pushed the offgraph-kv-native-export branch from 01fcf32 to 05200ca Compare November 5, 2026 22:29
kiymetakdemir
kiymetakdemir reviewed Nov 6, 2026
View reviewed changes
Comment thread backends/native/examples/hf_llm/export_hf_llm.py
directory (or only its config, with random weights), wraps it with the cache
adapter from ``transformers.integrations.executorch``, and exports the
``(input_ids [1, T], cache_position [T])`` signature. Reports which operators
the native partitioner leaves behind. ``--gguf`` takes the weights from a GGUF

@kiymetakdemir kiymetakdemir Nov 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: sentence is cut in half here

@metascroy metascroy Nov 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed: this PR had accidentally dropped "the native partitioner leaves behind." from that sentence; restored it.

kiymetakdemir
kiymetakdemir approved these changes Nov 6, 2026
View reviewed changes
Comment thread backends/native/examples/hf_llm/export_hf_llm.py Outdated
config = model.config.get_text_config()
# The cache op attends every position it holds, so a windowed layer
# would silently attend past its window.
if "sliding_attention" in (getattr(config, "layer_types", None) or ()):

@kiymetakdemir kiymetakdemir Nov 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude says MistralConfig has layer_types=None and a sliding_window (4096 by default), so this check lets it through. Could it also reject when layer_types is unset but sliding_window is set?
https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json

@metascroy metascroy Nov 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, thanks. --offgraph_cache now also rejects configs with no layer_types and a non-None sliding_window; Mistral, Mixtral, Phi3 and Starcoder2 window every layer in that case. When layer_types is set it stays authoritative, since Qwen2/3 and SmolLM3 can carry a sliding_window value they don't use.

@metascroy
metascroy force-pushed the offgraph-kv-native-export branch from 05200ca to 655cb68 Compare November 11, 2026 18:05
@metascroy
metascroy force-pushed the offgraph-kv-hf-attention branch from 9cc9b4d to 57123b7 Compare November 11, 2026 18:05
This was referenced Nov 12, 2026
Draft
Draft
Draft
Draft
Draft
Draft
Draft
Draft
Draft
Closed
Closed
Draft
Draft
Draft
Draft
Draft
Draft
Draft
Draft
Draft
@metascroy
Export HF LLMs for the native backend with an off-graph KV cache
c3d3ba5 
Add --offgraph_cache to the native HF LLM export. Attention routes through
kvcache::update_and_attend, so the exported graph holds no KV cache buffers;
the runtime supplies the cache. The native partitioner accepts the op, and
golden samples come from the reference sequence cache.

Also add the missing __main__ guard, without which python -m on the export
script silently did nothing.
@metascroy
metascroy force-pushed the offgraph-kv-native-export branch from 655cb68 to c3d3ba5 Compare November 14, 2026 00:18

This branch had an error being deployed

1 failed deployment
cadence — c3d3ba51 Deployed Nov 14, 2026 by metascroy via hifi-build / hifi4 #38322
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kiymetakdemir kiymetakdemir kiymetakdemir approved these changes

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@metascroy @kiymetakdemir