first commit

2026-09-23 21:51:44 +02:00
commit 91f98f0260
7 changed files with 427 additions and 0 deletions
@@ -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
@@ -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!
@@ -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
@@ -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
@@ -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
@@ -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}')
@@ -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())"