Skip to content

[https://nvbugs/6329227][fix] Use pkgutil.extend_path to merge the two flash_attn distributions before… #15498

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
tensorrt-cicd wants to merge 4 commits into from
Closed

[https://nvbugs/6329227][fix] Use pkgutil.extend_path to merge the two flash_attn distributions before… #15498

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ce97e9a
[nvbugs/6329227][fix] Propagate parent sys.path to mp.spawn child
tensorrt-cicd Jun 19, 2026
a776fc5
[nvbugs/6329227][fix] Resolve flash_attn.cute and REPO_ROOT after Vis…
tensorrt-cicd Jun 19, 2026
adea537
[nvbugs/6329227][chore] Remove stale waiver after fix
tensorrt-cicd Jun 25, 2026
0ffa4c2
[nvbugs/6329227][fix] Wire _visual_gen_deps into multi-GPU LPIPS tests
trtllm-agent Jul 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 15 additions & 0 deletions tensorrt_llm/_torch/visual_gen/attention_backend/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down
10 changes: 7 additions & 3 deletions tests/integration/defs/examples/visual_gen/test_visual_gen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
17 changes: 12 additions & 5 deletions tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
"""Multi-GPU integration tests for VisualGen LPIPS quality checks."""

import os
import sys
from typing import Callable

import pytest
Expand All @@ -37,6 +38,7 @@
_run_lpips_eval,
_run_wan_lpips_pipeline,
_save_lpips_video_mp4,
_visual_gen_deps,
)

try:
Expand Down Expand Up @@ -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(
Comment on lines +113 to 118

@coderabbitai coderabbitai Bot Jun 19, 2026

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.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Restore PYTHONPATH after spawning workers.

Line 116 mutates process-global environment state and never restores it; later tests can inherit this synthetic value and become order-dependent/flaky.

🔧 Proposed fix 
-    os.environ["PYTHONPATH"] = os.pathsep.join(filter(None, sys.path))
-    mp.spawn(
-        _distributed_worker,
-        args=(world_size, backend, test_fn, port, kwargs),
-        nprocs=world_size,
-        join=True,
-    )
+    previous_pythonpath = os.environ.get("PYTHONPATH")
+    os.environ["PYTHONPATH"] = os.pathsep.join(map(str, filter(None, sys.path)))
+    try:
+        mp.spawn(
+            _distributed_worker,
+            args=(world_size, backend, test_fn, port, kwargs),
+            nprocs=world_size,
+            join=True,
+        )
+    finally:
+        if previous_pythonpath is None:
+            os.environ.pop("PYTHONPATH", None)
+        else:
+            os.environ["PYTHONPATH"] = previous_pythonpath
🤖 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/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py`
around lines 112 - 117, The code modifies os.environ["PYTHONPATH"] on line 115
to set up the environment for the spawned workers but never restores the
original value after mp.spawn() completes, which can cause subsequent tests to
inherit this modified state and become order-dependent. Save the original
PYTHONPATH value before setting it (handle the case where it may not exist),
then restore it after the mp.spawn() call finishes to ensure the environment is
clean for subsequent tests.

_distributed_worker,
args=(world_size, backend, test_fn, port, kwargs),
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -200,14 +207,14 @@ 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(
"variant_name,parallel",
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)
3 changes: 0 additions & 3 deletions tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading