Instructions to use AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP with MLX:

# Make sure mlx-vlm is installed
# pip install --upgrade mlx-vlm

from mlx_vlm import load, generate
from mlx_vlm.prompt_utils import apply_chat_template
from mlx_vlm.utils import load_config

# Load the model
model, processor = load("AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP")
config = load_config("AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP")

# Prepare input
image = ["http://images.cocodataset.org/val2017/000000039769.jpg"]
prompt = "Describe this image."

# Apply chat template
formatted_prompt = apply_chat_template(
    processor, config, prompt, num_images=1
)

# Generate output
output = generate(model, processor, formatted_prompt, image)
print(output)

Notebooks
Google Colab
Kaggle
Local Apps Settings
LM Studio

How to use AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP with Pi:

Start the MLX server

# Install MLX LM:
uv tool install mlx-lm
# Start a local OpenAI-compatible server:
mlx_lm.server --model "AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP"

Configure the model in Pi

# Install Pi:
npm install -g @mariozechner/pi-coding-agent
# Add to ~/.pi/agent/models.json:
{
  "providers": {
    "mlx-lm": {
      "baseUrl": "http://localhost:8080/v1",
      "api": "openai-completions",
      "apiKey": "none",
      "models": [
        {
          "id": "AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP"
        }
      ]
    }
  }
}

Run Pi

# Start Pi in your project directory:
pi

Hermes Agent new

How to use AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP with Hermes Agent:

Start the MLX server

# Install MLX LM:
uv tool install mlx-lm
# Start a local OpenAI-compatible server:
mlx_lm.server --model "AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP"

Configure Hermes

# Install Hermes:
curl -fsSL https://hermes-agent.nousresearch.com/install.sh | bash
hermes setup
# Point Hermes at the local server:
hermes config set model.provider custom
hermes config set model.base_url http://127.0.0.1:8080/v1
hermes config set model.default AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP

Run Hermes

hermes

OpenClaw new

How to use AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP with OpenClaw:

Start the MLX server

# Install MLX LM:
uv tool install mlx-lm
# Start a local OpenAI-compatible server:
mlx_lm.server --model "AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP"

Configure OpenClaw

# Install OpenClaw:
npm install -g openclaw@latest
# Register the local server and set it as the default model:
openclaw onboard --non-interactive --mode local \
  --auth-choice custom-api-key \
  --custom-base-url http://127.0.0.1:8080/v1 \
  --custom-model-id "AutomatosX/AX-Gemma-4-31B-IT-MLX-QAT-4bit-Assistant-MTP" \
  --custom-provider-id mlx-lm \
  --custom-compatibility openai \
  --custom-text-input \
  --accept-risk \
  --skip-health

Run OpenClaw

openclaw agent --local --agent main --message "Hello from Hugging Face"

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

Parameter count: approximately 31.27B logical target parameters (31B class). 4-bit is the target quantization precision, not a 4B 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 quantization-aware-trained Gemma 4 31B instruction target with its exact paired 4-bit assistant for AX Engine multi-token prediction (MTP) / speculative decoding.

The target verifies every drafted token. The assistant improves decode speed without replacing the target model. This repository does not contain PyTorch, GGUF, or the unquantized Google QAT weights.

Model details

Official QAT base: google/gemma-4-31B-it-qat-q4_0-unquantized
MLX target: mlx-community/gemma-4-31B-it-qat-4bit, revision 83b1a12d1da4b07afe435bd2cc310308bfc13984
MLX assistant: mlx-community/gemma-4-31B-it-qat-assistant-4bit, revision 5234fd588403c9b68f3bd20a140b7e61700cb7e2
Format: MLX Safetensors target plus an AX Engine assistant subtree
Target quantization: 4-bit affine, group size 64; FFN projections retain 8-bit overrides
Assistant quantization: 4-bit affine, group size 64
Configured target context length: 262,144 tokens
MTP pairing: exact
Maximum packaged draft depth: 1
Intended hardware: Apple Silicon

QAT means that the upstream checkpoint was optimized during training for its target quantization scheme before the MLX conversion. It is distinct from a post-training-only 4-bit conversion.

Download

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

The repository is approximately 29.1 GB.

Use with AX Engine MTP

Install AX Engine, then serve the downloaded repository directory:

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

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

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

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 target source
Assistant weight: byte-exact against the pinned MLX 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 target and assistant

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