[diffusion] feat: data-parallel serving (--dp-size)

[mickqian]

Motivation

--data-parallel-size parsed, the process-group layout reserved the dimension ("tp-sp-pp-cfg-dp", dp outermost), and validation rejected it: dp_size > 1 raised "DP is not yet supported", because nothing routed requests to replicas. This implements the missing routing, so N replicas serve N requests concurrently on one launch.

Modifications

The design leans on two properties the codebase already had. Because dp is the outermost axis, every replica is a contiguous rank block whose sp/cfg/tp subgroups never cross replicas; and request fan-out already travels a broadcast_pyobj relay over exactly those subgroups (Scheduler.recv_reqs). So per-replica serving needs only a per-replica ingress:

• Scheduler: each replica's first rank is its driver and binds its own ZMQ endpoint (one settled scheduler port per replica). The existing relay distributes requests within the replica; non-driver ranks already no-op on result return; request-based warmup keys on per-scheduler state, so each replica warms itself.
• Clients (sync + async): generation requests round-robin across replica endpoints; realtime sessions stick to one replica by session-id hash, since their GPU state lives there; control ops that mutate replica state (weights, LoRA, memory occupation, shutdown) fan out to every replica with first-error-wins merging — call sites unchanged; ping reports ready only when every replica answers.
• Reply-rank fix: output materialization, file saving, frame transport, peak-memory recording and request metrics were gated on global rank == 0, which is the replying rank only at dp=1. On any other replica the driver nulled the output without saving it, so every request served by replica ≥ 1 returned "Model generation returned no output". These gates now key on is_output_rank (first rank of the replica).
• Deletes dp_degree, a dead twin of dp_size that nothing read.
• DP stays monolithic-only: disaggregated roles scale by adding role instances, so dp_size > 1 with a disagg role still raises.

Test

• test/unit/test_dp_routing.py: round-robin coverage, dp=1 degenerate case, session stickiness, control-op classification, fan-out merge semantics, per-replica endpoint derivation.
• test/single_test_file/test_dp_serving_2_gpu.py (registered in the 2-gpu suite): launches --dp-size 2 (one GPU per replica), asserts both drivers bind their own ingress, a fixed seed returns identical bytes through two round-robined replicas, and a concurrent pair completes in < 1.6x one request's wall time — two requests serialized on one replica take ~2x, so the bound separates the behaviors. The no-output bug above was caught by exactly this concurrency check: a single offline batch only ever exercises replica 0.
• 2xH200, Qwen-Image-2512 (1024x1024, 12 steps, one GPU per replica), manual e2e:
  • both drivers bind distinct endpoints (dp replica 0 ... :5575, dp replica 1 ... :5576)
  • same seed under dp=2 is bitwise identical to dp=1 (md5 match)
  • concurrent pair: both 200 OK, created in the same second on different replicas
  • concurrent-pair wall time: dp=2: 3.31 s for two requests vs 3.28 s for one (1.01x); dp=1 control: 6.64 s vs 3.32 s (2.00x, perfect serialization) -- the pair truly runs in parallel

Checklist

• Format your code according to the Code Formatting with Pre-Commit.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #31059198106
Latest PR Test (Extra): ❌ Run #31059197889

[mickqian]

/tag-and-rerun-ci

[mickqian]

The multimodal-gen-unit-test failure is pre-existing on main, not from this PR. Evidence rather than assertion:

• This diff touches no GLM files (git diff --stat origin/main...HEAD has no glm/ar entries).
• The exact two tests fail identically on pristine upstream main (acaab22d0, current HEAD):

FAILED test_glm_image_ar.py::TestGlmImageARSrtBackend::test_srt_ar_forward_aggregates_usage
FAILED test_glm_image_multi_output.py::test_ar_stage_generates_one_prior_per_requested_output
2 failed, 3 passed in 0.17s

Failure signatures — RuntimeError: GLM-Image AR batch returned an unexpected response: expected 2 outputs, got dict. (glm_image.py:505) and ValueError: not enough values to unpack (expected 3, got 2) — point at a return-contract mismatch between the two recent GLM changes, #33378 (usage report) and #30683 (AR batching), whose paths these tests cover.

The other 946 unit tests pass, including this PR's test_dp_routing.py.

[mickqian]

Update: the unit-test breakage healed on main — 5424d20 [CI] Fix GLM-Image usage unit tests (#33731) landed inside today's acaab22d..beabc59 window, and both files now pass there (10 passed, verified on a clean checkout). Re-running the failed job should clear this PR's only red check.