[Perf][qwen-image-edit] Skip attention-mask to avoid varlen path

[kTorp]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

This PR builds on top of [Perf][Wan2.2] Skip attention mask for zero-padded SP sequences to avoid varlen path and similarly avoids the varlen path for a 14-15 % latency improvement for qwen-image-edit.

NOTE: this should be merged after #3763 since it depends on the mask_sp_padding boolean in vllm_omni/diffusion/data.py introduced in that PR.

Test Plan

• Before and after results, with commands, shown below.

Test Result

Commands used

VLLM_LOGGING_LEVEL=DEBUG vllm serve Qwen/Qwen-Image-Edit --omni --port 8098 --ulysses-degree 8
curl -s -X POST http://localhost:8098/v1/images/edits \
  -F "model=Qwen/Qwen-Image-Edit" \
  -F "image=@/app/data/flux_cat.png" \
  -F "prompt=Add a cool hat to the cat." \
  -F "negative_prompt=' '" \
  -F "size=2048x2048" \
  -F "num_inference_steps=50" \
  -F "guidance_scale=4.0" \
  -F "seed=42" \
  -F "response_format=b64_json" \
  | jq -r '.data[0].b64_json' | base64 -d > results/output.png

VLLM_LOGGING_LEVEL=DEBUG vllm serve Qwen/Qwen-Image --omni --port 8098 --ulysses-degree 8
curl -s -X POST http://localhost:8098/v1/images/generations \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen-Image",
    "prompt": "A cat holding a sign that says hello world.",
    "negative_prompt": " ",
    "size": "2048x2048",
    "num_inference_steps": 50,
    "guidance_scale": 0.0,
    "seed": 42,
    "response_format": "b64_json"
  }' | jq -r '.data[0].b64_json' | base64 -d > results/output.png

Qwen-Image-Edit

baseline: DiffusionEngine.step breakdown: preprocess=132.42 ms, add_req_and_wait=34167.64 ms, postprocess=136.11 ms, total=34439.35 ms

PR: DiffusionEngine.step breakdown: preprocess=129.74 ms, add_req_and_wait=29107.34 ms, postprocess=123.19 ms, total=29364.35

Qwen-Image

baseline: DiffusionEngine.step breakdown: preprocess=0.00 ms, add_req_and_wait=18037.45 ms, postprocess=145.08 ms, total=18183.86 ms

PR: DiffusionEngine.step breakdown: preprocess=0.00 ms, add_req_and_wait=17911.60 ms, postprocess=131.42 ms, total=18044.44 ms

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[hsliuustc0106]

Follows the same pattern as PR #3763, 14-15% latency improvement for qwen-image-edit.

[hsliuustc0106]

can you provide a perf comparison table for different models with different input sizes

[fhfuih]

Could you also provide quantitative metrics for the image similarities? Like SSIM, PSNR, pixel similarity. Thanks

[NumberWan]

Thank you.

At 2048×2048 with --ulysses-degree 8, the packed image sequence length appears to be divisible by 8 for both Qwen-Image (size=2048×2048> 16,384 token) and Qwen-Image-Edit ( tokens, noise + source latents =16,384 * 2 token ).
In that case _shard_with_auto_pad should not pad, and hidden_states_mask should stay None even on baseline.
As I understand, SP padding only occurs when the pixel like "1328*1328"

Could you paste the DEBUG log lines from sequence_parallel.py showing the actual seq_len, pad_size, and world_size during your benchmark? That would make it clear whether the ~14% Edit gain is from SP padding → varlen avoidance, or if there is another contributing factor.

[kTorp]

Hi, thanks for the reviews and follow-ups 😄

Addressing your comments @NumberWan for qwen-image-edit the size of the input image I used was 1024 × 704 so the 2k command resolution adds in extra details around the subject to get to the desired size. Therefore some padding is present in this case, I should have added those details

DEBUG 06-11 08:32:02 [sequence_parallel.py:449] Auto-padded sequence from 20596 to 20600 (pad_size=4, world_size=8, dim=1)

However, your comment did make me realize an oversight on my part and that is that there's no padding for qwen-image when using 2024x2024 so the given numbers are not very representative. I re-ran the command using your suggested 1328x1328 resolution and get a much larger difference

baseline:  DiffusionEngine.step breakdown: preprocess=0.00 ms, add_req_and_wait=20265.65 ms, postprocess=80.57 ms, total=20347.94 ms

PR: DiffusionEngine.step breakdown: preprocess=0.00 ms, add_req_and_wait=11220.27 ms, postprocess=81.33 ms, total=11303.20 ms

@fhfuih I collected your suggested metrics using PeakSignalNoiseRatio and StructuralSimilarityIndexMeasure from torchmetrics.image, at least on visual inspection I cannot tell the images apart

qwen-image-edit
    SSIM:          0.981868    
    PSNR:          36.8550 dB 
    mean_abs_diff: 3.507649e-03
    
qwen-image (1328x1328)
    SSIM:          0.977995
    PSNR:          30.9831 dB
    mean_abs_diff: 6.109676e-03

@hsliuustc0106 do you mean for other models than qwen?

[fhfuih]

Thanks for the follow up with data.

Any idea/guesses why it is not identical (SSIM>=0.99 or 1, and PSNR ~ inf)? 🤔 Because if this branch is redundant and the images are meant to be the same, we should be able to observe higher values.

EDIT: seems like it is mentioned here #3763 (comment)

---

And I can also add some more perf test data on my side:

With 4 GPU H200 and 1789*1789 resolution (picked a random number not divisible by 4):

• USP 4: 29.43 -> 23.89 (-19%)
• USP 4 + VAE patch parallel 4: 29.82 -> 24.97 (-16%)
• USP 4 + CacheDiT: 12.925 -> 11.158 (-14%)

---

In your # Purpose, could you add an explicit notice that we should wait for #3763 to merge first? (Lest the "build on top" is not clear enough and someone merged it to break the code 😁

[kTorp]

@hsliuustc0106

I got numbers for Flux2 comparing 1328x1328 before and after skipping varlen-mask

1328x1328: Auto-padded sequence from 6889 to 6896 (pad_size=7, world_size=8, dim=1)

| Model      | Resolution | Baseline total (s) | PR total (s) | Change  |
|------------|------------|--------------------|--------------|---------|
| FLUX.2-dev | 1328x1328  |       7.41         |    5.86      |  20.9%  |
| FLUX.2-dev | 2048x2048  |       12.13        |    12.11     |  0.2%   |

So there's a pretty good boost in the padded case while the 2048x204 remains unchanged as expected

Outputs:

1328x1328 baseline

1328x1328 skip_mask

[hsliuustc0106]

@jjuvonen-amd PTAL

[jjuvonen-amd]

@kTorp, thank you for adding this for Qwen and Flux. Were you planning on enabling and testing it also with HunyuanVideo 1.5 and GLM Image?

Would it make sense to include a one-time warning message when mask_sp_padding=False and SP adds padding tokens? Similar to what was added for Wan2.2 models:

logger.warning_once(
                "SP auto-padding applied %d token(s) (seq_len=%d, ulysses_degree=%d). "
                "Padding tokens are not masked from attention (mask_sp_padding=False), "
                "which avoids the varlen attention path but may produce minor numerical differences. "
                "Set parallel_config.mask_sp_padding=True to restore strict masking.",
                ctx.sp_padding_size,
                ctx.sp_original_seq_len,
                parallel_config.sequence_parallel_size,
            )

[kTorp]

@jjuvonen-amd thanks for the original PR and for feedback, I'll add in the logger warnings as done for Wan2.2.

Yes, I'll enable and test it for HunyuanVideo1.5 and GLM image as well. Is it alright if I do so in a separate PR @hsliuustc0106, so that it doesn't clutter this one?

[kTorp]

@hsliuustc0106 sorry for the delay, I've now had time to run and test the remaining models and opened this #4645

Flux2 and HunyuanVideo1.5 shows a good boost for padded sequence lengths similarly to the previous models. Unfortunately I had a lot of troubles gettings results for GLMImage. I got this setup to work for ulysses-degree=8

VLLM_LOGGING_LEVEL=DEBUG vllm serve zai-org/GLM-Image --omni --port 8098 \
  --deploy-config /app/results/glm_image_sp8.yaml \
  2>&1 | tee /app/results/glm_serve.log

curl -s http://localhost:8098/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [{"role": "user", "content": "A small cat"}],
    "height": 1024,
    "width": 1024,
    "seed": 42
  }' \
  | jq -r '.choices[0].message.content[0].image_url.url' \
  | cut -d',' -f2- \
  | base64 -d > results/glm_image_baseline_ulysses_8_1024x1024.png

where glm_image_sp8.yaml (adapted from vllm-omni/deploy/glm_image.yaml) is

stages:
  # Stage 0: AR model (single GPU, co-located on GPU 0).
  - stage_id: 0
    max_num_seqs: 1
    gpu_memory_utilization: 0.6
    enforce_eager: false
    trust_remote_code: true
    enable_prefix_caching: false
    max_num_batched_tokens: 32768
    devices: "0"
    default_sampling_params:
      temperature: 0.9
      top_p: 0.75
      top_k: 16512
      stop_token_ids: [16385]
      max_tokens: 4353
      seed: 42
      detokenize: false

  # Stage 1: Diffusion (DiT + VAE) with Ulysses SP across 8 GPUs.
  - stage_id: 1
    max_num_batched_tokens: 32768
    max_num_seqs: 1
    enforce_eager: true
    trust_remote_code: true
    enable_prefix_caching: false
    devices: "0,1,2,3,4,5,6,7"
    parallel_config:
      tensor_parallel_size: 1
      sequence_parallel_size: 8
      ulysses_degree: 8
      ring_degree: 1
    default_sampling_params:
      seed: 42
      num_inference_steps: 25
      guidance_scale: 1.5
      height: 1024
      width: 1024

gave nice results

but as soon as I tried changing to a different resolution to trigger the automatic-padding I could only produce these strange tilings:

1056x1056 single-gpu

1056x1056 ulysses-8

1088x1088 ulysses-8

I'm probably just running the model erroneously 😅 any feedback / instructions would be welcome. In the meanwhile I've left it out of the PR so it's not a blocker.