Updates to support HF dynamic cache export

[metascroy]

As titled.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• linux.dgx.b200 runner queuing exceeds 4 hours and dgxb200-05 "CUDA not available"

❌ 4 New Failures, 2 Unrelated Failures

As of commit 683d734 with merge base 9b91b43 ():

NEW FAILURES - The following jobs have failed:

• MLX / test-mlx-llm (unsloth/Qwen3-0.6B, qwen3-0.6b, false, 4w, macos-14-xlarge) / test-mlx-llm-qwen3-0.6b-4w (gh)
  ##[error]fatal: Fetched in submodule path 'third-party/pybind11', but it did not contain d03662f0984f652b60e7ddce53d3868002275197. Direct fetching of that commit failed.
• MLX / test-mlx-muse-glimmer-build / test-mlx-muse-glimmer-build (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• MLX / test-mlx-parakeet / test-mlx-parakeet (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• MLX / test-mlx-stories110m / test-mlx-stories110m (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Apple / build-benchmark-app / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• Apple / build-frameworks-ios / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1

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

[metascroy]

@claude review this code

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

[claude]

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

---

Code Review — Support dynamic indexing

• Read the diff (op_helpers.py, ops.py, test_ops.py)
• Analyze _index_handler gather/reshape logic and emit_shape offset
• Post review

---

Nice change. The core idea is sound: instead of specializing the post-gather reshape target through int(SymInt) (which would bake in a guard and break dynamic shapes), you now emit the output shape at runtime via emit_shape, and only the broadcast index dims — which are the only ones that can be dynamic — need SymSize reads. The dim_offset addition to emit_shape is a clean, backward-compatible extension.

I walked the offset mapping through all three test configs and it holds:

config	axes	leading_dims	dynamic out dim → gather axis
leading=False	[0,1]	0	R@0→0, C@1→1 ✓
trailing, leading=False	[0,1]	0	R@0→0, C@1→1 ✓
trailing, leading=True	[1,2]	1	R@1→0, C@2→1 ✓

Why it's safe: non-indexed dims are forced static by the slice_sizes loop (ops.py:2477-2481), so every dynamic output dim is a broadcast index dim. In the contiguous case those broadcast dims start at output position axes[0] and at gather position 0, so dim_offset=-leading_dims maps them correctly; in the non-contiguous case PyTorch moves broadcast dims to the front, leading_dims=0, and the offset is a no-op. That same static-guarantee is also what prevents a negative dim in SymSizeNode (a dynamic dim at output index < leading_dims would need a leading non-indexed dim to be dynamic, which the loop already rejects). Worth a one-line comment noting this invariant, since it's load-bearing but implicit.

Main thing to flag: reshape data-ordering for leading non-indexed dims > 1

The leading_dims machinery correctly handles the symbolic sizes, but the underlying gather→reshape only preserves correct element ordering when leading non-indexed dims are singletons (or absent). For contiguous indexed axes with a leading non-indexed dim of size A > 1:

• gather output (row-major) is [broadcast…, A, 1…, trailing]
• aten expects [A, broadcast…, trailing]

A reshape reinterprets the flat buffer in place — it can't move A past the broadcast dims; that needs a transpose. The new leading=True test only exercises A = 1 (shape = (1,) + …), where a singleton doesn't affect ordering, so this gap is masked.

This looks pre-existing (the prior int()-based path used the same ReshapeNode), so it's not a regression from this PR — but the new offset logic makes leading non-indexed dims look supported. I'd suggest either:

1. adding a test with a leading non-indexed dim > 1 (e.g. x shape (2, 6, 7), x[:, rows, columns]) to confirm/expose the behavior, or
2. an explicit guard/assert that leading non-indexed dims are singleton in the contiguous-with-leading_dims>0 path.

If I've misread the MLX gather output ordering (indices.shape + slice_sizes), disregard — but a A>1 test would settle it either way.

Minor

• ops.py:2506 — the contiguity check axes == list(range(axes[0], axes[-1] + 1)) is correct (axes is built ascending from enumerate); a short inline comment ("contiguous indexed axes ⇒ broadcast dims stay in place") would help the next reader parse the ternary.
• Test coverage is otherwise good — dynamic rows/columns, runtime shapes differing from export-time inputs, and the min/max Dim bounds all look right.

Overall: correct and well-documented for the supported cases; the only open question is the leading-dim > 1 ordering.
· branch hf-export

[metascroy]

@claude review this code

[claude]

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

---

Code Review — HF dynamic cache export support

• Read the full current diff (6 files)
• Review op_helpers.py + ops.py (dynamic indexing)
• Review passes.py (CollapseDtypeConversionPass)
• Review MLXInterpreter.h (SDPA masked rows)
• Review new tests

---

The PR grew since the last review — it now bundles three logically distinct changes. All three look correct and are well-tested. Notes below.

1. ops.py — dynamic multi-index + leading-dim transpose ✅

This resolves the leading-non-indexed-dim ordering gap I flagged earlier. The new TransposeNode reorders the gather result before reshape, and I walked the permutation through:

• gather layout: [broadcast(bn), slice(x_ndim)], where the slice block is [leading_nonidx, indexed(1s), trailing_nonidx]
• ATen contiguous layout: [leading_nonidx, broadcast, trailing_nonidx]
• perm range(bn, bn+ld) + range(bn) + range(bn+ld, bn+x_ndim) maps exactly that. ✓

The reshape-after-transpose is order-preserving because the stripped indexed dims are size-1, so flattening [leading, broadcast, 1s, trailing] equals [leading, broadcast, trailing]. The guard (slice_sizes[:leading_dims] has a non-singleton and the broadcast block is dynamic-or->1) correctly keeps the singleton cases reshape-only. AdvancedIndexLayoutTest covers contiguous/non-contiguous, singleton, scalar, and separated-axes configs with exact TransposeNode counts — nice.

One thing worth an explicit note: out_shape reads SymSize from gather_slot (the pre-transpose slot) via dim_offset=-leading_dims, while the reshape input is reshape_slot (post-transpose). This is correct — out_shape is just a size list in ATen order, and the offset maps each dynamic (broadcast) output index back to its original gather axis — but it's subtle enough that a one-line comment at ops.py:2508 ("sizes read from the pre-transpose slot; offset maps broadcast dims back to gather axes") would help. The negative-dim invariant in SymSizeNode still holds because the slice_sizes loop forces every non-indexed dim static, so dynamic dims only ever land at output index >= leading_dims.

2. passes.py — CollapseDtypeConversionPass lossy-cast guard ✅

Good correctness fix. The pass previously collapsed any pair of pure dtype casts, which drops an intermediate narrowing/truncating cast (e.g. fp32 → int32 → fp32 or fp32 → fp16 → fp32) and silently changes results. The new whitelist only collapses when the intermediate cast losslessly represents the source dtype, which is exactly the condition needed: if x.to(d1) == x value-for-value, then x.to(d1).to(d2) == x.to(d2). The final dtype d2 correctly doesn't matter.

The whitelist is conservative (e.g. int32 → fp64 is lossless but not listed, so it's skipped) — that's the safe direction, just flagging it's an intentional missed-optimization rather than a completeness bug. The added source_val is None and len(parent.users) == 1 guards, plus the test_multi_user_parent_not_collapsed / test_missing_source_metadata_not_collapsed / test_non_pure_cast_not_collapsed tests, cover the edges well.

3. MLXInterpreter.h — SDPA zero-row post-processing ✅

exec_sdpa now zeroes fully-masked query rows (bool mask all-False, or additive mask all -inf) to match PyTorch's fused-kernel zero-row semantics. Using where(has_key, out, 0) rather than a multiply is the right call — it avoids NaN * 0 = NaN from a row that softmaxed to NaN. Verified:

• partial rows (allowed[2, 1:] = False) keep has_key = true and are left to the kernel ✓
• the "finite" mask kind (finfo.min, not -inf) is not zeroed, since not_equal(M, -inf) stays true — this correctly avoids clobbering finite biases ✓
• has_key shape [..., seq_q, 1] broadcasts against out [b, h, seq_q, head_dim] for both shared and per-batch masks ✓

Two minor points:

• This adds not_equal + any + where to every masked SDPA, including plain causal-tensor masks that never have empty rows. Correctness need justifies it, but if masked SDPA is hot it may be worth gating (e.g. skip when the mask is statically known dense). Not blocking.
• The zero-row behavior matches PyTorch's fused (flash/mem-efficient) kernels; the eager math path returns NaN for fully-masked rows in some versions. Since SDPAMaskedRowsTest asserts against eager PyTorch with atol=0, CI green confirms the reference here returns 0 — just be aware the "correct" reference is backend/version-dependent.

Summary

Correct and thoroughly tested across all three changes; the earlier leading-dim ordering concern is fully resolved. Only non-blocking nits: a clarifying comment on the pre-transpose emit_shape read Fix this →, and the always-on SDPA mask post-processing overhead.
· branch hf-export