[AMD] CI: cut two setup cycles from the AMD multimodal-gen lanes

[michaelzhang-ai]

Motivation

AMD multimodal-gen jobs spend most of their runner time not testing. Measured on mi300 in run 31443692177, every job pays a container pull of 27–97 min plus a dependency install of 23–98 min, so 100–155 min of fixed cost before a single test executes. Across the eight multimodal jobs in that run that is ~15.8 runner-hours of setup against ~5.6 hours of tests — roughly 74% overhead, and because both matrices are max-parallel: 1 it is serialized into wall-clock. CUDA for comparison: Install dependencies takes 27 seconds and total fixed overhead is under a minute.

The two clearest wins are removing whole setup cycles.

Modifications

Fold the unit suite into a warm shard. multimodal-gen-unit-test-amd paid ~110 min of setup to run ~5 min of portable CPU-style tests. It now runs as a step of multimodal-gen-test-1-gpu-amd part 0 (and the ROCm 7.2 counterpart), on the container that job already warmed. The step runs after the diffusion tests and on failure as well, so a diffusion regression cannot mask a unit regression, and it carries its own timeout-minutes so it cannot eat the diffusion step's budget. Trigger conditions are unchanged — on both lanes the 1-GPU job already gates on exactly the inputs the unit job did (multimodal_gen on pr-test-amd, main_package || sgl_kernel on ROCm 7.2). The standalone job, its target_stage_select entry, and its needs entry are removed; the ci_coverage_report.py comment naming it is updated (the "unit": ("CUDA", "AMD") mapping itself still holds).

2-GPU: 3 shards → 2. Per-case timings from the same run were 14, 26 and 21 min of healthy work across the three shards; the third shard's other 96 min was flux_image_t2i_2_gpus (77.7 min) and flux2_modelopt_fp8_tp2_t2i (18.1 min) retrying, which is the known ROCm flux bug (#34351 / #34352), not real work. Two shards therefore land near 30 min of case time each — and even if one shard inherited both flux cases while they stay broken, ~126 min is still inside the 180 min step timeout.

Net effect: 8 AMD multimodal jobs → 5 per lane, i.e. ~4 fewer setup cycles across the two AMD workflows, on the order of 8 runner-hours per full PR run.

What this deliberately does not do

1-GPU stays at 4 shards. It has no headroom today: shard 0 already ends with The action 'Run diffusion server tests (1-GPU)' has timed out after 90 minutes, spent retrying the broken flux_image_t2i. Merging 1-GPU shards should wait until that case is fixed or skipped, and would likely want the step timeout raised to 150 min to match the 2-GPU lane.

The retry amplification and the image pull are untouched. Retrying a deterministically broken case seven times is what turns a 20-minute shard into a 127-minute one, and the ROCm image pull is the single largest line item in the whole lane — caching or pre-pulling it on the mi300 runners would beat any amount of matrix tuning. Both are worth separate work.

Validation

Dispatched on this branch against the ROCm 7.2 lane (the PR gate since #34204): run 31638833841.

The unit fold works, including the case it had to get right. On 1-GPU part 0 the diffusion step failed (7 passed, 1 setup error) and the folded Run diffusion unit tests step still ran and passed — 1360 passed, 10 skipped in 102 s. if: matrix.part == 0 && (success() || failure()) fires after a failed predecessor, and the step is correctly skipped on parts 1-3.

The 2 two-GPU shards cover the suite and fit the timeout, but with less margin than I estimated. They split into 13 cases + 3 standalone and 13 cases + 4 standalone. Shard 1 spent 45 min on cases plus 9 min on standalone. Shard 0 reached 149 min against the 180 min step cap — not because two shards is too coarse, but because flux_image_t2i_2_gpus fails and the harness retries it: 64 min for the case block, then repeated ~13 min retry passes. Both shards' remaining failures are the pre-existing ROCm set (flux #34351 / #34352, and test_ar_models failing warmup with Inference tensors do not track version counter).

So the split is safe today and drops to ~45-65 min per shard once the flux retries stop, but if a reviewer wants insurance while flux is broken, the options are to hold this until flux lands or to raise the 2-GPU step timeout. I have deliberately not raised it here: a larger cap means a wedged job holds a scarce mi300 runner longer.

Accuracy Tests

N/A — CI orchestration only.

Speed Tests and Profiling

Covered above: the change is motivated entirely by measured CI wall-clock, and the numbers come from the per-step and per-case timings of run 31443692177.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): 🚫 Run #32331156244
Latest PR Test (Extra): ❌ Run #32331156067
Latest PR Test (AMD ROCm 7.2): ❌ Run #32331156310

[michaelzhang-ai]

Dispatched the ROCm 7.2 lane (the PR gate after #34204) on this branch, targeting both changed stages: run 31638833841.

The job list already confirms both edits took effect — the matrix expanded to 2 multimodal-gen-test-2-gpu-amd-rocm720 shards instead of 3, and there is no multimodal-gen-unit-test-amd-rocm720 job. What the run still has to prove:

• the 2 two-GPU shards each stay inside the 180 min step timeout;
• the folded Run diffusion unit tests step executes on 1-GPU part 0 after the diffusion step fails — that shard currently ends in the 90 min timeout from flux_image_t2i retrying ([Bug][AMD][diffusion] FLUX server crashes on ROCm: Triton PTX inline-asm ('=f' constraint) in diffusion layernorm/numerics kernels fails to compile on gfx942 #34351 / [AMD] Disable PTX-inline-asm diffusion LayerNorm fast path on ROCm (fix FLUX warmup crash) #34352), so this is a direct test of the if: matrix.part == 0 && (success() || failure()) guard.

Not dispatching pr-test-amd.yml (ROCm 7.0): identical edits, and it is a daily shadow now, so it would only consume mi300 runners.

Reminder that this PR targets cursor/fix-mm-gen-2gpu-partition-guard-e992, not main — the 2-shard split needs #34328 first. I'll retarget once that merges.

[michaelzhang-ai]

Thanks — the gate reading is accurate (only amd is set, so all 7 red checks are pr-gate + the *-finish cascade, no executed test failures). Two corrections and the missing coverage now filled in.

The changed lanes have been exercised, just not by PR-attached checks. I dispatched the ROCm 7.2 lane on this branch earlier: run 31638833841. It answers all three of your confirm items:

1. Run diffusion unit tests fired on matrix.part == 0 after the diffusion step failed, and passed — 1360 passed / 10 skipped in 102 s. (Parts 1-3 correctly skip it.)
2. --total-partitions 2 covered the whole suite across the two shards — 13 cases + 3 standalone and 13 cases + 4 standalone = all 26 cases and 7 standalone files.
3. pr-test-amd-rocm720-finish executed, so the needs graph is intact — a dangling needs is a workflow-validation error and the run would not have started at all.

ROCm 7.0 cannot be covered by adding a label. Since #34204, pr-test-amd.yml has no pull_request trigger — only the daily cron, workflow_dispatch, and workflow_call (it is invoked by amd-aiter-scout.yml). So I dispatched it directly, targeted at the two changed stages: run 32208164744 — 2 two-GPU shards (confirming 3→2 there too) plus the 4 one-GPU shards. check-changes in that run flaked on fatal: shallow file has changed since we read it while deepening for the merge base; the targeted jobs run off target_stage and are unaffected.

On "confirm the shards pass": not attainable today for reasons this PR does not cause. The AMD multimodal shards are red on merged main too — flux (#34351 / #34352) and test_ar_models failing warmup with Inference tensors do not track version counter. #34328 itself merged with those same shards red. The workable criterion is: every item scheduled, each shard inside its step timeout, no new failure mode — all shown above, with the one caveat recorded in the description that the worst 2-GPU shard reached 149 of the 180 min cap because the broken flux cases retry.

[michaelzhang-ai]

Added the run-ci label as requested and re-ran the gated lanes. Result is worth recording, because it settles the question rather than the way either of us expected:

run-ci cannot give this PR coverage of the jobs it changes. On attempt 2 of PR Test ROCm 7.2 the gate and check-changes are both green — and multimodal-gen-test-1-gpu-amd-rocm720 / multimodal-gen-test-2-gpu-amd-rocm720 are skipped anyway. Those jobs require needs.check-changes.outputs.multimodal_gen == 'true', and that filter is:

multimodal_gen:
  - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)"

This PR touches only .github/workflows/pr-test-amd*.yml and scripts/ci/utils/ci_coverage_report.py, so the flag is false and the lanes skip regardless of the label. What the label did buy is ~6 unrelated sgl-kernel / jit-kernel / stage-a jobs now occupying mi300 runners. So a targeted dispatch is not a shortcut here, it is the only mechanism — unless we add the AMD workflow paths to that filter, which might be a reasonable follow-up in its own right.

ROCm 7.0 is now covered too — run 32208164744, targeted at the two changed stages:

• The fold works on 7.0 as well. 1-GPU part 0's diffusion step failed (7 passed, 1 error) and Run diffusion unit tests then ran and passed: 1360 passed / 10 skipped in 99.8 s.
• The 2-shard split covers the suite on 7.0. Shard 1 ran 13 cases (9 passed, 1 failed, 3 errors), then proceeded into its standalone files — test_ar_models.py passed (5:37) and test_disagg_server.py returned 1 passed / 2 errors, after which fast-fail stopped the remaining two.
• Remaining failures are the known pre-existing set: flux_image_t2i, flux2_modelopt_fp8_tp2_t2i ([Bug][AMD][diffusion] FLUX server crashes on ROCm: Triton PTX inline-asm ('=f' constraint) in diffusion layernorm/numerics kernels fails to compile on gfx942 #34351 / [AMD] Disable PTX-inline-asm diffusion LayerNorm fast path on ROCm (fix FLUX warmup crash) #34352) and joyai_image_edit_ti2i.

Two incidental findings from that run, both independent of this PR but useful:

1. test_ar_models.py passes on ROCm 7.0 and fails on 7.2 (Inference tensors do not track version counter), so that one is 7.2-specific rather than a general ROCm problem.
2. test_disagg_server.py shows 2 errors on 7.0 (TestDisaggZImage1Rank and TestDisaggZImageTracing, both a 500 during warmup). This is the first time that file has executed on AMD since the partition abort started on Aug 6, so it is new information rather than a new regression — worth its own issue if someone owns disagg on ROCm.