[MLX] Add off-graph KV cache export mode for HF models

[kiymetakdemir]

Summary

Adds --use-offgraph-cache, which exports a HuggingFace causal LM against kvcache::update_and_attend instead of an in-graph cache. The model runs with use_cache=False and past_key_values=None, so each attention layer emits one op fed only that step's k/v; history lives in a cache the runtime owns and binds by cache_key. KV-sharing layers address their donor's cache rather than one of their own, so gemma-4 E2B needs 15 caches for its 35 layers.

Files

• hf_attention.py — registers the mlx_offgraph attention implementation and a mask function returning None, since the op masks internally from the position it is given. OffGraphExportWrapper exposes the (input_ids, cache_position) signature a runner drives and passes position_ids explicitly, because HuggingFace otherwise derives them from past_key_values.get_seq_length(), which is
• export_llm_hf.py — adds the off-graph export path and publishes the layout as get_n_caches, get_kv_heads, get_head_dims, get_windows.
• cache.py — reshapes the HFStaticCache fallback's cache_position to 1-D. That path belongs to the in-graph cache, but it is the fix gemma-4 needs: gemma-4 calls update() without cache_kwargs, so the fallback reads a 0-dim cumulative_length and torch.export fails indexing it.

Test

Exported Llama-3.2-1B and gemma-3-1b; both partition into a single MLX subgraph, and the published layout matches each architecture — 16 flat caches for llama, 26 for gemma-3 with full attention at layers 5/11/17/23.

python -m executorch.backends.mlx.examples.llm.export_llm_hf
--model-id unsloth/gemma-3-1b-it --output gemma3_offgraph.pte
--use-offgraph-cache --max-seq-len 2048 --dtype fp32

python -m executorch.backends.mlx.examples.llm.export_llm_hf
--model-id unsloth/Llama-3.2-1B-Instruct --output llama_offgraph.pte
--use-offgraph-cache --max-seq-len 2048 --dtype bf16

Add --qlinear 4w for INT4.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 2e8cb61 with merge base 14326d9 ():
💚 Looks good so far! There are no failures yet. 💚

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 39s —— View job

---

Code Review — Off-graph KV cache export mode

Read the diff plus update_and_attend.py, reference_cache.py, resolve_hf_cache_layout, and the existing custom-components export path. Overall this is a clean, well-documented addition that composes nicely with the existing op. The _cache_id donor logic and the cumulative_length.reshape(1) fix are both nicely reasoned in the comments. A few things worth addressing:

Correctness

1. softcap is silently dropped (hf_attention.py:106). mlx_offgraph_attention_forward accepts softcap in its signature but never forwards it — and kvcache::update_and_attend has no softcap parameter at all. For the two tested models (Llama-3.2, Gemma-3) this is fine since neither uses attention logit softcapping, but a model that does (e.g. Gemma-2) would export "successfully" and silently produce wrong logits. Same applies to head_mask. Since the op can't honor them, consider asserting they're unset rather than dropping them:

assert softcap is None, "update_and_attend does not support attention logit softcapping"
assert head_mask is None, "update_and_attend does not support head_mask"

This turns a silent-wrong-output footgun into a clear export-time error. Fix this →

2. Sliding-window masking is published but not enforced in the eager reference. The export publishes get_windows for the runtime, but ContiguousReferenceCache/attend in reference_cache.py only apply causal/none masking — no window truncation. So if you validate a sliding-window model (Gemma-3 has sliding layers at all but 5/11/17/23) against the eager op past the window length, the eager oracle will attend to full history and diverge from the actual MLX runtime. Not introduced by this PR, but worth being aware of when verifying numerics. Consider a note in the PR description or a follow-up.

Consistency / minor

3. scale fallback may be wrong for models with custom pre-attn scaling. float(query.shape[-1]) ** -0.5 (hf_attention.py:118) is only correct when head_dim ** -0.5 is the intended scale. Gemma uses query_pre_attn_scalar ** -0.5, which differs from head_dim. In practice HF always passes scaling explicitly via module.scaling, so the fallback shouldn't trigger — but given it would be silently wrong if it did, it's arguably safer to assert scaling is not None and drop the fallback entirely. Matches the "fail loud" spirit of the existing position_ids assert right above it.

4. Module docstring is now stale (hf_attention.py:9-30). The top-of-file docstring describes only the "mlx" custom_sdpa implementation. Since this file now also hosts the off-graph path, a sentence pointing at register_mlx_offgraph_attention / OffGraphExportWrapper would help the next reader.

5. get_n_caches value. "get_n_caches": len(layer_types) (export_llm_hf.py:400) — worth a one-line comment that layer_types here is already the KV-shared-truncated list returned by resolve_hf_cache_layout (line 56-57 of cache.py drops the shared tail), which is exactly why it equals the donor count (15 for Gemma-4-E2B) rather than the layer count. Right now the reader has to cross-reference two files to confirm get_n_caches, get_windows, and _cache_id all agree on the same truncated indexing. They do — the invariant is just implicit.

Nits

• cumulative_length.reshape(1) (cache.py:530) is correct for a 0-dim scalar. Fine as-is.
• The _export_with_offgraph_cache body duplicates the dtype map, quantize call, and export/lower boilerplate from _export_with_custom_components. Not worth refactoring for two callers, but if a third export path lands, the shared torch.export + to_edge_transform_and_lower + to_executorch tail is a candidate to extract.

I did not run the export end-to-end (requires downloading the HF checkpoints), so the review is static only. The logic and the published layout look self-consistent.
· branch mlx-offgraph-hf-export

[kiymetakdemir]

@claude review this code

[claude]

Claude finished @kiymetakdemir's task in 1m 50s —— View job

---

Code Review — Off-graph KV cache export mode

Re-read the diff against origin/main plus the supporting op (update_and_attend.py), the eager oracle (reference_cache.py), and resolve_hf_cache_layout. Note first that this commit already folds in most of the earlier review round: softcap/head_mask now raise instead of silently dropping (both attention paths), scaling is asserted rather than falling back to head_dim ** -0.5, and the module docstring now describes both "mlx" and "mlx_offgraph". Nice.

The design is coherent and the tracer-facing contract (mutates_args=(), cache off-graph, layer_id/scale/out_dtype as node constants) lines up cleanly with what _export_with_offgraph_cache publishes. The _cache_id donor logic is consistent with the truncation in resolve_hf_cache_layout (cache.py:56-57 drops the KV-shared tail, so len(layer_types) == donor count == the index space _cache_id maps into). A few things worth a look:

Correctness / worth verifying

1. Does the dynamic seq dim admit q_len == 1 at decode? (export_llm_hf.py) The export declares seq_len_dim = torch.export.Dim("seq_length_dim", max=max_seq_len - 1) with an example seq_length = 3. A bare torch.export.Dim has historically carried an implicit min of 2. If that holds in the pinned torch version, a single-token decode step (cache_position of length 1) would violate the guard baked into the .pte at runtime. Since the runner drives one token per decode step, please confirm the traced program actually accepts q_len == 1 — either by exporting with Dim(..., min=1) (or min=0) or by verifying the guard doesn't reject it. Worth a one-line comment on why max_seq_len - 1 rather than max_seq_len as well. Check this →

2. Sliding-window masking is published but not enforced by the eager oracle. get_windows is emitted for the runtime, but ContiguousReferenceCache/attend in reference_cache.py only distinguish CAUSAL vs NONE — no window truncation. So if you validate a sliding-window model (Gemma-3's sliding layers) against the eager op past the window length, the oracle attends to full history and will diverge from the MLX runtime. Not introduced here, but it means the eager path can't be used as ground truth for sliding layers beyond sliding_window tokens. Worth a note in the PR description or a follow-up.

3. get_windows assumes a single uniform window size. cache_windows fills every sliding_attention layer with the one text_config.sliding_window. Fine for Gemma-3, but models with per-layer window sizes would be mis-published silently. Since the layout is per-cache already, this is a latent gap rather than a bug for the tested models.

Consistency / minor

4. Batch dim is implicitly 1. position_ids[0].reshape(-1, 1) (hf_attention.py) and cache_position.unsqueeze(0) in OffGraphExportWrapper both assume B == 1, which matches the single-sequence ContiguousReferenceCache scope. That's a reasonable constraint, but there's no assert — a B > 1 input would trace and then produce wrong placement rather than failing. A one-line assert in the wrapper would match the fail-loud spirit of the position_ids/scaling asserts.

5. get_n_caches indexing is implicit. "get_n_caches": len(layer_types) — a one-line comment that layer_types is the already-truncated list from resolve_hf_cache_layout (not the model's layer count) would save the next reader a cross-file hop to confirm get_n_caches, get_windows, and _cache_id all share the same donor-indexed space. The comment above already gestures at this; making the invariant explicit at the assignment would close it.

6. Boilerplate duplication. _export_with_offgraph_cache repeats the dtype map, quantize_model_ call, and the to_edge_transform_and_lower → to_executorch → _save_program tail from _export_with_custom_components. Not worth extracting for two callers, but a third export path would justify pulling out the shared lowering tail.

Nit

• cumulative_length.reshape(1) (cache.py) is correct for the 0-dim scalar case and the comment explains why gemma-4 needs it. Fine as-is. Note this touches the in-graph path, not the off-graph one — the comment already says so, good.

I did not run the export end-to-end (requires the HF checkpoints), so this is static review only. Items 2–6 are advisory; item 1 is the one I'd confirm before merge.
· branch mlx-offgraph-hf-export