FlashInfer on ROCm installation

FlashInfer on ROCm installation#

2026-03-03

4 min read time

Applies to Linux

System requirements#

To use FlashInfer 0.2.5, you need the following prerequisites:

  • ROCm version: 6.4.1

  • Operating system: Ubuntu 24.04

  • GPU platform: AMD Instinct™ MI300X

  • PyTorch: 2.7.1

  • Python: 3.12

Install FlashInfer#

To install FlashInfer on ROCm, you have the following options:

Use a prebuilt Docker image with FlashInfer pre-installed#

Docker is the recommended method to set up a FlashInfer environment, as it avoids potential installation issues. The tested, prebuilt image includes FlashInfer, PyTorch, ROCm, and other dependencies.

  1. Pull the Docker image:

    docker pull rocm/flashinfer:flashinfer-0.2.5_rocm6.4_ubuntu24.04_py3.12_pytorch2.7

  2. Launch and connect to the container:

    docker run -it --rm \
--privileged -v ./:/app \
--network=host --device=/dev/kfd \
--device=/dev/dri --group-add video \
--name=my_flashinfer --cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--ipc=host --shm-size 16G \
rocm/flashinfer:flashinfer-0.2.5_rocm6.4_ubuntu24.04_py3.12_pytorch2.7

Build your own Docker image#

FlashInfer on ROCm can be run directly by setting up a Docker container from scratch. A Dockerfile is provided in the ROCm/flashinfer repository to help you get started.

  1. Clone the ROCm/flashinfer repository:

    git clone https://github.com/ROCm/flashinfer.git

  2. Enter the directory and build the Dockerfile:

    cd flashinfer
docker build -t rocm/flashinfer:flashinfer-0.2.5_rocm6.4_ubuntu24.04_py3.12_pytorch2.7

  3. Run the Docker container:

    docker run -it --device=/dev/kfd --device=/dev/dri --group-add video rocm/flashinfer:flashinfer-0.2.5_rocm6.4_ubuntu24.04_py3.12_pytorch2.7

  4. The above step will create a Docker container with FlashInfer pre-installed. During this process, the Dockerfile will have pre-installed and configured a micromamba environment named flashinfer-py3.12-torch2.7.1-rocm6.4.1.

Install FlashInfer using pip#

Use a base PyTorch Docker image and follow these steps to install FlashInfer using pip.

  1. Pull the base ROCm PyTorch Docker image:

    docker pull rocm/pytorch:rocm6.4.1_ubuntu24.04_py3.12_pytorch_release_2.7.1

  2. Change the <container name> and then use the following command:

    docker run -it --privileged --network=host --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ipc=host --shm-size 128G --name=<container name> rocm/pytorch:rocm6.4.1_ubuntu24.04_py3.12_pytorch_release_2.7.1

  3. After setting up the container, install FlashInfer from the AMD-hosted PyPI repository.

    pip install flashinfer==0.2.5.post10 --index-url=https://pypi.amd.com/simple

Test the FlashInfer installation#

Once you have the Docker container running, start using FlashInfer by following these steps:

  1. Activate the micromamba environment:

    Note

    If you followed Install FlashInfer using pip, you do not need to activate the micromamba environment. If you followed Use a prebuilt Docker image with FlashInfer pre-installed or Build your own Docker image, don’t forget this step.

    micromamba activate flashinfer-py3.12-torch2.7.1-rocm6.4.1

  2. Enter the FlashInfer directory:

    Note

    If you followed Install FlashInfer using pip, ensure you git clone the repository first.

    git clone https://github.com/ROCm/flashinfer.git

    cd flashinfer/

  3. Run the example provided in the flashinfer/examples directory. This example runs Batch Decode and then verifies the output.

    python examples/test_batch_decode_example.py

  1. If FlashInfer was installed correctly, you should see the following output:

    PASS

  2. The above output indicates that FlashInfer is installed correctly. You can now use FlashInfer in your projects.

Contents