diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py index 532e86e91e08..312aa13796b4 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py @@ -26,9 +26,11 @@ from ...attention_backend.interface import PredefinedAttentionMask from .interface import AttentionBackend, AttentionTensorLayout +from .utils import _merge_flash_attn_namespace _flash_attn_fwd_import_error = None try: + _merge_flash_attn_namespace() from flash_attn.cute.interface import _flash_attn_fwd except (ImportError, OSError) as e: _flash_attn_fwd = None diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py index b9d5bac5e423..51091d3931ef 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -33,9 +33,11 @@ from ...attention_backend.interface import PredefinedAttentionMask from .interface import AttentionBackend, AttentionTensorLayout +from .utils import _merge_flash_attn_namespace _flash_attn_combine_import_error = None try: + _merge_flash_attn_namespace() from flash_attn.cute.interface import flash_attn_combine as _flash_attn_combine except (ImportError, OSError) as e: _flash_attn_combine = None diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index a3ce3df89cd9..8ac73770c8ec 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -30,6 +30,21 @@ from .interface import AttentionBackend +def _merge_flash_attn_namespace() -> None: + """Merge legacy `flash-attn` (v2) and `flash-attn-4` (cute-only) wheels. + + `flash-attn-4` ships only `flash_attn/cute/` and relies on namespace + extension into the co-installed legacy `flash-attn` regular package. + When both wheels coexist, Python pins `flash_attn.__path__` to the v2 + location and never sees `cute/`; re-extend the path to merge them. + """ + import pkgutil + + import flash_attn + + flash_attn.__path__ = pkgutil.extend_path(flash_attn.__path__, flash_attn.__name__) + + def get_visual_gen_attention_backend( backend_name: str, ) -> Type[AttentionBackend]: diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen.py b/tests/integration/defs/examples/visual_gen/test_visual_gen.py index 916518acb644..365b079fde77 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen.py @@ -235,9 +235,13 @@ def _visual_gen_deps(llm_venv): """Install av + diffusers + ffmpeg once per session (shared by all video-gen fixtures).""" llm_venv.run_cmd(["-m", "pip", "install", "av"]) llm_venv.run_cmd(["-m", "pip", "install", "diffusers>=0.37.0"]) - # Install ffmpeg system package required by save_video() for MP4 encoding - check_call(["apt-get", "update", "-y"], shell=False) - check_call(["apt-get", "install", "-y", "ffmpeg"], shell=False) + if shutil.which("ffmpeg"): + return + # apt-get needs root; add `sudo` when the test harness is running as an + # unprivileged user (e.g. LOCAL_USER=1 dev containers) but sudo is available. + prefix = [] if os.geteuid() == 0 else ["sudo", "-n"] + check_call([*prefix, "apt-get", "update", "-y"], shell=False) + check_call([*prefix, "apt-get", "install", "-y", "ffmpeg"], shell=False) @pytest.fixture(scope="session") diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index a60765fe600c..8973191ac66c 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -15,6 +15,7 @@ """Multi-GPU integration tests for VisualGen LPIPS quality checks.""" import os +import sys from typing import Callable import pytest @@ -37,6 +38,7 @@ _run_lpips_eval, _run_wan_lpips_pipeline, _save_lpips_video_mp4, + _visual_gen_deps, ) try: @@ -108,6 +110,11 @@ def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") backend = "nccl" if use_cuda else "gloo" port = get_free_port() + # mp.spawn starts a fresh interpreter that does not inherit the parent's + # runtime sys.path mutations; without this, the child cannot import + # tensorrt_llm.bindings.internal.process_group (needed by the tp_size>1 + # C++ allreduce path). + os.environ["PYTHONPATH"] = os.pathsep.join(filter(None, sys.path)) mp.spawn( _distributed_worker, args=(world_size, backend, test_fn, port, kwargs), @@ -158,7 +165,7 @@ def _wan22_lpips_distributed_worker(rank: int, world_size: int, **kwargs) -> Non dist.barrier() -def _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel): +def _run_wan22_t2v_lpips_case(_visual_gen_deps, tmp_path, variant_name, parallel): _skip_if_insufficient_gpus_for_parallel(parallel) parallel_cfg = ParallelConfig(**parallel) generated_path = tmp_path / f"wan22_t2v_generated_{variant_name}.mp4" @@ -200,8 +207,8 @@ def _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel): WAN22_LPIPS_MULTI_GPU_VARIANTS, ids=[name for name, _ in WAN22_LPIPS_MULTI_GPU_VARIANTS], ) -def test_wan22_t2v_lpips_against_golden_multi_gpu(tmp_path, variant_name, parallel): - _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel) +def test_wan22_t2v_lpips_against_golden_multi_gpu(_visual_gen_deps, tmp_path, variant_name, parallel): + _run_wan22_t2v_lpips_case(_visual_gen_deps, tmp_path, variant_name, parallel) @pytest.mark.parametrize( @@ -209,5 +216,5 @@ def test_wan22_t2v_lpips_against_golden_multi_gpu(tmp_path, variant_name, parall WAN22_LPIPS_TP_VARIANTS, ids=[name for name, _ in WAN22_LPIPS_TP_VARIANTS], ) -def test_wan22_t2v_lpips_against_golden_tp(tmp_path, variant_name, parallel): - _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel) +def test_wan22_t2v_lpips_against_golden_tp(_visual_gen_deps, tmp_path, variant_name, parallel): + _run_wan22_t2v_lpips_case(_visual_gen_deps, tmp_path, variant_name, parallel) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 11beee20e467..e42853fc2b9c 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -280,9 +280,6 @@ examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_g examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2] SKIP (https://nvbugs/6272644) examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2_attn2d_2x1] SKIP (https://nvbugs/6272644) examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[ulysses4] SKIP (https://nvbugs/6272644) -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_tp[cfg2_tp2] SKIP (https://nvbugs/6329227) -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_tp[tp2] SKIP (https://nvbugs/6329227) -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_tp[tp2_ulysses2] SKIP (https://nvbugs/6329227) full:A100/disaggregated/test_workers.py::test_workers_conditional_disaggregation_deepseek_v3_lite_bf16[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6329052) full:A100X/llmapi/test_llm_examples.py::test_llmapi_speculative_decoding_mtp SKIP (https://nvbugs/6287561) full:A100X/unittest/llmapi/test_llm_pytorch.py -m "part0" SKIP (https://nvbugs/6416249)