From b3a5bba49f26f1423d1ed462c1e870e60675c5b5 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sun, 1 Nov 2026 22:22:43 +0100 Subject: [PATCH] [None][fix] Enable the fused per-token AdaLN on SM103 (B300/GB300) 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 --- .../kernels/blackwell/pertoken_adaln/pertoken_adaln.py | 4 ++-- tensorrt_llm/_torch/visual_gen/models/wan/utils_wan.py | 2 +- tests/unittest/_torch/visual_gen/test_pertoken_adaln.py | 7 +++++-- .../_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py | 4 ++-- 4 files changed, 10 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/kernels/blackwell/pertoken_adaln/pertoken_adaln.py b/tensorrt_llm/_torch/visual_gen/kernels/blackwell/pertoken_adaln/pertoken_adaln.py index 99102b4c93cd..1f8cec860cce 100644 --- a/tensorrt_llm/_torch/visual_gen/kernels/blackwell/pertoken_adaln/pertoken_adaln.py +++ b/tensorrt_llm/_torch/visual_gen/kernels/blackwell/pertoken_adaln/pertoken_adaln.py @@ -408,8 +408,8 @@ def _launch( hash_key = (device, vec, *PerTokenAdaLN.make_hash_key(*torch_tensors)) compiled_fn = _COMPILE_CACHE.get(hash_key) if compiled_fn is None: - if torch.cuda.get_device_capability(device) != (10, 0): - raise ValueError("fused per-token AdaLN requires an SM100 GPU") + if torch.cuda.get_device_capability(device) not in ((10, 0), (10, 3)): + raise ValueError("fused per-token AdaLN requires an SM100 or SM103 GPU") kernel = PerTokenAdaLN(D, vec) fake_sig_args = [_to_fake_cute_arg(t) for t in torch_tensors] compiled_fn = cute.compile(kernel, *fake_sig_args, options="--enable-tvm-ffi") diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/utils_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/utils_wan.py index d787f1bb994c..d1d872bb8803 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/utils_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/utils_wan.py @@ -46,7 +46,7 @@ def _pertoken_adaln_arch_ok(device: torch.device) -> bool: if device.type != "cuda" or not torch.cuda.is_available(): return False try: - return torch.cuda.get_device_capability(device) == (10, 0) + return torch.cuda.get_device_capability(device) in ((10, 0), (10, 3)) except RuntimeError: return False diff --git a/tests/unittest/_torch/visual_gen/test_pertoken_adaln.py b/tests/unittest/_torch/visual_gen/test_pertoken_adaln.py index 1243755b9e52..8f4dc87dee79 100644 --- a/tests/unittest/_torch/visual_gen/test_pertoken_adaln.py +++ b/tests/unittest/_torch/visual_gen/test_pertoken_adaln.py @@ -34,8 +34,11 @@ def _require_sm100() -> torch.device: - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): - pytest.skip("The fused per-token AdaLN kernel requires an SM100 GPU.") + if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ( + (10, 0), + (10, 3), + ): + pytest.skip("The fused per-token AdaLN kernel requires an SM100 or SM103 GPU.") return torch.device("cuda", torch.cuda.current_device()) diff --git a/tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py index cc32027b37cf..9240ec80fcc5 100644 --- a/tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py @@ -195,7 +195,7 @@ def _assert_pipeline_matches_hf( assert trtllm_pipe.transformer_2 is None assert trtllm_pipe.expand_timesteps is True - is_sm100 = torch.cuda.get_device_capability() == (10, 0) + is_sm100 = torch.cuda.get_device_capability() in ((10, 0), (10, 3)) plain_call = None residual_call = None with ExitStack() as stack: @@ -443,7 +443,7 @@ def test_fp8_cache_dit_trtllm(self): assert t2v_result.video.dim() == 5 B, _T, H, W, C = t2v_result.video.shape assert B == 1 and H == 704 and W == 1280 and C == 3 - if torch.cuda.get_device_capability() == (10, 0): + if torch.cuda.get_device_capability() in ((10, 0), (10, 3)): assert all(block._pertoken_adaln._enabled for block in pipeline.transformer.blocks) assert pipeline.cache_accelerator is not None