I have created the following repository for this work on gitlab https://gitlab.wikimedia.org/repos/machine-learning/rocm-wheelhouse
I have identified the following issues which we should either resolve or decide on before proceeding with this work:

1. We need to use an image that has ROCm + PyTorch installed. The official images available on dockerhub use ubuntu and none of them uses python 3.11 which we use in debian stable (bookworm). We should either build a debian image with ROCm drivers installed for this purpose or we should reassess our production images and use ubuntu images for LLM workloads which would make things easier. There are no official ROCm packages for bookworm and likely trixie is going to support pytorch-rocm natively
2. access to 3rd party images from gitlab runners: our gitlab runners don't have access to all third party images available on dockerhub. More information can be found in the permission matrix on the wikitech page for gitlab runners. We could use the shared runners where we'd need to add the image we want in the allowed images (profile::gitlab::runner::allowed_images in operations/puppet/hieradata/cloud.yaml)

Since these python packages aren't going to be so frequent we should first make sure we have reproducible builds from a local environment and then reconsider if we need a CI/CD process or not.
The local env would consist of a docker image based on debian stable with rocm drivers installed + pytorch for rocm.
Then inside this image we can build flash attention and vllm (and any other python package we need that doesn't have a published wheel).

isaranto opened https://gitlab.wikimedia.org/repos/machine-learning/rocm-wheelhouse/-/merge_requests/1

Draft: Add Debian/ROCm/Pytorch dockerfiles

There is an official image now available for rocm/vllm for the MI300 ( blog | dockerhub images)
My understanding is that this image is built from the dockerfile on the vllm repo although this is not explicitly stated anywhere. However the 2 tags available for MI300 are ~22GB each. When I build the Dockerfile.rocm for MI210 the resulting image is ~36GB

DOCKER_BUILDKIT=1 docker build --build-arg FX_GFX_ARCHS="gfx90a" -f Dockerfile.rocm -t vllm-rocm

In order to run benchmarks(T382343#10437184) for awq, gptq, and fa2 in the ml-lab environment, I built wheels compatible with Debian 12 (bookworm), AMD ROCm 6.1, and torch 2.4.1+rocm6.1. Below is the process I used to build these wheels:

1.AutoGPTQ:

This process failed for both bitsandbytes(P71788) and vllm(P71890) wheels. The successfully built wheels can be found here: https://gitlab.wikimedia.org/repos/machine-learning/huggingface-optimum-benchmark-automation/-/tree/main/wheels

In T381859#10443645, @kevinbazira wrote:

This process failed for bitsandbytes(P71788) wheels.

The problem here is that hip header files get picked up from /usr/include/hip instead of /opt/rocm/include:

...
In file included from /usr/include/hip/hip_fp16.h:29:                           
/usr/include/hip/amd_detail/amd_hip_fp16.h:1520:21: error: use of undeclared identifier '__llvm_amdgcn_rcp_f16'
 1520 |                     __llvm_amdgcn_rcp_f16(static_cast<__half_raw>(x).data)};
      | 
...

These files are compatible with ROCm 5.2.x instead of 6.1.x (and the bundled llvm). The way I was able to build it was by running:

export CPLUS_INCLUDE_PATH=/opt/rocm-6.1.0/include

before running cmake ... and subsequent commands, so that this path gets searched before /usr/include for header files.
I haven't tried building vllm but looking at your logs, it looks like this is also at least part of the problem there:

/usr/include/hip/amd_detail/amd_hip_fp16.h:1520:21: error: use of undeclared identifier '__llvm_amdgcn_rcp_f16'
 1520 |                     __llvm_amdgcn_rcp_f16(static_cast<__half_raw>(x).data)};
      |                     ^

and so the above workaround should in theory also help to unblock us there.

In T381859#10448020, @MunizaA wrote:

In T381859#10443645, @kevinbazira wrote:

This process failed for bitsandbytes(P71788) wheels.

The problem here is that hip header files get picked up from /usr/include/hip instead of /opt/rocm/include

Thank you so much for this tip Muniza, after setting C++ headers location, the bitsandbytes wheel built successfully:

In T381859#10448020, @MunizaA wrote:

In T381859#10443645, @kevinbazira wrote:

This process failed for vllm(P71890) wheels.

The problem here is that hip header files get picked up from /usr/include/hip instead of /opt/rocm/include
I haven't tried building vllm but looking at your logs, it looks like this is also at least part of the problem there

Indeed the build process picking C++ headers from the wrong location was part of the problem. After setting C++ headers location to point to ROCm 6.1.x, the vllm wheel built successfully:

In an effort to build packages locally but with docker I've made the following attempt:

✅ use official rocm image based in ubuntu 22.04 ( rocm/dev-ubuntu-22.04:6.1-complete )
✅ install python3.11 
✅ create a virtual environment and install any required dependencies + torch-rocm
❌ build flash attention

I am getting the following error

python -m build --no-isolation --wheel .
* Getting build dependencies for wheel...

ERROR Backend subprocess exited when trying to invoke get_requires_for_build_wheel

and if I pass --skip-dependency-check it will again fail in the build_wheel step.
Of course there is something wrong with the docker image which can be found in the open MR.