From c8ad8521069479ea5aef62ae86ef7c0fe93baf31 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 29 Jan 2026 12:32:29 -0800 Subject: [PATCH 1/2] Remove all uses of nvFuser legacy bindings as a follow-up to https://github.com/Lightning-AI/lightning-thunder/pull/2812 --- dockers/ubuntu-cuda/Dockerfile | 2 +- docs/source/basic/inspecting_traces.rst | 2 +- thunder/executors/nvfuserex.py | 13 +------------ thunder/executors/nvfuserex_impl.py | 26 +++++++++---------------- thunder/tests/test_dynamo.py | 2 +- 5 files changed, 13 insertions(+), 32 deletions(-) diff --git a/dockers/ubuntu-cuda/Dockerfile b/dockers/ubuntu-cuda/Dockerfile index 6806a2c18c..9b839a0eed 100644 --- a/dockers/ubuntu-cuda/Dockerfile +++ b/dockers/ubuntu-cuda/Dockerfile @@ -158,5 +158,5 @@ RUN \ pip list && \ python -c "import sys; ver = sys.version_info ; assert f'{ver.major}.{ver.minor}' == '$PYTHON_VERSION', ver" && \ python -c "import torch; print(f'PyTorch=={torch.__version__} with {torch.cuda.device_count()} GPUs')" && \ - python -c "import nvfuser; print(f'nvFuser=={nvfuser.version()}')" && \ + python -c "import nvfuser_direct as nvfuser; print(f'nvFuser=={nvfuser.version()}')" && \ python -c "import triton; print(f'Triton=={triton.__version__}')" diff --git a/docs/source/basic/inspecting_traces.rst b/docs/source/basic/inspecting_traces.rst index ea91ad0f77..2e7ad77eae 100644 --- a/docs/source/basic/inspecting_traces.rst +++ b/docs/source/basic/inspecting_traces.rst @@ -206,7 +206,7 @@ This will print the following:: # cuda version: 12.1 # nvfuser version: 0.2.8 import torch - from nvfuser import FusionDefinition, DataType + from nvfuser_direct import FusionDefinition, DataType def nvfuser_fusion_id0(fd : FusionDefinition) -> None : T0 = fd.define_tensor(shape=[-1, -1], contiguity=[True, True], dtype=DataType.Float, is_cpu=False, stride_order=[1, 0]) diff --git a/thunder/executors/nvfuserex.py b/thunder/executors/nvfuserex.py index 93ddd29b5a..49b95ffd41 100644 --- a/thunder/executors/nvfuserex.py +++ b/thunder/executors/nvfuserex.py @@ -18,23 +18,12 @@ def nvfuser_version() -> LooseVersion | None: try: import nvfuser_direct except ImportError: - try: - import nvfuser - except ImportError: - pass - else: - if hasattr(nvfuser, "version"): - return LooseVersion(nvfuser.version()) - else: - # NOTE: This import of nvFuser may or may not have version info - return LooseVersion("0.0.0") + return None else: if hasattr(nvfuser_direct, "version"): return LooseVersion(nvfuser_direct.version()) else: return LooseVersion("0.0.0") - # NOTE This occurs when nvFuser couldn't be imported - return None def required_nvfuser_version() -> LooseVersion: diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py index 29a474e30a..cbcc045664 100644 --- a/thunder/executors/nvfuserex_impl.py +++ b/thunder/executors/nvfuserex_impl.py @@ -71,24 +71,16 @@ # NOTE This impl file is here because nvFuser may not be available, so it's imported conditionally # by nvfuserex.py when nvFuser is available. -DIRECT_BINDINGS_SUPPORTED_VERSION = LooseVersion("0.2.34") DTENSOR_SUPPORTED_VERSION = LooseVersion("0.2.28") -if nvfuser_version() >= DIRECT_BINDINGS_SUPPORTED_VERSION: - import nvfuser_direct as nvfuser - from nvfuser_direct import ( - DataType, - FusionDefinition, - multidevice, - ParallelType, - execute_with_dtensors, - compute_tensor_descriptor as nv_compute_td, - ) -else: - if nvfuser_version() >= DTENSOR_SUPPORTED_VERSION: - from nvfuser_direct import FusionDefinition as DirectFusionDefinition - from nvfuser_direct import multidevice, ParallelType, execute_with_dtensors - import nvfuser - from nvfuser import DataType, FusionDefinition, compute_tensor_descriptor as nv_compute_td +import nvfuser_direct as nvfuser +from nvfuser_direct import ( + DataType, + FusionDefinition, + multidevice, + ParallelType, + execute_with_dtensors, + compute_tensor_descriptor as nv_compute_td, +) # # Helper functions diff --git a/thunder/tests/test_dynamo.py b/thunder/tests/test_dynamo.py index a60bc50797..987dd1f6fe 100644 --- a/thunder/tests/test_dynamo.py +++ b/thunder/tests/test_dynamo.py @@ -1403,7 +1403,7 @@ def foo(x): @pytest.mark.skip(reason="https://github.com/Lightning-AI/lightning-thunder/issues/2546") @requiresCUDA def test_WallTime_KernelTime(): - from nvfuser import FusionDefinition, DataType + from nvfuser_direct import FusionDefinition, DataType def nvfuser_fusion_id2(fd: FusionDefinition) -> None: T0 = fd.define_tensor( From 373fb152cdd67a9e22b685f5381ae78f8cbc77c4 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 29 Jan 2026 15:48:41 -0800 Subject: [PATCH 2/2] Fix --- thunder/executors/nvfuserex_impl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py index cbcc045664..f084eabab1 100644 --- a/thunder/executors/nvfuserex_impl.py +++ b/thunder/executors/nvfuserex_impl.py @@ -386,7 +386,7 @@ def check_dtensor_tracing_and_runtime_metadata(inp): lambda: "nvfuser: Expected runtime and tracing metadata to be the same for DTensor.", ) - fd = FusionDefinition() if nvfuser_version() >= DIRECT_BINDINGS_SUPPORTED_VERSION else DirectFusionDefinition() + fd = FusionDefinition() # Device may be set in one of the "factory" methods like full, iota, or uniform # NOTE: This should be called before defining because a factory method may look-up at `_selected_device` while being defined. fd._selected_device = None