Repository navigation
Conversation
The WAN VisualGen fused per-token AdaLN was gated on exactly compute capability 10.0 at the kernel launch and in _pertoken_adaln_arch_ok, so B300/GB300 (10.3) always took the unfused eager path. The kernel uses only SIMT CuTe DSL constructs (warp shuffles, CTA reductions, vectorized copies) and is JIT-compiled for the running device, so admit 10.0 and 10.3 (but not 10.7), and update the test skips and fused-path assertions to match. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Runtime architecture checks tensorrt_llm/_torch/visual_gen/kernels/blackwell/pertoken_adaln/pertoken_adaln.py, tensorrt_llm/_torch/visual_gen/models/wan/utils_wan.py |
The per-token AdaLN launch and Wan eligibility check accept compute capabilities (10, 0) and (10, 3). Other capabilities remain unsupported by these checks. |
Architecture test checks tests/unittest/_torch/visual_gen/test_pertoken_adaln.py, tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py |
The per-token AdaLN unit and pipeline tests include compute capability (10, 3) in their architecture checks. The unit test still skips when CUDA is unavailable or the capability is unsupported. |
Estimated code review effort: 2 (Simple) | ~10 minutes
Change: Bug fix
Suggested reviewers: chang-l
Merge Risk: 
b3a5b
-
The change is mergeable based on the consistent runtime and test architecture checks; full pipeline validation remains environment-dependent.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Title check The title clearly and concisely states the main change: enabling fused per-token AdaLN on SM103. Description check The description includes the required Description, Test Coverage, and PR Checklist sections. It explains the issue and solution, reports test results and limitations, and marks the checklist as review… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Full details: Description check
-
Explanation
The description includes the required Description, Test Coverage, and PR Checklist sections. It explains the issue and solution, reports test results and limitations, and marks the checklist as reviewed.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 5 functions across 4 files.
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches
-
Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Description
The fused per-token AdaLN for WAN VisualGen (#17695) is gated on exactly compute capability 10.0, in two places:
pertoken_adaln.pyraisesValueError("fused per-token AdaLN requires an SM100 GPU");_pertoken_adaln_arch_okinutils_wan.pyreturnsFalse, so WAN takes the unfused path.As a result, B300/GB300 (10.3) always run the unfused eager path.
The kernel itself has nothing SM100-specific.
PerTokenAdaLNandreduce.pyuse onlythread_idx/block_idx,shuffle_sync_down,sync_threadsand vectorized copies. There is no tcgen05, TMA or TMEM, andcute.compileJIT-compiles it for the device it runs on.sol_attn/interface.pyalready relies on the same reasoning for SM103 ("no SM100-exclusive construct (measured identical on B200 and B300)").This PR admits 10.0 and 10.3, and still excludes 10.7, matching
is_blackwell_datacenter():pertoken_adaln.py: the launch check accepts(10, 0)and(10, 3).utils_wan.py:_pertoken_adaln_arch_okuses the same check.test_pertoken_adaln.py: the skip now admits 10.3.test_wan22_ti2v_5b_pipeline.py: the two fused-path engagement asserts now also expect the fused path on 10.3.Test Coverage
On one B300 (CC 10.3), with nvidia-cutlass-dsl 4.8.0, apache-tvm-ffi 0.1.13.post2 and PyTorch 2.14.0+cu130 as in
requirements.txt, I ran the kernel package (pertoken_adaln.pybyte-identical to this PR, and tomainat 309426d for the baseline) against the kernel-level tests oftest_pertoken_adaln.py. Those tests compare against an FP32 reference and cover torch.compile fullgraph and contract checks. I loaded the kernel package standalone and left out the WAN-module test, which needs the full package.mainValueError: fused per-token AdaLN requires an SM100 GPU), 2 passed (the contract-only tests)I did not run
utils_wan.pyortest_wan22_ti2v_5b_pipeline.py, which need the full package and the WAN 2.2 checkpoint, and did not measure the end-to-end speedup on B300. #17695 measured 4.3% end-to-end on B200.ruff checkandruff format --checkpass on the changed files.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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.
Dev Engineer Review
QA Engineer Review
test_pertoken_adaln.pynow permits tests on compute capability 10.3. The author reports 7 kernel-level tests passing on B300, including FP32-reference comparisons,torch.compilefullgraph coverage, and contract checks. The file is listed in thel0_b200.ymlCI list.test_wan22_ti2v_5b_pipeline.pynow checks the fused path on capabilities 10.0 and 10.3. It is listed in thel0_b200.ymlCI list and thellm_function_core.txtmanual-QA list. The author did not run this pipeline test because it requires the full package and WAN 2.2 checkpoint.Per-File QA Perspective
pertoken_adaln.py: Verify fused-kernel dispatch on SM100 and SM103, and confirm unsupported capabilities still raiseValueError.utils_wan.py: Verify the Wan model selects fused per-token AdaLN on SM100 and SM103 and keeps it disabled on other architectures.test_pertoken_adaln.py: Covers kernel behavior, compilation, and contracts on supported devices. It is listed in the CI test list.test_wan22_ti2v_5b_pipeline.py: Covers fused-path assertions and combined optimizations for the Wan 2.2 TI2V pipeline. It appears in both the CI and manual-QA lists; execution remains unreported.