[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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8d0517b with merge base 5b3da18 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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:

assertsoftcapisNone, "update_and_attend does not support attention logit softcapping"asserthead_maskisNone, "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

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 23s —— View job

---

Code Review — Off-graph KV cache export mode

Re-read the full diff (hf_attention.py, export_llm_hf.py, cache.py) plus the op (extension/llm/cache/update_and_attend.py) and resolve_hf_cache_layout. This commit already folds in the earlier two rounds — softcap/head_maskraise on both attention paths, scaling is asserted (no head_dim ** -0.5 fallback), and the module docstring now covers both "mlx" and "mlx_offgraph". The tracer contract (mutates_args=(), layer_id/scale/out_dtype as node constants, cache off-graph) lines up cleanly with what _export_with_offgraph_cache publishes, and _cache_id's donor logic matches 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, get_windows, get_kv_heads, get_head_dims all share). Solid.

One item I'd confirm before merge (still open from the last round), the rest advisory.

Worth confirming before merge

1. Does seq_length_dim admit q_len == 1 at decode? (export_llm_hf.py:415) The off-graph export declares torch.export.Dim("seq_length_dim", max=max_seq_len - 1) with example seq_length = 3. A bare torch.export.Dim has historically carried an implicit min of 2 — if that holds in the pinned torch (release/2.13), the .pte bakes in a seq >= 2 guard and a single-token decode step (cache_position of length 1) would be rejected at runtime. The runner drives one token per decode step, so please confirm the traced program accepts q_len == 1 — either export with Dim(..., min=1) or verify the guard doesn't reject 1. I couldn't run the export here (sandbox blocks Python + HF checkpoint download) to check empirically. Note the same Dim(max=... - 1) shape appears in _export_with_custom_components (:269), so whatever the resolution, keeping both consistent — plus a one-line comment on why max_seq_len - 1 rather than max_seq_len — would help. Check this →

Advisory (not blocking)

2. B == 1 is assumed but never asserted.OffGraphExportWrapper.forward does position_ids=cache_position.unsqueeze(0) (hf_attention.py:234) and mlx_offgraph_attention_forward does position_ids[0].reshape(-1, 1) (:156) — both collapse the batch dim to index 0. A B > 1 input would trace and then silently place all batches at batch-0's positions rather than failing. A one-line assert input_ids.shape[0] == 1 in the wrapper would match the fail-loud spirit of the position_ids/scaling asserts right below it. The op's own contract is [q_len, n_dims] positions, so single-sequence is the intended scope — just make it explicit.

3. Sliding-window masking is published but not enforced by the eager oracle.get_windows is emitted for the runtime (export_llm_hf.py:396-398), but attend/ContiguousReferenceCache in reference_cache.py only distinguish causal vs none — no window truncation. So the eager op can't serve as ground truth for Gemma-3's sliding layers past sliding_window tokens; it'll attend to full history and diverge from the MLX runtime. Not introduced here — worth a note in the PR description or a follow-up so numerics validation doesn't get misread.

4. get_windows assumes one uniform window size.cache_windows fills every sliding_attention layer with the single text_config.sliding_window (:396-398). Fine for Gemma-3, but a model with per-layer window sizes would be mis-published silently. Latent gap, not a bug for the tested models — the layout is already per-cache, so it's a small step to source per-layer if that ever lands.

5. get_n_caches indexing is implicit (:400)."get_n_caches": len(layer_types) — the comment at :391-392 already gestures at "these are indexed by cache," but making the invariant explicit at the assignment (that layer_types is the truncated donor list, not the model's layer count) would save the next reader a cross-file hop to confirm get_n_caches/get_windows/_cache_id share the donor-indexed space.

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:530) is correct for the 0-dim scalar and the comment explains the gemma-4 fallback. Fine as-is — and it correctly touches only the in-graph path, not the off-graph one.

Static review only — I could not run the export end-to-end (Python execution is gated in this environment and the HF checkpoints aren't available). Item 1 is the one I'd verify before merge; 2–6 are advisory.
· branch mlx-offgraph-hf-export

[metascroy]

Is it better to have **kwargs catch everything else (including softcap), and then raise if kwargs is non-empty? That way more unknown kwargs don't sneak in

[kiymetakdemir]

It makes sense for catching unknown ones, I added a sweep that pops the args we handle elsewhere (is_causal, use_cache, sliding_window, dropout) and raises on anything left. But softcap and head_mask still need to be named, models pass them often as None which still would raise.

[kiymetakdemir]

For claude review;

1. Dim's default min is 0, not 2 (torch/export/dynamic_shapes.py:180), so the constraint is in range [0, max_seq_len-1] and q_len == 1 is verified in practice, since the runner decodes one token per step.
2. Added assert input_ids.shape[0] == 1 in the export wrapper.
3. I'll add MaskKind.EXPLICIT in a separate PR.
4. The value is uniform because transformers exposes sliding_window as a scalar on the text config; layer_types only says which layers are sliding, so there's no per-layer window to read today.
5. Added a comment at the assignment stating that get_n_caches is the cache count rather than num_hidden_layers.
6. Extracting it would mean touching the other export paths, which are out of scope for this PR.

[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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8d0517b with merge base 5b3da18 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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:

assertsoftcapisNone, "update_and_attend does not support attention logit softcapping"asserthead_maskisNone, "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

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 23s —— View job

---

Code Review — Off-graph KV cache export mode

Re-read the full diff (hf_attention.py, export_llm_hf.py, cache.py) plus the op (extension/llm/cache/update_and_attend.py) and resolve_hf_cache_layout. This commit already folds in the earlier two rounds — softcap/head_maskraise on both attention paths, scaling is asserted (no head_dim ** -0.5 fallback), and the module docstring now covers both "mlx" and "mlx_offgraph". The tracer contract (mutates_args=(), layer_id/scale/out_dtype as node constants, cache off-graph) lines up cleanly with what _export_with_offgraph_cache publishes, and _cache_id's donor logic matches 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, get_windows, get_kv_heads, get_head_dims all share). Solid.

One item I'd confirm before merge (still open from the last round), the rest advisory.

Worth confirming before merge

1. Does seq_length_dim admit q_len == 1 at decode? (export_llm_hf.py:415) The off-graph export declares torch.export.Dim("seq_length_dim", max=max_seq_len - 1) with example seq_length = 3. A bare torch.export.Dim has historically carried an implicit min of 2 — if that holds in the pinned torch (release/2.13), the .pte bakes in a seq >= 2 guard and a single-token decode step (cache_position of length 1) would be rejected at runtime. The runner drives one token per decode step, so please confirm the traced program accepts q_len == 1 — either export with Dim(..., min=1) or verify the guard doesn't reject 1. I couldn't run the export here (sandbox blocks Python + HF checkpoint download) to check empirically. Note the same Dim(max=... - 1) shape appears in _export_with_custom_components (:269), so whatever the resolution, keeping both consistent — plus a one-line comment on why max_seq_len - 1 rather than max_seq_len — would help. Check this →

Advisory (not blocking)

2. B == 1 is assumed but never asserted.OffGraphExportWrapper.forward does position_ids=cache_position.unsqueeze(0) (hf_attention.py:234) and mlx_offgraph_attention_forward does position_ids[0].reshape(-1, 1) (:156) — both collapse the batch dim to index 0. A B > 1 input would trace and then silently place all batches at batch-0's positions rather than failing. A one-line assert input_ids.shape[0] == 1 in the wrapper would match the fail-loud spirit of the position_ids/scaling asserts right below it. The op's own contract is [q_len, n_dims] positions, so single-sequence is the intended scope — just make it explicit.

3. Sliding-window masking is published but not enforced by the eager oracle.get_windows is emitted for the runtime (export_llm_hf.py:396-398), but attend/ContiguousReferenceCache in reference_cache.py only distinguish causal vs none — no window truncation. So the eager op can't serve as ground truth for Gemma-3's sliding layers past sliding_window tokens; it'll attend to full history and diverge from the MLX runtime. Not introduced here — worth a note in the PR description or a follow-up so numerics validation doesn't get misread.

4. get_windows assumes one uniform window size.cache_windows fills every sliding_attention layer with the single text_config.sliding_window (:396-398). Fine for Gemma-3, but a model with per-layer window sizes would be mis-published silently. Latent gap, not a bug for the tested models — the layout is already per-cache, so it's a small step to source per-layer if that ever lands.

5. get_n_caches indexing is implicit (:400)."get_n_caches": len(layer_types) — the comment at :391-392 already gestures at "these are indexed by cache," but making the invariant explicit at the assignment (that layer_types is the truncated donor list, not the model's layer count) would save the next reader a cross-file hop to confirm get_n_caches/get_windows/_cache_id share the donor-indexed space.

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:530) is correct for the 0-dim scalar and the comment explains the gemma-4 fallback. Fine as-is — and it correctly touches only the in-graph path, not the off-graph one.

Static review only — I could not run the export end-to-end (Python execution is gated in this environment and the HF checkpoints aren't available). Item 1 is the one I'd verify before merge; 2–6 are advisory.
· branch mlx-offgraph-hf-export

[metascroy]

Is it better to have **kwargs catch everything else (including softcap), and then raise if kwargs is non-empty? That way more unknown kwargs don't sneak in

[kiymetakdemir]

It makes sense for catching unknown ones, I added a sweep that pops the args we handle elsewhere (is_causal, use_cache, sliding_window, dropout) and raises on anything left. But softcap and head_mask still need to be named, models pass them often as None which still would raise.

[kiymetakdemir]

For claude review;

1. Dim's default min is 0, not 2 (torch/export/dynamic_shapes.py:180), so the constraint is in range [0, max_seq_len-1] and q_len == 1 is verified in practice, since the runner decodes one token per step.
2. Added assert input_ids.shape[0] == 1 in the export wrapper.
3. I'll add MaskKind.EXPLICIT in a separate PR.
4. The value is uniform because transformers exposes sliding_window as a scalar on the text config; layer_types only says which layers are sliding, so there's no per-layer window to read today.
5. Added a comment at the assignment stating that get_n_caches is the cache count rather than num_hidden_layers.
6. Extracting it would mean touching the other export paths, which are out of scope for this PR.

[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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8d0517b with merge base 5b3da18 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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:

assertsoftcapisNone, "update_and_attend does not support attention logit softcapping"asserthead_maskisNone, "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

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 23s —— View job

---

Code Review — Off-graph KV cache export mode

Re-read the full diff (hf_attention.py, export_llm_hf.py, cache.py) plus the op (extension/llm/cache/update_and_attend.py) and resolve_hf_cache_layout. This commit already folds in the earlier two rounds — softcap/head_maskraise on both attention paths, scaling is asserted (no head_dim ** -0.5 fallback), and the module docstring now covers both "mlx" and "mlx_offgraph". The tracer contract (mutates_args=(), layer_id/scale/out_dtype as node constants, cache off-graph) lines up cleanly with what _export_with_offgraph_cache publishes, and _cache_id's donor logic matches 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, get_windows, get_kv_heads, get_head_dims all share). Solid.

One item I'd confirm before merge (still open from the last round), the rest advisory.

Worth confirming before merge

1. Does seq_length_dim admit q_len == 1 at decode? (export_llm_hf.py:415) The off-graph export declares torch.export.Dim("seq_length_dim", max=max_seq_len - 1) with example seq_length = 3. A bare torch.export.Dim has historically carried an implicit min of 2 — if that holds in the pinned torch (release/2.13), the .pte bakes in a seq >= 2 guard and a single-token decode step (cache_position of length 1) would be rejected at runtime. The runner drives one token per decode step, so please confirm the traced program accepts q_len == 1 — either export with Dim(..., min=1) or verify the guard doesn't reject 1. I couldn't run the export here (sandbox blocks Python + HF checkpoint download) to check empirically. Note the same Dim(max=... - 1) shape appears in _export_with_custom_components (:269), so whatever the resolution, keeping both consistent — plus a one-line comment on why max_seq_len - 1 rather than max_seq_len — would help. Check this →

Advisory (not blocking)

2. B == 1 is assumed but never asserted.OffGraphExportWrapper.forward does position_ids=cache_position.unsqueeze(0) (hf_attention.py:234) and mlx_offgraph_attention_forward does position_ids[0].reshape(-1, 1) (:156) — both collapse the batch dim to index 0. A B > 1 input would trace and then silently place all batches at batch-0's positions rather than failing. A one-line assert input_ids.shape[0] == 1 in the wrapper would match the fail-loud spirit of the position_ids/scaling asserts right below it. The op's own contract is [q_len, n_dims] positions, so single-sequence is the intended scope — just make it explicit.

3. Sliding-window masking is published but not enforced by the eager oracle.get_windows is emitted for the runtime (export_llm_hf.py:396-398), but attend/ContiguousReferenceCache in reference_cache.py only distinguish causal vs none — no window truncation. So the eager op can't serve as ground truth for Gemma-3's sliding layers past sliding_window tokens; it'll attend to full history and diverge from the MLX runtime. Not introduced here — worth a note in the PR description or a follow-up so numerics validation doesn't get misread.

4. get_windows assumes one uniform window size.cache_windows fills every sliding_attention layer with the single text_config.sliding_window (:396-398). Fine for Gemma-3, but a model with per-layer window sizes would be mis-published silently. Latent gap, not a bug for the tested models — the layout is already per-cache, so it's a small step to source per-layer if that ever lands.

5. get_n_caches indexing is implicit (:400)."get_n_caches": len(layer_types) — the comment at :391-392 already gestures at "these are indexed by cache," but making the invariant explicit at the assignment (that layer_types is the truncated donor list, not the model's layer count) would save the next reader a cross-file hop to confirm get_n_caches/get_windows/_cache_id share the donor-indexed space.

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:530) is correct for the 0-dim scalar and the comment explains the gemma-4 fallback. Fine as-is — and it correctly touches only the in-graph path, not the off-graph one.

Static review only — I could not run the export end-to-end (Python execution is gated in this environment and the HF checkpoints aren't available). Item 1 is the one I'd verify before merge; 2–6 are advisory.
· branch mlx-offgraph-hf-export

[metascroy]

Is it better to have **kwargs catch everything else (including softcap), and then raise if kwargs is non-empty? That way more unknown kwargs don't sneak in

[kiymetakdemir]

It makes sense for catching unknown ones, I added a sweep that pops the args we handle elsewhere (is_causal, use_cache, sliding_window, dropout) and raises on anything left. But softcap and head_mask still need to be named, models pass them often as None which still would raise.

[kiymetakdemir]

For claude review;

1. Dim's default min is 0, not 2 (torch/export/dynamic_shapes.py:180), so the constraint is in range [0, max_seq_len-1] and q_len == 1 is verified in practice, since the runner decodes one token per step.
2. Added assert input_ids.shape[0] == 1 in the export wrapper.
3. I'll add MaskKind.EXPLICIT in a separate PR.
4. The value is uniform because transformers exposes sliding_window as a scalar on the text config; layer_types only says which layers are sliding, so there's no per-layer window to read today.
5. Added a comment at the assignment stating that get_n_caches is the cache count rather than num_hidden_layers.
6. Extracting it would mean touching the other export paths, which are out of scope for this PR.

[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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8d0517b with merge base 5b3da18 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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:

assertsoftcapisNone, "update_and_attend does not support attention logit softcapping"asserthead_maskisNone, "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

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 23s —— View job

---

Code Review — Off-graph KV cache export mode

Re-read the full diff (hf_attention.py, export_llm_hf.py, cache.py) plus the op (extension/llm/cache/update_and_attend.py) and resolve_hf_cache_layout. This commit already folds in the earlier two rounds — softcap/head_maskraise on both attention paths, scaling is asserted (no head_dim ** -0.5 fallback), and the module docstring now covers both "mlx" and "mlx_offgraph". The tracer contract (mutates_args=(), layer_id/scale/out_dtype as node constants, cache off-graph) lines up cleanly with what _export_with_offgraph_cache publishes, and _cache_id's donor logic matches 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, get_windows, get_kv_heads, get_head_dims all share). Solid.

One item I'd confirm before merge (still open from the last round), the rest advisory.

Worth confirming before merge

1. Does seq_length_dim admit q_len == 1 at decode? (export_llm_hf.py:415) The off-graph export declares torch.export.Dim("seq_length_dim", max=max_seq_len - 1) with example seq_length = 3. A bare torch.export.Dim has historically carried an implicit min of 2 — if that holds in the pinned torch (release/2.13), the .pte bakes in a seq >= 2 guard and a single-token decode step (cache_position of length 1) would be rejected at runtime. The runner drives one token per decode step, so please confirm the traced program accepts q_len == 1 — either export with Dim(..., min=1) or verify the guard doesn't reject 1. I couldn't run the export here (sandbox blocks Python + HF checkpoint download) to check empirically. Note the same Dim(max=... - 1) shape appears in _export_with_custom_components (:269), so whatever the resolution, keeping both consistent — plus a one-line comment on why max_seq_len - 1 rather than max_seq_len — would help. Check this →

Advisory (not blocking)

2. B == 1 is assumed but never asserted.OffGraphExportWrapper.forward does position_ids=cache_position.unsqueeze(0) (hf_attention.py:234) and mlx_offgraph_attention_forward does position_ids[0].reshape(-1, 1) (:156) — both collapse the batch dim to index 0. A B > 1 input would trace and then silently place all batches at batch-0's positions rather than failing. A one-line assert input_ids.shape[0] == 1 in the wrapper would match the fail-loud spirit of the position_ids/scaling asserts right below it. The op's own contract is [q_len, n_dims] positions, so single-sequence is the intended scope — just make it explicit.

3. Sliding-window masking is published but not enforced by the eager oracle.get_windows is emitted for the runtime (export_llm_hf.py:396-398), but attend/ContiguousReferenceCache in reference_cache.py only distinguish causal vs none — no window truncation. So the eager op can't serve as ground truth for Gemma-3's sliding layers past sliding_window tokens; it'll attend to full history and diverge from the MLX runtime. Not introduced here — worth a note in the PR description or a follow-up so numerics validation doesn't get misread.

4. get_windows assumes one uniform window size.cache_windows fills every sliding_attention layer with the single text_config.sliding_window (:396-398). Fine for Gemma-3, but a model with per-layer window sizes would be mis-published silently. Latent gap, not a bug for the tested models — the layout is already per-cache, so it's a small step to source per-layer if that ever lands.

5. get_n_caches indexing is implicit (:400)."get_n_caches": len(layer_types) — the comment at :391-392 already gestures at "these are indexed by cache," but making the invariant explicit at the assignment (that layer_types is the truncated donor list, not the model's layer count) would save the next reader a cross-file hop to confirm get_n_caches/get_windows/_cache_id share the donor-indexed space.

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:530) is correct for the 0-dim scalar and the comment explains the gemma-4 fallback. Fine as-is — and it correctly touches only the in-graph path, not the off-graph one.

Static review only — I could not run the export end-to-end (Python execution is gated in this environment and the HF checkpoints aren't available). Item 1 is the one I'd verify before merge; 2–6 are advisory.
· branch mlx-offgraph-hf-export

[metascroy]

Is it better to have **kwargs catch everything else (including softcap), and then raise if kwargs is non-empty? That way more unknown kwargs don't sneak in

[kiymetakdemir]

It makes sense for catching unknown ones, I added a sweep that pops the args we handle elsewhere (is_causal, use_cache, sliding_window, dropout) and raises on anything left. But softcap and head_mask still need to be named, models pass them often as None which still would raise.

[kiymetakdemir]

For claude review;

1. Dim's default min is 0, not 2 (torch/export/dynamic_shapes.py:180), so the constraint is in range [0, max_seq_len-1] and q_len == 1 is verified in practice, since the runner decodes one token per step.
2. Added assert input_ids.shape[0] == 1 in the export wrapper.
3. I'll add MaskKind.EXPLICIT in a separate PR.
4. The value is uniform because transformers exposes sliding_window as a scalar on the text config; layer_types only says which layers are sliding, so there's no per-layer window to read today.
5. Added a comment at the assignment stating that get_n_caches is the cache count rather than num_hidden_layers.
6. Extracting it would mean touching the other export paths, which are out of scope for this PR.

[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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8d0517b with merge base 5b3da18 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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:

assertsoftcapisNone, "update_and_attend does not support attention logit softcapping"asserthead_maskisNone, "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

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 23s —— View job

---

Code Review — Off-graph KV cache export mode

Re-read the full diff (hf_attention.py, export_llm_hf.py, cache.py) plus the op (extension/llm/cache/update_and_attend.py) and resolve_hf_cache_layout. This commit already folds in the earlier two rounds — softcap/head_maskraise on both attention paths, scaling is asserted (no head_dim ** -0.5 fallback), and the module docstring now covers both "mlx" and "mlx_offgraph". The tracer contract (mutates_args=(), layer_id/scale/out_dtype as node constants, cache off-graph) lines up cleanly with what _export_with_offgraph_cache publishes, and _cache_id's donor logic matches 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, get_windows, get_kv_heads, get_head_dims all share). Solid.

One item I'd confirm before merge (still open from the last round), the rest advisory.

Worth confirming before merge

1. Does seq_length_dim admit q_len == 1 at decode? (export_llm_hf.py:415) The off-graph export declares torch.export.Dim("seq_length_dim", max=max_seq_len - 1) with example seq_length = 3. A bare torch.export.Dim has historically carried an implicit min of 2 — if that holds in the pinned torch (release/2.13), the .pte bakes in a seq >= 2 guard and a single-token decode step (cache_position of length 1) would be rejected at runtime. The runner drives one token per decode step, so please confirm the traced program accepts q_len == 1 — either export with Dim(..., min=1) or verify the guard doesn't reject 1. I couldn't run the export here (sandbox blocks Python + HF checkpoint download) to check empirically. Note the same Dim(max=... - 1) shape appears in _export_with_custom_components (:269), so whatever the resolution, keeping both consistent — plus a one-line comment on why max_seq_len - 1 rather than max_seq_len — would help. Check this →

Advisory (not blocking)

2. B == 1 is assumed but never asserted.OffGraphExportWrapper.forward does position_ids=cache_position.unsqueeze(0) (hf_attention.py:234) and mlx_offgraph_attention_forward does position_ids[0].reshape(-1, 1) (:156) — both collapse the batch dim to index 0. A B > 1 input would trace and then silently place all batches at batch-0's positions rather than failing. A one-line assert input_ids.shape[0] == 1 in the wrapper would match the fail-loud spirit of the position_ids/scaling asserts right below it. The op's own contract is [q_len, n_dims] positions, so single-sequence is the intended scope — just make it explicit.

3. Sliding-window masking is published but not enforced by the eager oracle.get_windows is emitted for the runtime (export_llm_hf.py:396-398), but attend/ContiguousReferenceCache in reference_cache.py only distinguish causal vs none — no window truncation. So the eager op can't serve as ground truth for Gemma-3's sliding layers past sliding_window tokens; it'll attend to full history and diverge from the MLX runtime. Not introduced here — worth a note in the PR description or a follow-up so numerics validation doesn't get misread.

4. get_windows assumes one uniform window size.cache_windows fills every sliding_attention layer with the single text_config.sliding_window (:396-398). Fine for Gemma-3, but a model with per-layer window sizes would be mis-published silently. Latent gap, not a bug for the tested models — the layout is already per-cache, so it's a small step to source per-layer if that ever lands.

5. get_n_caches indexing is implicit (:400)."get_n_caches": len(layer_types) — the comment at :391-392 already gestures at "these are indexed by cache," but making the invariant explicit at the assignment (that layer_types is the truncated donor list, not the model's layer count) would save the next reader a cross-file hop to confirm get_n_caches/get_windows/_cache_id share the donor-indexed space.

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:530) is correct for the 0-dim scalar and the comment explains the gemma-4 fallback. Fine as-is — and it correctly touches only the in-graph path, not the off-graph one.

Static review only — I could not run the export end-to-end (Python execution is gated in this environment and the HF checkpoints aren't available). Item 1 is the one I'd verify before merge; 2–6 are advisory.
· branch mlx-offgraph-hf-export

[metascroy]

Is it better to have **kwargs catch everything else (including softcap), and then raise if kwargs is non-empty? That way more unknown kwargs don't sneak in

[kiymetakdemir]

It makes sense for catching unknown ones, I added a sweep that pops the args we handle elsewhere (is_causal, use_cache, sliding_window, dropout) and raises on anything left. But softcap and head_mask still need to be named, models pass them often as None which still would raise.

[kiymetakdemir]

For claude review;

1. Dim's default min is 0, not 2 (torch/export/dynamic_shapes.py:180), so the constraint is in range [0, max_seq_len-1] and q_len == 1 is verified in practice, since the runner decodes one token per step.
2. Added assert input_ids.shape[0] == 1 in the export wrapper.
3. I'll add MaskKind.EXPLICIT in a separate PR.
4. The value is uniform because transformers exposes sliding_window as a scalar on the text config; layer_types only says which layers are sliding, so there's no per-layer window to read today.
5. Added a comment at the assignment stating that get_n_caches is the cache count rather than num_hidden_layers.
6. Extracting it would mean touching the other export paths, which are out of scope for this PR.

[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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8d0517b with merge base 5b3da18 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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:

assertsoftcapisNone, "update_and_attend does not support attention logit softcapping"asserthead_maskisNone, "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

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 23s —— View job

---

Code Review — Off-graph KV cache export mode

Re-read the full diff (hf_attention.py, export_llm_hf.py, cache.py) plus the op (extension/llm/cache/update_and_attend.py) and resolve_hf_cache_layout. This commit already folds in the earlier two rounds — softcap/head_maskraise on both attention paths, scaling is asserted (no head_dim ** -0.5 fallback), and the module docstring now covers both "mlx" and "mlx_offgraph". The tracer contract (mutates_args=(), layer_id/scale/out_dtype as node constants, cache off-graph) lines up cleanly with what _export_with_offgraph_cache publishes, and _cache_id's donor logic matches 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, get_windows, get_kv_heads, get_head_dims all share). Solid.

One item I'd confirm before merge (still open from the last round), the rest advisory.

Worth confirming before merge

1. Does seq_length_dim admit q_len == 1 at decode? (export_llm_hf.py:415) The off-graph export declares torch.export.Dim("seq_length_dim", max=max_seq_len - 1) with example seq_length = 3. A bare torch.export.Dim has historically carried an implicit min of 2 — if that holds in the pinned torch (release/2.13), the .pte bakes in a seq >= 2 guard and a single-token decode step (cache_position of length 1) would be rejected at runtime. The runner drives one token per decode step, so please confirm the traced program accepts q_len == 1 — either export with Dim(..., min=1) or verify the guard doesn't reject 1. I couldn't run the export here (sandbox blocks Python + HF checkpoint download) to check empirically. Note the same Dim(max=... - 1) shape appears in _export_with_custom_components (:269), so whatever the resolution, keeping both consistent — plus a one-line comment on why max_seq_len - 1 rather than max_seq_len — would help. Check this →

Advisory (not blocking)

2. B == 1 is assumed but never asserted.OffGraphExportWrapper.forward does position_ids=cache_position.unsqueeze(0) (hf_attention.py:234) and mlx_offgraph_attention_forward does position_ids[0].reshape(-1, 1) (:156) — both collapse the batch dim to index 0. A B > 1 input would trace and then silently place all batches at batch-0's positions rather than failing. A one-line assert input_ids.shape[0] == 1 in the wrapper would match the fail-loud spirit of the position_ids/scaling asserts right below it. The op's own contract is [q_len, n_dims] positions, so single-sequence is the intended scope — just make it explicit.

3. Sliding-window masking is published but not enforced by the eager oracle.get_windows is emitted for the runtime (export_llm_hf.py:396-398), but attend/ContiguousReferenceCache in reference_cache.py only distinguish causal vs none — no window truncation. So the eager op can't serve as ground truth for Gemma-3's sliding layers past sliding_window tokens; it'll attend to full history and diverge from the MLX runtime. Not introduced here — worth a note in the PR description or a follow-up so numerics validation doesn't get misread.

4. get_windows assumes one uniform window size.cache_windows fills every sliding_attention layer with the single text_config.sliding_window (:396-398). Fine for Gemma-3, but a model with per-layer window sizes would be mis-published silently. Latent gap, not a bug for the tested models — the layout is already per-cache, so it's a small step to source per-layer if that ever lands.

5. get_n_caches indexing is implicit (:400)."get_n_caches": len(layer_types) — the comment at :391-392 already gestures at "these are indexed by cache," but making the invariant explicit at the assignment (that layer_types is the truncated donor list, not the model's layer count) would save the next reader a cross-file hop to confirm get_n_caches/get_windows/_cache_id share the donor-indexed space.

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:530) is correct for the 0-dim scalar and the comment explains the gemma-4 fallback. Fine as-is — and it correctly touches only the in-graph path, not the off-graph one.

Static review only — I could not run the export end-to-end (Python execution is gated in this environment and the HF checkpoints aren't available). Item 1 is the one I'd verify before merge; 2–6 are advisory.
· branch mlx-offgraph-hf-export

[metascroy]

Is it better to have **kwargs catch everything else (including softcap), and then raise if kwargs is non-empty? That way more unknown kwargs don't sneak in

[kiymetakdemir]

It makes sense for catching unknown ones, I added a sweep that pops the args we handle elsewhere (is_causal, use_cache, sliding_window, dropout) and raises on anything left. But softcap and head_mask still need to be named, models pass them often as None which still would raise.

[kiymetakdemir]

For claude review;

1. Dim's default min is 0, not 2 (torch/export/dynamic_shapes.py:180), so the constraint is in range [0, max_seq_len-1] and q_len == 1 is verified in practice, since the runner decodes one token per step.
2. Added assert input_ids.shape[0] == 1 in the export wrapper.
3. I'll add MaskKind.EXPLICIT in a separate PR.
4. The value is uniform because transformers exposes sliding_window as a scalar on the text config; layer_types only says which layers are sliding, so there's no per-layer window to read today.
5. Added a comment at the assignment stating that get_n_caches is the cache count rather than num_hidden_layers.
6. Extracting it would mean touching the other export paths, which are out of scope for this PR.

[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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8d0517b with merge base 5b3da18 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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:

assertsoftcapisNone, "update_and_attend does not support attention logit softcapping"asserthead_maskisNone, "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

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 23s —— View job

---

Code Review — Off-graph KV cache export mode

Re-read the full diff (hf_attention.py, export_llm_hf.py, cache.py) plus the op (extension/llm/cache/update_and_attend.py) and resolve_hf_cache_layout. This commit already folds in the earlier two rounds — softcap/head_maskraise on both attention paths, scaling is asserted (no head_dim ** -0.5 fallback), and the module docstring now covers both "mlx" and "mlx_offgraph". The tracer contract (mutates_args=(), layer_id/scale/out_dtype as node constants, cache off-graph) lines up cleanly with what _export_with_offgraph_cache publishes, and _cache_id's donor logic matches 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, get_windows, get_kv_heads, get_head_dims all share). Solid.

One item I'd confirm before merge (still open from the last round), the rest advisory.

Worth confirming before merge

1. Does seq_length_dim admit q_len == 1 at decode? (export_llm_hf.py:415) The off-graph export declares torch.export.Dim("seq_length_dim", max=max_seq_len - 1) with example seq_length = 3. A bare torch.export.Dim has historically carried an implicit min of 2 — if that holds in the pinned torch (release/2.13), the .pte bakes in a seq >= 2 guard and a single-token decode step (cache_position of length 1) would be rejected at runtime. The runner drives one token per decode step, so please confirm the traced program accepts q_len == 1 — either export with Dim(..., min=1) or verify the guard doesn't reject 1. I couldn't run the export here (sandbox blocks Python + HF checkpoint download) to check empirically. Note the same Dim(max=... - 1) shape appears in _export_with_custom_components (:269), so whatever the resolution, keeping both consistent — plus a one-line comment on why max_seq_len - 1 rather than max_seq_len — would help. Check this →

Advisory (not blocking)

2. B == 1 is assumed but never asserted.OffGraphExportWrapper.forward does position_ids=cache_position.unsqueeze(0) (hf_attention.py:234) and mlx_offgraph_attention_forward does position_ids[0].reshape(-1, 1) (:156) — both collapse the batch dim to index 0. A B > 1 input would trace and then silently place all batches at batch-0's positions rather than failing. A one-line assert input_ids.shape[0] == 1 in the wrapper would match the fail-loud spirit of the position_ids/scaling asserts right below it. The op's own contract is [q_len, n_dims] positions, so single-sequence is the intended scope — just make it explicit.

3. Sliding-window masking is published but not enforced by the eager oracle.get_windows is emitted for the runtime (export_llm_hf.py:396-398), but attend/ContiguousReferenceCache in reference_cache.py only distinguish causal vs none — no window truncation. So the eager op can't serve as ground truth for Gemma-3's sliding layers past sliding_window tokens; it'll attend to full history and diverge from the MLX runtime. Not introduced here — worth a note in the PR description or a follow-up so numerics validation doesn't get misread.

4. get_windows assumes one uniform window size.cache_windows fills every sliding_attention layer with the single text_config.sliding_window (:396-398). Fine for Gemma-3, but a model with per-layer window sizes would be mis-published silently. Latent gap, not a bug for the tested models — the layout is already per-cache, so it's a small step to source per-layer if that ever lands.

5. get_n_caches indexing is implicit (:400)."get_n_caches": len(layer_types) — the comment at :391-392 already gestures at "these are indexed by cache," but making the invariant explicit at the assignment (that layer_types is the truncated donor list, not the model's layer count) would save the next reader a cross-file hop to confirm get_n_caches/get_windows/_cache_id share the donor-indexed space.

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:530) is correct for the 0-dim scalar and the comment explains the gemma-4 fallback. Fine as-is — and it correctly touches only the in-graph path, not the off-graph one.

Static review only — I could not run the export end-to-end (Python execution is gated in this environment and the HF checkpoints aren't available). Item 1 is the one I'd verify before merge; 2–6 are advisory.
· branch mlx-offgraph-hf-export

[metascroy]

Is it better to have **kwargs catch everything else (including softcap), and then raise if kwargs is non-empty? That way more unknown kwargs don't sneak in

[kiymetakdemir]

It makes sense for catching unknown ones, I added a sweep that pops the args we handle elsewhere (is_causal, use_cache, sliding_window, dropout) and raises on anything left. But softcap and head_mask still need to be named, models pass them often as None which still would raise.

[kiymetakdemir]

For claude review;

1. Dim's default min is 0, not 2 (torch/export/dynamic_shapes.py:180), so the constraint is in range [0, max_seq_len-1] and q_len == 1 is verified in practice, since the runner decodes one token per step.
2. Added assert input_ids.shape[0] == 1 in the export wrapper.
3. I'll add MaskKind.EXPLICIT in a separate PR.
4. The value is uniform because transformers exposes sliding_window as a scalar on the text config; layer_types only says which layers are sliding, so there's no per-layer window to read today.
5. Added a comment at the assignment stating that get_n_caches is the cache count rather than num_hidden_layers.
6. Extracting it would mean touching the other export paths, which are out of scope for this PR.

[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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8d0517b with merge base 5b3da18 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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:

assertsoftcapisNone, "update_and_attend does not support attention logit softcapping"asserthead_maskisNone, "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

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 23s —— View job

---

Code Review — Off-graph KV cache export mode

Re-read the full diff (hf_attention.py, export_llm_hf.py, cache.py) plus the op (extension/llm/cache/update_and_attend.py) and resolve_hf_cache_layout. This commit already folds in the earlier two rounds — softcap/head_maskraise on both attention paths, scaling is asserted (no head_dim ** -0.5 fallback), and the module docstring now covers both "mlx" and "mlx_offgraph". The tracer contract (mutates_args=(), layer_id/scale/out_dtype as node constants, cache off-graph) lines up cleanly with what _export_with_offgraph_cache publishes, and _cache_id's donor logic matches 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, get_windows, get_kv_heads, get_head_dims all share). Solid.

One item I'd confirm before merge (still open from the last round), the rest advisory.

Worth confirming before merge

1. Does seq_length_dim admit q_len == 1 at decode? (export_llm_hf.py:415) The off-graph export declares torch.export.Dim("seq_length_dim", max=max_seq_len - 1) with example seq_length = 3. A bare torch.export.Dim has historically carried an implicit min of 2 — if that holds in the pinned torch (release/2.13), the .pte bakes in a seq >= 2 guard and a single-token decode step (cache_position of length 1) would be rejected at runtime. The runner drives one token per decode step, so please confirm the traced program accepts q_len == 1 — either export with Dim(..., min=1) or verify the guard doesn't reject 1. I couldn't run the export here (sandbox blocks Python + HF checkpoint download) to check empirically. Note the same Dim(max=... - 1) shape appears in _export_with_custom_components (:269), so whatever the resolution, keeping both consistent — plus a one-line comment on why max_seq_len - 1 rather than max_seq_len — would help. Check this →

Advisory (not blocking)

2. B == 1 is assumed but never asserted.OffGraphExportWrapper.forward does position_ids=cache_position.unsqueeze(0) (hf_attention.py:234) and mlx_offgraph_attention_forward does position_ids[0].reshape(-1, 1) (:156) — both collapse the batch dim to index 0. A B > 1 input would trace and then silently place all batches at batch-0's positions rather than failing. A one-line assert input_ids.shape[0] == 1 in the wrapper would match the fail-loud spirit of the position_ids/scaling asserts right below it. The op's own contract is [q_len, n_dims] positions, so single-sequence is the intended scope — just make it explicit.

3. Sliding-window masking is published but not enforced by the eager oracle.get_windows is emitted for the runtime (export_llm_hf.py:396-398), but attend/ContiguousReferenceCache in reference_cache.py only distinguish causal vs none — no window truncation. So the eager op can't serve as ground truth for Gemma-3's sliding layers past sliding_window tokens; it'll attend to full history and diverge from the MLX runtime. Not introduced here — worth a note in the PR description or a follow-up so numerics validation doesn't get misread.

4. get_windows assumes one uniform window size.cache_windows fills every sliding_attention layer with the single text_config.sliding_window (:396-398). Fine for Gemma-3, but a model with per-layer window sizes would be mis-published silently. Latent gap, not a bug for the tested models — the layout is already per-cache, so it's a small step to source per-layer if that ever lands.

5. get_n_caches indexing is implicit (:400)."get_n_caches": len(layer_types) — the comment at :391-392 already gestures at "these are indexed by cache," but making the invariant explicit at the assignment (that layer_types is the truncated donor list, not the model's layer count) would save the next reader a cross-file hop to confirm get_n_caches/get_windows/_cache_id share the donor-indexed space.

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:530) is correct for the 0-dim scalar and the comment explains the gemma-4 fallback. Fine as-is — and it correctly touches only the in-graph path, not the off-graph one.

Static review only — I could not run the export end-to-end (Python execution is gated in this environment and the HF checkpoints aren't available). Item 1 is the one I'd verify before merge; 2–6 are advisory.
· branch mlx-offgraph-hf-export

[metascroy]

Is it better to have **kwargs catch everything else (including softcap), and then raise if kwargs is non-empty? That way more unknown kwargs don't sneak in

[kiymetakdemir]

It makes sense for catching unknown ones, I added a sweep that pops the args we handle elsewhere (is_causal, use_cache, sliding_window, dropout) and raises on anything left. But softcap and head_mask still need to be named, models pass them often as None which still would raise.

[kiymetakdemir]

For claude review;

1. Dim's default min is 0, not 2 (torch/export/dynamic_shapes.py:180), so the constraint is in range [0, max_seq_len-1] and q_len == 1 is verified in practice, since the runner decodes one token per step.
2. Added assert input_ids.shape[0] == 1 in the export wrapper.
3. I'll add MaskKind.EXPLICIT in a separate PR.
4. The value is uniform because transformers exposes sliding_window as a scalar on the text config; layer_types only says which layers are sliding, so there's no per-layer window to read today.
5. Added a comment at the assignment stating that get_n_caches is the cache count rather than num_hidden_layers.
6. Extracting it would mean touching the other export paths, which are out of scope for this PR.