Good First Issue: Enable Gemma 4 on MLX Backend

[metascroy]

Good First Issue: Enable Gemma 4 on MLX Backend

Summary

Add support for exporting and running Google's Gemma 4 model on the MLX delegate. Gemma 3 (1B and 4B) already works on MLX and is tested in CI. Gemma 4 should follow the same path.

Background

Gemma 3 is currently supported on MLX via two export paths:

1. HuggingFace path (export_llm_hf.py): Uses AutoModelForCausalLM from HuggingFace transformers. Gemma 3 1B is tested in CI with both 4w and nvfp4 quantization.
2. export_llm path (extension/llm/export/export_llm.py): Uses the internal Llama-based model definition with HF weight conversion. Gemma 3 has a dedicated examples/models/gemma3/ directory with convert_weights.py.

Gemma 4 likely uses a similar architecture to Gemma 3 (dense transformer with sliding window attention and GQA). If the architecture is compatible, enabling it may be as simple as testing it through the existing HF export path.

Steps

1. Test via the HuggingFace export path (try this first)

The HF path auto-detects model architecture, so Gemma 4 may work out of the box:

python -m executorch.backends.mlx.examples.llm.export_llm_hf \
 --model-id google/gemma-4-<VARIANT> \
 --output gemma4_mlx.pte \
 --qlinear 4w \
 --qembedding 4w \
 --use-custom-sdpa \
 --use-custom-kv-cache

If this fails, note the error. Common issues:

• Unsupported aten op: Add a handler in ops.py (see the "Add MLX Op Handler" issue template)
• Data-dependent control flow: May need source transformations to avoid torch.export graph breaks
• New attention pattern: May need updates to the sliding window / KV cache handling

2. Test via the export_llm pipeline

If the HF path works but you also want the export_llm path:

1. Create examples/models/gemma4/ with:
   • __init__.py
   • convert_weights.py (remap HF checkpoint keys to ExecuTorch transformer format)
   • config/ with model config JSON
2. Add gemma4 entries to HUGGING_FACE_REPO_IDS and the weight converter dispatch in examples/models/llama/export_llama_lib.py
3. Follow the pattern from examples/models/gemma3/

3. Run inference

python -m executorch.backends.mlx.examples.llm.run_llm_hf \
 --pte gemma4_mlx.pte \
 --model-id google/gemma-4-<VARIANT> \
 --prompt "What is the capital of France?" \
 --max-new-tokens 50

4. Add CI test

Add a Gemma 4 entry to the test-mlx-llm matrix in .github/workflows/mlx.yml:

matrix:
 model:
 - id: "google/gemma-4-<VARIANT>"
 name: "gemma4"

What might need changes

• If Gemma 4 has new architectural features (e.g., new attention type, new activation, MoE): these may need new op handlers or source transformations
• If HuggingFace transformers has export-incompatible code (like Phi-4's LongRoPE data-dependent branching): may need to wait for a transformers fix or use the export_llm path instead
• Sliding window configuration: Gemma models use sliding window attention. The MLX export path auto-detects this from the HF config and uses RingBufferKVCache

References

• Existing Gemma 3 MLX CI test (search for gemma3-1b)
• Gemma 3 model definition
• LLM HF export script
• LLM HF export README
• MLX Delegate README

[zeel2104]

I would like to work on it. Could you assign it to me? previously contributed to Nvidia, Microsoft, pytorch, apache etc

[metascroy]

I would like to work on it. Could you assign it to me? previously contributed to Nvidia, Microsoft, pytorch, apache etc

Done! Thanks for taking this on!

[zeel2104]

Hey @metascroy
I have opened a PR, a summary of what it does

I validated the HuggingFace MLX path for Gemma 4 on Apple Silicon with google/gemma-4-E2B-it.

Working configuration:

• --qlinear 4w
• --use-custom-sdpa
• --use-custom-kv-cache

What is now working:

• export_llm_hf.py export succeeds
• run_llm_hf.py inference succeeds
• the normal installed package path also works without needing PYTHONPATH

Smoke test prompt:

• What is the capital of France?

Output included:

• Paris

I also updated MLX docs and CI, but only for the exact Gemma 4 configuration that was validated.

What is still not covered in this PR:

• the internal extension.llm.export.export_llm / examples/models/gemma4/ path
• broader Gemma 4 quantization coverage such as qembedding 4w or nvfp4
• additional Gemma 4 variants beyond google/gemma-4-E2B-it

Reasoning:

• the issue asked to try the HuggingFace path first, and that path needed a series of Gemma 4-specific compatibility fixes before it worked reliably
• qembedding 4w and broader quantization combinations were not part of the final working configuration I could validate end to end
• the internal export_llm path is a separate chunk of work because it would require adding examples/models/gemma4/, a weight converter, config wiring, and validating that the internal Llama-style model path cleanly represents Gemma 4’s attention/cache layout

What I’ll do next:

• land the HuggingFace MLX support first
• follow up separately on the internal export_llm path
• expand testing later to additional quantization configs and, if applicable, more Gemma 4 variants

Please let me know if you want any modifications or any other change

[metascroy]

Awesome @zeel2104! Can you link your PR here, or tag me as a reviewer?

[zeel2104]

Hey @metascroy, I have opened the PR here: #19065