ci: compute matrix partition counts from est_time (#25193)

.github/workflows/_pr-test-check-changes.yml

@@ -32,12 +32,8 @@ on:
         value: ${{ jobs.run.outputs.jit_kernel }}
       multimodal_gen:
         value: ${{ jobs.run.outputs.multimodal_gen }}
-      max_parallel:
-        value: ${{ jobs.run.outputs.max_parallel }}
-      max_parallel_small:
-        value: ${{ jobs.run.outputs.max_parallel_small }}
-      max_parallel_2gpu:
-        value: ${{ jobs.run.outputs.max_parallel_2gpu }}
+      partitions:
+        value: ${{ jobs.run.outputs.partitions }}
       b200_runner:
         value: ${{ jobs.run.outputs.b200_runner }}
       enable_retry:
@@ -63,9 +59,7 @@ jobs:
       sgl_kernel_raw: ${{ steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel }}
       jit_kernel: ${{ steps.filter-api.outputs.jit_kernel || steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }}
       multimodal_gen: ${{ steps.filter-api.outputs.multimodal_gen || steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }}
-      max_parallel: ${{ steps.set-parallel.outputs.max_parallel }}
-      max_parallel_small: ${{ steps.set-parallel.outputs.max_parallel_small }}
-      max_parallel_2gpu: ${{ steps.set-parallel.outputs.max_parallel_2gpu }}
+      partitions: ${{ steps.partitions.outputs.partitions }}
       b200_runner: ${{ steps.set-runner.outputs.b200_runner }}
       enable_retry: ${{ steps.set-retry.outputs.enable_retry }}
       continue_on_error: ${{ steps.set-continue-on-error.outputs.continue_on_error }}
@@ -198,22 +192,29 @@ jobs:
             echo "multimodal_gen=false" >> $GITHUB_OUTPUT
           fi
 
-      - name: Set max-parallel based on run type
-        id: set-parallel
+      - name: Determine full-parallel mode
+        id: parallel-mode
         env:
           GH_TOKEN: ${{ github.token }}
         run: |
-          # Determine if this run gets full parallelism (scheduled / high priority)
+          # `full=true` lifts the matrix-fanout throttle so each suite's
+          # max_parallel = size. Conditions (matching the prior set-parallel
+          # step exactly):
+          #   1. Scheduled cron run.
+          #   2. pull_request event with the `high priority` label.
+          #   3. workflow_dispatch with target_stage set (i.e. /rerun-stage)
+          #      whose underlying PR carries `high priority`. The labels
+          #      aren't on the dispatch payload, so look them up via API:
+          #      try SHA -> /pulls (works for fork PRs), fall back to
+          #      branch name -> gh pr list (works for non-fork PRs).
           FULL=false
           if [[ "${{ github.event_name }}" == "schedule" ]]; then
             FULL=true
-            echo "Scheduled run detected, using full parallelism"
+            echo "Scheduled run -> full parallelism"
           elif [[ "${{ github.event_name }}" == "pull_request" && "${{ contains(github.event.pull_request.labels.*.name, 'high priority') }}" == "true" ]]; then
             FULL=true
-            echo "High priority PR detected, using full parallelism"
+            echo "high priority PR -> full parallelism"
           elif [[ -n "${{ inputs.target_stage }}" ]]; then
-            # /rerun-stage (workflow_dispatch): query PR labels via GitHub API
-            # Try SHA lookup first (fork PRs), fallback to branch name (non-fork PRs)
             LABELS=""
             PR_HEAD_SHA="${{ inputs.pr_head_sha }}"
             if [[ -n "$PR_HEAD_SHA" ]]; then
@@ -224,28 +225,25 @@ jobs:
               LABELS=$(gh pr list --head "${{ github.ref_name }}" --repo "${{ github.repository }}" \
                 --json labels --jq '.[0].labels[].name' 2>/dev/null || true)
             fi
-            echo "PR labels: ${LABELS:-"(none)"}"
+            echo "PR labels (via API): ${LABELS:-"(none)"}"
             if echo "$LABELS" | grep -Fxq "high priority"; then
               FULL=true
-              echo "High priority PR detected via API (/rerun-stage), using full parallelism"
+              echo "high priority PR (via API) -> full parallelism"
             fi
           fi
+          echo "full=$FULL" >> "$GITHUB_OUTPUT"
 
-          # Set max-parallel for each runner type
-          #   1-gpu-h100: 14 partitions, 1-gpu-5090: 8 partitions, 2-gpu-h100: 4 partitions
-          if [[ "$FULL" == "true" ]]; then
-            LEVEL=full
-            echo "max_parallel=14" >> $GITHUB_OUTPUT
-            echo "max_parallel_small=8" >> $GITHUB_OUTPUT
-            echo "max_parallel_2gpu=4" >> $GITHUB_OUTPUT
-          else
-            LEVEL=low
-            echo "max_parallel=3" >> $GITHUB_OUTPUT
-            echo "max_parallel_small=3" >> $GITHUB_OUTPUT
-            echo "max_parallel_2gpu=2" >> $GITHUB_OUTPUT
-          fi
-          echo "parallel_level=$LEVEL" >> $GITHUB_OUTPUT
-          echo "Parallelism level: $LEVEL"
+      - name: Compute partitions
+        id: partitions
+        run: |
+          # Emit a single JSON output `partitions` keyed by suite name with
+          # {size, arr, max_parallel} fields per suite. Replaces the prior
+          # full/low max-parallel presets; `--full-parallel` keeps the
+          # `high priority` PR / scheduled cron escape hatch.
+          # See scripts/ci/utils/compute_partitions.py.
+          python3 scripts/ci/utils/compute_partitions.py \
+            --full-parallel ${{ steps.parallel-mode.outputs.full }} \
+            >> "$GITHUB_OUTPUT"
 
       - name: Set B200 runner tag
         id: set-runner
@@ -318,7 +316,6 @@ jobs:
             echo "| multimodal_gen    | ${{ steps.filter-api.outputs.multimodal_gen || steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} |"
             echo "| target_stage      | ${{ inputs.target_stage || '(none)' }} |"
             echo "| detection_method  | ${{ inputs.target_stage && 'GitHub API' || 'dorny/paths-filter' }} |"
-            echo "| max_parallel      | ${{ steps.set-parallel.outputs.parallel_level }} (h100=${{ steps.set-parallel.outputs.max_parallel }}, 5090=${{ steps.set-parallel.outputs.max_parallel_small }}, 2gpu=${{ steps.set-parallel.outputs.max_parallel_2gpu }}) |"
             echo "| b200_runner       | ${{ steps.set-runner.outputs.b200_runner }} |"
             echo "| enable_retry      | ${{ steps.set-retry.outputs.enable_retry }} |"
             echo "| continue_on_error | ${{ steps.set-continue-on-error.outputs.continue_on_error }} |"

.github/workflows/pr-test.yml

@@ -122,7 +122,11 @@ jobs:
         id: wait
         with:
           stage-name: stage-a
-          jobs: '["stage-a-test-1-gpu-small", {"prefix": "stage-a-test-cpu", "expected_count": 4}]'
+          jobs: |
+            [
+              {"prefix": "stage-a-test-1-gpu-small", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-1-gpu-small'].size }}},
+              {"prefix": "stage-a-test-cpu", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-cpu'].size }}}
+            ]
           max-wait-minutes: '240'
 
   wait-for-stage-b:
@@ -150,10 +154,10 @@ jobs:
           stage-name: stage-b
           jobs: |
             [
-              {"prefix": "stage-b-test-1-gpu-small", "expected_count": 8},
-              {"prefix": "stage-b-test-1-gpu-large", "expected_count": 14},
-              {"prefix": "stage-b-test-2-gpu-large", "expected_count": 4},
-              {"prefix": "stage-b-test-4-gpu-b200", "expected_count": 1}
+              {"prefix": "stage-b-test-1-gpu-small", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-small'].size }}},
+              {"prefix": "stage-b-test-1-gpu-large", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-large'].size }}},
+              {"prefix": "stage-b-test-2-gpu-large", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-2-gpu-large'].size }}},
+              {"prefix": "stage-b-test-4-gpu-b200", "expected_count": ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-4-gpu-b200'].size }}}
             ]
           max-wait-minutes: '480'
 
@@ -285,6 +289,11 @@ jobs:
       )
     runs-on: 1-gpu-5090
     timeout-minutes: 240
+    strategy:
+      fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-1-gpu-small'].max_parallel }}
+      matrix:
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-1-gpu-small'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -314,7 +323,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test/
-          python3 run_suite.py --hw cuda --suite stage-a-test-1-gpu-small $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-a-test-1-gpu-small --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-1-gpu-small'].size }} $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
@@ -339,8 +348,9 @@ jobs:
     timeout-minutes: 240
     strategy:
       fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-cpu'].max_parallel }}
       matrix:
-        partition: [0, 1, 2, 3]
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-cpu'].arr }}
     steps:
       - name: Free disk space
         run: |
@@ -391,7 +401,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test/
-          python3 run_suite.py --hw cpu --suite stage-a-test-cpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 4 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cpu --suite stage-a-test-cpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-a-test-cpu'].size }} $CONTINUE_ON_ERROR_FLAG
 
   # Runs on 5090 (32GB, SM120)
   stage-b-test-1-gpu-small:
@@ -410,9 +420,9 @@ jobs:
     timeout-minutes: 240
     strategy:
       fail-fast: false
-      max-parallel: ${{ fromJson(needs.check-changes.outputs.max_parallel_small) }}
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-small'].max_parallel }}
       matrix:
-        partition: [0, 1, 2, 3, 4, 5, 6, 7]
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-small'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -442,7 +452,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test/
-          python3 run_suite.py --hw cuda --suite stage-b-test-1-gpu-small --auto-partition-id ${{ matrix.partition }} --auto-partition-size 8 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-b-test-1-gpu-small --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-small'].size }} $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
@@ -470,9 +480,9 @@ jobs:
     timeout-minutes: 240
     strategy:
       fail-fast: false
-      max-parallel: ${{ fromJson(needs.check-changes.outputs.max_parallel) }}
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-large'].max_parallel }}
       matrix:
-        partition: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-large'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -502,7 +512,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test/
-          python3 run_suite.py --hw cuda --suite stage-b-test-1-gpu-large --auto-partition-id ${{ matrix.partition }} --auto-partition-size 14 --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-b-test-1-gpu-large --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-1-gpu-large'].size }} --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
@@ -529,9 +539,9 @@ jobs:
     timeout-minutes: 240
     strategy:
       fail-fast: false
-      max-parallel: ${{ fromJson(needs.check-changes.outputs.max_parallel_2gpu) }}
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-2-gpu-large'].max_parallel }}
       matrix:
-        partition: [0, 1, 2, 3]
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-2-gpu-large'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -561,7 +571,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test/
-          python3 run_suite.py --hw cuda --suite stage-b-test-2-gpu-large --auto-partition-id ${{ matrix.partition }} --auto-partition-size 4 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-b-test-2-gpu-large --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-2-gpu-large'].size }} $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
@@ -588,6 +598,9 @@ jobs:
     timeout-minutes: 240
     strategy:
       fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-4-gpu-b200'].max_parallel }}
+      matrix:
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-4-gpu-b200'].arr }}
 
     steps:
       - name: Checkout code
@@ -618,7 +631,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test
-          python3 run_suite.py --hw cuda --suite stage-b-test-4-gpu-b200 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-b-test-4-gpu-b200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-b-test-4-gpu-b200'].size }} $CONTINUE_ON_ERROR_FLAG
 
       - name: Run FA4 jit_kernel tests (SM100+)
         timeout-minutes: 10
@@ -681,8 +694,9 @@ jobs:
     timeout-minutes: 240
     strategy:
       fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-h100'].max_parallel }}
       matrix:
-        part: [0, 1, 2]
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-h100'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -712,12 +726,12 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test
-          python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-h100 --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-h100 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-h100'].size }} $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
         with:
-          artifact-suffix: ${{ matrix.part }}
+          artifact-suffix: ${{ matrix.partition }}
 
       - name: Cleanup venv
         if: always()
@@ -739,8 +753,9 @@ jobs:
     timeout-minutes: 240
     strategy:
       fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h200'].max_parallel }}
       matrix:
-        part: [0, 1, 2, 3]
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h200'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -789,12 +804,12 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test
-          python3 run_suite.py --hw cuda --suite stage-c-test-8-gpu-h200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 4 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-c-test-8-gpu-h200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h200'].size }} $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
         with:
-          artifact-suffix: ${{ matrix.part }}
+          artifact-suffix: ${{ matrix.partition }}
 
       - name: Cleanup venv
         if: always()
@@ -816,6 +831,11 @@ jobs:
     timeout-minutes: 240
     env:
       SGLANG_CI_RDMA_ALL_DEVICES: "mlx5_1,mlx5_2,mlx5_3,mlx5_4"
+    strategy:
+      fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h20'].max_parallel }}
+      matrix:
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h20'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -845,7 +865,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test
-          python3 run_suite.py --hw cuda --suite stage-c-test-8-gpu-h20 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-c-test-8-gpu-h20 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-8-gpu-h20'].size }} $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
@@ -868,6 +888,11 @@ jobs:
       )
     runs-on: 4-gpu-h100
     timeout-minutes: 240
+    strategy:
+      fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-4-gpu-h100'].max_parallel }}
+      matrix:
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-4-gpu-h100'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -914,7 +939,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test
-          python3 run_suite.py --hw cuda --suite stage-c-test-deepep-4-gpu-h100 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-c-test-deepep-4-gpu-h100 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-4-gpu-h100'].size }} $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
@@ -937,6 +962,11 @@ jobs:
       )
     runs-on: 8-gpu-h200-deepep
     timeout-minutes: 240
+    strategy:
+      fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-8-gpu-h200'].max_parallel }}
+      matrix:
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-8-gpu-h200'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -984,7 +1014,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test
-          python3 run_suite.py --hw cuda --suite stage-c-test-deepep-8-gpu-h200 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-c-test-deepep-8-gpu-h200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-deepep-8-gpu-h200'].size }} $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
@@ -1009,8 +1039,9 @@ jobs:
     timeout-minutes: 240
     strategy:
       fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-b200'].max_parallel }}
       matrix:
-        part: [0, 1, 2, 3, 4, 5]
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-b200'].arr }}
 
     steps:
       - name: Checkout code
@@ -1041,12 +1072,12 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test
-          python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-b200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 6 --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-b200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-4-gpu-b200'].size }} --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
         with:
-          artifact-suffix: ${{ matrix.part }}
+          artifact-suffix: ${{ matrix.partition }}
 
       - name: Cleanup venv
         if: always()
@@ -1066,6 +1097,11 @@ jobs:
       )
     runs-on: ${{ needs.check-changes.outputs.b200_runner }}
     timeout-minutes: 240
+    strategy:
+      fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-4-gpu-b200'].max_parallel }}
+      matrix:
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-4-gpu-b200'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -1095,7 +1131,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test
-          python3 run_suite.py --hw cuda --suite stage-c-test-dsv4-4-gpu-b200 --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-c-test-dsv4-4-gpu-b200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-4-gpu-b200'].size }} --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()
@@ -1118,6 +1154,11 @@ jobs:
       )
     runs-on: 8-gpu-h200
     timeout-minutes: 240
+    strategy:
+      fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-8-gpu-h200'].max_parallel }}
+      matrix:
+        partition: ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-8-gpu-h200'].arr }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -1147,7 +1188,7 @@ jobs:
           CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
         run: |
           cd test
-          python3 run_suite.py --hw cuda --suite stage-c-test-dsv4-8-gpu-h200 --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cuda --suite stage-c-test-dsv4-8-gpu-h200 --auto-partition-id ${{ matrix.partition }} --auto-partition-size ${{ fromJson(needs.check-changes.outputs.partitions)['stage-c-test-dsv4-8-gpu-h200'].size }} --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG
 
       - uses: ./.github/actions/upload-cuda-coredumps
         if: failure()

scripts/ci/utils/compute_partitions.py

@@ -0,0 +1,163 @@
+"""Sum CIRegistry est_time per per-commit suite and emit one $GITHUB_OUTPUT line
+keyed by suite name. Consumed by pr-test.yml stage jobs as
+`fromJson(needs.check-changes.outputs.partitions)['<suite>']`.
+
+    partitions={"stage-b-test-1-gpu-small": {"size": 8, "arr": [0,...,7], "max_parallel": 2}, ...}
+"""
+
+import argparse
+import glob
+import importlib.util
+import json
+import math
+import os
+from collections import defaultdict
+
+REPO_ROOT = os.path.dirname(
+    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+)
+
+# Load ci_register.py directly: `import sglang.test...` pulls torch/orjson via
+# sglang.__init__ but check-changes runs on bare ubuntu-latest. ci_register
+# itself is stdlib-only (AST).
+_CI_REGISTER_PATH = os.path.join(
+    REPO_ROOT, "python", "sglang", "test", "ci", "ci_register.py"
+)
+_spec = importlib.util.spec_from_file_location("ci_register", _CI_REGISTER_PATH)
+_ci_register = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_ci_register)
+collect_tests = _ci_register.collect_tests
+HWBackend = _ci_register.HWBackend
+
+# pr-test-amd.yml / pr-test-npu.yml have their own dispatch.
+_TARGET_BACKENDS = {HWBackend.CUDA, HWBackend.CPU}
+
+# stage-a is the critical-path entry gate; pin its fanout to smoke-coverage
+# defaults instead of est_time. max_parallel = size (no throttle).
+_STAGE_A_OVERRIDES = {
+    "stage-a-test-cpu": 4,
+    "stage-a-test-1-gpu-small": 1,
+}
+
+# Per-partition wall-clock target + ceiling. Single knob for the whole
+# pipeline. ~17 min avg under perfect LPT (TARGET / LPT_SLOP), ~22 min under
+# worst-case LPT 4/3 imbalance, fail-fast above 30 min.
+TARGET_SECONDS = 20 * 60
+
+# LPT (Longest Processing Time first) worst case is 4/3 * OPT; pad ~15% so a
+# slightly-unlucky LPT result still fits inside MAX_PARTITION_SECONDS.
+LPT_SLOP = 1.15
+
+# Hard ceiling. Exceeded → raise, forcing the maintainer to split a slow file
+# or bump TARGET_SECONDS deliberately.
+MAX_PARTITION_SECONDS = 30 * 60
+
+
+def discover_files(repo_root: str) -> list[str]:
+    test_dir = os.path.join(repo_root, "test")
+    files = [
+        f
+        for f in glob.glob(
+            os.path.join(test_dir, "registered", "**", "*.py"), recursive=True
+        )
+        if not f.endswith("/conftest.py") and not f.endswith("/__init__.py")
+    ]
+    jit_kernel_dir = os.path.join(repo_root, "python", "sglang", "jit_kernel")
+    files += glob.glob(
+        os.path.join(jit_kernel_dir, "tests", "**", "test_*.py"), recursive=True
+    )
+    files += glob.glob(
+        os.path.join(jit_kernel_dir, "benchmark", "**", "bench_*.py"), recursive=True
+    )
+    return files
+
+
+def compute_max_parallel(size: int) -> int:
+    return max(size // 4, 1)
+
+
+def compute_partitions(tests, full_parallel=False):
+    """Group per-commit tests by suite and emit partition metadata.
+
+    `full_parallel=True` (scheduled cron or `high priority` PR) sets
+    max_parallel = size, lifting the matrix-fanout throttle.
+    """
+    suite_tests = defaultdict(list)
+    for t in tests:
+        if t.backend not in _TARGET_BACKENDS:
+            continue
+        if t.nightly or t.disabled is not None:
+            continue
+        suite_tests[t.suite].append(t)
+
+    result = {}
+    for suite, group in suite_tests.items():
+        total = sum(t.est_time for t in group)
+        if suite in _STAGE_A_OVERRIDES:
+            size = _STAGE_A_OVERRIDES[suite]
+            max_parallel = size
+        else:
+            size = max(1, math.ceil(total * LPT_SLOP / TARGET_SECONDS))
+            max_parallel = size if full_parallel else compute_max_parallel(size)
+        # Check naive average (total/size). LPT can be ~4/3 of that, but the
+        # ceil + LPT_SLOP padding above absorbs that slack.
+        if total / size > MAX_PARTITION_SECONDS:
+            raise RuntimeError(
+                f"Suite {suite!r}: total est_time {total:.0f}s / size {size} "
+                f"= {total / size:.0f}s exceeds MAX_PARTITION_SECONDS "
+                f"({MAX_PARTITION_SECONDS}s). Split a slow file or raise "
+                f"TARGET_SECONDS deliberately."
+            )
+        result[suite] = {
+            "size": size,
+            "arr": list(range(size)),
+            "max_parallel": max_parallel,
+        }
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--repo-root", default=REPO_ROOT)
+    parser.add_argument(
+        "--output-format",
+        choices=("gha", "json"),
+        default="gha",
+        help="`gha` emits `partitions=<json>` for $GITHUB_OUTPUT; `json` is raw",
+    )
+    parser.add_argument(
+        "--full-parallel",
+        choices=("true", "false"),
+        default="false",
+        help="Lift the max_parallel throttle (set by schedule / `high priority`)",
+    )
+    args = parser.parse_args()
+
+    files = discover_files(args.repo_root)
+    # Warn-not-fail on unregistered files: run_suite.py catches this at
+    # test-execution time with sanity_check=True; dispatch should keep going.
+    all_tests = collect_tests(files, sanity_check=False)
+
+    result = compute_partitions(all_tests, full_parallel=(args.full_parallel == "true"))
+    payload = json.dumps(result, separators=(",", ":"), sort_keys=True)
+    if args.output_format == "gha":
+        print(f"partitions={payload}")
+    else:
+        print(payload)
+
+    summary_path = os.environ.get("GITHUB_STEP_SUMMARY")
+    if summary_path:
+        with open(summary_path, "a") as f:
+            f.write("## Partitions\n\n")
+            f.write(
+                f"`full_parallel={args.full_parallel}` "
+                f"(`size//4` throttle is lifted when true)\n\n"
+            )
+            f.write("| Suite | size | max_parallel |\n")
+            f.write("|---|---:|---:|\n")
+            for suite, info in sorted(result.items()):
+                f.write(f"| `{suite}` | {info['size']} | {info['max_parallel']} |\n")
+
+
+if __name__ == "__main__":
+    main()