[MLX] Add quant regrouping to MLX lowering flow, and update quantize_and_save in Qwen to support MLX path

[metascroy]

For the MLX repacking and quantize_and_save work, we removed the redundant hand-rolled MLX regrouping/packing code in both gemma4 and qwen3.5 (deleting the old pack_mlx/export.py regroup paths and their tests).

We also added MLX path in qwen's quantize_and_save.py

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Merge Blocking SEVs

There is 1 active merge blocking SEVs. Please view them below:

• (merge blocking) [Maintenance] - Due to #190347 we're advancing a maintenance planned to Sunday to today

If you must merge, use @pytorchbot merge -f.

❌ 65 Pending, 1 Unclassified Failure

As of commit e0f870e with merge base 931ba24 ():

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/gemma-3-1b-it, gemma3-1b, true, 4w, macos-14-xlarge) / test-mlx-llm-gemma3-1b-custom-4w (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  The process '/opt/homebrew/bin/git' failed with exit code 128

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 10m 16s —— View job

---

Code Review

Reviewed the full diff against origin/main. This is a solid, well-documented refactor. The centralization of group_size regrouping into the MLX pattern handlers (regroup_affine_scales / mlx_affine_group_size) is a clean improvement over duplicating it in each model's pack_mlx.py, and the end-to-end nibble-packed int4 path (no int8 expansion during lowering) is a nice memory win. Comments and docstrings are unusually thorough — appreciated.

Findings below, roughly by importance.

1. _int4_packed is dead state

backends/mlx/llm/switch.py:83,120,155 — the attribute is initialized and set to True in both quantized branches but is never read anywhere in the repo (confirmed via grep -rn _int4_packed). The runtime dispatch in forward() keys off self._is_quantized + qdata.dtype, and the lowering keys off w_data.dtype == torch.uint8, so this flag is redundant. Either remove it, or wire it into pack()/forward() if it was meant to be the packing discriminator. Per the repo's "no unused state" style, I'd drop it. Fix this →

2. Verify the CUDA-bundle format change is backward-compatible

quantize_and_save.py replaces the hand-rolled _SUBCLASS_REGISTRY + __tensor_flatten__ header scheme with torchao's flatten_tensor_state_dict / unflatten_tensor_state_dict. This changes the on-disk safetensors layout for both backends, including the existing CUDA path (Int4TilePackedTo4dTensor). Two things worth confirming before merge:

• torchao's helpers actually support Int4TilePackedTo4dTensor (the CUDA dense-layer subclass), not just Int4Tensor / IntxUnpackedToInt8Tensor. If it silently falls back or errors on that type, the CUDA path regresses.
• Previously-saved CUDA bundles can still be loaded, or it's acceptable that they can't. The old load_quantized_state_dict read a "quantization" JSON header; the new one reads torchao's metadata. Old bundles will fail to reconstruct. If any are in circulation this is a breaking change worth a note in the PR description / README.

Since I can't execute torchao here, please confirm a CUDA round-trip (quantize_and_save.py default backend → export.py --prequantized) still works.

3. group_size=16 behavior change — confirm the fallback still lowers

The env-gated non-fused path (ET_MLX_ALLOW_NON_FUSED_QUANTIZED_OPS) is removed, and QuantizedLinearHandler / Int4QuantizedLinearHandler now return None when group_size < 32. That's reasonable (16 can't be regrouped coarser). The intent is that a group_size=16 linear falls through to the generic _dequantize_affine_handler + matmul. Worth a sentence in the PR confirming that path is actually reached — i.e. an unmatched fused-linear pattern doesn't leave the dequantize_affine node unlowered. The removed test_ops.py cases (group_size=16) previously covered this; there's now no test exercising a 16-group linear end-to-end on the MLX backend, so a regression here would be silent.

4. _load_remapped_subset scans the full weight map per layer

quantize_and_save.py — stream_quantize_and_save_mlx calls _load_remapped_subset once per decoder layer, and each call iterates the entire weight_map (for ckpt_key, shard in weight_map.items() + predicate). That's O(layers × total_keys). For a 40–60 layer MoE with many per-expert keys this is quadratic. It's a one-time offline script so it's not blocking, but a prefix bucketing (group keys by layer index once up front) would avoid the rescan if you find it slow in practice.

5. Minor

• save_quantized_tensors silently skips duplicate keys (if key in state_dict: continue). With tied lm_head/embed_tokens this relies on iteration order (params before buffers); fine as-is, but a one-line comment on why first-wins is intended would help.
• _open_checkpoint_shards caches safe_open handles and never closes them. Acceptable for a short-lived script, just noting.
• The prepacked assertion message in to_mlx_qparams ("packed cols must be a multiple of 4") is good; the accompanying comment (gs >= 32) is the right invariant.

Things that look correct

• The scale-layout reconciliation in switch.py (.t() on ExportableInt4Tensor scale/zero_point since it's [in//gs, out], vs. no transpose on IntxUnpackedToInt8Tensor which is already [out, in//gs]) is subtle and looks right — both branches produce [E, out, in//gs] for gather_qmm.
• The removed AddmmNode / TransposeNode / DequantizeNode imports in patterns.py are only referenced in docstrings now; the handlers use emit_quantized_gather, so removal is safe.
• regroup_affine_scales numerical equivalence is well-covered by test_quant_regroup.py::test_regroup_preserves_dequant.
• All symbols the new streaming/loading code imports from model.py and mlx_source_transformations.py exist (Block, GemmaRMSNorm, _process_checkpoint_key, _fuse_projection_weights, from_hf_config, _swap_rms_norm → _rms_weight).

Overall: good to merge once #2 (CUDA round-trip) and #3 (group_size=16 fallback) are confirmed, and #1 is cleaned up.
• mlx-quant-improvements

[Gasoonjia]

why repeat < 1 is legal?

[metascroy]

target_group_size >= old_group_size, so repeat can only be 1 in this case. A bit confusing, though. I'll fix it

[metascroy]

@claude review this code

[claude]

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

---

Code Review

Reviewed the current diff (2 commits) against the base. This is a clean, well-documented refactor — centralizing group_size regrouping into the MLX pattern handlers (regroup_affine_scales / mlx_affine_group_size) instead of duplicating it in each model's pack_mlx.py is a solid improvement, and the end-to-end nibble-packed int4 path (no int8 expansion during lowering) is a real memory win. The docstrings and comments are thorough.

A note on state since the last review: the _int4_packed dead-state finding is resolved — that attribute is gone from switch.py, and save_quantized_tensors now carries the first-wins dedup comment. Nice.

Findings below, roughly by importance.

1. group_size < 32 int4 linear/embedding has no fallback lowering

patterns.py:1204 (Int4QuantizedLinearHandler) and the affine handler at patterns.py:877 both return None when group_size < 32. For affine (dequantize_affine) that's safe: there's a standalone _dequantize_affine_handler (ops.py:4396) that catches the unfused node and now regroups internally (ops.py:4448), so a group_size=16 affine linear still lowers via generic dequant + matmul.

For int4 (dequantize_int4_tensor) there is no standalone handler — I grep'd and only the two fused pattern handlers reference it. So a dequantize_int4_tensor linear/embedding with an MLX-legal-but-<32 group (e.g. an original group_size of 16, which mlx_affine_group_size(16) returns as 16) would fail to match the fused handler and then hit no fallback → an unlowered node → export error, not a graceful slow path. In practice torchao Int4Tensor groups are ≥32 so this is an edge case, but the asymmetry with the affine path is worth either a one-line note or a guard. Fix this →

2. regroup_affine_scales repeat naming (Gasoonjia's question)

Answering @Gasoonjia's inline why is repeat < 1 legal? — in the current code (op_helpers.py:730-744) it can't be. The assert old_group_size >= target_group_size and old_group_size % target_group_size == 0 guarantees repeat = old_group_size // target_group_size >= 1, and if repeat == 1: return ... False handles the no-op. So repeat is always ≥ 1 and the ==1 early-out is exact. @metascroy already said he'd clean up the confusing phrasing — 👍, the logic itself is correct.

3. CUDA-bundle format change — confirm backward-compat / round-trip

quantize_and_save.py replaces the hand-rolled _SUBCLASS_REGISTRY + __tensor_flatten__ header scheme with torchao's flatten_tensor_state_dict / unflatten_tensor_state_dict for saving, and keeps a legacy reader for old bundles (_load_legacy_state_dict). Two things to confirm before merge:

• The modern saver (flatten_tensor_state_dict) handles Int4TilePackedTo4dTensor (the CUDA dense subclass), not just Int4Tensor / IntxUnpackedToInt8Tensor. The new CPU tests (test_quantize_roundtrip.py::TestSerializerFormatDispatch) only exercise Int4Tensor (min/max, CPU), so the CUDA subclass round-trip is still untested here. A quick CUDA quantize_and_save.py (default backend) → export.py --prequantized confirmation would close this.
• The saved on-disk layout for the CUDA path changed (torchao format, not the old quantization JSON header). Old bundles still load via the legacy reader, so this is read-compatible — just worth a line in the PR description noting new CUDA bundles are written in the new format.

4. _load_remapped_subset scans the full weight map per layer

quantize_and_save.py:250 — stream_quantize_and_save_mlx calls _load_remapped_subset once per decoder layer, and each call iterates the entire weight_map with a startswith(f"model.layers.{i}.") predicate → O(layers × total_keys). For a 40–60 layer MoE with many per-expert keys this is quadratic. It's a one-time offline script so not blocking, but bucketing keys by layer index once up front would avoid the rescan if it's slow in practice. Fix this →

5. Minor

• _open_checkpoint_shards caches safe_open handles and never closes them. Fine for a short-lived script; just noting.
• load_prequantized_model_mlx (export.py) infers max_seq_len by scanning for a .kv_cache.k_cache buffer and raising if absent — good defensive check, clear error message.
• The removed group_size=16 cases in test_ops.py (and the deleted ET_MLX_ALLOW_NON_FUSED_QUANTIZED_OPS env-gating) mean there's now no end-to-end test exercising a 16-group linear on the MLX backend. The generic _dequantize_affine_handler fallback for that case (item Add support for quantized LeakyReLU #1) is therefore only covered by the unit-level test_quant_regroup.py, not end-to-end. Consider keeping one group_size=16 affine linear in GatherQmm/QuantizedLinear e2e tests so a regression in the fallback isn't silent.

Things that look correct

• The scale-layout reconciliation in switch.py::pack — .t() on ExportableInt4Tensor scale/zero_point (it's [in//gs, out]) vs. no transpose on IntxUnpackedToInt8Tensor (already [out, in//gs]) — is subtle and consistent; both branches yield [E, out, in//gs] for gather_qmm.
• to_mlx_qparams(prepacked=True) fast path: the qdata.shape[-1] % 4 == 0 assert and view(torch.uint32) reinterpret are correct given in_features % 8 == 0 (holds for gs ≥ 32), and the _gather_qmm_handler dtype sniff (w_data.dtype == torch.uint8) correctly routes packed vs. unpacked experts.
• regroup_affine_scales numerical equivalence is well-covered by test_quant_regroup.py::test_regroup_preserves_dequant, and the mlx_affine_group_size divisor logic is covered by TestMlxAffineGroupSize.
• The pack_mlx.py simplification (drop _mlx_group_size, defer regrouping to export) is consistent with the updated test_pack_mlx.py expectations (block_size now passes through untouched).

Overall: good to merge once #1 (int4 <32 fallback — note or guard) is addressed and #3 (CUDA round-trip) is confirmed. #4/#5 are non-blocking polish.
• mlx-quant-improvements