Skip to content

[diffusion] add tp support for qwen-image and refactor some tests#830

Merged
SamitHuang merged 4 commits intovllm-project:mainfrom
ZJY0516:qwen-image-tp
Jan 19, 2026
Merged

[diffusion] add tp support for qwen-image and refactor some tests#830
SamitHuang merged 4 commits intovllm-project:mainfrom
ZJY0516:qwen-image-tp

Conversation

@ZJY0516
Copy link
Collaborator

@ZJY0516 ZJY0516 commented Jan 18, 2026

Purpose

  • tp support for qwen image
  • refator some tests: add a class for monitoring gpu peak mem

Test

Qwen Image

Image size: 1024x1024

tp size time generated image
1 63.7661s qwen_image_output
2 37.6579s qwen_image_output
4 23.6846s qwen_image_output

Qwen Image Edit

Input image size: (1242, 1483)

tp size time generated image
1 179.8156s output_image_edit
2 99.3778s output_image_edit
4 58.5206s output_image_edit

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, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
  • (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)

Signed-off-by: zjy0516 <[email protected]>
@ZJY0516 ZJY0516 requested a review from SamitHuang January 18, 2026 07:27
Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 35a98fb248

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment on lines +71 to +74
torch.cuda.empty_cache()
device_index = torch.cuda.current_device()
monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start()

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Reset CUDA peak stats before collecting TP memory

GPUMemoryMonitor.peak_used_mb falls back to torch.cuda.max_memory_allocated/reserved, which are process‑wide peaks and are not reset by empty_cache(). Since _run_zimage_generate is invoked twice in the same process, the TP=2 run will inherit the TP=1 peak and can never be lower even if it actually uses less memory, making the new assertion flaky. Consider calling torch.cuda.reset_peak_memory_stats(device_index) before starting the monitor (or dropping the max_memory fallback) so each run measures its own peak.

Useful? React with 👍 / 👎.

Signed-off-by: zjy0516 <[email protected]>
Signed-off-by: zjy0516 <[email protected]>
Signed-off-by: zjy0516 <[email protected]>
@ZJY0516 ZJY0516 added the ready label to trigger buildkite CI label Jan 19, 2026
Copy link
Collaborator

@hsliuustc0106 hsliuustc0106 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@SamitHuang SamitHuang merged commit 3fc4f98 into vllm-project:main Jan 19, 2026
7 checks passed
@ZJY0516 ZJY0516 mentioned this pull request Jan 19, 2026
41 tasks
with1015 pushed a commit to with1015/vllm-omni that referenced this pull request Jan 20, 2026
@ZJY0516 ZJY0516 mentioned this pull request Jan 20, 2026
5 tasks
@ZJY0516 ZJY0516 deleted the qwen-image-tp branch January 23, 2026 15:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready label to trigger buildkite CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants