Enable use_sdpa_with_kv_cache in the qwen3_5 example config (#22390) #125
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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/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_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 mlx_cell_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 | |
| ./cmake-out/backends/mlx/test/mlx_cell_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-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::" | |
| # 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::" | |
| # 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_ctx_len / get_prefill_chunk_size / 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_ctx_len', 'get_prefill_chunk_size', 'get_max_block_len', 'get_mask_token_id']}; want = {'get_max_ctx_len': ${MAX_CTX_LEN}, 'get_prefill_chunk_size': ${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 | |
| 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-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::" |