55# This source code is licensed under the BSD-style license found in the
66# LICENSE file in the root directory of this source tree.
77
8- # Test CUDA model end-to-end, need to run .ci/scripts/export_model_cuda_artifact .sh first
8+ # Test CUDA/Metal model end-to-end, need to run .ci/scripts/export_model_artifact .sh first
99
1010show_help () {
1111 cat << EOF
12- Usage: test_model_cuda_e2e .sh <hf_model> <quant_name> [model_dir]
12+ Usage: test_model_e2e .sh <device> <hf_model> <quant_name> [model_dir]
1313
14- Build and run end-to-end tests for CUDA models.
14+ Build and run end-to-end tests for CUDA/Metal models.
1515
1616Arguments:
17+ device cuda or metal (required)
18+
1719 hf_model HuggingFace model ID (required)
1820 Supported models:
1921 - mistralai/Voxtral-Mini-3B-2507
@@ -27,12 +29,12 @@ Arguments:
2729 - quantized-int4-weight-only
2830
2931 model_dir Directory containing model artifacts (optional, default: current directory)
30- Expected files: model.pte, aoti_cuda_blob.ptd
32+ Expected files: model.pte, aoti_cuda_blob.ptd/aoti_metal_blob.ptd
3133 Tokenizers and test files will be downloaded to this directory
3234
3335Examples:
34- test_model_cuda_e2e .sh "openai/whisper-small" "non-quantized"
35- test_model_cuda_e2e .sh "mistralai/Voxtral-Mini-3B-2507" "quantized-int4-tile-packed" "./model_output"
36+ test_model_e2e .sh metal "openai/whisper-small" "non-quantized"
37+ test_model_e2e .sh cuda "mistralai/Voxtral-Mini-3B-2507" "quantized-int4-tile-packed" "./model_output"
3638EOF
3739}
3840
5557
5658set -eux
5759
58- HF_MODEL=" $1 "
59- QUANT_NAME=" $2 "
60+ DEVICE=" $1 "
61+ HF_MODEL=" $2 "
62+ QUANT_NAME=" $3 "
6063# Download tokenizers, audio, and image files to this directory
61- MODEL_DIR=" ${3 :- .} "
64+ MODEL_DIR=" ${4 :- .} "
6265
6366echo " Testing model: $HF_MODEL (quantization: $QUANT_NAME )"
6467
65- # Make sure model.pte and aoti_cuda_blob .ptd exist
68+ # Make sure model.pte and aoti_${DEVICE}_blob .ptd exist
6669if [ ! -f " $MODEL_DIR /model.pte" ]; then
6770 echo " Error: model.pte not found in $MODEL_DIR "
6871 exit 1
6972fi
70- if [ ! -f " $MODEL_DIR /aoti_cuda_blob .ptd" ]; then
71- echo " Error: aoti_cuda_blob .ptd not found in $MODEL_DIR "
73+ if [ ! -f " $MODEL_DIR /aoti_ ${DEVICE} _blob .ptd" ]; then
74+ echo " Error: aoti_ ${DEVICE} _blob .ptd not found in $MODEL_DIR "
7275 exit 1
7376fi
7477# Locate EXECUTORCH_ROOT from the directory of this script
@@ -152,14 +155,24 @@ ls -al
152155echo " ::endgroup::"
153156
154157echo " ::group::Build $MODEL_NAME Runner"
158+
159+ if [ " $DEVICE " = " cuda" ]; then
160+ BUILD_BACKEND=" EXECUTORCH_BUILD_CUDA"
161+ elif [ " $DEVICE " = " metal" ]; then
162+ BUILD_BACKEND=" EXECUTORCH_BUILD_METAL"
163+ else
164+ echo " Error: Unsupported device '$DEVICE '. Must be 'cuda' or 'metal'."
165+ exit 1
166+ fi
167+
155168cmake --preset llm \
156- -DEXECUTORCH_BUILD_CUDA =ON \
169+ -D ${BUILD_BACKEND} =ON \
157170 -DCMAKE_INSTALL_PREFIX=cmake-out \
158171 -DCMAKE_BUILD_TYPE=Release \
159172 -Bcmake-out -S.
160173cmake --build cmake-out -j$( nproc) --target install --config Release
161174
162- cmake -DEXECUTORCH_BUILD_CUDA =ON \
175+ cmake -D ${BUILD_BACKEND} =ON \
163176 -DCMAKE_BUILD_TYPE=Release \
164177 -Sexamples/models/$RUNNER_PATH \
165178 -Bcmake-out/examples/models/$RUNNER_PATH /
@@ -168,11 +181,13 @@ echo "::endgroup::"
168181
169182echo " ::group::Run $MODEL_NAME Runner"
170183set +e
171- export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH
184+ if [ " $DEVICE " = " cuda" ]; then
185+ export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH
186+ fi
172187
173188# Build runner command with common arguments
174189RUNNER_BIN=" cmake-out/examples/models/$RUNNER_PATH /$RUNNER_TARGET "
175- RUNNER_ARGS=" --model_path ${MODEL_DIR} /model.pte --data_path ${MODEL_DIR} /aoti_cuda_blob .ptd --temperature 0"
190+ RUNNER_ARGS=" --model_path ${MODEL_DIR} /model.pte --data_path ${MODEL_DIR} /aoti_ ${DEVICE} _blob .ptd --temperature 0"
176191
177192# Add model-specific arguments
178193case " $MODEL_NAME " in
0 commit comments