[TRTLLM-14798][perf] Fuse Wan DupUp3D output mapping - #17003
Conversation
09933d5 to
9aacc3e
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:
WalkthroughAdds a validated Triton fused implementation for Wan’s ChangesWan DupUp3D acceleration
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Input as CUDA input
participant DupUp3D
participant can_implement_dup_up3d
participant TritonKernel
participant EagerPath
Input->>DupUp3D: call forward
DupUp3D->>can_implement_dup_up3d: validate fused support
alt supported non-empty output
DupUp3D->>TritonKernel: launch fused mapping
TritonKernel-->>DupUp3D: channels-last-3d output
else unsupported or empty input
DupUp3D->>EagerPath: run repeat and reshape fallback
EagerPath-->>DupUp3D: eager output
end
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/visual_gen/test_wan_dup_up3d.py (1)
1-80: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd the new visual-gen unit test to the test-db list
tests/unittest/_torch/visual_gen/test_wan_dup_up3d.pyis not listed in the integration test sets; add it totests/integration/test_lists/test-db/l0_b200.ymlso CI runs it.- Coverage summary:
test_fused_dup_up3d_matches_eager,test_fused_dup_up3d_int32_index_limit, andtest_fused_dup_up3d_falls_back_above_index_limitcover eager parity, the int32 indexing boundary, and the fallback path.🤖 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/test_wan_dup_up3d.py` around lines 1 - 80, Add tests/unittest/_torch/visual_gen/test_wan_dup_up3d.py to the appropriate test list in l0_b200.yml so CI includes all three DupUp3D tests: test_fused_dup_up3d_matches_eager, test_fused_dup_up3d_int32_index_limit, and test_fused_dup_up3d_falls_back_above_index_limit.Source: Path instructions
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/models/wan/dup_up3d.py (1)
80-96: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd Args docs to the public
dup_up3ddocstring.The one-line docstring doesn't document parameters (
x,output_channels,repeats,factor_t,factor_s,first_chunk) or return semantics beyond the summary. As per coding guidelines, "Prefer docstrings for external interfaces, use Google-style docstrings, document public function arguments, and include tensor dimensions and constrained dtypes where applicable."🤖 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/wan/dup_up3d.py` around lines 80 - 96, Expand the public dup_up3d docstring with Google-style Args entries for x, output_channels, repeats, factor_t, factor_s, and first_chunk, including tensor dimensions and applicable dtype constraints. Retain the existing summary and document the return behavior, including when None is returned because Triton cannot index the output.Source: Coding guidelines
🤖 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 `@tensorrt_llm/_torch/visual_gen/models/wan/dup_up3d.py`:
- Around line 29-42: Remove tl.constexpr annotations from the shape, factor,
crop, and stride parameters in the kernel signature, including output_channels
through stride_xw, so they remain runtime values and do not expand
specialization. Retain tl.constexpr only for block_size.
---
Outside diff comments:
In `@tests/unittest/_torch/visual_gen/test_wan_dup_up3d.py`:
- Around line 1-80: Add tests/unittest/_torch/visual_gen/test_wan_dup_up3d.py to
the appropriate test list in l0_b200.yml so CI includes all three DupUp3D tests:
test_fused_dup_up3d_matches_eager, test_fused_dup_up3d_int32_index_limit, and
test_fused_dup_up3d_falls_back_above_index_limit.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/wan/dup_up3d.py`:
- Around line 80-96: Expand the public dup_up3d docstring with Google-style Args
entries for x, output_channels, repeats, factor_t, factor_s, and first_chunk,
including tensor dimensions and applicable dtype constraints. Retain the
existing summary and document the return behavior, including when None is
returned because Triton cannot index the output.
🪄 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: 892296c3-9a85-487b-8fb9-63fefea33f4c
📒 Files selected for processing (3)
tensorrt_llm/_torch/visual_gen/models/wan/dup_up3d.pytensorrt_llm/_torch/visual_gen/models/wan/wan_vae.pytests/unittest/_torch/visual_gen/test_wan_dup_up3d.py
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
6c3f1fe to
b07d588
Compare
taianz-nv
changed the title
zhenhuaw-me
left a comment
zhenhuaw-me
left a comment
There was a problem hiding this comment.
Not part of the PR: Considering the Wan VAE is also used by Cosmos, maybe we can move the Wan VAE code to https://github.com/NVIDIA/TensorRT-LLM/tree/main/tensorrt_llm/_torch/visual_gen/modules/vae
| self.repeats = out_channels * self.factor // in_channels | ||
|
|
||
| def forward(self, x: torch.Tensor, first_chunk: bool = False) -> torch.Tensor: | ||
| if x.is_cuda: |
There was a problem hiding this comment.
Sorry, I didn't get it. I'm not sure why in eager pass, the tensor will be on CPU side. I was thinking that The tensor will rarely be on CPU side in the pipeilne forward function.
There was a problem hiding this comment.
Sorry, I misunderstood your question earlier. You're right that the pipeline input is expected to be on CUDA, and the CPU eager path was only retained as a test reference. Since can_implement_dup_up3d() already checks x.is_cuda, I'll remove this redundant outer check and rely on the capability check for the fallback.
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 `@tensorrt_llm/_torch/visual_gen/models/wan/dup_up3d.py`:
- Around line 124-170: Update the output-shape and launch logic in the fused
DupUp3D path around _supports_triton_indexing and _dup_up3d_kernel so empty
inputs preserve eager behavior: avoid negative output_frames when input_frames
is zero with first_chunk enabled, and return None for any zero-sized batch or
spatial/temporal dimension before allocating or launching Triton. Keep normal
non-empty shape calculation and kernel execution unchanged.
In `@tensorrt_llm/_torch/visual_gen/models/wan/wan_vae.py`:
- Around line 266-278: The fused dup_up3d dispatch in WanVAE must be skipped
when autograd is enabled and x.requires_grad is true; update the CUDA guard
around dup_up3d to preserve the eager differentiable path, while retaining fused
execution for inference or non-gradient inputs. In
tests/unittest/_torch/visual_gen/test_wan_dup_up3d.py lines 13-50, add a CUDA
regression test that runs module(x).sum().backward() and verifies x.grad is
populated.
In `@tests/unittest/_torch/visual_gen/test_wan_dup_up3d.py`:
- Around line 35-50: Add a genuinely non-contiguous, non-overlapping CUDA input
view in the DupUp3D test, such as a spatially sliced tensor with non-unit
strides, while preserving the existing dtype, device, and shape. Use this
strided input for both the eager CPU reference and CUDA execution, then retain
the exact output comparison and channels-last-contiguous assertion.
🪄 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: d4004206-7dc3-4c78-8b74-8fb70731aec4
📒 Files selected for processing (4)
tensorrt_llm/_torch/visual_gen/models/wan/dup_up3d.pytensorrt_llm/_torch/visual_gen/models/wan/wan_vae.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/visual_gen/test_wan_dup_up3d.py
Signed-off-by: Taian Zhang <taianz@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tensorrt_llm/_torch/visual_gen/models/wan/wan_vae.py`:
- Around line 267-285: Guard the fused dispatch in the Wan VAE upsampling path
around can_implement_dup_up3d and dup_up3d so it runs only when gradients are
not required: require not (torch.is_grad_enabled() and x.requires_grad),
preserving the eager path for autograd. In
tests/unittest/_torch/visual_gen/test_wan_dup_up3d.py lines 25-56, add a CUDA
regression that calls module(x).sum().backward() and verifies x.grad is
populated.
In `@tests/unittest/_torch/visual_gen/test_wan_dup_up3d.py`:
- Around line 80-81: Annotate the monkeypatched fail_if_launched function with
explicit parameter and return types, while preserving its existing pytest.fail
behavior and *args/**kwargs interface.
🪄 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: 7c8b85e6-0c92-4bdd-882c-5e5d564b05c3
📒 Files selected for processing (3)
tensorrt_llm/_torch/visual_gen/models/wan/dup_up3d.pytensorrt_llm/_torch/visual_gen/models/wan/wan_vae.pytests/unittest/_torch/visual_gen/test_wan_dup_up3d.py
Signed-off-by: Taian Zhang <taianz@nvidia.com>
|
/bot run |
|
PR_Github #62961 [ run ] triggered by Bot. Commit: |
|
PR_Github #62961 [ run ] completed with state |
|
/bot run |
|
PR_Github #63620 [ run ] triggered by Bot. Commit: |
|
PR_Github #63620 [ run ] completed with state |
Signed-off-by: Taian Zhang <taianz@nvidia.com>
|
/bot run |
|
PR_Github #64299 [ run ] triggered by Bot. Commit: |
|
PR_Github #64299 [ run ] completed with state |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Verified the kernel's inverse mapping against the eager path (repeat_interleave decomposition, permute(0,1,5,2,6,3,7,4) inversion, first-chunk crop) — it matches, and the inclusive <= 2^31 boundary is exactly right since the max linear offset is N−1. Test coverage is thorough. Comments below are touch-ups: a logger formatting bug in the fallback warning, one gap in the int32 guard for exotic strided inputs, and a small test-matrix addition. One heads-up, not a blocker: shapes and strides are all tl.constexpr, so each unique (shape, stride) combination triggers a Triton recompile — fine for fixed-resolution chunked VAE decode, but worth remembering if input shapes ever become dynamic per request.
| output_elements = math.prod(output_shape) | ||
| supported = _supports_triton_indexing(output_elements) | ||
| if not supported: | ||
| logger.warning_once( |
There was a problem hiding this comment.
Logger.log joins *msg with spaces (" ".join(map(str, msg))) — it does not do printf-style interpolation, so this prints a literal %d with the count appended after it. Use an f-string:
logger.warning_once(
f"Fused DupUp3D output has {output_elements} elements, exceeding the Triton "
"signed 32-bit indexing limit; falling back to the eager implementation.",
key="wan_dup_up3d_int32_index_fallback",
)There was a problem hiding this comment.
Fixed in 7651709. The warning now uses a single f-string, so the element counts are rendered correctly.
| first_chunk, | ||
| ) | ||
| output_elements = math.prod(output_shape) | ||
| supported = _supports_triton_indexing(output_elements) |
There was a problem hiding this comment.
This bounds only the output offsets, but the kernel also computes input_offset in 32-bit arithmetic from arbitrary strides. A strided view into a base tensor larger than 2^31 elements (e.g. big[..., ::2]) can pass this check while its max input linear offset sum((s - 1) * st for s, st in zip(x.shape, x.stride())) exceeds INT32_MAX, silently gathering from wrapped offsets. Unlikely in the VAE path (conv outputs are contiguous), but the docstring promises arbitrary strides — worth adding the input-extent bound to the same check.
There was a problem hiding this comment.
Agreed and fixed in 7651709. The eligibility check now bounds both the output size and the maximum relative input span, falling back before signed 32-bit indexing can overflow; a host-side regression test covers the large-stride case.
| A logical NCTHW tensor stored in PyTorch ``channels_last_3d`` physical | ||
| order, which corresponds to NTHWC. | ||
| """ | ||
| assert can_implement_dup_up3d( |
There was a problem hiding this comment.
This re-runs the full eligibility check the caller just performed (including the shape math), disappears under python -O, and hides a warning_once side effect inside an assert expression. Since forward already gates on can_implement_dup_up3d, consider dropping this to a cheap contract assert (or an explicit raise if you want direct callers protected under -O).
There was a problem hiding this comment.
Agreed and fixed in 7651709. The repeated full eligibility assert is replaced by the cheap internal contract validation; DupUp3D.forward remains the capability gate.
| @pytest.mark.parametrize( | ||
| ("in_channels", "out_channels", "factor_t"), | ||
| [ | ||
| pytest.param(4, 4, 2, id="temporal-repeat-8"), |
There was a problem hiding this comment.
All parity configs have repeats > 1, so eager's repeats == 1 branch (which skips repeat_interleave entirely, wan_vae.py:284) is never compared against the kernel's uniform // repeats path. Add one param where out_channels * factor == in_channels, e.g. pytest.param(32, 4, 2, id="repeats-1").
There was a problem hiding this comment.
Added the repeats == 1 parity configuration in 7651709.
Signed-off-by: Taian Zhang <taianz@nvidia.com>
|
/bot run |
|
PR_Github #64493 [ run ] triggered by Bot. Commit: |
| self.repeats = out_channels * self.factor // in_channels | ||
|
|
||
| def forward(self, x: torch.Tensor, first_chunk: bool = False) -> torch.Tensor: | ||
| from .dup_up3d import can_implement_dup_up3d, dup_up3d |
There was a problem hiding this comment.
maybe check if the code base has a convention about file/kernel/function naming of the triton kernels?
| def forward(self, x: torch.Tensor, first_chunk: bool = False) -> torch.Tensor: | ||
| from .dup_up3d import can_implement_dup_up3d, dup_up3d | ||
|
|
||
| if can_implement_dup_up3d( |
|
PR_Github #64493 [ run ] completed with state |
5 participants
Dev Engineer Review
DupUp3Doutput mapping intensorrt_llm/_torch/visual_gen/models/wan/dup_up3d.pythat performs channel repetition + temporal/spatial pixel-shuffle and writes directly to atorch.channels_last_3d(NTHWC) output using maskedtl.load/tl.store.can_implement_dup_up3d(...) -> boolto validate the internalDupUp3Dcontract, compute the fused output shape (includingfirst_chunktemporal cropping), and ensure the fused output element count fits Triton’s signed 32-bit indexing limit (_MAX_TRITON_INDEXED_ELEMENTS = 1<<31). For oversized outputs it emits a one-time warning (key="wan_dup_up3d_int32_index_fallback") and returnsFalse.dup_up3d(...)toassert can_implement_dup_up3d(...)before allocating/launching, removing any oversized-outputNonereturn path (now returnstorch.Tensor).DupUp3D.forwardintensorrt_llm/_torch/visual_gen/models/wan/wan_vae.pyto only take the CUDA fused path whencan_implement_dup_up3d(...)isTrue; otherwise it falls back to the existing eager tensor implementation.0 in x.shape), and (2) kernel-side masked offset handling (mask = offsets < output_elements).QA Engineer Review
tests/unittest/_torch/visual_gen/test_wan_dup_up3d.py:test_fused_dup_up3d_matches_eager(expanded): parameterized acrossinput_layout(contiguous,channels_last_3d,strided),first_chunk, and multiple(in_channels, out_channels, factor_t)configurations; asserts exact parity vs eager CPU and enforceschannels_last_3dcontiguity.test_fused_dup_up3d_empty_input_uses_eager: verifies empty CUDA inputs bypass fuseddup_up3dby monkeypatching the fused entrypoint to fail if launched.test_fused_dup_up3d_int32_index_limit: validates_supports_triton_indexingat the(1<<31)boundary.test_fused_dup_up3d_falls_back_above_index_limit: forces non-Triton dispatch by monkeypatching_MAX_TRITON_INDEXED_ELEMENTSand checks exact equality with eager CPU.tests/integration/test_lists/test-db/l0_b200.ymlto includeunittest/_torch/visual_gen/test_wan_dup_up3d.py.Description
Jira: TRTLLM-14798 (Task under epic TRTLLM-12775)
Wan 2.2's residual upsampling shortcut uses
DupUp3Dto repeat input channels,reshape and permute them into temporal/spatial pixel-shuffle positions,
materialize the permuted tensor, and optionally crop the first temporal chunk.
The eager path writes large repeated and permuted intermediates even though the
operation only rearranges values.
This PR adds one Triton kernel that maps each element of the final
channels-last output directly to its source input index. It fuses the channel
repeat, temporal/spatial pixel shuffle, final layout, and first-chunk temporal
crop without intermediate tensors.
The implementation:
channels_last_3d;output exceeds the Triton kernel's signed 32-bit indexing range
(
2^31elements); andThis is a shared VAE activation kernel, not an FP4 kernel. Native BF16 and FP4
Conv paths feed it the same BF16 activations and execute the same Triton source.
Performance
Matched four-GPU NSC A/B measurements show gains on both paths:
DupUp3DDupUp3DA later selected-stack FP4 campaign independently measured
368.843 -> 337.792 ms, a31.052 msreduction.Test Coverage
row-major, channels-last, and genuinely strided inputs, first/later chunks, temporal/spatial
factors, and production repeat counts 8, 4, and 2.
2^31signed-index limit and a forcedCUDA dispatch test that verifies oversized outputs take the eager fallback.
channels_last_3d.path selected automatically.
1553978task 4 passed 20 kernel matrix cases and thecheckpoint parity case.
1560274passed all 36 focused CUDA tests, including stridedaddressing and empty N/T/H/W eager-dispatch coverage.
git diff --checkpass.PR Checklist
FP4 integration report
This precision-independent optimization composes with the FP4 Wan VAE integration in #17262. The consolidated performance/accuracy report and dependency map are maintained there.