From f800e4be62c18d601e4c0acb60761202b9d1952f Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Tue, 20 Oct 2026 23:56:31 -0700 Subject: [PATCH] [executorch][ci] Split unittest-cuda into kernel and system shards unittest-cuda has timed out at 90 min on every main run since a04a9f9918 made the SDPA autotune tests ~6.5 min slower; the job was already at 83-89 min before that. Split it into a matrix: - kernel: Triton kernel and autotune tests (SDPA, tq4 SDPA, int4 matmul, int5 large GEMM, fused MoE, deltanet, topk, autotune inputs/launch params, mutation cache). ~42 min of pytest. - system: everything else under backends/cuda/tests and passes/tests (lowering passes, op dispatch, export/partitioner/runtime), plus the pybindings device tests, C++ shim tests, and Qwen 3.5 MoE / Gemma 4 31B model tests. ~28 min of pytest. The system shard runs the test directories with the kernel and quantized GEMM files ignored, so new tests default to it. Test plan: CI on this PR; both unittest-cuda shards should finish well under the 90 min timeout. --- .github/workflows/cuda.yml | 137 ++++++++++++++++++++++++------------- 1 file changed, 89 insertions(+), 48 deletions(-) diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 23ec43c78ce..2de810544c9 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -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. @@ -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 @@ -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 @@ -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