AX Gemma 4 31B IT MLX 6-bit Assistant MTP

Parameter count: approximately 31.27B logical target parameters (31B class). 6-bit is the target quantization precision, not a 6B model-size claim. The separately packaged assistant is an acceleration component and is not included in the target count.

This is a self-contained MLX model package for Apple Silicon. It combines the Gemma 4 31B instruction-tuned target with its exact paired Google assistant model for AX Engine multi-token prediction (MTP) / speculative decoding.

The target model verifies every drafted token. The assistant improves decode speed without replacing the target model or changing its output distribution. This repository does not contain PyTorch or GGUF target weights.

Model details

• Base model: google/gemma-4-31B-it
• MLX target: mlx-community/gemma-4-31b-it-6bit, revision 938d4fb4ebff2df7f6c8200977cf82a06d20f5b9
• Assistant: google/gemma-4-31B-it-assistant, revision 627c5ec1458b9086b841a91e0512fd31fd2fbbf1
• Format: MLX Safetensors target plus an AX Engine assistant subtree
• Target quantization: 6-bit affine, group size 64
• Assistant precision: upstream unquantized bfloat16 Safetensors
• Configured context length: 262,144 tokens
• MTP pairing: exact
• Maximum packaged draft depth: 1
• Intended hardware: Apple Silicon

The target weight shards and assistant weights are byte-identical to the pinned upstream revisions. AX Engine adds native manifests, an assistant pairing contract, an AX-compatible assistant configuration, and a target-tokenizer copy inside assistant/.

Download

hf download AutomatosX/AX-Gemma-4-31B-IT-MLX-6bit-Assistant-MTP \
  --local-dir ./AX-Gemma-4-31B-IT-MLX-6bit-Assistant-MTP

The repository is approximately 27.1 GB. Ensure that the destination has enough free space before downloading.

Use with AX Engine MTP

Install AX Engine, then serve the downloaded repository directory:

ax-engine serve ./AX-Gemma-4-31B-IT-MLX-6bit-Assistant-MTP --port 31418

Gemma assistant MTP is enabled by default. AX Engine reads ax_gemma4_assistant_mtp.json, loads the compact drafter from assistant/, and uses the 31B target to verify proposals. Do not load assistant/ by itself as a general-purpose chat model.

The target weights can also be used for direct MLX generation, but the nested assistant acceleration and pairing contract are AX Engine-specific.

Package layout

• model-*.safetensors: 6-bit MLX target weights
• model.safetensors.index.json: target shard index
• assistant/model.safetensors: paired bfloat16 assistant weights
• model-manifest.json and assistant/model-manifest.json: AX native manifests
• ax_gemma4_assistant_mtp.json: exact-pairing and draft-depth contract
• ax_provenance.json: pinned sources and checksums

Validation and provenance

Validated on macOS arm64 with AX Engine 6.9.0 on 2026-07-20:

• AX native artifact validation: ready, with no issues
• All target weight shards: byte-exact against the pinned MLX source
• Assistant weight: byte-exact against the pinned Google assistant source
• Assistant and target tokenizer: byte-identical inside the package
• Pairing contract: exact
• Canonical chat template: pinned from Google Gemma 4 and applied to the target

See ax_provenance.json for immutable source revisions and SHA-256 values.

License

Apache License 2.0. Review the Gemma terms and the original Google model cards for usage limitations and responsible-use guidance.