[CI] Partition stage-a-test-cpu into 4 matrix shards (#23208)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

.claude/skills/ci-workflow-guide/SKILL.md

@@ -270,7 +270,7 @@ Large suites are split across matrix jobs using the **LPT (Longest Processing Ti
 | Suite | Partitions | Runner | max_parallel |
 |-------|-----------|--------|-------------|
 | `stage-a-test-1-gpu-small` | 1 (no matrix) | `1-gpu-5090` | — |
-| `stage-a-test-cpu` | 1 (no matrix) | `ubuntu-latest` | — |
+| `stage-a-test-cpu` | 4 | `ubuntu-latest` | — |
 | `stage-b-test-1-gpu-small` | 8 | `1-gpu-5090` | 8 |
 | `stage-b-test-1-gpu-large` | 14 | `1-gpu-h100` | dynamic (3 or 14) |
 | `stage-b-test-2-gpu-large` | 4 | `2-gpu-h100` | — |

.github/workflows/pr-test.yml

@@ -361,7 +361,7 @@ jobs:
         id: wait
         with:
           stage-name: stage-a
-          jobs: '["stage-a-test-1-gpu-small", "stage-a-test-cpu"]'
+          jobs: '["stage-a-test-1-gpu-small", {"prefix": "stage-a-test-cpu", "expected_count": 4}]'
           max-wait-minutes: '240'
 
   wait-for-stage-b:
@@ -616,6 +616,10 @@ jobs:
       )
     runs-on: ubuntu-latest
     timeout-minutes: 240
+    strategy:
+      fail-fast: false
+      matrix:
+        partition: [0, 1, 2, 3]
     steps:
       - name: Free disk space
         run: |
@@ -661,7 +665,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 $CONTINUE_ON_ERROR_FLAG
+          python3 run_suite.py --hw cpu --suite stage-a-test-cpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 4 $CONTINUE_ON_ERROR_FLAG
 
   # Runs on 5090 (32GB, SM120)
   stage-b-test-1-gpu-small: