mlx.yml

name: MLX

on:
  push:
    branches:
      - main
      - release/*
    tags:
      - ciflow/mlx/*
  pull_request:
    paths:
      - .github/workflows/mlx.yml
      - backends/mlx/**
      - extension/llm/export/**
      - extension/audio/**
      - examples/models/gemma4_31b/**
      - examples/models/muse-glimmer/**
      - examples/models/parakeet/**
      - examples/models/voxtral_realtime/**
      - examples/models/qwen3_5_moe/**
  workflow_dispatch:

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
  cancel-in-progress: true

permissions:
  contents: read

jobs:
  # Emits is-full-run='true' for workflow_dispatch / ciflow tag /
  # sampled-push commits (every 4th main/release commit by depth).
  # Returns 'false' for pull_request events — PR jobs use the workflow-
  # level `paths:` filter (above) for path-based gating instead.
  run-decision:
    name: CI run decision
    uses: ./.github/workflows/_ci-run-decision.yml

  test-mlx:
    needs: run-decision
    if: |
      github.event_name == 'pull_request' ||
      needs.run-decision.outputs.is-full-run == 'true'
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    with:
      default-packages: ""
      job-name: test-mlx
      runner: macos-14-xlarge
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      timeout: 90
      script: |
        set -eux

        echo "::group::Install ExecuTorch and configure build"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        # The sanitizers fail on github VM runner, but pass on real device
        # TODO: figure out why
        ${CONDA_RUN} cmake --preset mlx-release -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_MLX_ENABLE_SANITIZERS=OFF
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Install Python test requirements"
        ${CONDA_RUN} pip install gguf
        echo "::endgroup::"

        echo "::group::Build test runners"
        ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_mutable_state_test mlx_sequence_cache_test -j$(( $(sysctl -n hw.ncpu) - 1 ))
        echo "::endgroup::"

        echo "::group::Run mutable-state (multi-session) unit test"
        ./cmake-out/backends/mlx/test/mlx_mutable_state_test
        echo "::endgroup::"

        echo "::group::Run off-graph KV-cache op test"
        ./cmake-out/backends/mlx/test/mlx_sequence_cache_test
        echo "::endgroup::"

        echo "::group::Run op unit tests"
        ${CONDA_RUN} python -m executorch.backends.mlx.test.run_all_tests -j4 --max-tasks-per-worker 10 --clean-after
        echo "::endgroup::"

        echo "::group::Run Python unit tests"
        ${CONDA_RUN} python -m pytest \
          backends/mlx/test/test_passes.py \
          backends/mlx/test/test_pattern_utils.py \
          backends/mlx/test/test_partitioner.py \
          backends/mlx/test/test_serialization_dedup.py \
          backends/mlx/test/test_slot_recycling.py \
          backends/mlx/test/test_sample.py \
          examples/models/gemma4_31b/tests/test_mlx_pipeline.py \
          examples/models/muse-glimmer/tests/test_mlx_pipeline.py \
          -v
        echo "::endgroup::"

        echo "::group::Run multi-thread stress test"
        ${CONDA_RUN} python backends/mlx/test/export_multi_thread_test_model.py /tmp/multi_thread_test_model.pte
        ET_TESTING_MODEL_PATH=/tmp/multi_thread_test_model.pte \
          ET_TESTING_NUM_THREADS=50 \
          ET_PREDICTIONS_PER_THREAD=100 \
          ./cmake-out/backends/mlx/test/multi_thread_test_runner
        echo "::endgroup::"

        echo "::group::Run custom_kernel_ops op tests"
        # Run every custom_kernel_ops/**/test/test_*.py via its OpTestCase `run`
        # CLI. Recurses into per-format subpackages (e.g. gguf/test), so adding a
        # new op test file requires no change here.
        set -e
        for t in $(find backends/mlx/custom_kernel_ops -path '*/test/test_*.py' | sort); do
          mod="executorch.$(echo "${t%.py}" | tr '/' '.')"
          echo "--- ${mod} ---"
          ${CONDA_RUN} python -m "${mod}" run -v
        done
        echo "::endgroup::"

  test-mlx-qwen35-moe:
    needs: run-decision
    if: |
      github.event_name == 'pull_request' ||
      needs.run-decision.outputs.is-full-run == 'true'
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    with:
      default-packages: ""
      job-name: test-mlx-qwen35-moe
      runner: macos-14-xlarge
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      timeout: 90
      script: |
        set -eux

        echo "::group::Install ExecuTorch"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Export Qwen 3.5 MoE (tiny model)"
        ${CONDA_RUN} python -m executorch.examples.models.qwen3_5_moe.export \
          --tiny-test \
          --backend mlx \
          --qlinear 4w \
          --qlinear-group-size 32 \
          --output-dir /tmp/qwen35_moe_mlx_tiny
        echo "::endgroup::"

        echo "::group::Check AsType node count"
        ASTYPE_COUNT=$(${CONDA_RUN} python -m executorch.backends.mlx.pte_inspector \
          /tmp/qwen35_moe_mlx_tiny/model.pte --mlx-instructions 2>&1 | grep -c "AsTypeNode" || true)
        echo "AsType nodes: ${ASTYPE_COUNT}"
        if [ "$ASTYPE_COUNT" -gt 23 ]; then
          echo "Failed: expected no more than 23 AsType nodes, got ${ASTYPE_COUNT}"
          exit 1
        fi
        echo "::endgroup::"

        echo "::group::Run Qwen 3.5 MoE inference"
        OUTPUT=$(${CONDA_RUN} python -m executorch.examples.models.qwen3_5_moe.run \
          --pte /tmp/qwen35_moe_mlx_tiny/model.pte \
          --prompt-len 4 \
          --max-new-tokens 5 2>&1)
        echo "$OUTPUT"
        if echo "$OUTPUT" | grep -q "Generated token ids: \[167, 180, 167, 180, 167\]"; then
          echo "Success: Qwen 3.5 MoE MLX export + inference completed with expected output"
        else
          echo "Failed: unexpected output (expected [167, 180, 167, 180, 167])"
          exit 1
        fi
        echo "::endgroup::"

        echo "::group::Verify chunked == unchunked prefill"
        QWEN_TINY_PTE=/tmp/qwen35_moe_mlx_tiny/model.pte \
          ${CONDA_RUN} python -m pytest \
          examples/models/qwen3_5_moe/test_chunked_prefill.py -v
        echo "::endgroup::"

        echo "::group::Build Qwen 3.5 MoE MLX C++ runner"
        # Validates the MLX C++ runner build wiring (compile + link + metallib).
        # The tiny model has no compatible tokenizer (vocab 256, random weights),
        # so we don't run C++ inference here — only confirm it builds.
        ${CONDA_RUN} make qwen3_5_moe-mlx
        RUNNER=cmake-out/examples/models/qwen3_5_moe/qwen3_5_moe_runner
        if [ ! -x "$RUNNER" ]; then
          echo "Failed: runner not found at $RUNNER"
          exit 1
        fi
        if [ ! -f "$(dirname "$RUNNER")/mlx.metallib" ]; then
          echo "Failed: mlx.metallib not copied next to runner"
          exit 1
        fi
        echo "Success: built $RUNNER"
        echo "::endgroup::"

  # Build-only: the Muse Glimmer GGUF checkpoints are far too large to export in
  # CI, so this covers the compile / link / metallib wiring of the runners --
  # which is what MLX backend changes are most likely to break.
  test-mlx-muse-glimmer-build:
    needs: run-decision
    if: |
      github.event_name == 'pull_request' ||
      needs.run-decision.outputs.is-full-run == 'true'
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    with:
      default-packages: ""
      job-name: test-mlx-muse-glimmer-build
      runner: macos-14-xlarge
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      timeout: 90
      script: |
        set -eux

        echo "::group::Install ExecuTorch"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Build Muse Glimmer runners and worker with MLX"
        ${CONDA_RUN} make muse-glimmer-mlx
        BIN_DIR=cmake-out/examples/models/muse-glimmer
        for bin in solo_runner dflash_runner muse_glimmer_worker; do
          if [ ! -x "${BIN_DIR}/${bin}" ]; then
            echo "Failed: ${bin} not found at ${BIN_DIR}/${bin}"
            exit 1
          fi
        done
        # Runtime kernel load falls back to an absolute build-tree path, so a
        # missing colocated metallib only breaks a relocated binary. Assert the
        # copy so the runners stay self-contained (needs EXECUTORCH_BUILD_MLX=ON
        # in the preset, which drives executorch_target_copy_mlx_metallib).
        if [ ! -f "${BIN_DIR}/mlx.metallib" ]; then
          echo "Failed: mlx.metallib not copied next to the runners"
          exit 1
        fi
        echo "Success: built solo_runner, dflash_runner and muse_glimmer_worker"
        echo "::endgroup::"

  backend-tester:
    needs: run-decision
    if: |
      github.event_name == 'pull_request' ||
      needs.run-decision.outputs.is-full-run == 'true'
    strategy:
      fail-fast: false
      matrix:
        suite: [models, operators]
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    with:
      default-packages: ""
      job-name: test-mlx-backend-${{ matrix.suite }}
      runner: macos-14-xlarge
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      timeout: 120
      script: |
        set -eux

        echo "::group::Install ExecuTorch"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Run backend test suite (${{ matrix.suite }})"
        ${CONDA_RUN} pytest -c /dev/null backends/test/suite/${{ matrix.suite }}/ -m flow_mlx -n auto 2>&1 | tee pytest_output.txt || true
        echo "::endgroup::"

        # Parse pytest summary and check failure threshold
        if grep -E "^=+ .* =+$" pytest_output.txt | tail -1 | grep -q "failed"; then
          FAILED=$(grep -E "^=+ .* =+$" pytest_output.txt | tail -1 | grep -oE "[0-9]+ failed" | grep -oE "[0-9]+")
        else
          FAILED=0
        fi

        if [ "${{ matrix.suite }}" = "operators" ]; then
          MAX_FAILURES=0
        else
          MAX_FAILURES=3
        fi

        echo "Failed tests: $FAILED (max allowed: $MAX_FAILURES)"
        if [ "$FAILED" -gt "$MAX_FAILURES" ]; then
          echo "::error::Too many test failures: $FAILED > $MAX_FAILURES"
          exit 1
        fi

  test-mlx-parakeet:
    needs: run-decision
    if: |
      github.event_name == 'pull_request' ||
      needs.run-decision.outputs.is-full-run == 'true'
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    with:
      default-packages: ""
      job-name: test-mlx-parakeet
      runner: macos-14-xlarge
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      timeout: 90
      script: |
        set -eux

        echo "::group::Install ExecuTorch"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        echo "::endgroup::"

        echo "::group::Install Parakeet requirements"
        ${CONDA_RUN} pip install -r examples/models/parakeet/install_requirements.txt
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Export Parakeet"
        ${CONDA_RUN} python -m executorch.examples.models.parakeet.export_parakeet_tdt \
          --backend mlx \
          --dtype bf16 \
          --qlinear_encoder 4w \
          --qlinear_encoder_group_size 128 \
          --qlinear 4w \
          --qlinear_group_size 128 \
          --output-dir /tmp/parakeet_mlx
        echo "::endgroup::"

        echo "::group::Build Parakeet MLX runner"
        ${CONDA_RUN} make parakeet-mlx
        echo "::endgroup::"

        echo "::group::Run Parakeet MLX runner"
        curl -L https://dldata-public.s3.us-east-2.amazonaws.com/2086-149220-0033.wav -o /tmp/test_audio.wav
        OUTPUT=$(./cmake-out/examples/models/parakeet/parakeet_runner \
          --model_path /tmp/parakeet_mlx/model.pte \
          --audio_path /tmp/test_audio.wav \
          --tokenizer_path /tmp/parakeet_mlx/tokenizer.model 2>&1)
        echo "Runner output:"
        echo "$OUTPUT"
        if echo "$OUTPUT" | grep -iq "Phoebe"; then
          echo "Success: 'Phoebe' found in output"
        else
          echo "Failed: Expected 'Phoebe' not found in output"
          exit 1
        fi
        echo "::endgroup::"

  test-mlx-voxtral:
    # Requires HuggingFace secrets — skip on fork PRs.
    # Maintainers can opt-in by applying the ciflow/mlx label, which
    # pushes a ciflow/mlx/<PR> tag that re-runs this workflow with secrets.
    needs: run-decision
    if: |
      (github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request') &&
      (github.event_name == 'pull_request' || needs.run-decision.outputs.is-full-run == 'true')
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    secrets: inherit
    with:
      default-packages: ""
      job-name: test-mlx-voxtral
      runner: macos-14-xlarge
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      secrets-env: EXECUTORCH_HF_TOKEN
      timeout: 90
      script: |
        set -eux

        echo "::group::Install ExecuTorch"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        echo "::endgroup::"

        echo "::group::Install Voxtral requirements"
        ${CONDA_RUN} pip install mistral_common librosa soundfile datasets
        OPTIMUM_ET_VERSION=$(cat .ci/docker/ci_commit_pins/optimum-executorch.txt)
        ${CONDA_RUN} pip install "optimum-executorch @ git+https://github.com/huggingface/optimum-executorch.git@${OPTIMUM_ET_VERSION}"
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Export Voxtral"
        ${CONDA_RUN} python -m executorch.backends.mlx.examples.voxtral.export_voxtral_hf \
          --output-dir /tmp/voxtral_mlx \
          --dtype bf16 \
          --qlinear 4w
        echo "::endgroup::"

        echo "::group::Build Voxtral MLX runner"
        ${CONDA_RUN} make voxtral-mlx
        echo "::endgroup::"

        echo "::group::Run Voxtral MLX runner"
        curl -L https://huggingface.co/mistralai/Voxtral-Mini-3B-2507/resolve/main/tekken.json -o /tmp/tekken.json
        curl -L https://github.com/voxserv/audio_quality_testing_samples/raw/refs/heads/master/testaudio/16000/test01_20s.wav -o /tmp/test_audio.wav
        OUTPUT=$(./cmake-out/examples/models/voxtral/voxtral_runner \
          --model_path /tmp/voxtral_mlx/model.pte \
          --tokenizer_path /tmp/tekken.json \
          --audio_path /tmp/test_audio.wav \
          --processor_path /tmp/voxtral_mlx/preprocessor.pte \
          --prompt "What is happening in this audio?" \
          --temperature 0 2>&1)
        echo "Runner output:"
        echo "$OUTPUT"
        if echo "$OUTPUT" | grep -iq "poem"; then
          echo "Success: 'poem' found in output"
        else
          echo "Failed: Expected 'poem' not found in output"
          exit 1
        fi
        echo "::endgroup::"

  test-mlx-voxtral-realtime:
    # Requires HuggingFace secrets — skip on fork PRs.
    # Maintainers can opt-in by applying the ciflow/mlx label.
    needs: run-decision
    if: |
      (github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request') &&
      (github.event_name == 'pull_request' || needs.run-decision.outputs.is-full-run == 'true')
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    secrets: inherit
    with:
      default-packages: ""
      job-name: test-mlx-voxtral-realtime
      runner: macos-14-xlarge
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      secrets-env: EXECUTORCH_HF_TOKEN
      timeout: 90
      script: |
        set -eux
        # Disable HF Xet storage to avoid stalled downloads on CI runners
        export HF_HUB_DISABLE_XET=1

        echo "::group::Install ExecuTorch"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        echo "::endgroup::"

        echo "::group::Install Voxtral Realtime requirements"
        ${CONDA_RUN} pip install safetensors
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Download model"
        HF_TOKEN=$SECRET_EXECUTORCH_HF_TOKEN ${CONDA_RUN} python -c "from huggingface_hub import snapshot_download; snapshot_download('mistralai/Voxtral-Mini-4B-Realtime-2602')"
        MODEL_PATH=$(HF_TOKEN=$SECRET_EXECUTORCH_HF_TOKEN ${CONDA_RUN} python -c "from huggingface_hub import snapshot_download; print(snapshot_download('mistralai/Voxtral-Mini-4B-Realtime-2602'))")
        echo "Model path: ${MODEL_PATH}"
        echo "::endgroup::"

        echo "::group::Export preprocessor"
        ${CONDA_RUN} python -m executorch.extension.audio.mel_spectrogram \
          --feature_size 128 \
          --streaming \
          --backend mlx \
          --output_file /tmp/voxtral_rt_mlx/preprocessor.pte
        echo "::endgroup::"

        echo "::group::Export Voxtral Realtime (streaming)"
        ${CONDA_RUN} python -m executorch.examples.models.voxtral_realtime.export_voxtral_rt \
          --model-path "${MODEL_PATH}" \
          --backend mlx \
          --streaming \
          --output-dir /tmp/voxtral_rt_mlx \
          --qlinear-encoder 4w \
          --qlinear 4w \
          --qembedding 8w \
          --qembedding-group-size 128
        echo "::endgroup::"

        echo "::group::Build Voxtral Realtime MLX runner"
        ${CONDA_RUN} make voxtral_realtime-mlx
        echo "::endgroup::"

        echo "::group::Run Voxtral Realtime MLX runner"
        curl -L https://dldata-public.s3.us-east-2.amazonaws.com/2086-149220-0033.wav -o /tmp/test_audio.wav
        OUTPUT=$(./cmake-out/examples/models/voxtral_realtime/voxtral_realtime_runner \
          --model_path /tmp/voxtral_rt_mlx/model.pte \
          --tokenizer_path "${MODEL_PATH}/tekken.json" \
          --preprocessor_path /tmp/voxtral_rt_mlx/preprocessor.pte \
          --audio_path /tmp/test_audio.wav \
          --streaming 2>&1)
        echo "Runner output:"
        echo "$OUTPUT"
        if echo "$OUTPUT" | grep -iq "Phoebe"; then
          echo "Success: 'Phoebe' found in output"
        else
          echo "Failed: Expected 'Phoebe' not found in output"
          exit 1
        fi
        echo "::endgroup::"

  test-mlx-whisper:
    # Requires HuggingFace secrets — skip on fork PRs.
    # Maintainers can opt-in by applying the ciflow/mlx label.
    needs: run-decision
    if: |
      (github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request') &&
      (github.event_name == 'pull_request' || needs.run-decision.outputs.is-full-run == 'true')
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    secrets: inherit
    with:
      default-packages: ""
      job-name: test-mlx-whisper
      runner: macos-14-xlarge
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      secrets-env: EXECUTORCH_HF_TOKEN
      timeout: 90
      script: |
        set -eux
        # Disable HF Xet storage to avoid stalled downloads on CI runners
        export HF_HUB_DISABLE_XET=1

        echo "::group::Install ExecuTorch and configure MLX build"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        echo "::endgroup::"

        echo "::group::Install Whisper requirements"
        ${CONDA_RUN} pip install -U "huggingface_hub[cli]<1.0"
        ${CONDA_RUN} huggingface-cli login --token $SECRET_EXECUTORCH_HF_TOKEN
        ${CONDA_RUN} pip install transformers soundfile datasets librosa
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Export Whisper"
        ${CONDA_RUN} python -m executorch.backends.mlx.examples.whisper.export_whisper \
          --model-id "openai/whisper-tiny" \
          --output-dir /tmp/whisper_mlx \
          --dtype bf16 \
          --qlinear 4w
        echo "::endgroup::"

        echo "::group::Run Whisper inference"
        OUTPUT=$( ${CONDA_RUN} python -m executorch.backends.mlx.examples.whisper.run_whisper \
          --model-dir /tmp/whisper_mlx \
          --use-sample-audio 2>&1)
        echo "$OUTPUT"
        if echo "$OUTPUT" | grep -iq "Mr. Quilter"; then
          echo "Success: 'Mr. Quilter' found in transcription"
        else
          echo "Failed: Expected 'Mr. Quilter' not found in transcription"
          exit 1
        fi
        echo "::endgroup::"


  test-mlx-stories110m:
    needs: run-decision
    if: |
      github.event_name == 'pull_request' ||
      needs.run-decision.outputs.is-full-run == 'true'
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    with:
      default-packages: ""
      job-name: test-mlx-stories110m
      runner: macos-14-xlarge
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      timeout: 90
      script: |
        set -eux

        echo "::group::Install ExecuTorch"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        echo "::endgroup::"

        echo "::group::Install Llama requirements"
        ${CONDA_RUN} sh examples/models/llama/install_requirements.sh
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Build ExecuTorch with MLX delegate"
        ${CONDA_RUN} cmake --workflow --preset mlx-release
        echo "::endgroup::"

        echo "::group::Build Llama runner with MLX"
        pushd examples/models/llama
        ${CONDA_RUN} cmake --workflow --preset llama-release
        popd
        echo "::endgroup::"

        echo "::group::Download stories110M artifacts"
        curl -Ls "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.pt" --output stories110M.pt
        curl -Ls "https://raw.githubusercontent.com/karpathy/llama2.c/master/tokenizer.model" --output tokenizer.model
        echo '{"dim": 768, "multiple_of": 32, "n_heads": 12, "n_layers": 12, "norm_eps": 1e-05, "vocab_size": 32000}' > params.json
        echo "::endgroup::"

        echo "::group::Create tokenizer.bin"
        ${CONDA_RUN} python -m pytorch_tokenizers.tools.llama2c.convert -t tokenizer.model -o tokenizer.bin
        echo "::endgroup::"

        echo "::group::Export stories110M with MLX backend via export_llama_lib"
        ${CONDA_RUN} python -m extension.llm.export.export_llm \
          base.checkpoint=stories110M.pt \
          base.params=params.json \
          model.use_kv_cache=true \
          model.dtype_override=fp32 \
          backend.mlx.enabled=true \
          quantization.qmode=4w \
          quantization.group_size=32 \
          export.output_name=/tmp/stories110m_mlx.pte
        echo "::endgroup::"

        echo "::group::Run inference with C++ llama runner"
        ./cmake-out/examples/models/llama/llama_main \
          --model_path=/tmp/stories110m_mlx.pte \
          --tokenizer_path=tokenizer.bin \
          --prompt="Once upon a time," \
          --temperature=0 \
          --seq_len=10
        echo "::endgroup::"

  test-mlx-llm:
    # Requires HuggingFace secrets — skip on fork PRs.
    # Maintainers can opt-in by applying the ciflow/mlx label.
    needs: run-decision
    if: |
      (github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request') &&
      (github.event_name == 'pull_request' || needs.run-decision.outputs.is-full-run == 'true')
    strategy:
      fail-fast: false
      matrix:
        model:
          - id: "unsloth/Llama-3.2-1B-Instruct"
            name: "llama-1b"
          - id: "unsloth/Qwen3-0.6B"
            name: "qwen3-0.6b"
          - id: "unsloth/gemma-3-1b-it"
            name: "gemma3-1b"
        use-custom: [false, true]
        qconfig: ["4w", "nvfp4"]
        runner: ["macos-14-xlarge"]
        include:
          - model:
              id: "google/gemma-4-E2B-it"
              name: "gemma4-e2b"
            use-custom: true
            qconfig: "4w"
            runner: "macos-15-xlarge"
          - model:
              id: "google/gemma-4-E2B-it"
              name: "gemma4-e2b"
            use-custom: false
            qconfig: "4w"
            runner: "macos-15-xlarge"
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    secrets: inherit
    with:
      default-packages: ""
      job-name: test-mlx-llm-${{ matrix.model.name }}${{ matrix.use-custom && '-custom' || '' }}-${{ matrix.qconfig }}
      runner: ${{ matrix.runner }}
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      secrets-env: EXECUTORCH_HF_TOKEN
      timeout: 90
      script: |
        set -eux
        # Disable HF Xet storage to avoid stalled downloads on CI runners
        export HF_HUB_DISABLE_XET=1

        MODEL_ID="${{ matrix.model.id }}"
        MODEL_NAME="${{ matrix.model.name }}"
        USE_CUSTOM="${{ matrix.use-custom }}"
        QCONFIG="${{ matrix.qconfig }}"
        CUSTOM_ARGS=""
        if [ "${USE_CUSTOM}" = "true" ]; then
          CUSTOM_ARGS="--use-custom-sdpa --use-custom-kv-cache"
        fi

        QEMBEDDING_ARGS="--qembedding ${QCONFIG}"
        if [ "${MODEL_ID}" = "google/gemma-4-E2B-it" ]; then
          QEMBEDDING_ARGS=""
        fi

        echo "::group::Install ExecuTorch and configure MLX build"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        ${CONDA_RUN} cmake --preset mlx-release
        echo "::endgroup::"

        echo "::group::Install LLM requirements"
        ${CONDA_RUN} pip install -U "huggingface_hub[cli]<1.0"
        ${CONDA_RUN} huggingface-cli login --token $SECRET_EXECUTORCH_HF_TOKEN
        OPTIMUM_ET_VERSION=$(cat .ci/docker/ci_commit_pins/optimum-executorch.txt)
        ${CONDA_RUN} pip install transformers "optimum-executorch @ git+https://github.com/huggingface/optimum-executorch.git@${OPTIMUM_ET_VERSION}"
        if [ "${MODEL_ID}" = "google/gemma-4-E2B-it" ]; then
          # Gemma 4 requires a newer Transformers build than the CI-wide
          # optimum-executorch pin currently brings in. Keep this pinned to the
          # locally validated commit instead of floating on Transformers HEAD.
          GEMMA4_TRANSFORMERS_COMMIT=61461a7bcb458db7cf6eeea49678b9ab776a7821
          ${CONDA_RUN} pip install -U "transformers @ git+https://github.com/huggingface/transformers.git@${GEMMA4_TRANSFORMERS_COMMIT}"
        fi
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Export ${MODEL_NAME}"
        ${CONDA_RUN} python -m executorch.backends.mlx.examples.llm.export_llm_hf \
          --model-id "${MODEL_ID}" \
          --output /tmp/${MODEL_NAME}.pte \
          --qlinear ${QCONFIG} \
          ${QEMBEDDING_ARGS} \
          ${CUSTOM_ARGS}
        echo "::endgroup::"

        echo "::group::Run ${MODEL_NAME} inference"
        OUTPUT=$(${CONDA_RUN} python -m executorch.backends.mlx.examples.llm.run_llm_hf \
          --pte /tmp/${MODEL_NAME}.pte \
          --model-id "${MODEL_ID}" \
          --prompt "What is the capital of France?" \
          --max-new-tokens 50 2>&1)
        echo "$OUTPUT"
        if echo "$OUTPUT" | grep -iq "Paris"; then
          echo "Success: 'Paris' found in output"
        else
          echo "Failed: Expected 'Paris' not found in output"
          exit 1
        fi
        echo "::endgroup::"

  # Off-graph KV cache: the cache is a runtime object, so this path can only be
  # exercised by the C++ runner (pybindings cannot bind a cache_key). Also the
  # only coverage that the layout the export publishes is actually consumable.
  test-mlx-llm-offgraph:
    # Requires HuggingFace secrets — skip on fork PRs.
    needs: run-decision
    if: |
      (github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request') &&
      (github.event_name == 'pull_request' || needs.run-decision.outputs.is-full-run == 'true')
    strategy:
      fail-fast: false
      matrix:
        model:
          - id: "unsloth/Llama-3.2-1B-Instruct"
            name: "llama-1b"
            chat: "llama3"
            runner: "macos-14-xlarge"
          - id: "unsloth/gemma-3-1b-it"
            name: "gemma3-1b"
            chat: "gemma"
            runner: "macos-14-xlarge"
          # Only model with KV sharing: 15 caches for 35 layers, and sliding and
          # full-attention caches with different head shapes.
          - id: "google/gemma-4-E2B-it"
            name: "gemma4-e2b"
            chat: "gemma4"
            runner: "macos-15-xlarge"
    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
    secrets: inherit
    with:
      default-packages: ""
      job-name: test-mlx-llm-offgraph-${{ matrix.model.name }}
      runner: ${{ matrix.model.runner }}
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      secrets-env: EXECUTORCH_HF_TOKEN
      # Higher than the sibling LLM job: this one also builds and installs
      # ExecuTorch so the standalone runner project can link against it.
      timeout: 120
      script: |
        set -eux
        export HF_HUB_DISABLE_XET=1

        MODEL_ID="${{ matrix.model.id }}"
        MODEL_NAME="${{ matrix.model.name }}"
        CHAT="${{ matrix.model.chat }}"

        echo "::group::Install ExecuTorch and build the MLX runtime"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        ${CONDA_RUN} cmake --preset mlx-release
        ${CONDA_RUN} cmake --build cmake-out --target install -j$(( $(sysctl -n hw.ncpu) - 1 ))
        echo "::endgroup::"

        echo "::group::Build the C++ runner"
        # Standalone find_package(executorch) project, so it needs the install above.
        ${CONDA_RUN} cmake -S backends/mlx/examples/llm \
          -B cmake-out/backends/mlx/examples/llm -DCMAKE_BUILD_TYPE=Release
        ${CONDA_RUN} cmake --build cmake-out/backends/mlx/examples/llm \
          -j$(( $(sysctl -n hw.ncpu) - 1 ))
        RUNNER=cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf
        if [ ! -x "${RUNNER}" ]; then
          echo "Failed: runner not found at ${RUNNER}"
          exit 1
        fi
        echo "::endgroup::"

        echo "::group::Install LLM requirements"
        ${CONDA_RUN} pip install -U "huggingface_hub[cli]<1.0"
        ${CONDA_RUN} huggingface-cli login --token $SECRET_EXECUTORCH_HF_TOKEN
        OPTIMUM_ET_VERSION=$(cat .ci/docker/ci_commit_pins/optimum-executorch.txt)
        ${CONDA_RUN} pip install transformers "optimum-executorch @ git+https://github.com/huggingface/optimum-executorch.git@${OPTIMUM_ET_VERSION}"
        if [ "${MODEL_ID}" = "google/gemma-4-E2B-it" ]; then
          # Gemma 4 needs a newer Transformers than the CI-wide pin. Keep this
          # on the same commit test-mlx-llm validated against.
          GEMMA4_TRANSFORMERS_COMMIT=61461a7bcb458db7cf6eeea49678b9ab776a7821
          ${CONDA_RUN} pip install -U "transformers @ git+https://github.com/huggingface/transformers.git@${GEMMA4_TRANSFORMERS_COMMIT}"
        fi
        echo "::endgroup::"

        echo "::group::Export ${MODEL_NAME} off-graph"
        ${CONDA_RUN} python -m executorch.backends.mlx.examples.llm.export_llm_hf \
          --model-id "${MODEL_ID}" \
          --output /tmp/${MODEL_NAME}_offgraph.pte \
          --use-offgraph-cache \
          --max-seq-len 1024 \
          --dtype bf16 \
          --qlinear 4w
        echo "::endgroup::"

        echo "::group::Run ${MODEL_NAME} off-graph inference"
        # The cache geometry comes from the .pte; only capacity is given here.
        TOKENIZER=$(${CONDA_RUN} python -c \
          "from huggingface_hub import hf_hub_download; print(hf_hub_download('${MODEL_ID}', 'tokenizer.json'))")
        OUTPUT=$(${RUNNER} \
          --pte /tmp/${MODEL_NAME}_offgraph.pte \
          --tokenizer "${TOKENIZER}" \
          --chat "${CHAT}" \
          --kv-max-capacity 1024 \
          --prompt "What is the capital of France?" \
          --max-new-tokens 50 2>&1)
        echo "$OUTPUT"
        if echo "$OUTPUT" | grep -iq "Paris"; then
          echo "Success: 'Paris' found in output"
        else
          echo "Failed: Expected 'Paris' not found in output"
          exit 1
        fi
        echo "::endgroup::"