Skip to content

[None][chore] 2 more WAN multi-gpu tests - #15223

Merged
chang-l merged 4 commits into
NVIDIA:mainfrom
NVShreyas:user/shreyasm/visual-gen-blog-tests
Jun 12, 2026
Merged

[None][chore] 2 more WAN multi-gpu tests#15223
chang-l merged 4 commits into
NVIDIA:mainfrom
NVShreyas:user/shreyasm/visual-gen-blog-tests

Conversation

@NVShreyas

@NVShreyas NVShreyas commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator

Summary by CodeRabbit

  • Tests
    • Enhanced multi-GPU test coverage for text-to-video pipeline configurations, including new attention and parallelism scenarios.
    • Improved module-availability detection for FlashAttention components.
    • Added distributed correctness validation tests for parallel transformer configurations.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

@NVShreyas
NVShreyas requested a review from a team as a code owner June 10, 2026 15:07
@NVShreyas

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53325 [ run ] triggered by Bot. Commit: c6a2491 Link to invocation

@coderabbitai

coderabbitai Bot commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

This PR adds Flash Attention 4 backend support to multi-GPU visual generation test suites. It detects FA4 and flash_attn_combine kernel availability, then adds new distributed test configurations and correctness checks against HuggingFace references for both T2V pipeline and transformer parallel scenarios.

Changes

FA4 Attention Backend Testing Support

Layer / File(s) Summary
Backend capability detection for FA4 and flash_attn_combine
tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py, test_wan_transformer_parallel.py
Both test modules conditionally import _flash_attn_combine kernel alongside existing FA4 forward imports and set _ATTN2D_AVAILABLE flag based on successful import of both required symbols.
Pipeline parallel FA4 configuration and distributed test
tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py
New _build_cfg2_attn2d2x1_ulysses2_pvae8_args() builder creates VisualGenArgs for FA4 with 2×1 attention, ulysses=2, and parallel VAE at size 8. Distributed worker _logic_wan_cfg2_attn2d2x1_ulysses2_pvae8() runs TRTLLM pipeline at world size 8, verifies rank-0 video output, synchronizes all ranks, runs HuggingFace reference on rank 0, and asserts cosine similarity above threshold. Test method skips when kernels unavailable.
Transformer parallel FA4 2×2 configuration test
tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py
New test_parallel_attn2d_2x2_ulysses2_vs_single_gpu_8gpu() method runs 8-GPU distributed comparison between parallel and single-GPU transformer outputs for 2×2 attention parallelism with ulysses=2, gated by FA4 kernel availability check.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name Status Explanation Resolution
Description check ⚠️ Warning The PR description is empty—it contains only the template with no actual description, implementation details, or test coverage information filled in. Fill in the Description and Test Coverage sections with details about what the tests do and why they were added.
Docstring Coverage ⚠️ Warning Docstring coverage is 72.73% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (3 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately describes the main change: adding two new multi-GPU tests for the WAN visual generation pipeline.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🧹 Nitpick comments (1)
tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py (1)

341-401: ⚡ Quick win

Consider extracting a parameterized worker helper to reduce duplication.

This function shares ~90% of its logic with _logic_wan_cfg_ulysses_pvae() (lines 275-338). The only differences are the world-size assertion, the args builder called, and the log message. A parameterized helper would reduce maintenance burden if the core workflow changes.

However, test clarity often outweighs DRY, so this refactor is optional.

🤖 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_wan_pipeline_parallel.py`
around lines 341 - 401, The two test functions
_logic_wan_cfg2_attn2d2x1_ulysses2_pvae8 and _logic_wan_cfg_ulysses_pvae share
almost identical flow; extract a small parameterized worker helper (e.g.,
_run_wan_pipeline_worker(rank, world_size, checkpoint_path, expected_world_size,
args_builder, description)) that performs the world_size assertion, loads the
PipelineLoader via the provided args_builder, captures trtllm_video and
hf_video, frees resources, does the barrier and rank-0 checks, computes cosine
similarity and asserts threshold, and then replace both functions to call this
helper with different expected_world_size, args builder functions
(_build_cfg2_attn2d2x1_ulysses2_pvae8_args vs the other builder) and a short
description string to preserve the unique log message. Ensure the helper uses
the same symbols currently referenced: PipelineLoader, _capture_trtllm_video,
_capture_hf_video, _free, dist.barrier, DiffusionPipeline.from_pretrained, and
_cosine_similarity so behavior remains identical.
🤖 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.

Nitpick comments:
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py`:
- Around line 341-401: The two test functions
_logic_wan_cfg2_attn2d2x1_ulysses2_pvae8 and _logic_wan_cfg_ulysses_pvae share
almost identical flow; extract a small parameterized worker helper (e.g.,
_run_wan_pipeline_worker(rank, world_size, checkpoint_path, expected_world_size,
args_builder, description)) that performs the world_size assertion, loads the
PipelineLoader via the provided args_builder, captures trtllm_video and
hf_video, frees resources, does the barrier and rank-0 checks, computes cosine
similarity and asserts threshold, and then replace both functions to call this
helper with different expected_world_size, args builder functions
(_build_cfg2_attn2d2x1_ulysses2_pvae8_args vs the other builder) and a short
description string to preserve the unique log message. Ensure the helper uses
the same symbols currently referenced: PipelineLoader, _capture_trtllm_video,
_capture_hf_video, _free, dist.barrier, DiffusionPipeline.from_pretrained, and
_cosine_similarity so behavior remains identical.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 0bff0349-7d9c-4c38-86e7-c588e9bead76

📥 Commits

Reviewing files that changed from the base of the PR and between 74d8a48 and c6a2491.

📒 Files selected for processing (2)
  • tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py
  • tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py

@NVShreyas

Copy link
Copy Markdown
Collaborator Author

/bot kill

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53358 [ kill ] triggered by Bot. Commit: c6e53eb Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53325 [ run ] completed with state ABORTED. Commit: c6a2491

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53358 [ kill ] completed with state SUCCESS. Commit: c6e53eb
Successfully killed previous jobs for commit c6e53eb

Link to invocation

@NVShreyas

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53364 [ run ] triggered by Bot. Commit: c6e53eb Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53364 [ run ] completed with state FAILURE. Commit: c6e53eb
/LLM/main/L0_MergeRequest_PR pipeline #42543 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@NVShreyas

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53602 [ run ] triggered by Bot. Commit: 8cacbae Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53602 [ run ] completed with state FAILURE. Commit: 8cacbae
/LLM/main/L0_MergeRequest_PR pipeline #42748 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@NVShreyas

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53614 [ run ] triggered by Bot. Commit: 36b9745 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53614 [ run ] completed with state SUCCESS. Commit: 36b9745
/LLM/main/L0_MergeRequest_PR pipeline #42760 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@NVShreyas

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53693 [ run ] triggered by Bot. Commit: 36b9745 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53693 [ run ] completed with state SUCCESS. Commit: 36b9745
/LLM/main/L0_MergeRequest_PR pipeline #42828 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
…2_ring2 waive

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
@NVShreyas
NVShreyas force-pushed the user/shreyasm/visual-gen-blog-tests branch from 36b9745 to 9afd049 Compare June 12, 2026 00:57
@NVShreyas

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53733 [ run ] triggered by Bot. Commit: 9afd049 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53733 [ run ] completed with state FAILURE. Commit: 9afd049
/LLM/main/L0_MergeRequest_PR pipeline #42859 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@NVShreyas

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53900 [ run ] triggered by Bot. Commit: 9afd049 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53900 [ run ] completed with state SUCCESS. Commit: 9afd049
/LLM/main/L0_MergeRequest_PR pipeline #42999 completed with status: 'SUCCESS'

CI Report

Link to invocation

@chang-l
chang-l merged commit be7e978 into NVIDIA:main Jun 12, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants