update zluda loader

Author: lshqqytiger

modules/launch_utils.py

@@ -593,7 +593,6 @@ def prepare_environment():
                     if device is not None and zluda_installer.get_blaslt_enabled():
                         print(f'ROCm hipBLASLt: arch={device.name} available={device.blaslt_supported}')
                         zluda_installer.set_blaslt_enabled(device.blaslt_supported)
-                    zluda_installer.make_copy()
                     zluda_installer.load()
                     torch_command = os.environ.get('TORCH_COMMAND', 'pip install torch==2.6.0 torchvision --index-url https://download.pytorch.org/whl/cu118')
                     print(f'Using ZLUDA in {zluda_installer.path}')

modules/rocm.py

@@ -172,12 +172,12 @@ def get_version() -> str:
         return f'{arr[0]}.{arr[1]}' if len(arr) >= 2 else None
 
     def get_agents() -> List[Agent]:
-        if is_wsl: # WSL does not have 'rocm_agent_enumerator'
-            agents = spawn("rocminfo").split("\n")
-            agents = [x.strip().split(" ")[-1] for x in agents if x.startswith('  Name:') and "CPU" not in x]
-        else:
+        try:
             agents = spawn("rocm_agent_enumerator").split("\n")
             agents = [x for x in agents if x and x != 'gfx000']
+        except Exception: # old version of ROCm WSL doesn't have rocm_agent_enumerator
+            agents = spawn("rocminfo").split("\n")
+            agents = [x.strip().split(" ")[-1] for x in agents if x.startswith('  Name:') and "CPU" not in x]
         return [Agent(x) for x in agents]
 
     def load_hsa_runtime() -> None:
@@ -204,7 +204,9 @@ def get_blaslt_enabled() -> bool:
     def get_flash_attention_command(agent: Agent):
         default = "git+https://github.com/ROCm/flash-attention"
         if agent.gfx_version >= 0x1100 and os.environ.get("FLASH_ATTENTION_USE_TRITON_ROCM", "false").lower() != "true":
-            default = "git+https://github.com/ROCm/flash-attention@howiejay/navi_support"
+            # use the navi_rotary_fix fork because the original doesn't support rotary_emb for transformers
+            # original: "git+https://github.com/ROCm/flash-attention@howiejay/navi_support"
+            default = "https://github.com/Disty0/flash-attention@navi_rotary_fix"
         return os.environ.get("FLASH_ATTENTION_PACKAGE", default)
 
     is_wsl: bool = os.environ.get('WSL_DISTRO_NAME', 'unknown' if spawn('wslpath -w /') else None) is not None

modules/zluda.py

@@ -33,8 +33,8 @@ def initialize_zluda():
     from modules.zluda_hijacks import do_hijack
     do_hijack()
 
-    torch.backends.cudnn.enabled = zluda_installer.MIOpen_available
-    if not zluda_installer.MIOpen_available:
+    torch.backends.cudnn.enabled = zluda_installer.MIOpen_enabled
+    if not zluda_installer.MIOpen_enabled:
         torch.backends.cuda.enable_cudnn_sdp(False)
         torch.backends.cuda.enable_cudnn_sdp = do_nothing
     torch.backends.cuda.enable_flash_sdp(False)

modules/zluda_hijacks.py

@@ -1,7 +1,7 @@
 from functools import wraps
 import torch
 import torch._dynamo.device_interface
-from modules import rocm, zluda, shared
+from modules import shared, zluda # pylint: disable=unused-import
 
 
 MEM_BUS_WIDTH = {
@@ -13,11 +13,33 @@
     "AMD Radeon RX 7900 GRE": 256,
     "AMD Radeon RX 7800 XT": 256,
     "AMD Radeon RX 7700 XT": 192,
+    "AMD Radeon RX 7700": 192,
+    "AMD Radeon RX 7650 GRE": 128,
     "AMD Radeon RX 7600 XT": 128,
     "AMD Radeon RX 7600": 128,
+    "AMD Radeon RX 7500 XT": 96,
+    "AMD Radeon RX 6950 XT": 256,
+    "AMD Radeon RX 6900 XT": 256,
+    "AMD Radeon RX 6800 XT": 256,
+    "AMD Radeon RX 6800": 256,
+    "AMD Radeon RX 6750 XT": 192,
+    "AMD Radeon RX 6700 XT": 192,
+    "AMD Radeon RX 6700": 160,
+    "AMD Radeon RX 6650 XT": 128,
+    "AMD Radeon RX 6600 XT": 128,
+    "AMD Radeon RX 6600": 128,
+    "AMD Radeon RX 6500 XT": 64,
+    "AMD Radeon RX 6400": 64,
 }
 
 
+_topk = torch.topk
+def topk(input: torch.Tensor, *args, **kwargs): # pylint: disable=redefined-builtin
+    device = input.device
+    values, indices = _topk(input.cpu(), *args, **kwargs)
+    return torch.return_types.topk((values.to(device), indices.to(device),))
+
+
 class DeviceProperties:
     PROPERTIES_OVERRIDE = {"regs_per_multiprocessor": 65535, "gcnArchName": "UNKNOWN ARCHITECTURE"}
     internal: torch._C._CudaDeviceProperties
@@ -42,8 +64,7 @@ def torch__C__cuda_getCurrentRawStream(device):
 
 
 def do_hijack():
-    torch.version.hip = rocm.version
-
+    torch.topk = topk
     if zluda.default_agent is not None:
         DeviceProperties.PROPERTIES_OVERRIDE["gcnArchName"] = zluda.default_agent.name
     torch.cuda._get_device_properties = torch_cuda__get_device_properties # pylint: disable=protected-access

modules/zluda_installer.py

@@ -7,6 +7,7 @@
 import urllib.request
 from typing import Union
 from modules import rocm
+from modules.launch_utils import args
 
 
 DLL_MAPPING = {
@@ -18,15 +19,12 @@
 }
 HIPSDK_TARGETS = ['rocblas.dll', 'rocsolver.dll', 'hipfft.dll',]
 
-hipBLASLt_available = False
-MIOpen_available = False
+MIOpen_enabled = False
 
 path = os.path.abspath(os.environ.get('ZLUDA', '.zluda'))
 default_agent: Union[rocm.Agent, None] = None
 hipBLASLt_enabled = False
 
-nightly = os.environ.get("ZLUDA_NIGHTLY", "0") == "1"
-
 
 class ZLUDAResult(ctypes.Structure):
     _fields_ = [
@@ -66,42 +64,25 @@ def get_nightly_flag(self) -> int:
 ml = None
 
 
-def load_core_modules():
-    global core, ml # pylint: disable=global-statement
-    core = Core(ctypes.windll.LoadLibrary(os.path.join(path, 'nvcuda.dll')))
-    ml = ZLUDALibrary(ctypes.windll.LoadLibrary(os.path.join(path, 'nvml.dll')))
-
-
 def set_default_agent(agent: rocm.Agent):
     global default_agent # pylint: disable=global-statement
     default_agent = agent
 
-    is_nightly = False
-    try:
-        load_core_modules()
-        is_nightly = core.get_nightly_flag() == 1
-    except Exception:
-        pass
-
-    global hipBLASLt_available, hipBLASLt_enabled # pylint: disable=global-statement
-    hipBLASLt_available = is_nightly and os.path.exists(rocm.blaslt_tensile_libpath)
-    hipBLASLt_enabled = hipBLASLt_available and os.path.exists(os.path.join(rocm.path, "bin", "hipblaslt.dll"))
-
-    global MIOpen_available # pylint: disable=global-statement
-    MIOpen_available = is_nightly and os.path.exists(os.path.join(rocm.path, "bin", "MIOpen.dll"))
-
 
 def is_reinstall_needed() -> bool: # ZLUDA<3.8.7
     return not os.path.exists(os.path.join(path, 'cufftw.dll'))
 
 
-def install() -> None:
+def install():
     if os.path.exists(path):
         return
 
     platform = "windows"
-    commit = os.environ.get("ZLUDA_HASH", "ae0540beb129ffd140226ce956b386619b38f84c")
-    if nightly:
+    commit = os.environ.get("ZLUDA_HASH", "dba64c0966df2c71e82255e942c96e2e1cea3a2d")
+    if os.environ.get("ZLUDA_NIGHTLY", "0") == "1":
+        print("Warning: Environment variable 'ZLUDA_NIGHTLY' will be removed. Please use command-line argument '--use-nightly' instead.")
+        args.use_nightly = True
+    if args.use_nightly:
         platform = "nightly-" + platform
     urllib.request.urlretrieve(f'https://github.com/lshqqytiger/ZLUDA/releases/download/rel.{commit}/ZLUDA-{platform}-rocm{rocm.version[0]}-amd64.zip', '_zluda')
     with zipfile.ZipFile('_zluda', 'r') as archive:
@@ -113,7 +94,7 @@ def install() -> None:
     os.remove('_zluda')
 
 
-def uninstall() -> None:
+def uninstall():
     if os.path.exists(path):
         shutil.rmtree(path)
 
@@ -137,23 +118,29 @@ def link_or_copy(src: os.PathLike, dst: os.PathLike):
             shutil.copyfile(src, dst)
 
 
-def make_copy() -> None:
+def load():
+    global core, ml, hipBLASLt_enabled, MIOpen_enabled # pylint: disable=global-statement
+    core = Core(ctypes.windll.LoadLibrary(os.path.join(path, 'nvcuda.dll')))
+    ml = ZLUDALibrary(ctypes.windll.LoadLibrary(os.path.join(path, 'nvml.dll')))
+    is_nightly = core.get_nightly_flag() == 1
+    hipBLASLt_enabled = is_nightly and os.path.exists(rocm.blaslt_tensile_libpath) and os.path.exists(os.path.join(rocm.path, "bin", "hipblaslt.dll"))
+    MIOpen_enabled = is_nightly and os.path.exists(os.path.join(rocm.path, "bin", "MIOpen.dll"))
+
     for k, v in DLL_MAPPING.items():
         if not os.path.exists(os.path.join(path, v)):
             link_or_copy(os.path.join(path, k), os.path.join(path, v))
 
     if hipBLASLt_enabled and not os.path.exists(os.path.join(path, 'cublasLt64_11.dll')):
         link_or_copy(os.path.join(path, 'cublasLt.dll'), os.path.join(path, 'cublasLt64_11.dll'))
 
-    if MIOpen_available and not os.path.exists(os.path.join(path, 'cudnn64_9.dll')):
+    if MIOpen_enabled and not os.path.exists(os.path.join(path, 'cudnn64_9.dll')):
         link_or_copy(os.path.join(path, 'cudnn.dll'), os.path.join(path, 'cudnn64_9.dll'))
 
+    print(f"ZLUDA load: path='{path}' nightly={bool(core.get_nightly_flag())}")
 
-def load() -> None:
     os.environ["ZLUDA_COMGR_LOG_LEVEL"] = "1"
     os.environ["ZLUDA_NVRTC_LIB"] = os.path.join([v for v in site.getsitepackages() if v.endswith("site-packages")][0], "torch", "lib", "nvrtc64_112_0.dll")
 
-    load_core_modules()
     for v in HIPSDK_TARGETS:
         ctypes.windll.LoadLibrary(os.path.join(rocm.path, 'bin', v))
     for v in DLL_MAPPING.values():
@@ -166,12 +153,13 @@ def load() -> None:
     else:
         os.environ["DISABLE_ADDMM_CUDA_LT"] = "1"
 
-    if MIOpen_available:
+    if MIOpen_enabled:
         ctypes.windll.LoadLibrary(os.path.join(rocm.path, 'bin', 'MIOpen.dll'))
         ctypes.windll.LoadLibrary(os.path.join(path, 'cudnn64_9.dll'))
 
     def conceal():
-        import torch # noqa: F401
+        import torch
+        torch.version.hip = rocm.version
         platform = sys.platform
         sys.platform = ""
         from torch.utils import cpp_extension