Support ROCm on Windows with ADLX VRAM readings

diffsynth/core/device/__init__.py

@@ -1,2 +1,3 @@
 from .npu_compatible_device import parse_device_type, parse_nccl_backend, get_available_device_type, get_device_name
 from .npu_compatible_device import IS_NPU_AVAILABLE, IS_CUDA_AVAILABLE
+from .rocm_windows import IS_WINDOWS_ROCM, mem_get_info

diffsynth/core/device/rocm_windows.py

@@ -0,0 +1,88 @@
+import atexit
+import os
+import sys
+import threading
+import time
+import torch
+
+
+IS_WINDOWS_ROCM = sys.platform == "win32" and torch.version.hip is not None
+
+# ADLX takes about half a millisecond per query and `check_free_vram` runs on every layer forward,
+# so the VRAM used by other processes is refreshed at most this often (seconds).
+OTHER_PROCESSES_REFRESH = 1.0
+
+if torch.version.hip is not None:
+    # Without it, PyTorch SDPA has no flash or memory-efficient kernel on the GPUs AOTriton still
+    # marks experimental and falls back to the math path. Read at the first SDPA call.
+    os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
+
+_adlx = None
+_adlx_lock = threading.Lock()
+_other_processes = {}
+_missing_adlx_tips_printed = False
+
+
+def _adlx_session():
+    global _adlx, _missing_adlx_tips_printed
+    with _adlx_lock:
+        if _adlx is None:
+            _adlx = {}
+            try:
+                from adlx import ADLX
+            except ImportError:
+                if not _missing_adlx_tips_printed:
+                    print("On ROCm for Windows, torch only reports the VRAM used by this process. Install `amd-adlx` so that `vram_limit` also accounts for other programs and the desktop.")
+                    _missing_adlx_tips_printed = True
+                return None
+            try:
+                helper = ADLX.ADLXHelper()
+                if helper.Initialize() == ADLX.ADLX_RESULT.ADLX_OK:
+                    system = helper.GetSystemServices()
+                    # UniqueId packs the PCI location as (bus << 8) | (device << 3) | function
+                    gpus = {(gpu.UniqueId() >> 8) & 0xFF: gpu for gpu in system.GetGPUs()}
+                    _adlx.update(helper=helper, perf=system.GetPerformanceMonitoringServices(), gpus=gpus)
+                    atexit.register(_adlx_close)
+            except Exception as e:
+                print(f"ADLX is unavailable ({e}), falling back to the VRAM torch reports for this process.")
+        return _adlx or None
+
+
+def _adlx_close():
+    # ADLX crashes at interpreter exit unless every interface is released before Terminate.
+    with _adlx_lock:
+        helper = _adlx.pop("helper")
+        _adlx.clear()
+        helper.Terminate()
+
+
+def _adlx_used_bytes(bus):
+    session = _adlx_session()
+    if session is None or bus not in session["gpus"]:
+        return None
+    with _adlx_lock:
+        metrics = session["perf"].GetCurrentGPUMetrics(session["gpus"][bus])
+        used = metrics.GPUVRAM()
+        del metrics
+    return used * 1024 ** 2
+
+
+def mem_get_info(device=None):
+    """`torch.cuda.mem_get_info`, but for the whole GPU on ROCm for Windows.
+
+    There torch only counts its own process, so the desktop and other programs look like free VRAM.
+    The VRAM they use comes from ADLX instead.
+    """
+    free, total = torch.cuda.mem_get_info(device)
+    if not IS_WINDOWS_ROCM:
+        return free, total
+    used_by_this_process = total - free
+    bus = torch.cuda.get_device_properties(device).pci_bus_id
+    now = time.monotonic()
+    cached = _other_processes.get(bus)
+    if cached is None or now - cached[0] > OTHER_PROCESSES_REFRESH:
+        used = _adlx_used_bytes(bus)
+        if used is None:
+            return free, total
+        cached = _other_processes[bus] = (now, max(used - used_by_this_process, 0))
+    return max(total - used_by_this_process - cached[1], 0), total

diffsynth/core/vram/layers.py

@@ -2,7 +2,7 @@
 from typing import Union
 from .initialization import skip_model_initialization
 from .disk_map import DiskMap
-from ..device import parse_device_type, get_device_name, IS_NPU_AVAILABLE
+from ..device import parse_device_type, get_device_name, IS_NPU_AVAILABLE, mem_get_info
 
 
 class AutoTorchModule(torch.nn.Module):
@@ -67,6 +67,8 @@ def check_free_vram(self):
             allocated = torch.mps.current_allocated_memory()
             recommended_max = torch.mps.recommended_max_memory()
             gpu_mem_state = (recommended_max - allocated, recommended_max)
+        elif self.computation_device_type == "cuda":
+            gpu_mem_state = mem_get_info(self.computation_device)
         else:
             device = self.computation_device if not IS_NPU_AVAILABLE else get_device_name()
             gpu_mem_state = getattr(torch, self.computation_device_type).mem_get_info(device)

docs/en/Pipeline_Usage/GPU_support.md

@@ -12,6 +12,8 @@ All sample code provided by this project supports NVIDIA GPUs by default, requir
 
 AMD provides PyTorch packages based on ROCm, so most models can run without code changes. A small number of models may not be compatible due to their reliance on CUDA-specific instructions.
 
+On Windows, `torch.cuda.mem_get_info` with ROCm only counts the VRAM used by the current process. When `amd-adlx` is installed, `vram_limit` uses the VRAM used by the whole GPU instead, as it does on NVIDIA GPUs. On ROCm, `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` is set by default so that PyTorch SDPA can use its flash and memory-efficient kernels on GPUs where AOTriton support is still experimental.
+
 ### Apple Silicon
 
 On Apple Silicon devices, since VRAM and memory are unified, replace all `"cuda"` in the code with `"mps"` or `"cpu"`.

docs/en/Pipeline_Usage/Setup.md

@@ -20,6 +20,7 @@ To keep the framework lightweight, the base installation only installs the neces
 * `[quant]`: For parameter quantization, enabling precisions such as NF4, INT8, NVFP4.
 * `[training]`: For distributed large-scale pretraining.
 * `[logger]`: To enable training loggers such as TensorBoard, SwanLab, etc.
+* `[amd]`: For AMD GPUs on Windows, reading the VRAM used by the whole GPU through ADLX.
 * `[npu]`: For Ascend NPU devices with x86 architecture.
 * `[npu_aarch64]`: For Ascend NPU devices with aarch64/ARM architecture.
 * Dependencies of specific models
@@ -45,6 +46,13 @@ You need to install the `torch` package with ROCm support. Taking ROCm 6.4 (as o
 pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm6.4
 ```
 
+On Windows, install AMD's ROCm 10 `torch` packages instead (as of October 30, 2026), replacing `gfx1201` with your GPU's architecture, or with `all` if you are not sure. Then install the `[amd]` dependencies, which include `amd-adlx`, so that `vram_limit` accounts for the VRAM used by other programs and the desktop:
+
+```shell
+pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ "torch[device-gfx1201]==2.14.0+rocm10.1.0" "torchvision[device-gfx1201]==0.29.0a0+rocm10.1.0"
+pip install -e ".[amd]"
+```
+
 ### Apple Silicon
 
 On Apple Silicon devices, no changes to the installation steps are needed. However, since VRAM and memory are unified, replace all `"cuda"` in the code with `"mps"` or `"cpu"`.

docs/zh/Pipeline_Usage/GPU_support.md

@@ -12,6 +12,8 @@
 
 AMD 提供了基于 ROCm 的 torch 包，所以大多数模型无需修改代码即可运行，少数模型由于依赖特定的 cuda 指令无法运行。
 
+在 Windows 系统上，ROCm 版本的 `torch.cuda.mem_get_info` 只统计当前进程占用的显存。安装 `amd-adlx` 后，`vram_limit` 会改为使用整个 GPU 的显存占用，与 NVIDIA GPU 上的行为一致。在 ROCm 上，框架默认设置 `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1`，使 PyTorch SDPA 在 AOTriton 支持仍为实验性的 GPU 上也能使用 flash 和 memory-efficient 内核。
+
 ### Apple Silicon
 
 在 Apple Silicon 设备上，由于显存和内存是统一的，请将代码中的 `"cuda"` 全部修改为 `"mps"` 或 `"cpu"`。

docs/zh/Pipeline_Usage/Setup.md

@@ -20,6 +20,7 @@ pip install diffsynth
 * `[quant]`: 用于参数量化，启用 NF4、INT8、NVFP4 等精度
 * `[training]`: 用于分布式大规模预训练
 * `[logger]`: 用于启用 TensorBoard、SwanLab 等训练日志记录器
+* `[amd]`: 用于 Windows 系统上的 AMD GPU，通过 ADLX 读取整个 GPU 的显存占用
 * `[npu]`: 用于 x86 架构的 Ascend NPU 设备
 * `[npu_aarch64]`: 用于 aarch64/ARM 架构的 Ascend NPU 设备
 * 特定模型的依赖
@@ -45,6 +46,13 @@ pip install diffsynth
 pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm6.4
 ```
 
+在 Windows 系统上，请改为安装 AMD 提供的 ROCm 10 `torch` 包（本文更新于 2026 年 10 月 30 日），并将 `gfx1201` 替换为你的 GPU 架构，不确定时可使用 `all`。然后安装 `[amd]` 依赖（包含 `amd-adlx`），使 `vram_limit` 能够计入其他程序和桌面占用的显存：
+
+```shell
+pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ "torch[device-gfx1201]==2.14.0+rocm10.1.0" "torchvision[device-gfx1201]==0.29.0a0+rocm10.1.0"
+pip install -e ".[amd]"
+```
+
 ### Apple Silicon
 
 在 Apple Silicon 设备上，无需修改安装步骤。但由于显存和内存是统一的，请将代码中的 `"cuda"` 全部修改为 `"mps"` 或 `"cpu"`。

pyproject.toml

@@ -57,6 +57,9 @@ logger = [
     "swanlab",
     "wandb"
 ]
+amd = [
+    "amd-adlx; sys_platform == 'win32'"
+]
 npu = [
     "torch==2.7.1+cpu",
     "torch-npu==2.7.1",