MLX Q5K/Q6K upgrades

[metascroy]

This PR adds an MLX-native repack lowering path for GGUF Q5_K and Q6_K weights, mirroring the existing Q4_K path: at export time the raw GGUF blob is unpacked and repacked into MLX affine qparams (bits=5 / bits=6) consumed by MLX's native QuantizedMatmul/quantized-gather kernels, gated by ET_MLX_EMIT_DIRECT_GGUF (default = native, falling back to the fused Metal kernels otherwise). It also introduces a lossless group-size upgrade in to_intx_unpacked_to_int8_tensor (_max_uniform_group_size) that merges adjacent sub-blocks with identical scale/min into the largest MLX-supported group size (32→64→128) when bit-exact, which the repack paths request via max_group_size=128; because Q6_K's native group size (16) isn't MLX-compatible, its native path only activates when this merge reaches ≥32 and otherwise transparently falls back to fused kernels. Q6_K's symmetric zero-point lets biases be computed in the init chain via emit_quantized_biases instead of being serialized, and a FakeTensor guard keeps the data-dependent merge from tripping the partitioner's support check. Coverage includes new unit tests for the merge logic in test_gguf.py and native/fused + group-merged op tests across Q4_K/Q5_K/Q6_K in test_linear.py and test_embedding.py.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Unrelated Failure

As of commit b1f3acb with merge base 5f6d1aa ():

NEW FAILURE - The following job has failed:

• pull / test-multimodal-linux (gemma3-4b) / linux-job (gh)
  RuntimeError: Command docker exec -t 6df6b89c29cbef8fd0e03e630c961582ab82d1c9a711e7500913383aa5b85a66 /exec failed with exit code 139

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

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

• pull / unittest-nxp-neutron / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

---

Code Review: MLX Q5K/Q6K upgrades

I read through the repack paths, native kernels, the _max_uniform_group_size merge logic, the dispatch wiring, and the new tests. Overall this is a clean, well-documented extension of the existing Q4_K pattern — the design decisions (symmetric-Q6_K bias-in-init, FakeTensor guard, lossless group merge, group-size gating for Q6_K) are sound and well-tested. A few findings below, none blocking.

Correctness — looks solid

• _max_uniform_group_size (gguf.py:246): the ascending factor search then descending uniformity check is correct. Because factors are powers of two, n_groups % factor_max == 0 implies divisibility for every smaller factor, so the step-down never lands on a non-dividing factor. The "largest uniform run ⇒ sub-runs uniform" early-return is valid. Merging consecutive columns matches the K-contiguous sub-block layout, and taking s[:, :, 0] as representative is bit-exact since all entries are equal.
• Q6_K native gating (q6k/repack_mlx.py:58) returns None before touching P, so a gs < 32 weight leaves no orphan constants and cleanly falls back to fused. Same for emit_linear/emit_embedding returning None.
• The asymmetric-vs-symmetric split is right: Q5_K/Q4_K serialize biases (non-zero zero-point), Q6_K routes through emit_quantized_biases (zero-point 0 ⇒ biases = -scale·2^(bits-1) in the init chain). to_mlx_qparams and emit_quantized_biases agree on the offset.
• int5 centering [0,31] → [-16,15] fits, and to_mlx_qparams bits=5 lands in the generic LSB-packing branch where (K·5)%32==0 holds (K is a multiple of 256).

Minor issues

1. Wasted export-time unpack on the Q6_K fallback. When a Q6_K weight can't merge to gs ≥ 32, repack_mlx (q6k/repack_mlx.py:54) fully unpacks the blob via to_intx_unpacked_to_int8_tensor(max_group_size=128) and only then discovers group_size < 32 and returns None, after which _emit_linear_fused re-does the work. For a large lm_head this unpack is non-trivial. Since the merge decision only needs scale/min uniformity, you could check group size before the full field decode — though this only costs export time, so it's low priority.

2. Stale group size in the Q5_K native docstring. q5k/linear_mlx_native.py:12-13 states "mode 'affine', group_size 32", but the merge can produce 32/64/128 and the code correctly passes the merged group_size. Same fixed "32" phrasing recurs in the module doc; worth softening to "32 by default, larger when sub-blocks merge." Fix this →

3. Terminology drift across the __init__.py docstrings. q4k/__init__.py and q5k/__init__.py call the repack path "legacy", while q6k/__init__.py calls the same mechanism the "MLX-native repack path" (and it's the default, ET_MLX_EMIT_DIRECT_GGUF=0). Calling the default path "legacy" is a bit confusing — consider aligning the wording.

One question

emit_quantized_biases on a shared weight. For a weight referenced by both a linear and an embedding (e.g. tied lm_head/embedding sharing the same GGUF blob), emit_quantized_biases (op_helpers.py:393) is invoked twice with the same zero_point_key. make_or_get_constant dedupes the dummy constant, but P.emit_init(MultiplyNode(...)) appears to run unconditionally each call — is a duplicate init-chain multiply into the same biases slot harmless here, or should this be guarded? Q4_K didn't hit this since it serialized biases directly. Not likely to matter for the current models, but worth confirming.

Tests

Coverage is good — native/fused × Q4_K/Q5_K/Q6_K for both linear and embedding, ragged shapes, dynamic-M IfNode branches, group-merge (uniform=True) fixtures, and the standalone _max_uniform_group_size unit tests including the base-16 (Q6_K) climb and the min-mismatch block. Two small gaps:

• test_to_intx_unpacked_matches_reference (test_gguf.py:126) covers only q4_k/q6_k; adding q5_k there would give the 5-bit affine round-trip a direct assertion outside the MLX kernel tests.
• No unit test asserts Q6_K's repack_mlx returns None for a non-uniform (native gs=16) weight — the fallback is exercised end-to-end but a direct assertion would pin the contract.

I did not run the MLX kernel tests (they require Metal/macOS and aren't runnable on this Linux CI checkout); the review is static.

---

Review authored with Claude.
· branch mlx-q5k

[metascroy]

@claude review this code

[claude]

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

---

Code Review: MLX Q5K/Q6K upgrades

I read the repack paths, the native linear/embedding lowerings, _max_uniform_group_size, to_mlx_qparams / emit_quantized_biases, the dispatch wiring in patterns.py + per-type linear.py/embedding.py, and the new tests. This is a clean, well-documented extension of the Q4_K pattern; the core design decisions all hold up under inspection. Findings below — none blocking.

Correctness — verified

• Group-merge math (gguf.py:246): the ascending power-of-two search then descending uniformity check is correct. For Q6_K (base=16, max=128) the factor climbs 1→2→4→8 giving candidate group sizes {16,32,64,128} — every reachable size is either MLX-invalid 16 (rejected by the < 32 gate) or in MLX's {32,64,128} set, so no un-representable group size can slip through. n_groups = K/16 is a multiple of 16 (K is a multiple of 256), so the divisibility guard never blocks a valid climb. The "largest uniform run ⇒ sub-runs uniform" early return is valid, and s[:, :, 0] is bit-exact since the run is all-equal.
• Bit-packing widths (op_helpers.py:556): the generic LSB branch requires (K·bits)%32==0; with K a multiple of 256 both 256·5=1280 and 256·6=1536 are %32==0. int5 [0,31] (offset 16) and int6 [0,63] (offset 32) fit their widths.
• Symmetric Q6_K bias (op_helpers.py:422): init-chain biases = scale·(−2^(bits−1)) = −32·scale matches to_mlx_qparams's B = −scale·(0+32). Consistent.
• Q6_K fallback is clean (q6k/repack_mlx.py:58): the group_size < 32 check returns None before to_mlx_qparams/make_or_get_constant, so no orphan constants are emitted, and emit_linear/emit_embedding propagate the None to the fused path.

Minor issues

1. Q6_K computes biases twice, discards one. repack_mlx (q6k/repack_mlx.py:61) calls to_mlx_qparams(..., compute_biases=True by default), which materializes B = −32·scale. But since Q6_K is symmetric, emit_quantized_biases (op_helpers.py:418) takes the is_scale_only branch and computes biases in the init chain instead — so the B you just computed is thrown away. Passing compute_biases=False and biases=None would skip the wasted allocation. Low priority (export-time only). Fix this →

2. Confirmed: shared Q6_K weight emits a duplicate init multiply. For a weight referenced by both a linear and an embedding, repack_mlx runs twice with the same zero_point_key. make_or_get_constant dedupes the packed/scales/dummy constants, but emit_quantized_biases (op_helpers.py:427) calls P.emit_init(MultiplyNode(...)) unconditionally, and emit_init (program_builder.py:170) just appends — no dedup. So the init chain gets two identical scale·(−32) → biases multiplies into the same slot. It's idempotent (harmless), just a redundant op. Q4_K/Q5_K don't hit this since they serialize biases as a deduped constant. Worth a guard only if shared GGUF blobs are expected. Fix this →

3. Wasted full unpack on the Q6_K fallback. repack_mlx (q6k/repack_mlx.py:54) fully decodes the blob via to_intx_unpacked_to_int8_tensor(max_group_size=128) before discovering group_size < 32 and returning None, after which _emit_linear_fused redoes the work. The merge decision only needs scale uniformity, so a pre-check on the sub-block scales alone would avoid the full field decode for large lm_heads. Export-time cost only.

4. Docstring / terminology drift.

• q5k/linear_mlx_native.py:13 and q4k's equivalent say "group_size 32", but the merge yields 32/64/128 and the code passes the merged value — soften to "32 by default, larger when sub-blocks merge."
• q4k/__init__.py / q5k/__init__.py describe the default path as the "repack path" while q6k/__init__.py calls it "MLX-native repack path"; patterns.py:47 still comments _LINEAR_TYPES/_EMBEDDING_TYPES as "both via fused custom Metal kernels" even though the default is now the repack path. Aligning the wording (and updating that stale patterns.py comment) would help.

Tests

Strong coverage: native/fused × Q4_K/Q5_K/Q6_K for linear and embedding, dynamic-M IfNode branches, the base-16 Q6_K climb, test_group_upgrade_noop_on_nonuniform (pins the Q6_K→None fallback contract via block_size 16), and lossless test_group_upgrade_is_lossless. One small gap: test_to_intx_unpacked_matches_reference (test_gguf.py:126) already includes q5_k — good; but there's no direct assertion that Q6_K's repack_mlx itself returns None (the block_size==16 check is a proxy, but the repack_mlx contract is only exercised end-to-end).

I did not run the MLX kernel tests — they need Metal/macOS and aren't runnable on this Linux checkout. This review is static.

---

Review authored with Claude.
· branch mlx-q5k

[metascroy]

@claude review this code

[claude]

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

---

Code Review: MLX Q5K/Q6K upgrades (pass 3)

I re-read the repack paths, the native linear/embedding lowerings, _max_uniform_group_size, to_mlx_qparams / emit_quantized_biases, the dispatch wiring in patterns.py + per-type linear.py/embedding.py, and the tests. Prior-pass docstring/terminology feedback is mostly addressed (the native docstrings now say "32 by default, or the merged 64/128"; the patterns.py header says "MLX-native repack path by default"). The core design holds up. Remaining findings below — none blocking.

Correctness — verified

• Group merge (gguf.py:246): ascending power-of-two search then descending uniformity check is correct. For Q6_K (base=16, max=128) candidate sizes are {16,32,64,128}; 16 is rejected by the < 32 gate and the rest are MLX-valid, so no un-representable size escapes. n_groups = K/16 (K a multiple of 256) keeps the divisibility guard from blocking a valid climb. s[:, :, 0] is bit-exact since the merged run is all-equal. The FakeTensor early return keeps the data-dependent merge from tripping the partitioner's support check.
• Bit-packing (op_helpers.py:562): the generic LSB branch's (cols·bits)%32==0 holds for K a multiple of 256 (256·5, 256·6 both %32==0); the int64 chunked scatter with straddle carry is right, and index_add_ == OR since per-column bit ranges are disjoint. int5 [0,31](offset 16) and int6 [0,63](offset 32) fit.
• Symmetric Q6_K bias (op_helpers.py:433): init-chain biases = scale·(−2^(bits−1)) = −32·scale matches to_mlx_qparams's B = −scale·(0+offset). Consistent.
• Q6_K fallback is clean (q6k/repack_mlx.py:58): the group_size < 32 check returns None before to_mlx_qparams/make_or_get_constant, so no orphan constants leak, and emit_linear/emit_embedding propagate the None to the fused path.

Minor issues

1. Stale comment in patterns.py:47. # Quant types each pattern can lower (both via fused custom Metal kernels). still describes the old default even though the repack path is now default. The module header two lines up was updated; this inline comment was missed. Fix this →

2. Q6_K computes biases twice, discards one. q6k/repack_mlx.py:63 calls to_mlx_qparams(...) with the default compute_biases=True, materializing B = −32·scale. But Q6_K is symmetric, so emit_quantized_biases (op_helpers.py:424) takes the is_scale_only branch and recomputes biases in the init chain — the B just allocated is thrown away. Passing compute_biases=False and biases=None skips the wasted allocation. Export-time only. Fix this →

3. Shared Q6_K weight → duplicate init multiply. For a weight referenced by both a linear and an embedding, repack_mlx runs twice with the same zero_point_key. make_or_get_constant dedupes the packed/scales/dummy constants, but emit_quantized_biases (op_helpers.py:433) calls P.emit_init(MultiplyNode(...)) unconditionally — no dedup — so the init chain gets two identical scale·(−32) → biases multiplies into the same slot. Idempotent (harmless), just redundant. Q4_K/Q5_K don't hit this since they serialize biases as a deduped constant. Worth a guard only if shared GGUF blobs are expected in the target models.

4. Wasted full unpack on the Q6_K fallback. q6k/repack_mlx.py:54 fully decodes the blob via to_intx_unpacked_to_int8_tensor(max_group_size=128) before discovering group_size < 32 and returning None, after which _emit_linear_fused redoes the work. The merge decision only needs the sub-block scales, so a pre-check on those alone would skip the full field decode for large lm_heads. Export-time cost only.

Tests

Strong coverage: native/fused × Q4_K/Q5_K/Q6_K for linear and embedding, dynamic-M IfNode branches, the base-16 Q6_K climb, test_group_upgrade_noop_on_nonuniform (pins the Q6_K→fallback contract via block_size 16), test_group_upgrade_is_lossless, and q5_k now included in test_to_intx_unpacked_matches_reference. One small gap remains: no test asserts Q6_K's repack_mlx itself returns None on a non-uniform weight — the block_size==16 check in test_group_upgrade_noop_on_nonuniform is a proxy, but the repack_mlx None contract is only exercised end-to-end (and end-to-end needs Metal).

I did not run the MLX kernel tests — they need Metal/macOS and aren't runnable on this Linux checkout. This review is static.

---

Review authored with Claude.
· branch mlx-q5k