[diffusion] fix Z-Image SP sharding for portrait and padded resolutions

[Ratish1]

Motivation

Fix Z-Image-Turbo sequence-parallel sharding for portrait and padded resolutions when Ulysses/SP is enabled. Fixes #21021

Z-Image currently works for some resolutions such as 1024x1024 and 1280x720, but produces corrupted results for others such as 720x1280 and 720x720 when sharding is enabled. This PR fixes the Z-Image-specific SP path so it preserves native image geometry during denoising.

Z-Image patchifies latents in native F/H/W order, but the current Z-Image SP path mutates the image geometry before
denoising:

• it swaps H/W so sharding always happens along the larger spatial axis
• it pads fake spatial rows so each rank gets the same local shard height

Modifications

• Remove the Z-Image H/W swap path and keep sharding in native geometry.
• Build a per-request SP plan that chooses either native height or native width splitting based on lower token-padding cost.
• Shard rectangular latent slices directly along the chosen native axis.
• Gather uneven local shards explicitly with all_gather, then crop each rank’s contribution back to its true size before
  concatenation.
• Pass the Z-Image batch through the SP gather path so final gather can use the per-request shard plan.
• Pad only the local patch-token sequence in the Z-Image DiT path to the aligned target length required by Ulysses.
• Keep the shared gather_latents_for_sp(..., batch=None) signature compatible for other pipelines.

Accuracy Tests

I tested with 2 commands:
sglang generate --model-path Tongyi-MAI/Z-Image-Turbo --seed 42 --prompt "A crowded beach" --height 720 --width 720 --num-inference-steps 9 --ulysses-degree 2 --num-gpus 2 --guidance-scale 4.0

sglang generate --model-path Tongyi-MAI/Z-Image-Turbo --seed 42 --prompt "A crowded beach" --height 720 --width 1280 --num-inference-steps 9 --ulysses-degree 2 --num-gpus 2 --guidance-scale 4.0

Before PR:

After PR:

Benchmarking and Profiling

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 Process

1. Ping Merge Oncalls to start the PR flow. 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.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[BBuf]

This PR fixes the portrait / padded-resolution / W-shard path, but I don't see a regression test that actually exercises that path. The current 2-GPU Z-Image test still uses the default square case, so it likely never hits the new branch. Could we add one portrait or padded-resolution SP test here?

[mickqian]

please wait for #20679

[Ratish1]

/tag-and-rerun-ci

[Ratish1]

/rerun-failed-ci