Remove all uses of nvFuser legacy bindings

as a follow-up to #2812

Author: wujingyue

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__}')"

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])

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:

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

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(