[None][feat] Qwen Image Ulysses perf enhancement - #17147
Conversation
d68cb4a to
7fb2c86
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughQwen-Image-Edit-2511 now has a 2-GPU Ulysses configuration, updated support declarations, revised prompt encoding, and distributed attention tests. The B200 test list includes the pipeline configuration test. ChangesQwen Image Edit Ulysses Support
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Pytest
participant DistributedWorker
participant NCCLProcessGroup
participant QwenJointAttention
Pytest->>DistributedWorker: run two-GPU Ulysses test
DistributedWorker->>NCCLProcessGroup: initialize distributed workers
DistributedWorker->>QwenJointAttention: execute Ulysses attention
QwenJointAttention-->>DistributedWorker: return finite output tensors
DistributedWorker-->>Pytest: validate output shapes and values
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py (1)
14-14: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse union syntax for
prompt_masks.Replace
Optional[torch.Tensor]withtorch.Tensor | None. RemoveOptionalfrom Line 14.Proposed change
-from typing import Any, Optional +from typing import Any ... - *prompt_masks: Optional[torch.Tensor], + *prompt_masks: torch.Tensor | None,As per coding guidelines, use built-in generic types and
|.Also applies to: 262-266
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py` at line 14, Update the prompt_masks type annotations in the affected definitions to use torch.Tensor | None instead of Optional[torch.Tensor], and remove the unused Optional import while preserving the existing annotation behavior.Source: Coding guidelines
tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py (3)
39-40: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueComplete the callable annotations.
Add a generator return type to
_cleanup_mpi_env. ParameterizeCallableasCallable[[int, int], None]. Add-> Nonetotest_qwen_image_edit_ulysses_attention_2gpu.As per coding guidelines, “Annotate every function” and use precise
Callabletypes.Also applies to: 61-72, 148-149
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py` around lines 39 - 40, Complete annotations for _cleanup_mpi_env by adding its generator return type and annotate callback parameters with Callable[[int, int], None]. Add -> None to test_qwen_image_edit_ulysses_attention_2gpu and the other affected test functions, preserving their existing behavior.Source: Coding guidelines
61-69: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the broad exception handler.
The
finallyblock already runs cleanup for success and failure. Removeexcept Exceptionand let the original worker exception propagate.As per coding guidelines, “Catch specific exceptions instead of broad or bare exception handling such as
except:.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py` around lines 61 - 69, Remove the broad except Exception handler from _distributed_worker and its error print, allowing test_fn or initialization exceptions to propagate unchanged while retaining the finally block’s cleanup_distributed() call for both success and failure.Source: Coding guidelines
99-149: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd pipeline-level Ulysses coverage and register both tests.
test_qwen_image_edit_ulysses_attention_2gpuandtest_qwen_image_edit_ulysses_rejects_masked_prompt_conditioningdo not callQwenImageEditPlusPipeline.forward. Add a test that setsulysses_size=2, returns a non-Noneprompt mask, and assertsforwardraises before_prepare_edit_latentsruns.Test coverage summary:
- Added:
test_qwen_image_edit_ulysses_attention_2gpu.- Changed:
test_qwen_image_edit_ulysses_rejects_masked_prompt_conditioning.- CI and QA membership: neither test is listed.
- Add the multi-GPU test to
tests/integration/test_lists/test-db/l0_dgx_b200.ymland the pipeline-config test totests/integration/test_lists/test-db/l0_b200.yml.- Coverage verdict: insufficient.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py` around lines 99 - 149, Add pipeline-level Ulysses coverage alongside _test_qwen_image_edit_ulysses_attention: create or update test_qwen_image_edit_ulysses_rejects_masked_prompt_conditioning with ulysses_size=2, a non-None prompt mask, and an assertion that QwenImageEditPlusPipeline.forward raises before _prepare_edit_latents executes. Register test_qwen_image_edit_ulysses_attention_2gpu in l0_dgx_b200.yml and the pipeline-config rejection test in l0_b200.yml.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py`:
- Line 8: Preserve the test runner’s original TLLM_DISABLE_MPI value before the
assignment near the module setup, then restore that value during teardown
instead of always removing the variable. Update the teardown logic around the
existing cleanup at lines 39-42 so an originally set value is reinstated and an
originally absent value is removed.
---
Nitpick comments:
In
`@tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py`:
- Line 14: Update the prompt_masks type annotations in the affected definitions
to use torch.Tensor | None instead of Optional[torch.Tensor], and remove the
unused Optional import while preserving the existing annotation behavior.
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py`:
- Around line 39-40: Complete annotations for _cleanup_mpi_env by adding its
generator return type and annotate callback parameters with Callable[[int, int],
None]. Add -> None to test_qwen_image_edit_ulysses_attention_2gpu and the other
affected test functions, preserving their existing behavior.
- Around line 61-69: Remove the broad except Exception handler from
_distributed_worker and its error print, allowing test_fn or initialization
exceptions to propagate unchanged while retaining the finally block’s
cleanup_distributed() call for both success and failure.
- Around line 99-149: Add pipeline-level Ulysses coverage alongside
_test_qwen_image_edit_ulysses_attention: create or update
test_qwen_image_edit_ulysses_rejects_masked_prompt_conditioning with
ulysses_size=2, a non-None prompt mask, and an assertion that
QwenImageEditPlusPipeline.forward raises before _prepare_edit_latents executes.
Register test_qwen_image_edit_ulysses_attention_2gpu in l0_dgx_b200.yml and the
pipeline-config rejection test in l0_b200.yml.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f1f68866-877e-4229-8b97-ca1802cf4d54
📒 Files selected for processing (6)
docs/source/models/supported-models.mddocs/source/models/visual-generation.mdexamples/visual_gen/configs/qwen-image-edit-2511-fp8-2gpu-ulysses.yamltensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.pytests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.pytests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py`:
- Around line 149-150: Register both changed test modules in the appropriate
test lists: add
tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.py to
l0_dgx_b200.yml and
tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py to
l0_b200.yml, or equivalent applicable lists, so
test_qwen_image_edit_ulysses_attention_2gpu and the configuration test coverage
are included.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c2c82e34-fca1-4493-8148-ab2a20a20d30
📒 Files selected for processing (4)
tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.pytests/integration/defs/accuracy/test_cli_flow.pytests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.pytests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py
617d9bb to
2b878d4
Compare
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
a6049b8 to
40b75ea
Compare
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
|
/bot run |
|
PR_Github #70495 [ run ] triggered by Bot. Commit: |
|
PR_Github #70495 [ run ] completed with state
|
|
/bot run |
|
PR_Github #70537 [ run ] triggered by Bot. Commit: |
|
PR_Github #70537 [ run ] completed with state
|
|
/bot run |
|
PR_Github #70568 [ run ] triggered by Bot. Commit: |
|
PR_Github #70568 [ run ] completed with state
|
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot run |
|
PR_Github #70799 [ run ] triggered by Bot. Commit: |
|
PR_Github #70799 [ run ] completed with state
|
8 participants
Summary
This PR improves Qwen Image Ulysses execution and documents Qwen-Image-Edit-2511 Ulysses support.
U > 1);U = 1stays on the original local path.VANILLAattention, because padded text/image streams require key-padding-mask support.VANILLAUlysses consumes it, and non-mask-capable sequence-parallel backends raise from the transformer layer.trtllm-servesupport matrix entries from [TRTLLM-14730][feat] Add image edit serving endpoint for visual generation models #16987 while enabling Ulysses support.Implementation Notes
Qwen Image Ulysses gathers sequence shards and shards heads before attention. The previous path converted post-all-to-all tensors back to backend layout through eager PyTorch layout operations for Q, K, and V. The new native op fuses that post-unscatter layout conversion and removes launch/layout overhead on the multi-rank Ulysses path.
Runtime behavior:
U > 1: BF16 Ulysses attention uses the fused post-unscatter op when Q/K/V are same-shaped BF16 CUDA tensors and the kernel shape constraints are satisfied.U = 1: the fused op is not used, because there is no multi-rank all-to-all/post-unscatter work.Validation
Tests/coverage added or updated:
tests/unittest/_torch/thop/parallel_hw_agnostic/test_ulysses_post_unscatter.pyulysses_packed_qkv_post_unscatterD % 8, and oversized-block rejection teststests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_edit_ulysses.pygpu2mark, distributed barrier, and MPI-env restorationtests/integration/test_lists/test-db/l0_dgx_b200.ymlManual image-quality check:
BF16 Qwen Image U=2
Prompt:
LPIPS score: 0
Performance Notes
Manual B200 runs with the Python VisualGen scripts showed the fused post-unscatter op helps the multi-rank Ulysses path. The final branch gates the op to
U > 1;U = 1is intentionally baseline-only.BF16
PR Checklist