[solderzzc]

Proposed changes

QAT 4-bit Gemma 4 MTP assistants (e.g. mlx-community/gemma-4-26B-A4B-it-qat-assistant-4bit) ship pre_projection / post_projection as quantized .weight + .scales / .biases. Gemma4AssistantModel.sanitize kept the raw packed weight, so loading failed with unhandled keys and the packed weight would have been used as a dense matrix.

• sanitize now dequantizes both projections when .scales is present, using the checkpoint's quantization config.
• The gemma4_assistant factory entry passes perLayerQuantization from config.json to the model.
• Unquantized (bf16) assistants are unchanged.

This is a load fix only. MTP output is still not identical to plain decoding at temperature 0, with both the bf16 and the QAT assistant. That is a separate, pre-existing issue.

Test plan

• New test testAssistantSanitizeDequantizesProjections: fails without the fix (shapes [64,16] / [64,8]), passes with it
• All 24 Gemma4Tests pass
• End-to-end on M6: the QAT assistant now loads ("MTP assistant ready"). Before, it failed with unhandled keys.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

AI usage

• I have read this PR description in full and approve it as my own, and it
  accurately describes the code changes.
• AI usage disclosure: Code, tests, and this description were drafted with Claude Code. I reviewed the diff.

[solderzzc]

Cross-review (M6 session; AI-assisted: Claude Code with the code-review skill; approved by the repo owner in conversation)

Verdict: clean. No findings.

Checked

• denseProjectionWeight covers three cases:
  • Dense checkpoint: .weight only, the same behaviour as before.
  • Quantized checkpoint with its config: dequantized with groupSize/bits/mode. biases and global_scale are optional, so affine, mxfp4 and nvfp4 all work. All four keys are removed afterwards.
  • Quantized checkpoint without a config: keys are left in place, so loading fails with a clear unhandled-keys error instead of matmul'ing packed uint32 weights.
• The later quantize(model:) in loadWeights only turns modules into quantized ones when \(path).scales is still in the weights. The projections aren't modules and their keys are already gone, so nothing collides.
• The factory decodes BaseConfiguration from the same config.json. perLayerQuantization.quantization(layer:) returns nil for skipped layers, and that falls into the loud-failure case, as intended.
• The new test round-trips a 4-bit quantize and checks the shapes, the error bound, and that no *projection* keys are left over.

Verified on a Mac mini M6 (32 GB)

Tested with SwiftLM and --mtp, using gemma-4-26B-A4B-it-4bit + gemma-4-26B-A4B-it-qat-assistant-4bit. Before this PR it failed with unhandled keys (pre_projection.scales/biases, post_projection.scales/biases). With it, the model loads and generates coherent output.

Decode tok/s at 3 prompt lengths:

	short	mid	long
QAT assistant (this PR)	47.4	42.0	29.5
bf16 assistant	45.2	35.6	30.4
no MTP	53.0	50.8	46.0

I saw one PromptCache.save "expected a non-empty mlx_array" crash across the runs and could not reproduce it. It's probably unrelated to this PR, and I'll file it separately if it comes back. Gemma MTP being slower than no-MTP on M6 is tracked in SwiftLM ml-explore#184; the per-layer rollback fix PR is coming next.