[ExecuTorch][MLX] lowering memory/export improvements

[metascroy]

Muse Glimmer dflash + image (16 GB Q4_K target + 1.5 GB draft + 1.3 GB mmproj):

	before	after
export time	4.2 min	2.8 min (−32%)
peak memory	66.6 GB	48.3 GB (−27%)

What changed

1. Support checks no longer repack weights. The partitioner decided "can MLX
   lower this?" by running each handler, which for quantized ops means fully
   dequantizing and repacking the weight. exir asks several times per export, so
   every weight was repacked 4x. Handlers can now answer from node metadata via an
   optional supported() / register_support_check; anything without one keeps
   the old run-and-catch, so nothing regresses.

2. Raw blobs are freed as they're repacked. repack_mlx left the original
   GGUF bytes in the ExportedProgram until a sweep that only ran after the whole
   graph was built, so raw and repacked weights were live simultaneously (+16 GB).
   It now releases each blob on the spot, memoized per weight so tied
   embedding/lm_head repacks once.

3. Two guards, since a wrongly-released weight used to fail silently:
   serialization now raises naming the constant instead of emitting a .pte with
   a missing weight, and release is a no-op during support checks, where it would
   corrupt the caller's own program.

4. ET_MLX_MEM_PROFILE=1 adds phase-attributed peak-footprint logging (off by
   default); it's what located the regression above.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 4 Pending, 1 Unrelated Failure, 1 Unclassified Failure

As of commit 3e84925 with merge base 3b5bf0f ():

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

• MLX / test-mlx-llm (unsloth/Llama-3.2-1B-Instruct, llama-1b, true, nvfp4, macos-14-xlarge) / test-mlx-llm-llama-1b-custom-nvfp4 (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  ##[error]fatal: unable to access 'https://github.com/mreineck/pocketfft/': Could not resolve host: github.com

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) (detected as infra flaky with no log or failing log classifier)

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 4m 53s —— View job

---

Code review

• Read changed files
• Analyze correctness & design
• Post review

I read the full diff across all 12 files. This is a careful, well-documented change and I couldn't find a correctness bug that would produce a wrong .pte. The design is sound: the two guards (serialization raises on missing data; release is a no-op under _check_only) turn what used to be silent corruption into loud failures. Notes below, ordered by significance.

---

Design: the supported() / __call__ agreement is the load-bearing invariant

The whole speedup rests on supported() (and register_support_check) agreeing with what __call__ would decide — and that agreement is checked only indirectly (op tests + "false positive fails export, false negative → missing segment", per the docstring at program_builder.py:57). That's a reasonable call given the cost being avoided, and the mirror helpers are conscientious:

• mlx_qparams_supported (op_helpers.py:537) explicitly mirrors the to_mlx_qparams / regroup_affine_scales asserts.
• _blob_lowers (gguf/patterns.py:62) is deliberately shape-only so it holds for both the fused and MLX-native paths.
• _gather_qmm_supported (ops.py:1866) branches on uint8 vs int8 exactly as the handler's prepacked/non-prepacked paths do.

One thing worth confirming: _affine_qparams_lower returns False when metadata is missing or dynamic (qdata is None, non-int shapes). In __call__ the same weight is read concretely and would lower. So a fully-dynamic-but-lowerable quantized weight would now be dropped from ops_to_not_decompose where it previously (via run-and-catch) succeeded. In practice quantized weights are static constants, so this is almost certainly a non-issue — but it's the one place supported() is strictly more conservative than __call__, and it's silent (missing segment) rather than loud. Worth a sanity check that no supported model has a symbolic quantized-weight dim.

Minor: Q6_K re-unpacks twice for a tied small-group weight

In q6k/repack_mlx.py, when group_size < _MIN_MLX_GROUP_SIZE the function returns None before populating repack_cache. That's correct (the fused fallback needs the raw bytes, so it must not be released or cached). But it means a tied Q6_K weight whose sub-blocks don't merge to ≥32 runs the expensive from_raw(...).to_intx_unpacked_to_int8_tensor(...) on every consumer instead of once. The q4_k/q5_k paths don't have this because they always cache. Edge case (tied + small group + Q6_K), and it only costs time, not memory or correctness — flagging only because the memoization is otherwise thorough.

Nit: release_placeholder_tensor docstring vs. the guard it describes

The program_builder.py:312 docstring is excellent and explains why release is skipped under _check_only (state dict still needed by _unlift_exported_program_lifted_states). Just confirming the read: because _check_only is set in _process_nodes and never reset, any second build()-style pass on the same builder would carry the last value. In practice check_support_only() and build() are called on separate builder instances (partitioner.py constructs a fresh MLXProgramBuilder each time, and build() sets check_only=False), so this is fine — but the flag is process-wide mutable state on the builder rather than scoped to the pass, so keep that construction discipline.

_memprofile.py

Clean and correctly inert off-macOS: enabled() short-circuits on the env var and caches _self_check(), which fails gracefully via the (AttributeError, OSError, ValueError) catch when proc_pid_rusage is absent. The ri_proc_exit_abstime == 0 anchor for validating the hardcoded struct offsets is a nice touch. Since preprocess.py and partitioner.py now import it unconditionally, the module must stay import-clean on all platforms — it is (only stdlib ctypes), so no CI risk. if raised: is safe since peak footprint is monotonic (delta ≥ 0).

Behavior preservation I checked

• In _process_nodes, pattern body nodes (e.g. the dequantize_gguf node) still get marked supported via the deferred_handler no-op in both the old run-and-catch path and the new support-check path, so partition membership is unchanged when a head is supported. When a head is unsupported, body-supported/head-unsupported matches the old run-and-throw outcome too — no new partitioning regression.
• The mark_unsupported local (creating slots when meta['val'] exists) faithfully consolidates the three previously-duplicated blocks.
• _not_decompose_cache keys on weakref.ref(ep) is ep identity with the filter callback always None in both call paths — safe.

Overall: LGTM. The only item I'd actively verify before merge is the symbolic-shape edge in _affine_qparams_lower (first point); everything else is nits.
· mlx-mem-improvement

[Gasoonjia]

i feel like the preprocessing part for q4k/q5k/q6k are pretty much the same. Can we unify the logic by using utliy function or a common parent class?

[Gasoonjia]

will the memory block here impact the exportation speed in real world production?