From 91f98f026048e10e22eef2b999d8c7207a2eb7ef Mon Sep 17 00:00:00 2001 From: lullu Date: Wed, 23 Sep 2026 21:51:44 +0200 Subject: [PATCH] first commit --- Dockerfile | 50 +++++++ README.md | 196 +++++++++++++++++++++++++ docker-compose.yml | 28 ++++ docker-run.sh | 16 ++ scripts/check-comfyui.sh | 9 ++ scripts/test-pytorch-flashattention.py | 120 +++++++++++++++ scripts/test-pytorch.sh | 8 + 7 files changed, 427 insertions(+) create mode 100644 Dockerfile create mode 100644 README.md create mode 100644 docker-compose.yml create mode 100755 docker-run.sh create mode 100644 scripts/check-comfyui.sh create mode 100644 scripts/test-pytorch-flashattention.py create mode 100644 scripts/test-pytorch.sh diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..b1f335c --- /dev/null +++ b/Dockerfile @@ -0,0 +1,50 @@ +FROM rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.9.1 + +# "Installing" flash-attention +# - Not sure if branch `main_perf` is actually better, +# but everyone seems to be using it, so ¯\_(ツ)_/¯ +# - As far as I understood, we're not actually installing-installing +# flash-attention, but we're telling it to use the Triton backend/implementation +# which is already included in the rocm/pytorch image (built by amd), +# so the lines below are executed pretty fast. +# - That ^ is what the env variable seems to be for. + +ENV FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE + +RUN cd /opt && \ + git clone https://github.com/ROCm/flash-attention.git && \ + cd flash-attention && \ + git checkout main_perf && \ + python setup.py install + +# Cloning and installing ComfyUI in case the user doesn't provide their own +# - Nothing unusual here afaik + +RUN cd /opt && \ + git clone https://github.com/comfyanonymous/ComfyUI ComfyUI-pre-cloned && \ + cd ComfyUI-pre-cloned && \ + pip3 install -r requirements.txt && \ + pip3 install -r manager_requirements.txt + +# Some utilities to make life/debugging easier +# Feel free to remove these if you're building from scratch locally. + +RUN mkdir -p /opt/comfyui-gfx1151-utils + +WORKDIR /opt/comfyui-gfx1151-utils + +ADD scripts/check-comfyui.sh check-comfyui.sh +RUN chmod +x check-comfyui.sh + +ADD scripts/test-pytorch.sh test-pytorch.sh +RUN chmod +x test-pytorch.sh + +ADD scripts/test-pytorch-flashattention.py test-pytorch-flashattention.py +RUN chmod +x test-pytorch-flashattention.py + +# Run ComfyUI +# - Also nothing unusual here afaik + +EXPOSE 8188 + +CMD /opt/comfyui-gfx1151-utils/check-comfyui.sh && python3 /opt/ComfyUI/main.py --listen 0.0.0.0 --use-flash-attention --gpu-only --enable-manager diff --git a/README.md b/README.md new file mode 100644 index 0000000..649344e --- /dev/null +++ b/README.md @@ -0,0 +1,196 @@ +# ComfyUI for gfx1151 (Ryzen AI MAX) + +Dockerized ComfyUI with PyTorch & flash-attention for gfx1151 (AMD Strix Halo, Ryzen AI Max+ 395), +relying on AMD's pre-built and pre-configured environment (no custom wheels). + +Versions used: +* ROCm: 7.2 +* PyTorch: 2.9.1 +* Python: 3.12 +* ComfyUI (built-in): v0.15.0 + +**Last updated & tested**: Feb 25, 2026, on 6.18.9 (ArchLinux), AMD RYZEN AI MAX+ 395 (Framework Desktop), +with [opencl-amd](https://aur.archlinux.org/packages/opencl-amd) packages (7.2.0-1). + +> [!CAUTION] +> I kinda understand what's going on here, but not fully. It took me most of the day to figure out how to run +> ComfyUI on my Framework Desktop without it crashing (which is absurd for a CPU with _AI_ in the name), and +> the final working solution turned out to be much simpler than what I was able to find initially. +> +> That being said, it works (as of Feb 25, 2026), but I'm not sure if there's an even better / more correct way of +> achieving the same thing. Same goes for the environment variables that are supposedly making ComfyUI +> faster / resource efficient. +> +> I just want to share this solution to save someone else a couple of hours ¯\_(ツ)_/¯ + +## Get started now + +The Docker image is published to [Docker Hub](https://hub.docker.com/r/ignatberesnev/comfyui-gfx1151), +so you can, but don't have to build it yourself. + +There are two options: + +* Copy [docker-compose.yml](docker-compose.yml) and run `docker compose up -d` +* Copy [docker-run.sh](docker-run.sh) and run `./docker-run.sh`. After the first run, use `docker start comfyui-gfx1151`. + +ComfyUI will be available at http://localhost:8188. + +The starter templates should generate images without any issues. + +Once you've verified that it works, feel free to use this repository as the foundation for your own setup or workflow. + +#### Parameters + +Both options have the same pre-configured parameters, which are: + +* Allocate 8GB of shared memory (`shm_size`) for internal PyTorch / ComfyUI shenanigans, this should be plenty, + feel free to lower it. This should NOT be > than available RAM. This is NOT allocating VRAM. +* Mount `./ComfyUI` for the root of [ComfyUI](https://github.com/comfyanonymous/ComfyUI). If the directory is empty + when the container starts, it will copy a pre-cloned (baked in) version of ComfyUI. If it's not empty, it will be + used to run ComfyUI located in it. You can update this directory manually to use newer version of ComfyUI without + having to re-download the image +* Expose port `8188` for ComfyUI +* Add video + rendering devices and groups. While this just works on Arch, it might require some pre-requisite steps + on Ubuntu, I haven't checked. + +There are a couple of scripts that can check that both PyTorch and flash-attention work, you can find them below. + +#### Updating ComfyUI / dependencies + +If you need to install custom nodes or refresh ComfyUI dependencies after a manual update, +you can do it from within the container (until #4 is resolved): + +```bash +docker exec -it comfyui-gfx1151 /bin/bash + +cd /opt/ComfyUI + +pip install -r requirements.txt +``` + +## What's inside / how to replicate + +This image is based on AMD's [rocm/pytorch](https://hub.docker.com/r/rocm/pytorch) image that has Ubuntu 24.04, +ROCm 7.2, Python 3.12 and PyTorch 2.9.1, in which everything is configured to work together and it just works. +You can find out more about this image in +[AMD's ROCm documentation](https://rocm.docs.amd.com/projects/radeon-ryzen/en/latest/docs/install/installryz/native_linux/install-pytorch.html#use-docker-image-with-pre-installed-pytorch). + +There are only two missing pieces which this image adds: [flash-attention](https://github.com/ROCm/flash-attention/) +and, well, ComfyUI. + +There's nothing specific about the ComfyUI installation, you can actually bring your own, it should work. + +**flash-attention**, however, "doesn't work" out of the box if you run AMD's image. I'm saying "doesn't work" because, +as far as I understand, it doesn't have the frontend for it (the APIs), but it does have the backend: **Triton**. +So flash-attention can be "installed" with a special env variable `FLASH_ATTENTION_TRITON_AMD_ENABLE`, which makes +ComfyUI and other tools using flash-attention think that flash-attention is installed and works (even though it's +triton under the hood, which is actually doing the job). You can see the lines that install it in +[Dockerfile](Dockerfile), and if you try to do it yourself, you'll notice that it executes very fast +(because flash-attention isn't actually built in full). + +It's worth noting that flash-attention is cloned from a specific branch `main_perf` -- I'm not sure why exactly, +I haven't checked, but I assume it's because it has (stable?) support for Triton which is not yet in the main branch, +see [this issue](https://github.com/ROCm/flash-attention/issues/27). I basically copy-pasted this part from other +installations ([vLLM](https://community.frame.work/t/compiling-vllm-from-source-on-strix-halo/77241) and repos by +[kyuz0](https://github.com/kyuz0)), so I hope they know what they're doing :D + +In [scripts](scripts) there are two scripts that can check if PyTorch and flash-attention work as expected and utilize +the iGPU. I used these when looking for a solution, they proved to be helpful, so I'm adding them to the image in case +something breaks or doesn't work as expected, maybe they'll help debug the problem or something. + +With that knowledge, you should be able to take [Dockerfile](Dockerfile) and build an image yourself. + +If any of this makes more sense to you than it does to me and you know how to improve something or can add a helpful +comment with additional context, please do! + +## What I tried that didn't work + +The majority of other solutions seem to rely on custom-built wheels, such as the image by +[pccr10001/comfyui-gfx1151-fa](https://github.com/pccr10001/comfyui-gfx1151-fa). + +I never managed to make these custom wheels work, presumably because of non-locked dependencies (pulling in newer +version of rocm/etc with old wheels). However, they made me begin to understand what was happening and how to move +forward, so huge thanks to everyone who left any comments on the topic. + +Some other solutions also relied on the image `ghcr.io/rocm/therock_pytorch_dev_ubuntu_24_04_gfx1151`, which is no +longer published, so I never got that working either. The image I'm referencing +([rocm/pytorch](https://hub.docker.com/r/rocm/pytorch)) seems like a replacement for it though? + +Initially, I [copied over](https://github.com/pccr10001/comfyui-gfx1151-fa/blob/e6e59be08ff439ab5f9799aa2161f70709fcd975/README.md?plain=1#L33) +some environment variables that were supposed to speed up ComfyUI / PyTorch and make it more resource efficient: +`PYTORCH_TUNABLEOP_ENABLED`, `MIOPEN_FIND_MODE` and `ROCBLAS_USE_HIPBLASLT` (not adding them as a codeblock to avoid +someone copy-pasting them). However, at least one of them not only made it worse when it comes to the speed, but I +believe it would crash my display server (X11) every now and then when running stable diffusion models. Apparently, +this is relatively common to see with AMD drivers in general, so I'm not entirely sure that those env variables were +100% responsible for the crashes (might've been something else), but removing all of them helped (at least for now), +so I've removed them from this repo's scripts too. If you also experience display server crashes, let me know. + +## Tests + +There are two scripts that you can use to test if everything works correctly + +#### Test PyTorch + +While the container is running, running + +```bash +docker exec -it comfyui-gfx1151 /bin/bash /opt/comfyui-gfx1151-utils/test-pytorch.sh +``` + +should produce NO errors. The output should be something like: + +```text +GPU: AMD Radeon Graphics | FlashAttn: True +Mean: -0.026233481243252754 +``` + +#### Test flash-attention + +While the container is working, running + +```bash +docker exec -it comfyui-gfx1151 python3 /opt/comfyui-gfx1151-utils/test-pytorch-flashattention.py +``` + +should produce NO errors. The output should be something like: + +```text +=== PyTorch Installation Check === +PyTorch version: 2.9.1+rocm7.1.1.git351ff442 +PyTorch ROCm version: 7.1.52802-26aae437f6 +CUDA available: True +Device count: 1 +Device name: AMD Radeon Graphics + +=== Flash Attention Support Check === +/usr/lib/python3.12/contextlib.py:105: FutureWarning: `torch.backends.cuda.sdp_kernel()` is deprecated. In the future, this context manager will be removed. Please see `torch.nn.attention.sdpa_kernel()` for the new context manager, with updated signature. + self.gen = func(*args, **kwds) +Available SDP backends: +Flash Attention backend enabled +Test tensors created on cuda +Flash Attention test successful! Output shape: torch.Size([2, 8, 128, 64]) + +=== AOTriton Check === +AOTriton not available: No module named 'pyaotriton' + +=== Environment Variables === +ROCM_PATH: Not set +HIP_PATH: Not set +HIP_PLATFORM: Not set +HIP_ARCH: Not set +HSA_OVERRIDE_GFX_VERSION: Not set + +=== Testing GFX Version Override === +Set HSA_OVERRIDE_GFX_VERSION=11.0.0 to test gfx110x mapping +✓ Flash Attention worked with GFX override! + +=== Flash Attention Backend Detection === +✓ flash backend works +✓ mem_efficient backend works +✓ math backend works +``` + +## Acknowledgements + +Big thanks to [pccr10001](https://github.com/pccr10001), [lhl](https://github.com/lhl) and [kyuz0](https://github.com/kyuz0) for setting me on the right path! + diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..51a3608 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,28 @@ +services: + comfyui: + image: comfyui-gfx1151:v0.8 + container_name: comfyui + # Shared memory allocation, should not be > than available RAM, feel free to lower it + #shm_size: 8G + volumes: + - /data/@sd/ComfyUI:/opt/ComfyUI + ports: + - "8188:8188" + environment: + - TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 + devices: + - /dev/kfd:/dev/kfd + - /dev/dri:/dev/dri + group_add: + - video + cap_add: + - SYS_PTRACE + security_opt: + - seccomp:unconfined + networks: + - proxy-network + +networks: + proxy-network: + name: proxy-network + external: true diff --git a/docker-run.sh b/docker-run.sh new file mode 100755 index 0000000..a629a99 --- /dev/null +++ b/docker-run.sh @@ -0,0 +1,16 @@ +#!/bin/bash + +docker run -it \ + --cap-add=SYS_PTRACE \ + --security-opt seccomp=unconfined \ + --device=/dev/kfd \ + --device=/dev/dri \ + --group-add video \ + --group-add render \ + --ipc=host \ + -e TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 \ + -p 8188:8188 \ + -v $(pwd)/ComfyUI:/opt/ComfyUI \ + --shm-size 8G \ + --name comfyui-gfx1151 \ + ignatberesnev/comfyui-gfx1151:v0.2 diff --git a/scripts/check-comfyui.sh b/scripts/check-comfyui.sh new file mode 100644 index 0000000..ad00013 --- /dev/null +++ b/scripts/check-comfyui.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +# If the user mounts /opt/ComfyUI, but their directory doesn't have ComfyUI cloned, +# we're gonna copy the pre-cloned one into the mounted directory. + +if [ ! -e /opt/ComfyUI/requirements.txt ]; then + echo "No ComfyUI detected, copying a built-in (pre-cloned) one..." + cp -r /opt/ComfyUI-pre-cloned/{.,}* /opt/ComfyUI/ +fi diff --git a/scripts/test-pytorch-flashattention.py b/scripts/test-pytorch-flashattention.py new file mode 100644 index 0000000..71dcb17 --- /dev/null +++ b/scripts/test-pytorch-flashattention.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 + +######################################################################## +### Script copied from https://github.com/lhl/strix-halo-testing ### +### All credit goes to https://github.com/lhl, big thanks ### +######################################################################## + +import torch +import os + +print('=== PyTorch Installation Check ===') +print(f'PyTorch version: {torch.__version__}') +print(f'PyTorch ROCm version: {torch.version.hip}') +print(f'CUDA available: {torch.cuda.is_available()}') +if torch.cuda.is_available(): + print(f'Device count: {torch.cuda.device_count()}') + print(f'Device name: {torch.cuda.get_device_name()}') + +print() +print('=== Flash Attention Support Check ===') +try: + # Check if flash attention is available + from torch.nn.attention import SDPBackend + import torch.nn.functional as F + + # Check available backends + backends = torch.backends.cuda.sdp_kernel() + print(f'Available SDP backends: {backends}') + + # Try to use Flash Attention + with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False): + print('Flash Attention backend enabled') + + # Create test tensors + batch_size, seq_len, head_dim = 2, 128, 64 + device = 'cuda' if torch.cuda.is_available() else 'cpu' + + q = torch.randn(batch_size, 8, seq_len, head_dim, device=device, dtype=torch.float16) + k = torch.randn(batch_size, 8, seq_len, head_dim, device=device, dtype=torch.float16) + v = torch.randn(batch_size, 8, seq_len, head_dim, device=device, dtype=torch.float16) + + print(f'Test tensors created on {device}') + + # Run scaled dot product attention + output = F.scaled_dot_product_attention(q, k, v) + print(f'Flash Attention test successful! Output shape: {output.shape}') + +except Exception as e: + print(f'Flash Attention test failed: {e}') + +print() +print('=== AOTriton Check ===') +try: + import pyaotriton + print(f'AOTriton version: {pyaotriton.__version__}') + print('AOTriton imported successfully') +except ImportError as e: + print(f'AOTriton not available: {e}') +except Exception as e: + print(f'AOTriton error: {e}') + +print() +print('=== Environment Variables ===') +rocm_vars = ['ROCM_PATH', 'HIP_PATH', 'HIP_PLATFORM', 'HIP_ARCH', 'HSA_OVERRIDE_GFX_VERSION'] +for var in rocm_vars: + print(f'{var}: {os.environ.get(var, "Not set")}') + +print() +print('=== Testing GFX Version Override ===') +# Try setting HSA_OVERRIDE_GFX_VERSION to map gfx1151 to gfx1100 +os.environ['HSA_OVERRIDE_GFX_VERSION'] = '11.0.0' +print('Set HSA_OVERRIDE_GFX_VERSION=11.0.0 to test gfx110x mapping') + +try: + # Test with override + device = 'cuda' if torch.cuda.is_available() else 'cpu' + if device == 'cuda': + q = torch.randn(1, 1, 32, 64, device=device, dtype=torch.float16) + k = torch.randn(1, 1, 32, 64, device=device, dtype=torch.float16) + v = torch.randn(1, 1, 32, 64, device=device, dtype=torch.float16) + + with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False): + output = torch.nn.functional.scaled_dot_product_attention(q, k, v) + print('✓ Flash Attention worked with GFX override!') +except Exception as e: + print(f'✗ Flash Attention still failed with override: {e}') + +# Reset the environment variable +del os.environ['HSA_OVERRIDE_GFX_VERSION'] + +print() +print('=== Flash Attention Backend Detection ===') +try: + # More detailed backend checking + device = 'cuda' if torch.cuda.is_available() else 'cpu' + if device == 'cuda': + # Check what backends are actually available for this device + q = torch.randn(1, 1, 16, 32, device=device, dtype=torch.float16) + k = torch.randn(1, 1, 16, 32, device=device, dtype=torch.float16) + v = torch.randn(1, 1, 16, 32, device=device, dtype=torch.float16) + + # Test each backend individually + backends_to_test = [ + ('flash', {'enable_flash': True, 'enable_math': False, 'enable_mem_efficient': False}), + ('mem_efficient', {'enable_flash': False, 'enable_math': False, 'enable_mem_efficient': True}), + ('math', {'enable_flash': False, 'enable_math': True, 'enable_mem_efficient': False}) + ] + + for backend_name, kwargs in backends_to_test: + try: + with torch.backends.cuda.sdp_kernel(**kwargs): + output = torch.nn.functional.scaled_dot_product_attention(q, k, v) + print(f'✓ {backend_name} backend works') + except Exception as e: + print(f'✗ {backend_name} backend failed: {e}') + else: + print('CUDA not available, skipping backend tests') + +except Exception as e: + print(f'Backend detection failed: {e}') diff --git a/scripts/test-pytorch.sh b/scripts/test-pytorch.sh new file mode 100644 index 0000000..db2645d --- /dev/null +++ b/scripts/test-pytorch.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +python3 -c "import torch, importlib.util; print(f'GPU: {torch.cuda.get_device_name(0)} | FlashAttn: {importlib.util.find_spec('flash_attn') is not None}')" + +PYTORCH_TUNABLEOP_ENABLED=1 \ +MIOPEN_FIND_MODE=FAST \ +ROCBLAS_USE_HIPBLASLT=1 \ +python3 -c "import torch; x = torch.randn(2048, 2048, device='cuda'); y = torch.mm(x, x); print('Mean:', y.mean().item())"