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:

HuggingFace path (export_llm_hf.py): Uses AutoModelForCausalLM from HuggingFace transformers. Gemma 3 1B is tested in CI with both 4w and nvfp4 quantization.
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:

Create examples/models/gemma4/ with:
- __init__.py
- convert_weights.py (remap HF checkpoint keys to ExecuTorch transformer format)
- config/ with model config JSON
Add gemma4 entries to HUGGING_FACE_REPO_IDS and the weight converter dispatch in examples/models/llama/export_llama_lib.py
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

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

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!

@metascroy

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

@zeel2104

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

@metascroy

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

Navigation Menu

Repository navigation

Good First Issue: Enable Gemma 4 on MLX Backend #18928

Good First Issue: Enable Gemma 4 on MLX Backend

Summary

Background

Steps

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

2. Test via the export_llm pipeline

3. Run inference

4. Add CI test

What might need changes

References

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Good First Issue: Enable Gemma 4 on MLX Backend #18928

Description

Good First Issue: Enable Gemma 4 on MLX Backend

Summary

Background

Steps

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

2. Test via the export_llm pipeline

3. Run inference

4. Add CI test

What might need changes

References

Activity

zeel2104 commented on Apr 16, 2026

metascroy commented on Apr 16, 2026

zeel2104 commented on Apr 23, 2026

metascroy commented on Apr 23, 2026

zeel2104 commented on Apr 23, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions