Skip to content

[executorch][ci] Split unittest-cuda into kernel and system shards #24047

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Gasoonjia wants to merge 1 commit into from
+89 −48
Open

[executorch][ci] Split unittest-cuda into kernel and system shards #24047

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
137 changes: 89 additions & 48 deletions .github/workflows/cuda.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -467,7 +467,7 @@ jobs:
"${RUNNER_ARTIFACT_DIR}/native-sm86-on-sm80.log"

unittest-cuda:
name: unittest-cuda
name: unittest-cuda (${{ matrix.shard }})
needs: [changed-files, run-decision]
# This job runs the pybindings device test, so it also has to fire on the
# two pybindings paths the workflow filter lists.
Expand All @@ -482,6 +482,13 @@ jobs:
contains(needs.changed-files.outputs.changed-files, 'extension/pybindings') ||
contains(needs.changed-files.outputs.changed-files, 'runtime/__init__.py') ||
needs.run-decision.outputs.is-full-run == 'true'
# kernel: Triton kernels and their autotuning. system: lowering passes, op
# dispatch, export/runtime, and model tests. Together they do not fit the
# timeout.
strategy:
fail-fast: false
matrix:
shard: [kernel, system]
uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main
permissions:
id-token: write
Expand All @@ -495,6 +502,37 @@ jobs:
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
script: |
set -eux
SHARD="${{ matrix.shard }}"

# The system shard runs everything under backends/cuda/tests that is not
# listed here, so a new test lands there unless it is added to this list.
KERNEL_TESTS=(
backends/cuda/tests/test_triton_sdpa.py
backends/cuda/tests/test_triton_sdpa_splitk.py
backends/cuda/tests/test_triton_sdpa_nan.py
backends/cuda/tests/test_tq4_sdpa.py
backends/cuda/tests/test_int4_matmul.py
backends/cuda/tests/test_int5_large_quantized_gemm.py
backends/cuda/tests/test_quantized_gemm_family.py
backends/cuda/tests/test_fused_moe.py
backends/cuda/tests/test_fused_deltanet_decode.py
backends/cuda/tests/test_chunk_gated_delta_rule.py
backends/cuda/tests/test_topk.py
backends/cuda/tests/test_autotune_inputs.py
backends/cuda/tests/test_autotune_launch_params.py
backends/cuda/tests/test_triton_mutation_cache.py
)
# These run in unittest-cuda-quantized-gemm.
QUANTIZED_GEMM_TESTS=(
backends/cuda/tests/test_int4_quantized_gemm.py
backends/cuda/tests/test_int4_large_quantized_gemm.py
backends/cuda/tests/test_int5_quantized_gemm.py
backends/cuda/tests/test_int6_quantized_gemm.py
backends/cuda/tests/test_int6_large_quantized_gemm.py
backends/cuda/tests/test_int8_quantized_gemm.py
backends/cuda/tests/test_cuda_graph_timing.py
)

# Install executorch in editable mode so custom op libs land in-tree
bash ./install_executorch.sh

Expand All @@ -503,64 +541,67 @@ jobs:
conda install -y -c conda-forge 'libstdcxx-ng>=12'
export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH

# The pybindings loader only asks for device memory on a build that has
# a device allocator registered, so these tests cannot run in the CPU
# jobs. They need the install above and nothing built below, so run them
# here and fail before the long builds.
#
# unittest exits 0 when a test skips, and these tests skip when the CUDA
# backend is missing or no GPU is visible. Both are the prerequisites
# this job exists to provide, so treat a skip as a failure rather than
# let a packaging regression pass silently. `tee` would otherwise hide
# the real exit status behind its own, so read it back explicitly.
set +e
python -m unittest -v \
executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_device_planned_method_allocates_on_the_device \
executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_unplanned_device_output_is_device_memory \
2>&1 | tee /tmp/pybindings_device_test.log
pybindings_device_status=${PIPESTATUS[0]}
set -e
if [ "${pybindings_device_status}" -ne 0 ]; then
echo "::error::a pybindings device test failed"
exit "${pybindings_device_status}"
fi
if grep -qE "^OK \(skipped=|skipped=[1-9]" /tmp/pybindings_device_test.log; then
echo "::error::a pybindings device test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing"
exit 1
if [ "${SHARD}" = "system" ]; then
# The pybindings loader only asks for device memory on a build that has
# a device allocator registered, so these tests cannot run in the CPU
# jobs. They need the install above and nothing built below, so run them
# here and fail before the long builds.
#
# unittest exits 0 when a test skips, and these tests skip when the CUDA
# backend is missing or no GPU is visible. Both are the prerequisites
# this job exists to provide, so treat a skip as a failure rather than
# let a packaging regression pass silently. `tee` would otherwise hide
# the real exit status behind its own, so read it back explicitly.
set +e
python -m unittest -v \
executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_device_planned_method_allocates_on_the_device \
executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_unplanned_device_output_is_device_memory \
2>&1 | tee /tmp/pybindings_device_test.log
pybindings_device_status=${PIPESTATUS[0]}
set -e
if [ "${pybindings_device_status}" -ne 0 ]; then
echo "::error::a pybindings device test failed"
exit "${pybindings_device_status}"
fi
if grep -qE "^OK \(skipped=|skipped=[1-9]" /tmp/pybindings_device_test.log; then
echo "::error::a pybindings device test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing"
exit 1
fi
fi

# Build ExecuTorch with CUDA support
cmake --workflow --preset llm-release-cuda

# Build and run CUDA shim tests (C++)
pushd backends/cuda/runtime/shims/tests
cmake --workflow --preset default
popd
if [ "${SHARD}" = "system" ]; then
# Build and run CUDA shim tests (C++)
pushd backends/cuda/runtime/shims/tests
cmake --workflow --preset default
popd
fi

# Install flash-linear-attention for chunk_gated_delta_rule triton kernel tests
pip install "flash-linear-attention==0.4.2"

# Build executor_runner (needed by CUDA backend e2e tests)
# Build executor_runner (needed by CUDA backend e2e tests in both shards)
cmake --build cmake-out --target executor_runner

# Run CUDA backend Python tests. The quantized GEMM kernel suites run
# in unittest-cuda-quantized-gemm; together they do not fit this
# job's timeout.
python -m pytest backends/cuda/tests backends/cuda/passes/tests -v -o "addopts=" \
--ignore=backends/cuda/tests/test_int4_quantized_gemm.py \
--ignore=backends/cuda/tests/test_int4_large_quantized_gemm.py \
--ignore=backends/cuda/tests/test_int5_quantized_gemm.py \
--ignore=backends/cuda/tests/test_int6_quantized_gemm.py \
--ignore=backends/cuda/tests/test_int6_large_quantized_gemm.py \
--ignore=backends/cuda/tests/test_int8_quantized_gemm.py \
--ignore=backends/cuda/tests/test_cuda_graph_timing.py

# Run Qwen 3.5 MoE tests (quantize roundtrip + TurboQuant KV cache + sampler)
python -m pytest examples/models/qwen3_5_moe/test_quantize_roundtrip.py examples/models/qwen3_5_moe/test_turboquant.py examples/models/qwen3_5_moe/test_sampler.py -v -o "addopts="

# Run Gemma 4 31B tests (quant unit tests + pipeline integration tests)
pip install gguf
python -m pytest examples/models/gemma4_31b/tests/ --ignore=examples/models/gemma4_31b/tests/test_mlx_pipeline.py -v -o "addopts="
if [ "${SHARD}" = "kernel" ]; then
python -m pytest "${KERNEL_TESTS[@]}" -v -o "addopts=" --durations=20
else
IGNORES=()
for t in "${KERNEL_TESTS[@]}" "${QUANTIZED_GEMM_TESTS[@]}"; do
IGNORES+=("--ignore=${t}")
done
python -m pytest backends/cuda/tests backends/cuda/passes/tests -v -o "addopts=" --durations=20 \
"${IGNORES[@]}"

# Run Qwen 3.5 MoE tests (quantize roundtrip + TurboQuant KV cache + sampler)
python -m pytest examples/models/qwen3_5_moe/test_quantize_roundtrip.py examples/models/qwen3_5_moe/test_turboquant.py examples/models/qwen3_5_moe/test_sampler.py -v -o "addopts="

# Run Gemma 4 31B tests (quant unit tests + pipeline integration tests)
pip install gguf
python -m pytest examples/models/gemma4_31b/tests/ --ignore=examples/models/gemma4_31b/tests/test_mlx_pipeline.py -v -o "addopts="
fi

unittest-cuda-muse-glimmer-batching:
name: unittest-cuda-muse-glimmer-batching
Expand Down
Loading