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

unittest-cuda has timed out at 90 min on every main run since a04a9f9
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.

Author: Gasoonjia

.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