first commit

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

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: <contextlib._GeneratorContextManager object at 0x7f9f4568c1d0>
+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!
+

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

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

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

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

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())"