mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-10-12 02:16:43 +00:00
feat(deps): install ROCm 10 from AMD's index on Linux and Windows
The rocm extra moves from PyTorch's 2.13.0+rocm7.2 (Linux only) to AMD's 2.13.0+rocm10.0.0 for Linux x86_64 and Windows. The lock carries the GPU kernel packages of every target AMD builds, so the launcher can leave out the ones a machine does not need. Linux loses gfx900/gfx906, which ROCm 10 dropped. bitsandbytes moves to 0.50.2, the first release with a Windows ROCm binary. A new setting, rocm_aotriton_experimental (auto/on/off), sets TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL before anything runs attention. auto turns AOTriton's experimental fused kernels on only for a ROCm 10 build on GPUs where they were measured correct (gfx1200 so far); under torch 2.12+rocm7.14 the same switch made every fused call fail on that GPU. Builds on the dependency and documentation changes of #9606. Co-authored-by: 0xDELUXA <[email protected]>
This commit is contained in:
Alexander Eichhorn
co-authored by
0xDELUXA
0xDELUXA
parent
6c3277771d
commit
7001bf3780
@@ -4,7 +4,7 @@ sidebar:
|
||||
order: 2
|
||||
---
|
||||
|
||||
As of v5.6.0, Invoke has a low-VRAM mode. It works on systems with dedicated GPUs (Nvidia GPUs on Windows/Linux and AMD GPUs on Linux).
|
||||
As of v5.6.0, Invoke has a low-VRAM mode. It works on systems with dedicated GPUs (Nvidia and AMD GPUs on Windows/Linux).
|
||||
|
||||
This allows you to generate even if your GPU doesn't have enough VRAM to hold full models. Most users should be able to run even the beefiest models - like the ~24GB unquantised FLUX dev model.
|
||||
|
||||
@@ -225,7 +225,7 @@ It is strongly suggested to disable this feature:
|
||||
|
||||
### AMD GPUs on Windows (ROCm)
|
||||
|
||||
Invoke does not officially support ROCm on Windows (see the [System Requirements](/start-here/system-requirements/)). If you run a ROCm build of PyTorch there anyway, Windows behaves as if the Nvidia sysmem fallback were always on: an allocation that does not fit is moved into much slower system memory instead of failing, and it stays there. Task Manager shows it as shared GPU memory. Invoke works around this where it can:
|
||||
With a ROCm build of PyTorch (see the [System Requirements](/start-here/system-requirements/)), Windows behaves as if the Nvidia sysmem fallback were always on: an allocation that does not fit is moved into much slower system memory instead of failing, and it stays there. Task Manager shows it as shared GPU memory. Invoke works around this where it can:
|
||||
|
||||
- The model cache plans against the video-memory budget Windows grants Invoke, not against the card's total VRAM.
|
||||
- The PyTorch allocator defaults to `expandable_segments:True` (see [PyTorch CUDA allocator config](#pytorch-cuda-allocator-config)), so fragmented VRAM does not force allocations out.
|
||||
|
||||
@@ -111,6 +111,13 @@ The following commands vary depending on the version of Invoke being installed a
|
||||
--torch-backend=xpu
|
||||
```
|
||||
</TabItem>
|
||||
<TabItem label="AMD GPU">
|
||||
Do not use a torch backend. AMD publishes its ROCm wheels on its own index, with the GPU kernels in a separate package per GPU target. In the next step, pin them for your GPU's target and point `uv` at that index instead:
|
||||
```sh
|
||||
uv pip install <PACKAGE_SPECIFIER>==<VERSION> "torch[device-<TARGET>]==2.13.0+rocm10.0.0" "torchvision[device-<TARGET>]==0.28.0+rocm10.0.0" --index https://stable.repo.amd.com/rocm/whl-next/ --index-strategy unsafe-best-match --python 3.12 --python-preference only-managed --force-reinstall
|
||||
```
|
||||
`<TARGET>` is `gfx1201` for the RX 9070 series and Radeon AI PRO R9700, `gfx1200` for the RX 9060 series, `gfx1100` for the RX 7900 series, `gfx1101` for the RX 7700 and 7800 series, `gfx1102` for the RX 7600 series, and `gfx1151` for Ryzen AI Max. For other GPUs, see AMD's [GPU target table](https://github.com/ROCm/TheRock/blob/main/RELEASES.md#gfx-target-lookup-table).
|
||||
</TabItem>
|
||||
<TabItem label="Other / no GPU">
|
||||
Do not use a torch backend.
|
||||
</TabItem>
|
||||
@@ -133,7 +140,13 @@ The following commands vary depending on the version of Invoke being installed a
|
||||
On ARM64 (`aarch64`, e.g. Raspberry Pi 5), do **not** use a torch backend. The default PyPI wheels are CPU-only on ARM64 and work out of the box.
|
||||
</TabItem>
|
||||
<TabItem label="Linux + AMD GPU">
|
||||
Use:
|
||||
Do not use a torch backend. Install AMD's ROCm 10 wheels from AMD's index, as the launcher does, with the GPU kernels pinned for your GPU's target:
|
||||
```sh
|
||||
uv pip install <PACKAGE_SPECIFIER>==<VERSION> "torch[device-<TARGET>]==2.13.0+rocm10.0.0" "torchvision[device-<TARGET>]==0.28.0+rocm10.0.0" --index https://stable.repo.amd.com/rocm/whl-next/ --index-strategy unsafe-best-match --python 3.12 --python-preference only-managed --force-reinstall
|
||||
```
|
||||
`<TARGET>` is `gfx1201` for the RX 9070 series and Radeon AI PRO R9700, `gfx1200` for the RX 9060 series, `gfx1100` for the RX 7900 series, `gfx1101` for the RX 7700 and 7800 series, `gfx1102` for the RX 7600 series, and `gfx1151` for Ryzen AI Max. For other GPUs, see AMD's [GPU target table](https://github.com/ROCm/TheRock/blob/main/RELEASES.md#gfx-target-lookup-table).
|
||||
|
||||
ROCm 10 has no kernels for Vega-based cards (gfx900/gfx906). For those, use PyTorch's ROCm 7.2 build instead:
|
||||
```sh
|
||||
--torch-backend=rocm7.2
|
||||
```
|
||||
|
||||
@@ -17,7 +17,7 @@ The requirements below are rough guidelines for best performance. GPUs with less
|
||||
|
||||
- All Apple Silicon (M1, M2, etc) Macs work, but 16GB+ memory is recommended.
|
||||
- Nvidia GPUs need compute capability 7.5 or newer (GTX 16xx, RTX 20xx and everything after) and a driver from the R580 series or newer. Maxwell, Pascal and Volta cards (GTX 9xx/10xx, Titan V, Tesla P40/P100/V100) are not supported by Invoke's CUDA build: stay on the previous Invoke release, or do a [manual install](../manual) with `--torch-backend=cu126`, whose PyTorch build still includes these GPUs (not tested by the Invoke team). With a driver that is too old, Invoke starts on the CPU and logs an error saying so; on an unsupported GPU it logs an error at startup and generation fails.
|
||||
- AMD GPUs are supported on Linux only. The VRAM requirements are the same as Nvidia GPUs.
|
||||
- AMD GPUs are supported on Windows and Linux `x86_64` through ROCm 10 (see [AMD](#amd) below for the GPUs and drivers). On Linux, Vega-based cards (gfx900/gfx906) are no longer supported. The VRAM requirements are the same as Nvidia GPUs.
|
||||
- Intel Arc GPUs (Alchemist, Battlemage and newer) are supported on Windows and Linux `x86_64`. The VRAM requirements are the same as Nvidia GPUs.
|
||||
- Linux ARM64 (`aarch64`) devices — e.g. Raspberry Pi 5, other SBCs, ARM servers — are supported in CPU-only mode. Local generation is slow without a GPU, but API-backed models (e.g. GPT Image, Gemini) work well.
|
||||
|
||||
@@ -145,8 +145,15 @@ If, after restarting the app, this doesn't improve your performance, either rest
|
||||
|
||||
### AMD
|
||||
|
||||
:::tip[Linux Only]{icon="linux"}
|
||||
AMD GPUs are supported on Linux only, due to ROCm (the AMD equivalent of CUDA) support being Linux only.
|
||||
Invoke installs AMD's ROCm 10 build of PyTorch on Windows and Linux. It brings the ROCm runtime with it, so
|
||||
there is no separate ROCm or HIP SDK to install, but the GPU needs an up-to-date driver: [AMD Software: Adrenalin
|
||||
Edition] on Windows, AMD's kernel driver on Linux (see [Linux](#linux) below). AMD's [ROCm compatibility matrix] lists the GPUs and operating
|
||||
systems ROCm 10 supports. AMD lists Windows 11; the Invoke team has run it on Windows 10 with an RX 9060 XT.
|
||||
|
||||
:::caution[Vega-based cards on Linux]
|
||||
ROCm 10 has no kernels for gfx900 and gfx906 (Radeon RX Vega 56/64, Radeon VII, Instinct MI25/MI50/MI60).
|
||||
Stay on the previous Invoke release, or do a [manual install](../manual) with `--torch-backend=rocm7.2`,
|
||||
whose PyTorch build still includes them (not tested by the Invoke team).
|
||||
:::
|
||||
|
||||
:::caution[Bumps Ahead]
|
||||
@@ -160,11 +167,22 @@ If, after restarting the app, this doesn't improve your performance, either rest
|
||||
(for example `MIOPEN_FIND_MODE=NORMAL` to let MIOpen re-tune its kernel database).
|
||||
:::
|
||||
|
||||
:::note[Fused attention on newer Radeon GPUs]
|
||||
On some GPUs - among them the RX 7600/7700/7800 series, Ryzen AI 300/Max and the RX 9060 series -
|
||||
PyTorch uses its fused attention kernels only when `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL` is set.
|
||||
Without them, attention runs on the slower math kernel, which also needs far more memory. The
|
||||
`rocm_aotriton_experimental` setting decides. Its default, `auto`, turns the kernels on where they
|
||||
were measured correct with ROCm 10: so far the RX 9060 series (gfx1200), where Z-Image at 1024px
|
||||
takes 17 s instead of 36 s. Set it to `on` to try them on another GPU, and back to `auto` if images
|
||||
come out wrong or generation fails. The startup log says whether they are on.
|
||||
:::
|
||||
|
||||
:::note[Wide-head attention runs on the math kernel]
|
||||
On ROCm builds Invoke routes attention whose head dimension exceeds 256 to PyTorch's math kernel.
|
||||
The fused kernels return wrong output there on torch 2.13 / ROCm 7.2 (measured on a W7900), which
|
||||
showed up as black or noisy images from the image VAEs - their mid-block attention is one 512-wide
|
||||
head. The math kernel builds the whole attention matrix (about 4 GB for a 1024px VAE decode, 21 GB
|
||||
head. With ROCm 10 they were correct on an RX 9060 XT, but RDNA3 cards have not been measured
|
||||
again, so the threshold stays. The math kernel builds the whole attention matrix (about 4 GB for a 1024px VAE decode, 21 GB
|
||||
at 1536px), so Invoke computes it in pieces of at most 1 GB. The result can differ from a single
|
||||
math call in the last bit of a bf16 value, and each such call is roughly 10% slower. Export
|
||||
`INVOKE_ROCM_FUSED_SDPA_MAX_HEAD_DIM` to raise the threshold on a build whose fused kernels are
|
||||
@@ -187,11 +205,11 @@ If, after restarting the app, this doesn't improve your performance, either rest
|
||||
`rocprof`/ROCTracer-based kernel profiling, which this disables.
|
||||
:::
|
||||
|
||||
Run `rocm-smi` on your system's command line verify that drivers and ROCm are installed. If this command fails, or doesn't report versions, you will need to install them.
|
||||
#### Linux
|
||||
|
||||
Go to the [ROCm Documentation] and carefully follow the instructions for your system to get everything installed.
|
||||
|
||||
Confirm that `rocm-smi` displays driver and CUDA versions after installation.
|
||||
The ROCm libraries come with Invoke's PyTorch packages, so ROCm itself does not need to be installed. The GPU needs
|
||||
AMD's `amdgpu` kernel driver; on the distributions in AMD's [ROCm compatibility matrix], the driver that comes with
|
||||
the kernel is enough. For other distributions, follow the [ROCm Documentation] to install the driver.
|
||||
|
||||
#### Linux - via Docker Container
|
||||
|
||||
@@ -209,6 +227,8 @@ An alternative to installing ROCm locally is to use a [ROCm docker container] to
|
||||
[Intel Arc & Iris Xe Graphics driver]: https://www.intel.com/content/www/us/en/download/785597/intel-arc-iris-xe-graphics-windows.html
|
||||
[Intel client GPU install guide]: https://dgpu-docs.intel.com/driver/client/overview.html
|
||||
[ROCm Documentation]: https://rocmdocs.amd.com
|
||||
[ROCm compatibility matrix]: https://rocm.docs.amd.com/en/docs-10.0.0/compatibility/compatibility-matrix.html
|
||||
[AMD Software: Adrenalin Edition]: https://www.amd.com/en/support/download/drivers.html
|
||||
[ROCm docker container]: https://rocmdocs.amd.com/en/latest/Deep_learning/Deep_learning.html#docker-containers
|
||||
[ROCm/ROCm#6522]: https://github.com/ROCm/ROCm/issues/6522
|
||||
[ROCm/TheRock#7051]: https://github.com/ROCm/TheRock/issues/7051
|
||||
|
||||
@@ -641,6 +641,21 @@
|
||||
"type": "typing.Literal['auto', 'float16', 'bfloat16', 'float32']",
|
||||
"validation": {}
|
||||
},
|
||||
{
|
||||
"category": "DEVICE",
|
||||
"default": "auto",
|
||||
"description": "Use AOTriton's fused (flash and memory-efficient) attention kernels on AMD GPUs that PyTorch marks experimental for them, such as the RX 7600/7700/7800 series, Ryzen AI 300/Max and the RX 9060 series, by setting TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL. Without them, attention on those GPUs runs on the slower math kernel, which also needs far more memory. `auto` turns them on only with a ROCm 10 build of PyTorch, and only when every GPU used for generation is one they were measured correct on (so far gfx1200, the RX 9060 series); `on` and `off` decide for any build and GPU. A TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL value already set in the environment takes precedence. Has no effect on other GPUs.",
|
||||
"env_var": "INVOKEAI_ROCM_AOTRITON_EXPERIMENTAL",
|
||||
"literal_values": [
|
||||
"auto",
|
||||
"on",
|
||||
"off"
|
||||
],
|
||||
"name": "rocm_aotriton_experimental",
|
||||
"required": false,
|
||||
"type": "typing.Literal['auto', 'on', 'off']",
|
||||
"validation": {}
|
||||
},
|
||||
{
|
||||
"category": "GENERATION",
|
||||
"default": false,
|
||||
|
||||
@@ -76,6 +76,13 @@ def run_app() -> None:
|
||||
else:
|
||||
apply_rocm_windows_allocator_default(logger)
|
||||
|
||||
# Decide on AOTriton's experimental attention kernels before anything can run attention: torch reads the switch
|
||||
# once, at its first fused-kernel check, and the server already runs attention while starting up. This imports
|
||||
# torch, so it comes after the allocator configuration.
|
||||
from invokeai.app.util.rocm_aotriton import apply_rocm_aotriton_setting
|
||||
|
||||
apply_rocm_aotriton_setting(app_config.rocm_aotriton_experimental, app_config.generation_devices, logger)
|
||||
|
||||
# This import must happen after configure_torch_cuda_allocator() is called, because the module imports torch.
|
||||
from invokeai.app.invocations.baseinvocation import InvocationRegistry
|
||||
from invokeai.app.invocations.load_custom_nodes import load_custom_nodes
|
||||
|
||||
@@ -29,6 +29,7 @@ LEGACY_INIT_FILE = Path("invokeai.init")
|
||||
PRECISION = Literal["auto", "float16", "bfloat16", "float32"]
|
||||
ATTENTION_TYPE = Literal["auto", "normal", "xformers", "sliced", "torch-sdp"]
|
||||
NOISE_DTYPE = Literal["float32", "float16"]
|
||||
ROCM_AOTRITON = Literal["auto", "on", "off"]
|
||||
ATTENTION_SLICE_SIZE = Literal["auto", "balanced", "max", 1, 2, 3, 4, 5, 6, 7, 8]
|
||||
LOG_FORMAT = Literal["plain", "color", "syslog", "legacy"]
|
||||
LOG_LEVEL = Literal["debug", "info", "warning", "error", "critical"]
|
||||
@@ -120,6 +121,7 @@ class InvokeAIAppConfig(BaseSettings):
|
||||
pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, "expandable_segments:True" is used.
|
||||
device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)
|
||||
precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.<br>Valid values: `auto`, `float16`, `bfloat16`, `float32`
|
||||
rocm_aotriton_experimental: Use AOTriton's fused (flash and memory-efficient) attention kernels on AMD GPUs that PyTorch marks experimental for them, such as the RX 7600/7700/7800 series, Ryzen AI 300/Max and the RX 9060 series, by setting TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL. Without them, attention on those GPUs runs on the slower math kernel, which also needs far more memory. `auto` turns them on only with a ROCm 10 build of PyTorch, and only when every GPU used for generation is one they were measured correct on (so far gfx1200, the RX 9060 series); `on` and `off` decide for any build and GPU. A TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL value already set in the environment takes precedence. Has no effect on other GPUs.<br>Valid values: `auto`, `on`, `off`
|
||||
sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.
|
||||
noise_dtype: The dtype seeded noise is drawn in for SD1.5/SDXL, SD3, FLUX.1, FLUX.2, CogView4 and the Z-Image seed variance enhancer. `float32` draws the same noise on every platform. `float16` is the half-precision draw of earlier versions: on macOS it keeps the images your seeds gave before this update; on Windows and Linux both values give (nearly) the same images, and images made before the update cannot be reproduced there.<br>Valid values: `float32`, `float16`
|
||||
wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.
|
||||
@@ -244,6 +246,7 @@ class InvokeAIAppConfig(BaseSettings):
|
||||
generation_devices: Union[Literal["auto"], list[str]] = Field(default="auto", description="Devices to use for parallel generation. `auto` (the default) uses every available GPU, running one generation session per GPU concurrently and distributing jobs fairly across users — unless the legacy `device` setting is pinned to a specific device, in which case `auto` uses only that device (preserving configs that pinned `device` before multi-GPU support existed). Provide an explicit list (e.g. `[cuda:0, cuda:1]`) to use specific devices regardless of `device`, or a single-device list (e.g. `[cuda:0]`) to run serially. On systems without a GPU, `auto` resolves to the single `cpu`/`mps` device.<br>Valid values: `auto`, or a list whose entries are each `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, or `xpu:N` (where N is a device number)")
|
||||
offload_text_encoders_to_idle_gpus: bool = Field(default=True, description="When running on multiple GPUs, load text encoders onto a currently-idle GPU instead of the one running the denoise pipeline. This avoids churning the denoise model in and out of VRAM to make room for the encoder, and lets a cached encoder be reused across generations. Has no effect unless at least two `generation_devices` are configured and a GPU is idle; under full load encoders run on the session's own GPU as before.")
|
||||
precision: PRECISION = Field(default="auto", description="Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.")
|
||||
rocm_aotriton_experimental: ROCM_AOTRITON = Field(default="auto", description="Use AOTriton's fused (flash and memory-efficient) attention kernels on AMD GPUs that PyTorch marks experimental for them, such as the RX 7600/7700/7800 series, Ryzen AI 300/Max and the RX 9060 series, by setting TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL. Without them, attention on those GPUs runs on the slower math kernel, which also needs far more memory. `auto` turns them on only with a ROCm 10 build of PyTorch, and only when every GPU used for generation is one they were measured correct on (so far gfx1200, the RX 9060 series); `on` and `off` decide for any build and GPU. A TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL value already set in the environment takes precedence. Has no effect on other GPUs.")
|
||||
|
||||
# GENERATION
|
||||
sequential_guidance: bool = Field(default=False, description="Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.")
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
"""The `rocm_aotriton_experimental` setting: whether AOTriton's experimental fused SDPA kernels run on AMD GPUs.
|
||||
|
||||
PyTorch's ROCm builds ship flash and memory-efficient attention through AOTriton, which marks some GPU architectures
|
||||
"experimental" (gfx1101, gfx1102, gfx1150, gfx1151 and gfx1200 among them). On those, torch only uses the fused
|
||||
kernels when `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL` is set; otherwise every attention call falls back to the math
|
||||
kernel, which is slower and materializes the whole score matrix.
|
||||
|
||||
Whether the kernels are correct depends on the build as much as on the GPU. On an RX 9060 XT (gfx1200) under torch
|
||||
2.12+rocm7.14.1, setting the variable made every flash and memory-efficient call fail with `hipErrorInvalidValue`, and
|
||||
the CLIP text encoder hung on exit. Under torch 2.13.0+rocm10.0.0 every fused case matched the math reference, and
|
||||
Z-Image at 1024px denoised in 16.7 s instead of 36 s. So `auto` turns the kernels on only for a ROCm 10 build on
|
||||
architectures they were measured on, and `on`/`off` decide for any build and GPU.
|
||||
|
||||
torch reads the variable once per process, at the first fused-kernel check on an experimental architecture. The
|
||||
server runs attention during startup, so the setting is applied before the app is imported.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Iterable
|
||||
from typing import Optional, Union
|
||||
|
||||
from invokeai.app.services.config.config_default import ROCM_AOTRITON
|
||||
|
||||
AOTRITON_EXPERIMENTAL_ENV = "TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL"
|
||||
|
||||
# Architectures on which the experimental kernels were measured correct under `MIN_AUTO_ROCM_MAJOR`, with the
|
||||
# measurement: gfx1200 (RX 9060 XT), Windows, torch 2.13.0+rocm10.0.0 -- CLIP, Z-Image, FLUX and VAE attention shapes
|
||||
# all within the math kernel's own error, Z-Image 1024px 16.7 s instead of 36 s. Add an architecture only with a
|
||||
# measurement of its own.
|
||||
AUTO_MEASURED_ARCHS = frozenset({"gfx1200"})
|
||||
MIN_AUTO_ROCM_MAJOR = 10
|
||||
|
||||
# The values `c10::utils::check_env` reads as true.
|
||||
_TRUTHY = {"1", "on", "yes", "true", "y"}
|
||||
|
||||
|
||||
def rocm_major_version(torch_version: str) -> Optional[int]:
|
||||
"""The ROCm major version in a torch version's local label (`2.13.0+rocm10.0.0` -> 10), or None without one.
|
||||
|
||||
The label rather than `torch.version.hip`: AMD's ROCm 10 wheels report HIP 7.15. A locally built torch has no
|
||||
label, so `auto` leaves it alone.
|
||||
"""
|
||||
match = re.search(r"\+rocm(\d+)", torch_version)
|
||||
return int(match.group(1)) if match else None
|
||||
|
||||
|
||||
def resolve_aotriton_experimental(
|
||||
setting: ROCM_AOTRITON,
|
||||
exported: Optional[str],
|
||||
rocm_major: Optional[int],
|
||||
archs: Iterable[str],
|
||||
) -> tuple[Optional[str], str]:
|
||||
"""Decide what `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL` should be.
|
||||
|
||||
Returns the value to export (None: leave the environment as it is) and the reason, for the log.
|
||||
|
||||
Args:
|
||||
setting: The `rocm_aotriton_experimental` setting.
|
||||
exported: The variable's value already in the environment, if any. It always wins.
|
||||
rocm_major: The ROCm major version of the torch build (see `rocm_major_version`).
|
||||
archs: The architectures of the GPUs generation runs on, e.g. `{"gfx1200"}`.
|
||||
"""
|
||||
if exported is not None:
|
||||
return None, f"{AOTRITON_EXPERIMENTAL_ENV}={exported} is set in the environment"
|
||||
if setting == "on":
|
||||
return "1", "rocm_aotriton_experimental is 'on'"
|
||||
if setting == "off":
|
||||
return "0", "rocm_aotriton_experimental is 'off'"
|
||||
archs = sorted(set(archs))
|
||||
if rocm_major is None or rocm_major < MIN_AUTO_ROCM_MAJOR:
|
||||
build = "a ROCm build without a version label" if rocm_major is None else f"ROCm {rocm_major}"
|
||||
return None, f"auto: not measured on {build}"
|
||||
if not archs:
|
||||
return None, "auto: no ROCm GPU to generate on"
|
||||
unmeasured = [arch for arch in archs if arch not in AUTO_MEASURED_ARCHS]
|
||||
if unmeasured:
|
||||
return None, f"auto: not measured on {', '.join(unmeasured)}"
|
||||
return "1", f"auto: measured on {', '.join(archs)} with ROCm {rocm_major}"
|
||||
|
||||
|
||||
def _arch_name(gcn_arch_name: str) -> str:
|
||||
"""`gfx90a:sramecc+:xnack-` -> `gfx90a`."""
|
||||
return gcn_arch_name.split(":", 1)[0]
|
||||
|
||||
|
||||
def apply_rocm_aotriton_setting(
|
||||
setting: ROCM_AOTRITON,
|
||||
generation_devices: Union[str, list[str]],
|
||||
logger: logging.Logger,
|
||||
) -> None:
|
||||
"""Export `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL` as the setting decides, and log the effective value.
|
||||
|
||||
Must run after the allocator is configured (this imports torch) and before anything runs attention. A no-op on
|
||||
builds other than ROCm, apart from a warning for an explicit `on`.
|
||||
"""
|
||||
import torch
|
||||
|
||||
if torch.version.hip is None:
|
||||
if setting == "on":
|
||||
logger.warning("rocm_aotriton_experimental is 'on', but this PyTorch is not a ROCm build; ignoring it.")
|
||||
return
|
||||
|
||||
from invokeai.backend.util.devices import TorchDevice
|
||||
|
||||
archs: set[str] = set()
|
||||
try:
|
||||
devices = TorchDevice.get_generation_devices(generation_devices) or [TorchDevice.choose_torch_device()]
|
||||
for device in devices:
|
||||
if device.type == "cuda":
|
||||
archs.add(_arch_name(torch.cuda.get_device_properties(device).gcnArchName))
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not read the GPU architecture for rocm_aotriton_experimental: {e}")
|
||||
archs.clear()
|
||||
|
||||
exported = os.environ.get(AOTRITON_EXPERIMENTAL_ENV)
|
||||
value, reason = resolve_aotriton_experimental(setting, exported, rocm_major_version(torch.__version__), archs)
|
||||
if exported is not None and setting != "auto" and (exported.strip().lower() in _TRUTHY) != (setting == "on"):
|
||||
logger.warning(
|
||||
f"rocm_aotriton_experimental is '{setting}', but {AOTRITON_EXPERIMENTAL_ENV}={exported} is set in the "
|
||||
"environment and takes precedence."
|
||||
)
|
||||
if value is not None:
|
||||
os.environ[AOTRITON_EXPERIMENTAL_ENV] = value
|
||||
|
||||
effective = os.environ.get(AOTRITON_EXPERIMENTAL_ENV)
|
||||
enabled = effective is not None and effective.strip().lower() in _TRUTHY
|
||||
gpus = ", ".join(sorted(archs)) or "unknown GPU"
|
||||
if enabled:
|
||||
logger.info(f"ROCm ({gpus}): AOTriton's experimental fused attention kernels are on ({reason}).")
|
||||
else:
|
||||
logger.info(
|
||||
f"ROCm ({gpus}): AOTriton's experimental fused attention kernels are off ({reason}). On a GPU that "
|
||||
"AOTriton marks experimental, attention then runs on the slower math kernel; set "
|
||||
"rocm_aotriton_experimental: on to try the fused kernels."
|
||||
)
|
||||
File diff suppressed because one or more lines are too long
@@ -22204,6 +22204,7 @@ export type components = {
|
||||
* pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, "expandable_segments:True" is used.
|
||||
* device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)
|
||||
* precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.<br>Valid values: `auto`, `float16`, `bfloat16`, `float32`
|
||||
* rocm_aotriton_experimental: Use AOTriton's fused (flash and memory-efficient) attention kernels on AMD GPUs that PyTorch marks experimental for them, such as the RX 7600/7700/7800 series, Ryzen AI 300/Max and the RX 9060 series, by setting TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL. Without them, attention on those GPUs runs on the slower math kernel, which also needs far more memory. `auto` turns them on only with a ROCm 10 build of PyTorch, and only when every GPU used for generation is one they were measured correct on (so far gfx1200, the RX 9060 series); `on` and `off` decide for any build and GPU. A TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL value already set in the environment takes precedence. Has no effect on other GPUs.<br>Valid values: `auto`, `on`, `off`
|
||||
* sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.
|
||||
* noise_dtype: The dtype seeded noise is drawn in for SD1.5/SDXL, SD3, FLUX.1, FLUX.2, CogView4 and the Z-Image seed variance enhancer. `float32` draws the same noise on every platform. `float16` is the half-precision draw of earlier versions: on macOS it keeps the images your seeds gave before this update; on Windows and Linux both values give (nearly) the same images, and images made before the update cannot be reproduced there.<br>Valid values: `float32`, `float16`
|
||||
* wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.
|
||||
@@ -22602,6 +22603,13 @@ export type components = {
|
||||
* @enum {string}
|
||||
*/
|
||||
precision?: "auto" | "float16" | "bfloat16" | "float32";
|
||||
/**
|
||||
* Rocm Aotriton Experimental
|
||||
* @description Use AOTriton's fused (flash and memory-efficient) attention kernels on AMD GPUs that PyTorch marks experimental for them, such as the RX 7600/7700/7800 series, Ryzen AI 300/Max and the RX 9060 series, by setting TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL. Without them, attention on those GPUs runs on the slower math kernel, which also needs far more memory. `auto` turns them on only with a ROCm 10 build of PyTorch, and only when every GPU used for generation is one they were measured correct on (so far gfx1200, the RX 9060 series); `on` and `off` decide for any build and GPU. A TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL value already set in the environment takes precedence. Has no effect on other GPUs.
|
||||
* @default auto
|
||||
* @enum {string}
|
||||
*/
|
||||
rocm_aotriton_experimental?: "auto" | "on" | "off";
|
||||
/**
|
||||
* Sequential Guidance
|
||||
* @description Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
},
|
||||
"linux": {
|
||||
"cpu": "https://download.pytorch.org/whl/cpu",
|
||||
"rocm": "https://download.pytorch.org/whl/rocm7.2",
|
||||
"rocm": "https://stable.repo.amd.com/rocm/whl-next/",
|
||||
"cuda": "https://download.pytorch.org/whl/cu130",
|
||||
"xpu": "https://download.pytorch.org/whl/xpu"
|
||||
},
|
||||
|
||||
+226
-23
@@ -60,8 +60,7 @@ dependencies = [
|
||||
# `--torch-backend=<x>` with no extra (see docs/start-here/manual), so a blanket cap would also
|
||||
# reject torch>=2.12 on Windows/Linux CUDA, CPU and ARM64, where 2.12 has no known problem.
|
||||
# The ROCm path is constrained where it can actually be targeted: the `rocm` extra below pins an
|
||||
# exact version — 2.13.0+rocm7.2, validated end-to-end on gfx1100 (W7900), where its bundled
|
||||
# AOTriton 0.12b also triples flash-attention throughput vs the 2.10.0+rocm7.1 pin.
|
||||
# exact version, AMD's 2.13.0+rocm10.0.0.
|
||||
"torch>=2.7.0,<3.0; sys_platform != 'darwin'",
|
||||
"torch>=2.7.0,<2.8.0; sys_platform == 'darwin'",
|
||||
"torchsde", # diffusers needs this for SDE solvers, but it is not an explicit dep of diffusers
|
||||
@@ -125,9 +124,9 @@ dependencies = [
|
||||
# torch 2.4+cu carries its own triton dependency
|
||||
]
|
||||
|
||||
# The +cpu/+cu130/+rocm7.2 pins live on PyTorch's WHL indexes and are gated to the
|
||||
# platforms the project resolves and tests them on: Windows and Linux x86_64 (rocm: Linux
|
||||
# only). macOS is capped below torch 2.8 (see the base requirement), and linux_aarch64 stays
|
||||
# The +cpu/+cu130/+xpu pins live on PyTorch's WHL indexes, the +rocm10.0.0 pins on AMD's, and
|
||||
# all are gated to the platforms the project resolves and tests them on: Windows and Linux
|
||||
# x86_64. macOS is capped below torch 2.8 (see the base requirement), and linux_aarch64 stays
|
||||
# on the CPU-only PyPI 2.7.1 it was set up with (an aarch64 install requesting `cuda` gets
|
||||
# CPU torch), so both fall back to the base `torch` / `torchvision` from PyPI. The index has
|
||||
# carried aarch64 torchvision since 0.28, so moving aarch64 onto the pinned pairs -- and real
|
||||
@@ -153,17 +152,107 @@ dependencies = [
|
||||
"torch==2.7.1; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')",
|
||||
"torchvision==0.22.1; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')",
|
||||
]
|
||||
# ROCm 10 for Linux x86_64 and Windows, from AMD's index (PyTorch publishes no ROCm wheels for Windows, and none
|
||||
# built against ROCm 10). AMD splits the build: torch/torchvision and the ROCm runtime (`rocm-sdk-core`,
|
||||
# `rocm-sdk-libraries`) are shared, the GPU kernels come in one package per target and library. uv applies an
|
||||
# index source to direct requirements only, so every package that exists only on AMD's index is listed here.
|
||||
# The extra installs the kernels of every target AMD builds; the launcher leaves out the targets the machine
|
||||
# does not have (`uv sync --no-install-package`). ROCm 10 dropped gfx900/gfx906 (Vega 10/20, MI25/MI50), which
|
||||
# the previous 2.13.0+rocm7.2 pin from PyTorch's index still covered.
|
||||
"rocm" = [
|
||||
# ROCm wheels are x86_64-linux-only; gate by platform so macOS/win/linux-aarch64 resolution is unaffected.
|
||||
# 2.13.0+rocm7.2 bundles AOTriton 0.12b, whose regenerated gfx1100 tuning database triples
|
||||
# flash-attention throughput on RDNA3 vs the previous 2.10.0+rocm7.1 pin (measured 17.9 ->
|
||||
# 57.3 TFLOPS effective at the MiniMax H3 attention shape on a W7900).
|
||||
"torch==2.13.0+rocm7.2; sys_platform == 'linux' and platform_machine != 'aarch64'",
|
||||
"torchvision==0.28.0+rocm7.2; sys_platform == 'linux' and platform_machine != 'aarch64'",
|
||||
"triton-rocm==3.7.1; sys_platform == 'linux' and platform_machine != 'aarch64'",
|
||||
# linux_aarch64 fallback, see note under "cpu"
|
||||
"torch==2.7.1; sys_platform == 'linux' and platform_machine == 'aarch64'",
|
||||
"torchvision==0.22.1; sys_platform == 'linux' and platform_machine == 'aarch64'",
|
||||
"torch==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"torchvision==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
# The Linux torch wheel requires AMD's triton build; the Windows wheel requires none.
|
||||
"triton==3.8.0+git4cff872c.rocm10.0.0; sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
"rocm==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-core==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-libraries==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
# The first release with a Windows ROCm binary.
|
||||
"bitsandbytes>=0.50.2; sys_platform == 'win32'",
|
||||
# GPU kernels: per target, torch's, torchvision's and the ROCm runtime's; gfx110x, gfx115x and gfx12-0
|
||||
# are torch's kernels shared by a whole family.
|
||||
"amd-torch-device-gfx1010==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1010==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1010==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1011==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1011==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1011==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1012==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1012==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1012==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1030==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1030==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1030==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1031==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1031==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1031==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1032==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1032==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1032==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1033==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1033==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1033==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1034==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1034==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1034==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1035==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1035==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1035==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1036==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1036==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1036==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx110x==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1100==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1100==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1100==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1101==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1101==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1101==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1102==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1102==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1102==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1103==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1103==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1103==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx115x==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1150==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1150==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1150==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1151==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1151==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1151==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1152==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1152==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1152==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1153==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1153==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1153==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx12-0==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1200==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1200==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1200==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx1201==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx1201==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx1201==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx908==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx908==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx908==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torch-device-gfx90a==2.13.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"amd-torchvision-device-gfx90a==0.28.0+rocm10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"rocm-sdk-device-gfx90a==10.0.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
# Targets AMD builds for Linux only.
|
||||
"amd-torch-device-gfx1250==2.13.0+rocm10.0.0; sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
"amd-torchvision-device-gfx1250==0.28.0+rocm10.0.0; sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
"rocm-sdk-device-gfx1250==10.0.0; sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
"amd-torch-device-gfx942==2.13.0+rocm10.0.0; sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
"amd-torchvision-device-gfx942==0.28.0+rocm10.0.0; sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
"rocm-sdk-device-gfx942==10.0.0; sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
"amd-torch-device-gfx950==2.13.0+rocm10.0.0; sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
"amd-torchvision-device-gfx950==0.28.0+rocm10.0.0; sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
"rocm-sdk-device-gfx950==10.0.0; sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
# macOS / linux_aarch64 fallback, see note under "cpu"
|
||||
"torch==2.7.1; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')",
|
||||
"torchvision==0.22.1; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')",
|
||||
]
|
||||
# Intel XPU (Arc / Battlemage) wheels are published only for linux-x86_64 and
|
||||
# windows-amd64 on PyTorch's WHL index. Gate the +xpu pins to those platforms; the
|
||||
@@ -220,7 +309,7 @@ environments = ["sys_platform == 'win32' or sys_platform == 'darwin' or (sys_pla
|
||||
# there, or move it to the newest PyPI release, independently of the extras. `tool.uv` never reaches the
|
||||
# published metadata, so installing the released package with `--torch-backend` keeps the open range of the base
|
||||
# requirement; only `uv pip install` inside a checkout is held to it.
|
||||
# `==2.13.0` also matches the local builds (+cpu, +cu130, +rocm7.2, +xpu).
|
||||
# `==2.13.0` also matches the local builds (+cpu, +cu130, +rocm10.0.0, +xpu).
|
||||
constraint-dependencies = [
|
||||
"torch==2.13.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
"torchvision==0.28.0; sys_platform == 'win32' or (sys_platform == 'linux' and platform_machine == 'x86_64')",
|
||||
@@ -238,26 +327,139 @@ constraint-dependencies = [
|
||||
# lockfile regenerated by a uv too old to honor this.
|
||||
exclude-dependencies = ["notebook"]
|
||||
|
||||
# The Linux and Windows wheels of this build share a version but not their requirements: only the Linux wheel
|
||||
# requires (AMD's) triton, without a marker. uv reads one wheel's metadata for both, so restate the union here.
|
||||
[[tool.uv.dependency-metadata]]
|
||||
name = "torch"
|
||||
version = "2.13.0+rocm10.0.0"
|
||||
requires-dist = [
|
||||
"filelock",
|
||||
"typing-extensions>=4.10.0",
|
||||
"setuptools>=77.0.3",
|
||||
"sympy>=1.13.3",
|
||||
"networkx>=2.5.1",
|
||||
"jinja2",
|
||||
"fsspec>=0.8.5",
|
||||
"rocm[libraries]==10.0.0",
|
||||
"rocm-bootstrap",
|
||||
"triton==3.8.0+git4cff872c.rocm10.0.0; sys_platform == 'linux'",
|
||||
]
|
||||
|
||||
# `rocm` is published only as an sdist whose build asks the machine's GPU which kernel package to require, so the
|
||||
# lock would depend on the machine that wrote it. Its target-independent requirements, from rocm_sdk._dist_info;
|
||||
# the kernel packages are listed in the rocm extra instead.
|
||||
[[tool.uv.dependency-metadata]]
|
||||
name = "rocm"
|
||||
version = "10.0.0"
|
||||
requires-dist = [
|
||||
"rocm-sdk-core==10.0.0",
|
||||
"rocm-sdk-libraries==10.0.0; extra == 'libraries'",
|
||||
]
|
||||
provides-extras = ["libraries"]
|
||||
|
||||
[tool.uv.sources]
|
||||
# Keep linux_aarch64 (and, for cpu/cuda, macOS) off the PyTorch WHL indexes so that uv
|
||||
# resolves their 2.7.1 fallback pins from PyPI; see the note above the extras.
|
||||
# Keep linux_aarch64 and macOS off the PyTorch and AMD indexes so that uv resolves their
|
||||
# 2.7.1 fallback pins from PyPI; see the note above the extras.
|
||||
torch = [
|
||||
{ index = "torch-cpu", extra = "cpu", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
|
||||
{ index = "torch-cuda", extra = "cuda", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
|
||||
{ index = "torch-rocm", extra = "rocm", marker = "sys_platform != 'linux' or platform_machine != 'aarch64'" },
|
||||
{ index = "torch-rocm", extra = "rocm", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
|
||||
{ index = "torch-xpu", extra = "xpu", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
|
||||
{ index = "pypi", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
|
||||
]
|
||||
torchvision = [
|
||||
{ index = "torch-cpu", extra = "cpu", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
|
||||
{ index = "torch-cuda", extra = "cuda", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
|
||||
{ index = "torch-rocm", extra = "rocm", marker = "sys_platform != 'linux' or platform_machine != 'aarch64'" },
|
||||
{ index = "torch-rocm", extra = "rocm", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
|
||||
{ index = "torch-xpu", extra = "xpu", marker = "(sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'win32'" },
|
||||
{ index = "pypi", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
|
||||
]
|
||||
triton-rocm = [
|
||||
{ index = "torch-rocm", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
|
||||
# Only the rocm extra takes AMD's triton; cuda and the no-extra resolution keep PyPI's.
|
||||
triton = [
|
||||
{ index = "torch-rocm", extra = "rocm", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
|
||||
]
|
||||
# Packages that exist only on AMD's index.
|
||||
rocm = { index = "torch-rocm" }
|
||||
rocm-sdk-core = { index = "torch-rocm" }
|
||||
rocm-sdk-libraries = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1010 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1010 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1010 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1011 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1011 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1011 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1012 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1012 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1012 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1030 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1030 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1030 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1031 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1031 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1031 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1032 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1032 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1032 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1033 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1033 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1033 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1034 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1034 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1034 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1035 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1035 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1035 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1036 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1036 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1036 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx110x = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1100 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1100 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1100 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1101 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1101 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1101 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1102 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1102 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1102 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1103 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1103 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1103 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx115x = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1150 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1150 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1150 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1151 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1151 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1151 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1152 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1152 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1152 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1153 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1153 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1153 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx12-0 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1200 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1200 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1200 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1201 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1201 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1201 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx908 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx908 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx908 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx90a = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx90a = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx90a = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx1250 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx1250 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx1250 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx942 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx942 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx942 = { index = "torch-rocm" }
|
||||
amd-torch-device-gfx950 = { index = "torch-rocm" }
|
||||
amd-torchvision-device-gfx950 = { index = "torch-rocm" }
|
||||
rocm-sdk-device-gfx950 = { index = "torch-rocm" }
|
||||
# PyPI's xformers wheels are built against CUDA 12.8 and cannot load their extension next to a CUDA 13 torch (every
|
||||
# memory-efficient attention call then fails); PyTorch's cu130 index carries the CUDA 13 build of the same release.
|
||||
# Not scoped to the `cuda` extra: uv only allows that for packages the extra itself lists, and the CUDA 13 torch that
|
||||
@@ -283,9 +485,10 @@ name = "torch-cuda"
|
||||
url = "https://download.pytorch.org/whl/cu130"
|
||||
explicit = true
|
||||
|
||||
# AMD's index also mirrors PyPI packages (numpy, pillow, sympy, ...) without hashes, so it must stay explicit.
|
||||
[[tool.uv.index]]
|
||||
name = "torch-rocm"
|
||||
url = "https://download.pytorch.org/whl/rocm7.2"
|
||||
url = "https://stable.repo.amd.com/rocm/whl-next/"
|
||||
explicit = true
|
||||
|
||||
[[tool.uv.index]]
|
||||
|
||||
@@ -59,8 +59,9 @@ REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||
#
|
||||
# - darwin is deliberately empty: macOS uses MPS, for which PyTorch publishes no
|
||||
# separate index, so the launcher installs the default PyPI wheels.
|
||||
# - win32 has no rocm entry: PyTorch publishes no ROCm wheels for Windows, and
|
||||
# the `rocm` extra in pyproject.toml is marked `sys_platform == 'linux'`.
|
||||
# - win32 has no rocm entry: legacy (pre-6.14.0) installs do not support ROCm on
|
||||
# Windows. Current releases install the `rocm` extra from the lockfile, which
|
||||
# covers Windows through AMD's index.
|
||||
# - xpu is on both win32 and linux: PyTorch's XPU index publishes win_amd64 and
|
||||
# linux-x86_64 wheels, which is exactly what the `xpu` extra's markers allow.
|
||||
REQUIRED_BACKENDS: dict[str, set[str]] = {
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
"""The `rocm_aotriton_experimental` setting: which builds and GPUs get AOTriton's experimental attention kernels."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from invokeai.app.util import rocm_aotriton
|
||||
from invokeai.app.util.rocm_aotriton import (
|
||||
AOTRITON_EXPERIMENTAL_ENV,
|
||||
apply_rocm_aotriton_setting,
|
||||
resolve_aotriton_experimental,
|
||||
rocm_major_version,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("version", "expected"),
|
||||
[
|
||||
("2.13.0+rocm10.0.0", 10),
|
||||
("2.13.0+rocm7.2", 7),
|
||||
("2.10.0+rocm7.1", 7),
|
||||
("2.13.0+cu130", None),
|
||||
("2.13.0", None),
|
||||
# A locally built ROCm torch: no label, so `auto` cannot tell which ROCm it is.
|
||||
("2.8.0a0+gitfc14c65", None),
|
||||
],
|
||||
)
|
||||
def test_rocm_major_version_reads_the_local_label(version: str, expected: int | None):
|
||||
assert rocm_major_version(version) == expected
|
||||
|
||||
|
||||
def test_auto_turns_the_kernels_on_for_rocm_10_on_gfx1200():
|
||||
"""The measured case: RX 9060 XT, torch 2.13.0+rocm10.0.0."""
|
||||
value, reason = resolve_aotriton_experimental("auto", None, 10, {"gfx1200"})
|
||||
assert value == "1"
|
||||
assert "gfx1200" in reason
|
||||
|
||||
|
||||
def test_auto_leaves_rocm_7_alone_where_the_switch_broke_every_fused_call():
|
||||
"""torch 2.12+rocm7.14.1 on the same gfx1200: every flash/mem-efficient call failed with the switch on."""
|
||||
assert resolve_aotriton_experimental("auto", None, 7, {"gfx1200"})[0] is None
|
||||
|
||||
|
||||
def test_auto_leaves_an_unlabelled_build_alone():
|
||||
assert resolve_aotriton_experimental("auto", None, None, {"gfx1200"})[0] is None
|
||||
|
||||
|
||||
def test_auto_needs_every_generation_gpu_to_be_measured():
|
||||
"""The switch is process-wide, so one unmeasured GPU keeps it off for all."""
|
||||
value, reason = resolve_aotriton_experimental("auto", None, 10, {"gfx1200", "gfx1101"})
|
||||
assert value is None
|
||||
assert "gfx1101" in reason
|
||||
|
||||
|
||||
def test_auto_without_a_rocm_gpu_does_nothing():
|
||||
assert resolve_aotriton_experimental("auto", None, 10, set())[0] is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("setting", "expected"), [("on", "1"), ("off", "0")])
|
||||
@pytest.mark.parametrize("rocm_major", [None, 7, 10])
|
||||
def test_on_and_off_decide_for_any_build_and_gpu(setting, expected, rocm_major):
|
||||
assert resolve_aotriton_experimental(setting, None, rocm_major, {"gfx1101"})[0] == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setting", ["auto", "on", "off"])
|
||||
def test_a_value_in_the_environment_always_wins(setting):
|
||||
assert resolve_aotriton_experimental(setting, "0", 10, {"gfx1200"})[0] is None
|
||||
assert resolve_aotriton_experimental(setting, "1", 7, {"gfx1101"})[0] is None
|
||||
|
||||
|
||||
# ===== applying it at startup ===============================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_exported_switch(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.delenv(AOTRITON_EXPERIMENTAL_ENV, raising=False)
|
||||
yield
|
||||
# apply_rocm_aotriton_setting writes os.environ directly; monkeypatch restores it on teardown.
|
||||
monkeypatch.delenv(AOTRITON_EXPERIMENTAL_ENV, raising=False)
|
||||
|
||||
|
||||
def _rocm_build(version: str = "2.13.0+rocm10.0.0", arch: str = "gfx1200"):
|
||||
"""Patches standing in for a ROCm torch with one GPU of `arch` to generate on."""
|
||||
return (
|
||||
patch("torch.version.hip", "7.15.26333"),
|
||||
patch.object(torch, "__version__", version),
|
||||
patch(
|
||||
"invokeai.backend.util.devices.TorchDevice.get_generation_devices",
|
||||
return_value=[torch.device("cuda", 0)],
|
||||
),
|
||||
patch("torch.cuda.get_device_properties", return_value=SimpleNamespace(gcnArchName=f"{arch}:xnack-")),
|
||||
)
|
||||
|
||||
|
||||
def _apply(setting, patches, logger=None):
|
||||
logger = logger or MagicMock(spec=logging.Logger)
|
||||
with patches[0], patches[1], patches[2], patches[3]:
|
||||
apply_rocm_aotriton_setting(setting, "auto", logger)
|
||||
return logger
|
||||
|
||||
|
||||
def test_apply_exports_the_switch_for_the_measured_case(no_exported_switch):
|
||||
logger = _apply("auto", _rocm_build())
|
||||
assert os.environ[AOTRITON_EXPERIMENTAL_ENV] == "1"
|
||||
assert "are on" in logger.info.call_args.args[0]
|
||||
|
||||
|
||||
def test_apply_strips_the_feature_suffix_from_the_arch_name(no_exported_switch):
|
||||
"""`gcnArchName` carries target features (`gfx1200:xnack-`); the measured list holds bare names."""
|
||||
_apply("auto", _rocm_build(arch="gfx1200"))
|
||||
assert os.environ.get(AOTRITON_EXPERIMENTAL_ENV) == "1"
|
||||
|
||||
|
||||
def test_apply_leaves_an_unmeasured_gpu_unset_and_says_how_to_opt_in(no_exported_switch):
|
||||
logger = _apply("auto", _rocm_build(arch="gfx1151"))
|
||||
assert AOTRITON_EXPERIMENTAL_ENV not in os.environ
|
||||
message = logger.info.call_args.args[0]
|
||||
assert "are off" in message and "rocm_aotriton_experimental: on" in message
|
||||
|
||||
|
||||
def test_apply_keeps_an_exported_value_and_warns_when_the_setting_disagrees(monkeypatch):
|
||||
monkeypatch.setenv(AOTRITON_EXPERIMENTAL_ENV, "0")
|
||||
logger = _apply("on", _rocm_build())
|
||||
assert os.environ[AOTRITON_EXPERIMENTAL_ENV] == "0"
|
||||
assert AOTRITON_EXPERIMENTAL_ENV in logger.warning.call_args.args[0]
|
||||
|
||||
|
||||
def test_apply_does_nothing_on_a_non_rocm_build(no_exported_switch):
|
||||
logger = MagicMock(spec=logging.Logger)
|
||||
with patch("torch.version.hip", None):
|
||||
apply_rocm_aotriton_setting("auto", "auto", logger)
|
||||
apply_rocm_aotriton_setting("on", "auto", logger)
|
||||
assert AOTRITON_EXPERIMENTAL_ENV not in os.environ
|
||||
logger.info.assert_not_called()
|
||||
# Only the explicit 'on' is worth a word.
|
||||
assert logger.warning.call_count == 1
|
||||
|
||||
|
||||
def test_the_measured_list_only_holds_measured_architectures():
|
||||
"""Widening `auto` needs a measurement on that GPU; this pins the list so a change is a deliberate edit."""
|
||||
assert rocm_aotriton.AUTO_MEASURED_ARCHS == frozenset({"gfx1200"})
|
||||
assert rocm_aotriton.MIN_AUTO_ROCM_MAJOR == 10
|
||||
@@ -136,9 +136,9 @@ def test_rocm_url_regression_is_reported(repo_copy: Path, capsys: pytest.Capture
|
||||
|
||||
|
||||
def test_unsupported_backend_fails(repo_copy: Path):
|
||||
"""ROCm on Windows has no wheels; pinning an index for it must be rejected."""
|
||||
"""Legacy installs do not support ROCm on Windows; pinning an index for it must be rejected."""
|
||||
pins = _read_pins(repo_copy)
|
||||
pins["torchIndexUrl"]["win32"]["rocm"] = "https://download.pytorch.org/whl/rocm7.1"
|
||||
pins["torchIndexUrl"]["win32"]["rocm"] = "https://stable.repo.amd.com/rocm/whl-next/"
|
||||
_write_pins(repo_copy, pins)
|
||||
|
||||
assert check_pins.main(repo_copy) == 1
|
||||
|
||||
Reference in New Issue
Block a user