mlx.yml

name: MLX

on:
  push:
    branches:
      - main
      - release/*
    tags:
      - ciflow/mlx/*
  pull_request:
    paths:
      - .github/workflows/mlx.yml
      - Makefile
      - backends/mlx/**
      - extension/llm/batching/**
      - extension/llm/cache/**
      - extension/llm/export/**
      - extension/llm/runner/**
      - extension/llm/sampler/**
      - extension/llm/serving/**
      - extension/audio/**
      - examples/llm_server/cpp/**
      - examples/models/gemma4_31b/**
      - examples/models/muse-glimmer/**
      - examples/models/nemotron3_diarization/**
      - examples/models/parakeet/**
      - examples/models/supertonic/**
      - 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_BUILD_EXTENSION_RUNNER_UTIL=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"
        # Build all configured targets so newly registered tests are built too.
        ${CONDA_RUN} cmake --build cmake-out -j$(( $(sysctl -n hw.ncpu) - 1 ))
        echo "::endgroup::"

        echo "::group::Check MLX artifact sizes"
        size_check_failed=0
        check_artifact_size() {
          local path="$1"
          local max_size="$2"
          local baseline_size="$3"

          if [[ ! -f "${path}" ]]; then
            echo "::error file=${path}::Expected MLX artifact not found"
            size_check_failed=1
            return
          fi

          local actual_size
          actual_size=$(wc -c < "${path}" | tr -d '[:space:]')
          echo "${path}: ${actual_size} bytes (baseline ${baseline_size}, limit ${max_size})"
          if (( actual_size > max_size )); then
            echo "::error file=${path}::MLX artifact is ${actual_size} bytes, exceeding the ${max_size}-byte limit (baseline ${baseline_size})"
            size_check_failed=1
          fi
        }

        # Baselines recorded from the MLX v0.32.2 Release build on 2026-09-02.
        check_artifact_size cmake-out/backends/mlx/libmlxdelegate.a 5000000 4159552
        check_artifact_size cmake-out/backends/mlx/mlx/libmlx.a 14000000 11458544
        check_artifact_size cmake-out/backends/mlx/mlx/mlx/backend/metal/kernels/mlx.metallib 2000000 1460840
        echo "::endgroup::"
        if (( size_check_failed )); then
          exit 1
        fi

        echo "::group::Run registered MLX C++ tests"
        ${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 \
          ${CONDA_RUN} ctest --test-dir cmake-out/backends/mlx/test --output-on-failure --no-tests=error
        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 \
          examples/models/muse-glimmer/tests/test_mlx_offgraph.py \
          examples/models/muse-glimmer/tests/test_checkpoint_context.py \
          -v
        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-supertonic:
    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-supertonic
      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 Supertonic requirements"
        ${CONDA_RUN} python install_executorch.py > /dev/null
        ${CONDA_RUN} pip install -r examples/models/supertonic/requirements.txt
        echo "::endgroup::"

        echo "::group::Run asset-independent Supertonic Python tests"
        ${CONDA_RUN} python -m pytest examples/models/supertonic/tests -v
        echo "::endgroup::"

        echo "::group::Download pinned Supertonic assets"
        SUPERTONIC_ASSETS=/tmp/supertonic-3
        ${CONDA_RUN} python - <<'PY'
        from huggingface_hub import snapshot_download

        snapshot_download(
            "Supertone/supertonic-3",
            revision="3cadd1ee6394adea1bd021217a0e650ede09a323",
            allow_patterns=(
                "onnx/tts.json",
                "onnx/duration_predictor.onnx",
                "onnx/text_encoder.onnx",
                "onnx/vector_estimator.onnx",
                "onnx/vocoder.onnx",
                "onnx/unicode_indexer.json",
                "voice_styles/F1.json",
                "LICENSE",
            ),
            local_dir="/tmp/supertonic-3",
        )
        PY
        echo "::endgroup::"

        echo "::group::Run published-model parity tests"
        SUPERTONIC_MODEL_DIR="${SUPERTONIC_ASSETS}" \
          ${CONDA_RUN} python -m pytest \
            examples/models/supertonic/tests/test_checkpoint_loader.py \
            examples/models/supertonic/tests/test_stage_parity.py \
            -v
        echo "::endgroup::"

        echo "::group::Export the Supertonic PTE"
        SUPERTONIC_PTE="${SUPERTONIC_ASSETS}/supertonic_fp16_mlx.pte"
        ${CONDA_RUN} python -m examples.models.supertonic.export.export_supertonic \
          --asset-dir "${SUPERTONIC_ASSETS}" \
          --output "${SUPERTONIC_PTE}" \
          --max-text-length 512 \
          --max-latent-length 512 \
          --flow-steps 5
        echo "::endgroup::"

        echo "::group::Build ExecuTorch and Supertonic with MLX"
        ${CONDA_RUN} make supertonic-mlx
        echo "::endgroup::"

        echo "::group::Run Supertonic native tests"
        pushd examples/models/supertonic
        ${CONDA_RUN} cmake --preset supertonic-mlx \
          -DSUPERTONIC_INTEGRATION_PTE="${SUPERTONIC_PTE}" \
          -DSUPERTONIC_INTEGRATION_ASSET_DIR="${SUPERTONIC_ASSETS}" \
          -DSUPERTONIC_INTEGRATION_STYLE="${SUPERTONIC_ASSETS}/voice_styles/F1.json"
        ${CONDA_RUN} ctest --preset supertonic-mlx
        popd
        echo "::endgroup::"

        echo "::group::Verify the Supertonic runner and MLX metallib"
        RUNNER=cmake-out/examples/models/supertonic/supertonic_runner
        test -x "${RUNNER}"
        test -f "$(dirname "${RUNNER}")/mlx.metallib"
        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::"

  # No full-model export: Muse Glimmer GGUF checkpoints are too large for CI.
  # Cover runner compile/link/metallib wiring and asset-independent preparation.
  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::"

        echo "::group::Build and test Muse Glimmer batching with MLX"
        ${CONDA_RUN} make muse-glimmer-batching-mlx
        BATCHING_BIN_DIR="${PWD}/cmake-out/examples/models/muse-glimmer-batching"
        ${CONDA_RUN} cmake -S "${PWD}/examples/models/muse-glimmer" -B "${BATCHING_BIN_DIR}" \
          -DMUSE_GLIMMER_BUILD_BATCHING_TESTS=ON
        ${CONDA_RUN} cmake --build "${BATCHING_BIN_DIR}" --target muse_glimmer_batching_test --parallel
        ${CONDA_RUN} ctest --test-dir "${BATCHING_BIN_DIR}" \
          -R '^muse_glimmer_batching$' --output-on-failure --no-tests=error
        test -x "${BATCHING_BIN_DIR}/muse_glimmer_batching_runner"
        test -x "${BATCHING_BIN_DIR}/muse_glimmer_batching_worker"
        test -f "${BATCHING_BIN_DIR}/mlx.metallib"
        echo "::endgroup::"

  test-mlx-nemotron3-diarization-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-nemotron3-diarization-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
        ${CONDA_RUN} python install_executorch.py > /dev/null
        ${CONDA_RUN} make nemotron3-diarization-mlx
        BIN_DIR=cmake-out/examples/models/nemotron3_diarization
        test -x "${BIN_DIR}/nemotron3_diarization_runner"
        test -f "${BIN_DIR}/mlx.metallib"
        # Ensure the shared install still supplies the speech/LLM runner targets.
        test -f cmake-out/lib/libextension_asr_runner.a
        test -f cmake-out/lib/libextension_llm_runner.a

  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::"

  # DFlash speculative decoding: target and draft exported as two methods of one
  # .pte. Also the only coverage that the constants the export publishes
  # (get_max_context_len / get_max_seq_len / get_max_block_len /
  # get_mask_token_id) match what
  # the runner derives from them -- the runner takes no capacity or block flags.
  test-mlx-dflash:
    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-dflash
      runner: macos-14-xlarge
      python-version: "3.12"
      submodules: recursive
      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      # The 4B target is the smallest checkpoint DFlash drafts are published
      # for, so this job downloads and quantizes far more than the 1B LLM jobs.
      timeout: 120
      script: |
        set -eux
        export HF_HUB_DISABLE_XET=1

        MAX_CTX_LEN=1024
        PREFILL_CHUNK_SIZE=128
        N_DRAFT=3
        DRAFT_MODEL=z-lab/Qwen3-4B-DFlash-b16
        TOKENIZER_MODEL=Qwen/Qwen3-4B
        PTE=/tmp/qwen3_4b_dflash.pte

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

        echo "::group::Install DFlash requirements"
        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}"
        echo "::endgroup::"

        ${CONDA_RUN} pip list

        echo "::group::Export Qwen3-4B DFlash (target + draft)"
        # No --block-size: it must come from the draft checkpoint's config.
        ${CONDA_RUN} python -m executorch.backends.mlx.examples.llm.dflash.export \
          --output ${PTE} \
          --draft-model ${DRAFT_MODEL} \
          --target-model ${TOKENIZER_MODEL} \
          --max-ctx-len ${MAX_CTX_LEN} \
          --prefill-chunk-size ${PREFILL_CHUNK_SIZE}
        echo "::endgroup::"

        echo "::group::Cache the tokenizer"
        # The runner loads the tokenizer with local_files_only, so fetch it here
        # rather than letting an unauthenticated request happen mid-run. Also
        # confirms the chat template is present -- a template that fails to load
        # renders to nothing and the prompt tokenizes to zero tokens.
        TOKENIZER_MODEL=${TOKENIZER_MODEL} ${CONDA_RUN} python -c "import os; from transformers import AutoTokenizer; m = os.environ['TOKENIZER_MODEL']; t = AutoTokenizer.from_pretrained(m); assert t.chat_template, f'{m} has no chat_template'; ids = t.apply_chat_template([{'role': 'user', 'content': 'hi'}], add_generation_prompt=True, tokenize=True, return_tensors='pt'); ids = getattr(ids, 'input_ids', ids); assert ids.numel() > 0, 'chat template rendered no tokens'; print(f'tokenizer cached, chat template renders {ids.numel()} tokens')"
        echo "::endgroup::"

        echo "::group::Check published constant methods"
        # The runner reads these instead of taking flags, so a mismatch here is
        # a silently broken runner rather than a failed export. Kept on one line:
        # a YAML block scalar indents every line, which python -c would reject.
        BLOCK_SIZE=$(${CONDA_RUN} python -c "from huggingface_hub import snapshot_download; from executorch.backends.mlx.examples.llm.dflash.model import load_dflash_config; print(load_dflash_config(snapshot_download('${DRAFT_MODEL}', allow_patterns=['*.json'])).block_size)")
        MASK_TOKEN_ID=$(${CONDA_RUN} python -c "from huggingface_hub import snapshot_download; from executorch.backends.mlx.examples.llm.dflash.model import load_dflash_config; print(load_dflash_config(snapshot_download('${DRAFT_MODEL}', allow_patterns=['*.json'])).mask_token_id)")
        echo "draft checkpoint block_size: ${BLOCK_SIZE} mask_token_id: ${MASK_TOKEN_ID}"
        ${CONDA_RUN} python -c "from executorch.runtime import Runtime, Verification; from executorch.backends.mlx.examples.llm.runtime_meta import read_const_int as r; p = Runtime.get().load_program('${PTE}', verification=Verification.Minimal); got = {n: r(p, n) for n in ['get_max_context_len', 'get_max_seq_len', 'get_max_block_len', 'get_mask_token_id']}; want = {'get_max_context_len': ${MAX_CTX_LEN}, 'get_max_seq_len': ${PREFILL_CHUNK_SIZE}, 'get_max_block_len': ${BLOCK_SIZE}, 'get_mask_token_id': ${MASK_TOKEN_ID}}; print('published:', got); assert got == want, f'mismatch: {got} != {want}'; assert {'draft', 'target'} <= set(p.method_names), p.method_names; print('Success: constants match and both methods present')"
        echo "::endgroup::"

        echo "::group::Run DFlash speculative decoding"
        # --n-draft below the exported maximum, so a regression that silently
        # pins the block to BLOCK_SIZE gets caught.
        OUTPUT=$(${CONDA_RUN} python -m executorch.backends.mlx.examples.llm.dflash.run \
          --pte ${PTE} \
          --tokenizer ${TOKENIZER_MODEL} \
          --n-draft ${N_DRAFT} \
          --prompt "What is the capital of France?" \
          --max-new-tokens 32 \
          --verbose 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
        # A special token in the decoded text means skip_special_tokens regressed.
        if echo "$OUTPUT" | grep -q "<|im_end|>"; then
          echo "Failed: special tokens leaked into the decoded output"
          exit 1
        fi
        echo "::endgroup::"

        echo "::group::Check speculation actually happened"
        # tau <= 1.0 means every round fell back to a target-only step, i.e. the
        # draft was never accepted and DFlash degenerated to plain decoding.
        # tau also cannot exceed n_draft + 1, the emitted tokens per round.
        TAU=$(echo "$OUTPUT" | grep "avg emitted/round" | grep -oE "[0-9]+\.[0-9]+")
        echo "avg emitted/round (tau): ${TAU}"
        awk "BEGIN { exit !(${TAU} > 1.0 && ${TAU} <= ${N_DRAFT} + 1) }" || {
          echo "Failed: tau ${TAU} outside (1.0, $((N_DRAFT + 1))]"
          exit 1
        }
        echo "Success: tau ${TAU} in (1.0, $((N_DRAFT + 1))]"
        echo "::endgroup::"

        echo "::group::Reject an out-of-range --n-draft"
        # block_len is n_draft + 1 and the graph traced block_len <= BLOCK_SIZE,
        # so BLOCK_SIZE itself must be refused before reaching the runtime.
        if ${CONDA_RUN} python -m executorch.backends.mlx.examples.llm.dflash.run \
          --pte ${PTE} \
          --tokenizer ${TOKENIZER_MODEL} \
          --n-draft ${BLOCK_SIZE} \
          --max-new-tokens 1 > /tmp/n_draft_err.txt 2>&1; then
          echo "Failed: --n-draft ${BLOCK_SIZE} should have been rejected"
          cat /tmp/n_draft_err.txt
          exit 1
        fi
        cat /tmp/n_draft_err.txt
        grep -q "must be in \[1, $((BLOCK_SIZE - 1))\]" /tmp/n_draft_err.txt
        echo "Success: out-of-range --n-draft rejected"
        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
        BATCHED_RUNNER=cmake-out/backends/mlx/examples/llm/mlx_run_llm_batched
        for binary in "${RUNNER}" "${BATCHED_RUNNER}"; do
          if [ ! -x "${binary}" ]; then
            echo "Failed: runner not found at ${binary}"
            exit 1
          fi
        done
        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 \
          --logits-to-keep selected \
          --max-ctx-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::"

        echo "::group::Run ${MODEL_NAME} batched off-graph inference"
        ${BATCHED_RUNNER} \
          --pte /tmp/${MODEL_NAME}_offgraph.pte \
          --tokenizer "${TOKENIZER}" \
          --chat "${CHAT}" \
          --max-session-tokens 1024 \
          --max-new-tokens 50 \
          --max-decode-sequences 2 \
          --out-prefix /tmp/${MODEL_NAME}_batched \
          "What is the capital of France?" \
          "What color is grass?"
        for check in "0:Paris" "1:green"; do
          index="${check%%:*}"
          expected="${check#*:}"
          output="/tmp/${MODEL_NAME}_batched_${index}.txt"
          if grep -iq "${expected}" "${output}"; then
            echo "Success: '${expected}' found in ${output}"
          else
            echo "Failed: Expected '${expected}' not found in ${output}"
            cat "${output}"
            exit 1
          fi
        done
        echo "::endgroup::"