diff --git a/.claude/skills/debug-rocm-crash/SKILL.md b/.claude/skills/debug-rocm-crash/SKILL.md index 21ded8c7fd..7645aafda1 100644 --- a/.claude/skills/debug-rocm-crash/SKILL.md +++ b/.claude/skills/debug-rocm-crash/SKILL.md @@ -31,7 +31,7 @@ For an in-script view of what's being passed, wrap the suspect call with `print( | `Memory access fault by GPU node-N` / `hipErrorIllegalAddress` / "CUDA error: illegal memory access" (PyTorch's ROCm reports HIP errors as "CUDA" errors) | Run with the env combo above. Print tensor shapes/dtypes/strides just before the call. Verify: `is_contiguous()` where required, all tensors on the same `cuda:N`, `kv_indices` within `[0, num_pages)`, `head_dim_qk` matches between Q and KV. | | `backend="aiter"` `ValueError` before launch | `kv_layout != "NHD"` (only NHD is allowed — raised in the prefill wrapper's `plan()`, e.g. [`prefill_rocm.py:1978`](../../../flashinfer/prefill_rocm.py)). | | `backend="aiter"` `RuntimeError` | Non-gfx942/gfx950 GPU. | -| `backend="aiter"` `ImportError` | `amd-aiter` not installed (`pip install amd-aiter --index-url https://pypi.amd.com/simple/`). | +| `backend="aiter"` `ImportError` | `amd-aiter` not installed — see the AITER wheel section in `README.md` for the pinned version and its index. | | `backend="aiter"` hard GPU fault mid-kernel | `amd-aiter` version mismatch vs. ROCm. Reinstall matching your ROCm version. Try the default HIP backend to confirm the bug is in AITER, not our side. | | NaN / Inf in outputs | Insert `torch.isnan(t).any()` / `torch.isinf(t).any()` checks around the call. On CDNA3/4: `_fnuz` FP8 has different representable range than NVIDIA OCP FP8 — scale factors calibrated against NVIDIA refs overflow. Or `-inf` from a previous op fed into `exp`. Or `torch.empty` vs `torch.zeros`. | | `HIP out of memory` | `rocm-smi --showmeminfo vram --showpids` — kill zombies. JIT-compile spike → `MAX_JOBS=1`. Other tenant → `HIP_VISIBLE_DEVICES=N`. | diff --git a/.devcontainer/rocm/Dockerfile b/.devcontainer/rocm/Dockerfile index 34ac74ab58..6cc456ed1b 100644 --- a/.devcontainer/rocm/Dockerfile +++ b/.devcontainer/rocm/Dockerfile @@ -1,16 +1,45 @@ -ARG ROCM_VERSION=7.2 - -FROM mambaorg/micromamba:2.1.1 AS micromamba +# ROCm development container for amd-flashinfer. +# +# The base image already provides ROCm, Python and a HIP build of PyTorch inside +# a venv at /opt/venv, so this file adds only developer tooling and AITER. +# +# Two things that are NOT obvious and are load-bearing: +# +# 1. No conda/micromamba layer. Torch lives in the base image's interpreter, and +# a separate conda environment would shadow it with a torch-less Python. +# There is also no pip-installable torch for this ROCm release -- +# repo.radeon.com/rocm/manylinux/rocm-rel-7.14/ returns 404 -- so the bundled +# torch is the only one available. +# 2. Nothing may `import aiter` at build time. It runs rocminfo for arch +# detection *and* imports triton, which fails with +# "0 active drivers ([]). There should only be one." when no GPU is attached, +# and that has no env-var escape. Install it here; verify it at run time. +ARG ROCM_VERSION=7.14 +ARG UBUNTU_VERSION=26.04 +ARG PY_VERSION=3.14 +ARG TORCH_VERSION=2.12.0 + +# Header donor. The rocm/pytorch base is runtime-only: it carries the pip ROCm +# SDK (_rocm_sdk_core) rather than /opt/rocm, and that SDK ships no math or +# utility headers -- no thrust, rocPRIM, rocBLAS or hipBLASLt. torch 2.12's +# complex.h needs thrust/complex.h, so the JIT cannot compile without them. +FROM rocm/dev-ubuntu-${UBUNTU_VERSION}:${ROCM_VERSION}.0-full AS rocmdev + +FROM rocm/pytorch:rocm${ROCM_VERSION}_ubuntu${UBUNTU_VERSION}_py${PY_VERSION}_pytorch_release_${TORCH_VERSION} +ARG ROCM_VERSION +ARG PY_VERSION -FROM rocm/dev-ubuntu-24.04:${ROCM_VERSION}-complete +# AITER is pinned to an exact build including the local version segment: this is +# currently the only cp314 wheel published anywhere, and pip will not select a +# local version from a loose specifier. +ARG AITER_VERSION=0.1.16.post3.dev0+g620287969.d20260725 +ARG AITER_INDEX=https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/vllm-cdna/ -ARG ROCM_VERSION -ARG PY_VERSION=3.12 -ARG TORCH_VERSION=2.9.1 -ARG AITER_VERSION=0.1.10 -ARG AITER_ROCM_VERSION=7.1.1 +# Fail fast if the base image ever stops shipping a HIP-enabled torch. +RUN python3 -c "import torch, sys; hip = torch.version.hip; print(f'torch {torch.__version__} (hip={hip})'); sys.exit('ERROR: installed torch has no ROCm/HIP build (hip is None)' if hip is None else 0)" -# Update package lists and install system dependencies +# System dependencies. libstdc++-14-dev and clangd-19 are both still available on +# Ubuntu 26.04, so these names carry over from the 24.04 image unchanged. RUN apt-get update && apt-get install -y --no-install-recommends \ clang-format \ clangd-19 \ @@ -25,75 +54,91 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ zsh \ && rm -rf /var/lib/apt/lists/* -# Create a non-root user ARG USERNAME=devuser ARG USER_UID=1003 ARG USER_GID=$USER_UID -ARG MAMBA_USER_ID=${USER_UID} -ARG MAMBA_USER_GID=${USER_GID} -ENV MAMBA_USER=$USERNAME -ENV MAMBA_ROOT_PREFIX="/opt/conda" -ENV MAMBA_EXE="/bin/micromamba" - -# Silence warning where the UID and GID are out of the usual range, -# create a non-root user. + +# Ubuntu 26.04 ships a default `ubuntu` user at UID/GID 1000, which is the most +# common host UID and what the README's $(id -u) yields -- so it has to go +# *before* devuser is created, or groupadd fails with "GID '1000' already +# exists" and the build dies here. +RUN rm -rf /home/ubuntu && \ + if grep -q '^ubuntu:x:1000:1000' /etc/passwd; then userdel -f -r ubuntu; fi + +# Silence the warning about out-of-range UID/GID, then create the user. RUN sed -i 's/^\(UID_MAX\s*\).*$/\11000000000/' /etc/login.defs && \ sed -i 's/^\(GID_MAX\s*\).*$/\11000000000/' /etc/login.defs && \ groupadd --gid $USER_GID $USERNAME && \ useradd --uid $USER_UID --gid $USER_GID -m $USERNAME && \ echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME && \ - chmod 0440 /etc/sudoers.d/$USERNAME && \ - usermod -a -G render $USERNAME + chmod 0440 /etc/sudoers.d/$USERNAME -RUN echo "set-option -g default-command \"/bin/bash -i\"" >> /home/$USERNAME/.tmux.conf - -# Remove default 'ubuntu' user (UID 1000) to prevent devcontainer permission conflicts -RUN rm -rf /home/ubuntu && if grep ubuntu:x:1000:1000 /etc/passwd >/dev/null; then userdel -f -r ubuntu; fi +# The base image has a `video` group but no `render` group, while /dev/dri/render* +# is render-owned on the hosts. Create it so the user can be a member; callers +# must still pass the host's numeric render GID via --group-add, since the +# in-image GID will not generally match the host's. +RUN (getent group render >/dev/null || groupadd -r render) && \ + usermod -a -G render,video $USERNAME -# Adding micromamba functionality to an existing Docker image requires copying over needed files from the micromamba image -# Refer: https://micromamba-docker.readthedocs.io/en/latest/advanced_usage.html -COPY --from=micromamba "$MAMBA_EXE" "$MAMBA_EXE" -COPY --from=micromamba /usr/local/bin/_activate_current_env.sh /usr/local/bin/_activate_current_env.sh -COPY --from=micromamba /usr/local/bin/_dockerfile_shell.sh /usr/local/bin/_dockerfile_shell.sh -COPY --from=micromamba /usr/local/bin/_entrypoint.sh /usr/local/bin/_entrypoint.sh -COPY --from=micromamba /usr/local/bin/_dockerfile_initialize_user_accounts.sh /usr/local/bin/_dockerfile_initialize_user_accounts.sh -COPY --from=micromamba /usr/local/bin/_dockerfile_setup_root_prefix.sh /usr/local/bin/_dockerfile_setup_root_prefix.sh +RUN echo "set-option -g default-command \"/bin/bash -i\"" >> /home/$USERNAME/.tmux.conf -RUN /usr/local/bin/_dockerfile_initialize_user_accounts.sh && \ - /usr/local/bin/_dockerfile_setup_root_prefix.sh +# System python3.14 is PEP 668 externally-managed, so everything goes into the +# base image's existing venv -- that is where torch already lives, and creating a +# second one would hide it. +ENV VIRTUAL_ENV=/opt/venv +ENV PATH="$VIRTUAL_ENV/bin:/usr/lib/llvm-19/bin:$PATH" + +# Close the three gaps between the pip ROCm SDK and a /opt/rocm install, without +# which every JIT compile dies on `/opt/rocm/bin/amdclang++: not found`. +ENV ROCM_SDK="/opt/venv/lib/python${PY_VERSION}/site-packages/_rocm_sdk_core" +COPY --from=rocmdev /opt/rocm/core-${ROCM_VERSION}/include /opt/rocm-headers/include + +# 1. Headers, from the donor stage above. +ENV CPATH="/opt/rocm-headers/include${CPATH:+:$CPATH}" +# 2. Link names: the SDK ships only versioned libs (libamdhip64.so.7), so +# -lamdhip64 cannot resolve without an unversioned symlink. +# 3. Compiler name: flashinfer invokes $ROCM_HOME/bin/amdclang++, and the SDK +# ships bin/hipcc plus lib/llvm/bin/clang++ under a different name. The +# device bitcode also sits at lib/llvm/amdgcn rather than $ROCM_PATH/amdgcn. +RUN set -eu; \ + for so in "$ROCM_SDK"/lib/lib*.so.[0-9]*; do \ + base="${so%%.so.*}.so"; [ -e "$base" ] || ln -s "$(basename "$so")" "$base"; \ + done; \ + ln -sf ../lib/llvm/bin/clang++ "$ROCM_SDK/bin/amdclang++"; \ + ln -sf ../lib/llvm/bin/clang "$ROCM_SDK/bin/amdclang"; \ + [ -e "$ROCM_SDK/amdgcn" ] || ln -s lib/llvm/amdgcn "$ROCM_SDK/amdgcn" + +# Point at the SDK root, NOT at the grafted header tree: hipcc resolves its own +# clang via $ROCM_PATH/lib/llvm/bin, which only exists under the SDK. +ENV ROCM_PATH="$ROCM_SDK" +ENV ROCM_HOME="$ROCM_SDK" +# Appended rather than assigned so a future base image that sets these for +# torch's own bundled ROCm libs is not clobbered. Today's base sets none of them. +ENV LIBRARY_PATH="${LIBRARY_PATH:+$LIBRARY_PATH:}$ROCM_SDK/lib" +ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH:+$LD_LIBRARY_PATH:}$ROCM_SDK/lib" + +# setuptools and setuptools-scm are the in-tree backend's build requirements. +# They must be in the image because the documented install passes +# --no-build-isolation, which is itself required: isolation resolves torch from +# PyPI and pulls a non-ROCm wheel. +RUN pip install --no-cache-dir \ + ninja \ + "setuptools>=80" \ + "setuptools-scm>=9.2" \ + "packaging>=24" \ + pre-commit \ + numpy \ + pytest pytest-cov pytest-xdist pytest-rerunfailures \ + pybind11 \ + ruff \ + filelock && \ + pip install --no-cache-dir "amd-aiter==${AITER_VERSION}" --extra-index-url "${AITER_INDEX}" && \ + python3 -c "import importlib.metadata as m; print('amd-aiter', m.version('amd-aiter'))" + +# Editable installs write into the venv, so hand it to the dev user. +RUN chown -R $USER_UID:$USER_GID $VIRTUAL_ENV -# Switch to non-root user USER $USERNAME -# Set home directory to the user's home directory WORKDIR /home/$USERNAME -# Set clangd path -ENV PATH="/usr/lib/llvm-19/bin:$PATH" - -# Create the new micromamba environment -ARG MAMBA_ENV_NAME=flashinfer-py${PY_VERSION}-torch${TORCH_VERSION}-rocm${ROCM_VERSION} -ENV MAMBA_ENV_NAME=${MAMBA_ENV_NAME} - -RUN \ - sed -i 's/^#\(force_color_prompt=yes\)/\1/' /home/$USERNAME/.bashrc && \ - echo '' >> /home/$USERNAME/.bashrc && \ - echo 'eval "$(micromamba shell hook --shell bash)"' >> /home/$USERNAME/.bashrc && \ - echo "micromamba activate ${MAMBA_ENV_NAME}" >> /home/$USERNAME/.bashrc - -# Create a new micromamba env and install needed packages for flashinfer development. -# torch uses -f/--find-links (not --index-url) because the radeon repo is a flat -# wheel listing, not a PEP 503 simple index — with --index-url pip requests the -# per-package path (.../torch/), which 404s, and the install fails with "No -# matching distribution found". After installing, assert torch is a ROCm/HIP build -# (torch.version.hip is not None) so the build fails fast if -f ever resolves a -# PyPI CPU/CUDA wheel instead of the radeon ROCm wheel. -RUN /bin/micromamba create -n ${MAMBA_ENV_NAME} python=${PY_VERSION} gtest gmock bash-completion -c conda-forge && \ - /bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --no-cache-dir ninja build "setuptools>=80" "setuptools-scm>=9.2" "packaging>=24" pre-commit numpy pytest pytest-cov pytest-xdist pytest-rerunfailures pybind11 ruff && \ - /bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --no-cache-dir torch==${TORCH_VERSION} -f https://repo.radeon.com/rocm/manylinux/rocm-rel-${ROCM_VERSION}/ && \ - /bin/micromamba run -n ${MAMBA_ENV_NAME} python -c "import torch, sys; hip = torch.version.hip; print(f'torch {torch.__version__} (hip={hip})'); sys.exit('ERROR: installed torch has no ROCm/HIP build (hip is None)' if hip is None else 0)" && \ - /bin/micromamba run -n ${MAMBA_ENV_NAME} pip install amd_aiter==${AITER_VERSION} --extra-index-url https://pypi.amd.com/rocm-${AITER_ROCM_VERSION}/simple -SHELL ["/usr/local/bin/_dockerfile_shell.sh"] - -ENTRYPOINT ["/usr/local/bin/_entrypoint.sh"] - CMD ["/bin/bash"] diff --git a/CLAUDE.md b/CLAUDE.md index e32775a933..920ac6430b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -57,11 +57,15 @@ Details: `pr-workflow` skill. ## Installing Torch -Torch must come from AMD's ROCm repo via `--index-url` (not `-f`, which can -silently install a CPU-only wheel from PyPI). See the +Torch must come from AMD's ROCm repo, via `-f` — `repo.radeon.com` is a flat +wheel listing, not a PEP 503 index, so `--index-url` fails outright. See the [GPU, ROCm, and PyTorch Support](README.md#gpu-rocm-and-pytorch-support) table in `README.md` for the version and command. +ROCm 7.14 has no `rocm-rel-7.14/` directory on `repo.radeon.com` at all, so +there is no pip recipe for it. The devcontainer gets torch 2.12 from its +`rocm/pytorch:rocm7.14_*` base image instead. + ## Non-Obvious Gotchas **JIT build.ninja caching**: `JitSpec.build()` only writes `build.ninja` when @@ -86,26 +90,34 @@ count to avoid HSA/HIPBLAS flakiness under concurrent load. The `slow` marker gates 1M-trial sampling and 4 GB tensor tests — exclude with `-m "not slow"` for fast iteration. -**AITER is a separate install, and the version matters**: The AITER backend -(used by prefill attention on gfx942) is not bundled. Install the pinned wheel: +**AITER is version-pinned, and the pin depends on the interpreter**: the +devcontainer bundles the wheel, so no separate install is needed there. The +pin differs by channel because the channels carry different builds. + +On CPython 3.12 (the CI image, ROCm 7.1.1): ```bash pip install amd-aiter==0.1.10 --extra-index-url https://pypi.amd.com/rocm-7.1.1/simple ``` -`0.1.10` is what this repo is built and tested against — -`prefill_rocm.py` records it as `_AITER_LAST_VALIDATED`, and the README's -[Install AITER wheel package](README.md#install-aiter-wheel-package) section -explains the index choice. Note `amd-aiter` is **not** on the top-level -`pypi.amd.com/simple` index, and it must be `--extra-index-url` rather than -`--index-url` so AITER's own dependencies still resolve from PyPI. - -**Only cp310 and cp312 wheels exist** on that channel (verified 2026-08-20). -On any other interpreter — 3.11, 3.13, 3.14 — the command fails with -`No matching distribution found`, and there is no pinned-version fallback: -public PyPI tops out at a stale `0.1.7.post2.dev18`, and the nightlies index -only carries `>= 0.1.16`. Use CPython 3.12 unless you are prepared to run an -unvalidated AITER. +`amd-aiter` is **not** on the top-level `pypi.amd.com/simple` index, and it +must be `--extra-index-url` rather than `--index-url` so AITER's own +dependencies still resolve from PyPI. **Only cp310 and cp312 wheels exist** on +that channel (verified 2026-08-20); on 3.11, 3.13 or 3.14 it fails with +`No matching distribution found`, and public PyPI tops out at a stale +`0.1.7.post2.dev18`. + +The devcontainer runs CPython 3.14, for which the nightlies index carries the +only wheel that exists: + +```bash +pip install amd-aiter==0.1.16.post3.dev0+g620287969.d20260725 \ + --extra-index-url https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/vllm-cdna/ +``` + +Spell the version out in full including the local `+g...` segment; pip will +not select a local version from a loose specifier. `prefill_rocm.py` records +whatever is validated as `_AITER_LAST_VALIDATED`. A source build (`git clone --recursive https://github.com/ROCm/aiter.git && cd aiter && python3 setup.py develop`) tracks master, which is **many releases diff --git a/README.md b/README.md index 7b1e5807f8..29f3f85b9f 100644 --- a/README.md +++ b/README.md @@ -135,11 +135,11 @@ Measured on: -**Supported ROCm versions:** 7.0.2, 7.1.1, 7.2. +**Supported ROCm versions:** 7.0.2, 7.1.1, 7.2, 7.14. -**Supported PyTorch+ROCm versions:** 2.8.0, 2.9.1. +**Supported PyTorch+ROCm versions:** 2.8.0, 2.9.1, 2.12.0. -Install the matching ROCm-enabled PyTorch wheel from +Up to ROCm 7.2, install the matching ROCm-enabled PyTorch wheel from : ```bash @@ -151,6 +151,11 @@ ROCm version you need; refer to for available wheels. +ROCm 7.14 is the exception: `repo.radeon.com` publishes no `rocm-rel-7.14/` +directory, so there is no `pip install torch` recipe for it. Take torch from +the `rocm/pytorch:rocm7.14_ubuntu26.04_py3.14_pytorch_release_2.12.0` image +instead, as the development container does. + ## Getting Started ### Option 1: Get a Pre-built Docker Image @@ -224,13 +229,14 @@ Build the development Docker image with the repository's Dockerfile: ```bash docker build \ - --build-arg ROCM_VERSION=7.2 \ - --build-arg PY_VERSION=3.12 \ - --build-arg TORCH_VERSION=2.9.1 \ + --build-arg ROCM_VERSION=7.14 \ + --build-arg UBUNTU_VERSION=26.04 \ + --build-arg PY_VERSION=3.14 \ + --build-arg TORCH_VERSION=2.12.0 \ --build-arg USERNAME=$USER \ --build-arg USER_UID=$(id -u) \ --build-arg USER_GID=$(id -g) \ - -t flashinfer-0.5.3.amd1_rocm7.2_ubuntu24.04_py3.12_pytorch2.9.1 \ + -t flashinfer-0.5.3.amd1_rocm7.14_ubuntu26.04_py3.14_pytorch2.12.0 \ -f .devcontainer/rocm/Dockerfile . ``` @@ -238,9 +244,16 @@ docker build \
Build argument descriptions -* `ROCM_VERSION`: ROCm version (default: 7.2) -* `PY_VERSION`: Python version (default: 3.12) -* `TORCH_VERSION`: PyTorch version (default: 2.9.1) +The four version arguments select the `rocm/pytorch` base image tag, so they +must name a tag that exists on Docker Hub — they are not independent knobs. + +* `ROCM_VERSION`: ROCm version (default: 7.14) +* `UBUNTU_VERSION`: Ubuntu version (default: 26.04) +* `PY_VERSION`: Python version (default: 3.14) +* `TORCH_VERSION`: PyTorch version (default: 2.12.0) +* `AITER_VERSION`: AITER wheel version (default: + `0.1.16.post3.dev0+g620287969.d20260725`) +* `AITER_INDEX`: index serving that wheel (default: the vllm-cdna nightlies) * `USERNAME`: Username inside container (default: devuser) * `USER_UID`: User ID for matching host permissions * `USER_GID`: Group ID for matching host permissions @@ -255,10 +268,10 @@ docker run -it \ --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ --ipc=host --privileged --shm-size=128G --network=host \ --device=/dev/kfd --device=/dev/dri \ - --group-add video --group-add render \ + --group-add video --group-add "$(getent group render | cut -d: -f3)" \ -v $PWD:/workspace \ --name flashinfer-dev-container \ - flashinfer-0.5.3.amd1_rocm7.2_ubuntu24.04_py3.12_pytorch2.9.1 + flashinfer-0.5.3.amd1_rocm7.14_ubuntu26.04_py3.14_pytorch2.12.0 ``` @@ -272,7 +285,11 @@ docker run -it \ * `--shm-size=128G`: Shared memory size (adjust as needed) * `--network=host`: Uses host networking * `--device=/dev/kfd --device=/dev/dri`: Exposes AMD GPU devices -* `--group-add video --group-add render`: GPU access groups +* `--group-add video --group-add "$(getent group render | cut -d: -f3)"`: GPU + access groups. `render` must be the **host's numeric GID** — passing the name + resolves against the image's own `render` group, whose GID is assigned at + build time and will not match `/dev/dri/renderD*`, leaving the container + user unable to open the device ("No CUDA GPUs are available"). * `-v :`: Mounts source code
@@ -447,13 +464,25 @@ pip install amd-aiter==0.1.10 --extra-index-url https://pypi.amd.com/rocm-7.1.1/ Use `--extra-index-url`, not `--index-url`, so that AITER's own dependencies still resolve from PyPI. -Pin `0.1.10`: it is the version this repo is built and tested against, and -currently the only one published on that channel. FlashInfer links AITER's C++ -symbols by mangled name (`flashinfer/csrc_rocm/aiter_loader.cc`) and vendors its -argument structs (`include/flashinfer/attention/aiter/`), so a different AITER -build can fail at `dlsym` — or, if a struct layout changed, misinterpret kernel -arguments silently. Both the CI image (`docker/Dockerfile.rocm_ci`) and the -devcontainer install this exact version. +FlashInfer links AITER's C++ symbols by mangled name +(`flashinfer/csrc_rocm/aiter_loader.cc`) and vendors its argument structs +(`include/flashinfer/attention/aiter/`), so a different AITER build can fail at +`dlsym` — or, if a struct layout changed, misinterpret kernel arguments +silently. Always install a pinned version, never a floating one. + +The two images pin different versions, because the channels carry different +builds. The CI image (`docker/Dockerfile.rocm_ci`, ROCm 7.1.1 / Python 3.12) +uses `0.1.10` from the command above. The devcontainer is on Python 3.14, for +which `pypi.amd.com` publishes nothing, so it takes the only cp314 wheel that +exists: + +```bash +pip install amd-aiter==0.1.16.post3.dev0+g620287969.d20260725 \ + --extra-index-url https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/vllm-cdna/ +``` + +The version must be spelled out in full including the local `+g...` segment; +pip will not select a local version from a loose specifier. ### Known Limitations @@ -480,6 +509,17 @@ none of the ignored kwargs below are parameters of it. * batch decode: `use_tensor_cores=True` * the `aiter` Python package is not importable +**Known-wrong on the AITER path:** + +* `logits_soft_cap > 0` on **causal** single prefill with `head_dim=128` and + `kv_len >= 512`. A non-zero soft cap disables every AITER assembly path, + leaving `mha_varlen_fwd`'s CK kernel, which applies the cap incorrectly — + output is off by ~0.17 against an fp32 reference while `logits_soft_cap=0` + is exact. Non-causal soft-capped prefill is unaffected. This is an AITER + defect, reproducible with no FlashInfer in the call path, and is present + through at least aiter 0.1.21. Use `backend="fa2"` for causal soft-capped + models (Gemma-2, Grok-1). + **Features silently ignored on the AITER path** (kwargs are accepted by the FlashInfer wrapper but not forwarded to AITER, which can produce wrong results): diff --git a/amd-flashinfer-jit-cache/pyproject.toml b/amd-flashinfer-jit-cache/pyproject.toml index 9323bd0530..3ccffaab3b 100644 --- a/amd-flashinfer-jit-cache/pyproject.toml +++ b/amd-flashinfer-jit-cache/pyproject.toml @@ -1,7 +1,7 @@ [build-system] # NOTE: torch is intentionally NOT listed here because ROCm users must install # torch from AMD's ROCm repository, not from PyPI. For example, -# pip install torch==2.7.1 --index-url https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/ +# pip install torch==2.9.1 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/ # See README.md for complete installation instructions. requires = ["setuptools>=80", "setuptools_scm>=8", "packaging>=24", "wheel", "ninja"] build-backend = "build_backend" @@ -25,12 +25,13 @@ classifiers = [ "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.14", "Topic :: Scientific/Engineering :: Artificial Intelligence", ] # Runtime dependencies # NOTE: torch is intentionally NOT listed here because ROCm users must install # torch from AMD's ROCm repository, not from PyPI. For example, -# pip install torch==2.7.1 --index-url https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/ +# pip install torch==2.9.1 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/ # See README.md for complete installation instructions. dependencies = [] diff --git a/flashinfer/aiter_utils.py b/flashinfer/aiter_utils.py index ccd2834b6e..0335b285c9 100644 --- a/flashinfer/aiter_utils.py +++ b/flashinfer/aiter_utils.py @@ -3,6 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 import functools +import os +from typing import Optional import torch @@ -22,6 +24,59 @@ def is_aiter_supported(device: torch.device) -> bool: return arch in FLASHINFER_SUPPORTED_ROCM_ARCHS +def _ensure_aiter_gpu_archs() -> None: + """Give AITER's JIT a GPU_ARCHS, since from 0.1.16 it requires one. + + Unset reaches AITER's validator as ``['']`` and every JIT build asserts. Our + own shim build sets this for its own scope, but AITER's Python ops (decode, + paged-append, fused MoE) build outside it. Only fills a missing value, so an + operator-set GPU_ARCHS still wins. + """ + if os.environ.get("GPU_ARCHS"): + return + # Imported lazily: flashinfer.jit pulls in the compilation context, and + # importing it at module scope here would be circular. + from .jit.aiter_source import resolve_aiter_build_arch + + arch = resolve_aiter_build_arch() + if arch: + os.environ["GPU_ARCHS"] = arch + + +# The vendored structs in include/flashinfer/attention/aiter/ follow the 0.1.16 +# layout. They travel by value through dlsym'd pointers, so an older AITER +# mismatches offsets silently instead of failing to load -- hence a hard floor +# rather than a warning. +AITER_MIN_VERSION = "0.1.16" + + +def _aiter_installed_version() -> Optional[str]: + """The installed amd-aiter version, or None when it is not installed.""" + try: + import importlib.metadata as _md + + return _md.version("amd-aiter") + except Exception: + return None + + +def _aiter_version_supported() -> bool: + """True when amd-aiter is installed and new enough for our vendored ABI.""" + installed = _aiter_installed_version() + if installed is None: + return False + try: + from packaging.version import Version + + # Compare on base_version: the nightly wheels carry a local '+g' + # segment, and a '.dev0' pre-release segment that PEP 440 sorts *below* + # the release it is built from -- 0.1.16.post3.dev0 must still pass a + # 0.1.16 floor. Re-wrap in Version; base_version is a str. + return Version(Version(installed).base_version) >= Version(AITER_MIN_VERSION) + except Exception: + return False + + @functools.lru_cache(maxsize=1) def _aiter_importable() -> bool: """True when the AITER packages needed for the C++ backends actually import. @@ -29,15 +84,17 @@ def _aiter_importable() -> bool: Uses a real import (not ``find_spec``) so a broken or partially-installed AITER — where the spec exists but importing the compiled extension fails, e.g. missing ROCm deps — is reported as unavailable rather than routing ``auto`` into a path - that raises at build/load time. + that raises at build/load time. Too-old AITER counts as unavailable for the same + reason: ``auto`` must not route into an ABI it would corrupt. """ try: + _ensure_aiter_gpu_archs() import aiter # noqa: F401 import aiter_meta # noqa: F401 from aiter.jit import core as _core # noqa: F401 except Exception: return False - return True + return _aiter_version_supported() def is_aiter_available(device: torch.device, op: str) -> bool: @@ -75,6 +132,14 @@ def require_aiter(device: torch.device, op: str) -> None: """ require_capability(device, op, "aiter") if not _aiter_importable(): + installed = _aiter_installed_version() + if installed is not None and not _aiter_version_supported(): + raise ValueError( + f"backend='aiter' for {op} requires amd-aiter >= {AITER_MIN_VERSION}, " + f"but {installed} is installed; the vendored struct layouts do not " + "match older releases and would corrupt arguments silently. Upgrade " + "amd-aiter or use backend='native'." + ) raise ValueError( f"backend='aiter' for {op} requires the aiter package, which is not " "installed or failed to import. Install it (see the AITER Support " diff --git a/flashinfer/compilation_context_hip.py b/flashinfer/compilation_context_hip.py index ee8cdd74f0..f6f461ccdf 100644 --- a/flashinfer/compilation_context_hip.py +++ b/flashinfer/compilation_context_hip.py @@ -34,6 +34,10 @@ class CompilationContext: "-DFLASHINFER_ENABLE_FP8_E4M3", "-DFLASHINFER_ENABLE_FP8_E5M2", "-DHIP_ENABLE_WARP_SYNC_BUILTINS=1", + # Required from torch 2.12 on: c10/hip/HIPStream.h gates the c10::hip + # compat namespace behind `#ifdef USE_ROCM`, so without this every + # *_aiter.cu shim fails on c10::hip::getCurrentHIPStream(). + "-DUSE_ROCM=1", ] def __init__(self): diff --git a/flashinfer/csrc_rocm/activation_aiter.cu b/flashinfer/csrc_rocm/activation_aiter.cu index d4e6f28e53..e106e47dc0 100644 --- a/flashinfer/csrc_rocm/activation_aiter.cu +++ b/flashinfer/csrc_rocm/activation_aiter.cu @@ -9,15 +9,28 @@ #include #include -// AITER's public header (activation.h) pulls in → full -// pybind11, which clashes with FlashInfer's -DPy_LIMITED_API. torch::Tensor is -// at::Tensor, so forward-declare the entry point; the linker resolves it against -// the symbol-visible AITER .so. -namespace aiter { -void silu_and_mul(at::Tensor& out, at::Tensor& input); -} // namespace aiter +// AITER's real header, not a forward declaration: a signature change must be a +// compile error, not a load-time `undefined symbol`. Includable since 0.1.16, +// which dropped and with it the pybind11 clash against +// FlashInfer's -DPy_LIMITED_API. +#include +#include + +#include "aiter_tensor_compat.h" void silu_and_mul_aiter(at::Tensor out, at::Tensor input) { const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(input.device()); - aiter::silu_and_mul(out, input); + + // The kernel indexes linearly, so strides in aiter_tensor_t are not honoured. + // AITER's torch entry point used to reject this; the POD API cannot. + TORCH_CHECK(input.is_contiguous(), "silu_and_mul: input must be contiguous"); + TORCH_CHECK(out.is_contiguous(), "silu_and_mul: out must be contiguous"); + + const aiter_tensor_t out_a = flashinfer::aiter_compat::to_aiter(out); + const aiter_tensor_t in_a = flashinfer::aiter_compat::to_aiter(input); + + const flashinfer::aiter_compat::StreamGuard stream_guard(at::hip::getCurrentHIPStream()); + // `limit` (new in 0.1.16) gates an optional clamp; 0.0f is AITER's declared + // default and preserves the previous behaviour. + aiter::silu_and_mul(out_a, in_a, /*limit=*/0.0f); } diff --git a/flashinfer/csrc_rocm/aiter_tensor_compat.h b/flashinfer/csrc_rocm/aiter_tensor_compat.h new file mode 100644 index 0000000000..2741853291 --- /dev/null +++ b/flashinfer/csrc_rocm/aiter_tensor_compat.h @@ -0,0 +1,86 @@ +// SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +// SPDX-License-Identifier: Apache-2.0 +// +// at::Tensor -> aiter_tensor_t adapter, for AITER's POD C++ API (0.1.16+). +// +// Include AITER's real aiter_tensor.h, never a vendored copy: a layout change +// must be a compile error, not wrong numbers. It is safe under +// -DPy_LIMITED_API because it pulls no pybind11. +#pragma once + +#include +#include +#include +#include + +namespace flashinfer::aiter_compat { + +inline AiterDtype to_aiter_dtype(at::ScalarType t) { + switch (t) { + case at::kHalf: + return AITER_DTYPE_fp16; + case at::kBFloat16: + return AITER_DTYPE_bf16; + case at::kFloat: + return AITER_DTYPE_fp32; + // OCP (gfx950) and FNUZ (gfx942) e4m3 differ in exponent bias and NaN + // encoding, but AITER exposes one fp8 enum; the arch picks the meaning. + case at::kFloat8_e4m3fn: + case at::kFloat8_e4m3fnuz: + return AITER_DTYPE_fp8; + case at::kInt: + return AITER_DTYPE_i32; + case at::kShort: + return AITER_DTYPE_i16; + case at::kChar: + return AITER_DTYPE_i8; + case at::kByte: + return AITER_DTYPE_u8; + case at::kLong: + return AITER_DTYPE_i64; + default: + TORCH_CHECK(false, "no aiter_tensor_t dtype for at::ScalarType ", t); + } +} + +// aiter_tensor_t carries fixed shape[8]/strides[8] arrays, matching PyTorch's +// own dimension limit; anything deeper cannot be represented. +inline aiter_tensor_t to_aiter(const at::Tensor& t) { + TORCH_CHECK(t.dim() <= 8, "aiter_tensor_t supports at most 8 dims, got ", t.dim()); + + aiter_tensor_t out{}; + out.ptr = const_cast(t.data_ptr()); + out.numel_ = static_cast(t.numel()); + out.ndim = static_cast(t.dim()); + for (int i = 0; i < out.ndim; ++i) { + out.shape[i] = t.size(i); + out.strides[i] = t.stride(i); + } + out.dtype_ = to_aiter_dtype(t.scalar_type()); + // is_gpu() keys off device_id >= 0, and AITER kernels require device memory. + out.device_id = t.is_cpu() ? -1 : static_cast(t.device().index()); + return out; +} + +// Point AITER's thread_local stream at ours for the duration of a call. +// +// The POD entry points launch on aiter::getCurrentHIPStream(), which defaults +// to nullptr and is otherwise set only by AITER's Python layer. Scoped, because +// the value outlives the call otherwise: a caller inside a temporary +// torch.cuda.Stream would strand a freed handle for the next AITER call on this +// thread. c10's device guard restores the device, not the stream. +class StreamGuard { + public: + explicit StreamGuard(hipStream_t stream) : prev_(aiter::getCurrentHIPStream()) { + aiter::setCurrentHIPStream(stream); + } + ~StreamGuard() { aiter::setCurrentHIPStream(prev_); } + + StreamGuard(const StreamGuard&) = delete; + StreamGuard& operator=(const StreamGuard&) = delete; + + private: + hipStream_t prev_; +}; + +} // namespace flashinfer::aiter_compat diff --git a/flashinfer/csrc_rocm/fused_moe_aiter.cu b/flashinfer/csrc_rocm/fused_moe_aiter.cu index 0535c3aaa9..767829f7b1 100644 --- a/flashinfer/csrc_rocm/fused_moe_aiter.cu +++ b/flashinfer/csrc_rocm/fused_moe_aiter.cu @@ -21,6 +21,10 @@ // full pybind11, which clashes with FlashInfer's -DPy_LIMITED_API. torch::Tensor // is at::Tensor, so forward-declare the entry points; the linker resolves them // against the symbol-visible AITER .so. Both are at global namespace. +// +// These must track moe_ck.h exactly: a missing trailing parameter still compiles +// and then fails at dlopen with a mangled-name mismatch. 0.1.16 added +// `is_shuffled` to both stages. void moe_sorting_fwd(at::Tensor& topk_ids, at::Tensor& topk_weights, at::Tensor& sorted_token_ids, at::Tensor& sorted_weights, at::Tensor& sorted_expert_ids, at::Tensor& num_valid_ids, at::Tensor& moe_buf, int num_experts, int unit_size, @@ -33,7 +37,7 @@ void ck_moe_stage1(at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2, std::optional w1_scale, std::optional a1_scale, std::optional block_m, std::optional sorted_weights, int quant_type, int activation, std::optional splitk, bool nt, - std::optional dst_type); + std::optional dst_type, bool is_shuffled); void ck_moe_stage2(at::Tensor& inter_states, at::Tensor& w1, at::Tensor& w2, at::Tensor& sorted_token_ids, at::Tensor& sorted_expert_ids, @@ -41,7 +45,7 @@ void ck_moe_stage2(at::Tensor& inter_states, at::Tensor& w1, at::Tensor& w2, std::optional w2_scale, std::optional a2_scale, std::optional block_m, std::optional sorted_weights, int quant_type, int activation, std::optional splitk, bool nt, - std::optional dst_type); + std::optional dst_type, bool is_shuffled); namespace { @@ -193,10 +197,10 @@ void fused_moe_aiter(at::Tensor out, at::Tensor hidden_states, at::Tensor w1, at inter_states, topk_i32, kernel_name, /*w1_scale=*/std::nullopt, /*a1_scale=*/std::nullopt, block_m_i32, /*sorted_weights=*/std::nullopt, kQuantNone, activation_i32, /*splitk=*/1, /*nt=*/false, - /*dst_type=*/std::nullopt); + /*dst_type=*/std::nullopt, /*is_shuffled=*/true); ck_moe_stage2(inter_states, w1, w2, sorted_token_ids, sorted_expert_ids, num_valid_ids, out, topk_i32, kernel_name, /*w2_scale=*/std::nullopt, /*a2_scale=*/std::nullopt, block_m_i32, sorted_weights, kQuantNone, activation_i32, /*splitk=*/1, - /*nt=*/false, /*dst_type=*/std::nullopt); + /*nt=*/false, /*dst_type=*/std::nullopt, /*is_shuffled=*/true); } diff --git a/flashinfer/csrc_rocm/page_aiter.cu b/flashinfer/csrc_rocm/page_aiter.cu index 11f7bdbfe2..f7b2b5137b 100644 --- a/flashinfer/csrc_rocm/page_aiter.cu +++ b/flashinfer/csrc_rocm/page_aiter.cu @@ -13,15 +13,14 @@ #include -// Forward-declared rather than included: AITER's cache.h pulls in -// , which clashes with -DPy_LIMITED_API. Unlike the -// norm/rope/activation entry points this one lives in namespace aiter. -namespace aiter { -void reshape_and_cache_flash(at::Tensor& key, at::Tensor& value, at::Tensor& key_cache, - at::Tensor& value_cache, at::Tensor& slot_mapping, - const std::string& kv_cache_dtype, at::Tensor& k_scale, - at::Tensor& v_scale); -} // namespace aiter +// AITER's real header, not a forward declaration: a signature change must be a +// compile error, not a load-time `undefined symbol`. Includable since 0.1.16, +// which moved cache.h off and onto the POD aiter_tensor_t, +// removing the pybind11 clash against -DPy_LIMITED_API. +#include +#include + +#include "aiter_tensor_compat.h" namespace { @@ -172,6 +171,19 @@ void append_paged_kv_cache_aiter(at::Tensor append_key, at::Tensor append_value, // "auto" selects the no-quantization path; the scales are ignored but required. const std::string kv_cache_dtype = "auto"; - aiter::reshape_and_cache_flash(append_key, append_value, paged_k_cache, paged_v_cache, - slot_mapping, kv_cache_dtype, k_scale, v_scale); + + namespace compat = flashinfer::aiter_compat; + aiter_tensor_t key_a = compat::to_aiter(append_key); + aiter_tensor_t value_a = compat::to_aiter(append_value); + aiter_tensor_t key_cache_a = compat::to_aiter(paged_k_cache); + aiter_tensor_t value_cache_a = compat::to_aiter(paged_v_cache); + aiter_tensor_t slot_mapping_a = compat::to_aiter(slot_mapping); + aiter_tensor_t k_scale_a = compat::to_aiter(k_scale); + aiter_tensor_t v_scale_a = compat::to_aiter(v_scale); + + // The POD API launches on AITER's thread_local stream, which only its Python + // layer otherwise sets; scoped so the value does not outlive this call. + const flashinfer::aiter_compat::StreamGuard stream_guard(stream); + aiter::reshape_and_cache_flash(key_a, value_a, key_cache_a, value_cache_a, slot_mapping_a, + kv_cache_dtype, k_scale_a, v_scale_a); } diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 634f70d59d..293699a81a 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -343,7 +343,7 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: # Add new tuple for adding a new version group _ROCM_ARCH_GROUPS = [ ( - ["7.13", "7.12", "7.11", "7.3", "7.2", "7.1", "7.0"], + ["7.14", "7.13", "7.12", "7.11", "7.3", "7.2", "7.1", "7.0"], [ "gfx950", "gfx1201", diff --git a/flashinfer/jit/aiter_source.py b/flashinfer/jit/aiter_source.py index 12f92ce6c2..09a61ee376 100644 --- a/flashinfer/jit/aiter_source.py +++ b/flashinfer/jit/aiter_source.py @@ -5,10 +5,11 @@ FlashInfer wraps AITER kernels by compiling a small ``csrc_rocm/*_aiter.cu`` shim that calls AITER's C++ entry point directly and links the symbol-visible AITER -``.so``. The shim forward-declares the entry point rather than ``#include``-ing -AITER's public header, because that header pulls in pybind11, which clashes with -FlashInfer's ``-DPy_LIMITED_API`` build; ``torch::Tensor`` is ``at::Tensor``, so -the linker resolves the symbol from the AITER ``.so``. +``.so``. Prefer ``#include``-ing AITER's real header, so a signature change is a +compile error rather than a load-time ``undefined symbol``. Fall back to a +forward declaration only for headers that still pull in pybind11, which clashes +with FlashInfer's ``-DPy_LIMITED_API`` build (``rope.h``, ``rmsnorm.h`` as of +0.1.16); there ``torch::Tensor`` is ``at::Tensor``, so the linker still resolves. AITER's installed wheel builds its modules with ``-fvisibility=hidden``, so the kernel symbols (e.g. ``rope_cached_positions_2c_fwd_impl``) are not linkable. This @@ -23,6 +24,7 @@ """ import functools +import inspect import os import re import shutil @@ -392,20 +394,30 @@ def _build_aiter_lib( md_name, os.environ["GPU_ARCHS"], ) - build_module( - md_name=md_name, - srcs=a["srcs"], - flags_extra_cc=a["flags_extra_cc"], - flags_extra_hip=flags_extra_hip, - blob_gen_cmd=blob_gen_cmd, - extra_include=a["extra_include"], - extra_ldflags=a["extra_ldflags"], - verbose=os.environ.get("FLASHINFER_JIT_VERBOSE", "0") == "1", - is_python_module=a["is_python_module"], - is_standalone=a["is_standalone"], - torch_exclude=a["torch_exclude"], - hipify=a.get("hipify", False), - ) + kwargs = { + "md_name": md_name, + "srcs": a["srcs"], + "flags_extra_cc": a["flags_extra_cc"], + "flags_extra_hip": flags_extra_hip, + "blob_gen_cmd": blob_gen_cmd, + "extra_include": a["extra_include"], + "extra_ldflags": a["extra_ldflags"], + "verbose": os.environ.get("FLASHINFER_JIT_VERBOSE", "0") == "1", + "is_python_module": a["is_python_module"], + "is_standalone": a["is_standalone"], + "torch_exclude": a["torch_exclude"], + "hipify": a.get("hipify", False), + # Added after 0.1.10; required (no default) from 0.1.16 on, where it + # selects third-party sources AITER clones per build (CK, + # HipKittens). get_args_of_build supplies the right value. + "third_party": a.get("third_party"), + } + # AITER's build_module signature moves between releases, so pass only what + # the installed one accepts: 0.1.10 has no `third_party` and would raise + # TypeError on it, while 0.1.16+ makes it a required positional. Filtering + # here keeps a single code path working across both. + accepted = inspect.signature(build_module).parameters + build_module(**{k: v for k, v in kwargs.items() if k in accepted}) # AITER decides the output dir from a module-level `bd_dir` global that is # frozen at import time, so the .so does not reliably land in @@ -418,7 +430,11 @@ def _build_aiter_lib( finally: for k, v in prev.items(): if v is None: - os.environ.pop(k, None) + # GPU_ARCHS is the exception: AITER's own Python ops build + # outside this scope and assert on an unset value, and this is + # the arch _ensure_aiter_gpu_archs would resolve anyway. + if k != "GPU_ARCHS": + os.environ.pop(k, None) else: os.environ[k] = v diff --git a/flashinfer/jit/rope.py b/flashinfer/jit/rope.py index 1f8fea4de0..7f2d32361a 100644 --- a/flashinfer/jit/rope.py +++ b/flashinfer/jit/rope.py @@ -31,7 +31,15 @@ def gen_rope_module() -> JitSpec: def gen_rope_aiter_module() -> JitSpec: from .aiter_source import aiter_jitspec_flags, refresh_aiter_jitspec - extra_include_paths, extra_ldflags = aiter_jitspec_flags("module_rope_pos_fwd") + # AITER split the monolithic rope module by variant. Through 0.1.10 the whole + # forward path lived in "module_rope_pos_fwd"; from 0.1.16 that name is not + # registered at all, and the entry point this shim calls + # (rope_cached_positions_2c_fwd_impl) is built by the 2c cached-positions + # module. Asking for the old name does not error usefully -- AITER hands back + # an empty source list and the JIT dies on `assert len(sources) > 0`. + extra_include_paths, extra_ldflags = aiter_jitspec_flags( + "module_rope_2c_cached_positions_fwd" + ) return refresh_aiter_jitspec( gen_jit_spec( "rope_aiter", diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index 5e3077a5b2..cfa07cc1d3 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -61,7 +61,10 @@ # sizes we try. The second is the newest release we have actually validated against; # bumping it must not silently move the support boundary. _AITER_NATIVE_PAGING_SINCE = "0.1.10" -_AITER_LAST_VALIDATED = "0.1.10" +_AITER_LAST_VALIDATED = "0.1.16.post3.dev0+g620287969.d20260725" +# Newest AITER carrying the mha_varlen_fwd soft-cap defect. Bump only after +# re-measuring against an fp32 reference; the wrong answer is silent. +_AITER_SOFTCAP_DEFECT_THROUGH = "0.1.21" @functools.cache @@ -300,6 +303,13 @@ def _aiter_noop_plan(*args, **kwargs): @functools.cache def _aiter_ops_importable() -> bool: try: + # AITER 0.1.16+ freezes arch state at import, so GPU_ARCHS has to be set + # before this import, not before the first build. This is the second + # entry point that imports aiter; aiter_utils._aiter_importable is the + # other, and prefill/decode reach this one first. + from .aiter_utils import _ensure_aiter_gpu_archs + + _ensure_aiter_gpu_archs() import aiter.ops # noqa: F401 return True @@ -328,6 +338,28 @@ def _require_aiter_runtime(device: torch.device, op: str = "batch_prefill") -> N _aiter_auto_warned: set[tuple[torch.device, str]] = set() +def _aiter_softcap_defect( + causal: bool, + logits_soft_cap: Optional[float], + head_dim: int, + kv_len: Optional[int], +) -> bool: + """Would this call hit AITER's miscomputed soft cap? + + A non-zero cap disables AITER's asm paths, leaving mha_varlen_fwd's CK + kernel, which applies the cap wrongly for causal head_dim=128 with + kv_len >= 512. Non-causal is exact. kv_len=None means the caller does not + know it and the fallback is declined. + + Deliberately not version-gated: auto-expiring on an AITER newer than + _AITER_SOFTCAP_DEFECT_THROUGH would silently re-enable a wrong-answer path + on a nightly bump. Re-measure, then widen the constant by hand. + """ + if not (causal and logits_soft_cap and logits_soft_cap > 0): + return False + return head_dim == 128 and kv_len is not None and kv_len >= 512 + + def _auto_select_prefill_backend( device: torch.device, *, @@ -339,6 +371,9 @@ def _auto_select_prefill_backend( head_dim_vo: int, pos_encoding_mode: str = "NONE", op: str = "batch_prefill", + causal: bool = False, + logits_soft_cap: Optional[float] = None, + kv_len: Optional[int] = None, ) -> Tuple[str, Optional[str]]: """Return ``(backend, reason)``: 'aiter' when the GPU and call parameters satisfy AITER's constraints, else 'fa2' plus the reason AITER was declined. @@ -376,6 +411,11 @@ def _auto_select_prefill_backend( reason = ( f"pos_encoding_mode={pos_encoding_mode!r} (AITER only supports NONE)" ) + elif _aiter_softcap_defect(causal, logits_soft_cap, head_dim_qk, kv_len): + reason = ( + f"logits_soft_cap={logits_soft_cap} with causal head_dim={head_dim_qk} " + "(AITER mha_varlen_fwd computes the soft cap incorrectly)" + ) if reason is not None: key = (device, reason) @@ -1436,6 +1476,8 @@ def single_prefill_with_kv_cache( if scale_v is None: scale_v = torch.ones(v.shape[1], dtype=torch.float32, device=q.device) + kv_len = k.shape[0] if kv_layout == "NHD" else k.shape[1] + if backend == "auto": backend, _ = _auto_select_prefill_backend( q.device, @@ -1447,15 +1489,29 @@ def single_prefill_with_kv_cache( head_dim_vo=v.shape[-1], pos_encoding_mode=pos_encoding_mode, op="single_prefill", + causal=causal, + logits_soft_cap=logits_soft_cap, + kv_len=kv_len, ) if backend == "aiter": _require_aiter_runtime(q.device, "single_prefill") + if _aiter_softcap_defect(causal, logits_soft_cap, q.shape[-1], kv_len): + raise ValueError( + "AITER miscomputes logits_soft_cap for causal head_dim=128 prefill " + f"with kv_len >= 512 (through amd-aiter {_AITER_SOFTCAP_DEFECT_THROUGH}); " + "use backend='fa2' or backend='auto' instead." + ) if pos_encoding_mode != "NONE": raise ValueError( f"AITER backend does not support pos_encoding_mode={pos_encoding_mode!r}; " "use backend='fa2' or backend='auto' instead." ) + if kv_layout != "NHD": + raise ValueError( + f"AITER backend only supports kv_layout='NHD'; got {kv_layout!r}. " + "use backend='fa2' or backend='auto' instead." + ) # logits_soft_cap > 0 forces the varlen .so (mha_fwd template has no _logits # arm); the logits .so is split by causality (mask vs nmask) and neither is # pre-shipped by AITER, so bootstrap the variant matching the request. @@ -2134,6 +2190,16 @@ def plan( if self._jit_module is not None: self._cached_module = self._jit_module else: + # Only the flat-gather route carries the soft-cap defect; native + # paging uses mha_batch_prefill, which is exact. A page size outside + # the native set forces flat-gather, so that is the case we can rule + # out up front. When native paging is merely *claimed*, the run-time + # probe may still fall back to flat-gather -- see plan()'s + # use_native_paging. Shared by the auto route and the explicit-aiter + # guard below so the two cannot disagree about the same call. + softcap_kv_len = ( + None if page_size in _aiter_native_page_sizes() else self._max_kv_len + ) if self._backend == "auto": self._backend, self._backend_fallback_reason = ( _auto_select_prefill_backend( @@ -2146,8 +2212,19 @@ def plan( head_dim_vo=head_dim_vo, pos_encoding_mode=pos_encoding_mode, op="batch_prefill", + causal=causal, + logits_soft_cap=logits_soft_cap, + kv_len=softcap_kv_len, ) ) + if self._backend == "aiter" and _aiter_softcap_defect( + causal, logits_soft_cap, head_dim_qk, softcap_kv_len + ): + raise ValueError( + "AITER miscomputes logits_soft_cap for causal head_dim=128 prefill " + f"with kv_len >= 512 (through amd-aiter {_AITER_SOFTCAP_DEFECT_THROUGH}); " + "use backend='fa2' or backend='auto' instead." + ) if self._backend == "aiter" and pos_encoding_mode != "NONE": raise ValueError( f"AITER backend does not support pos_encoding_mode={pos_encoding_mode!r}; " @@ -3122,8 +3199,21 @@ def plan( head_dim_vo=head_dim_vo, pos_encoding_mode=pos_encoding_mode, op="batch_prefill", + # Ragged always dispatches through mha_varlen_fwd, so it + # carries the soft-cap defect exactly as single prefill does. + causal=causal, + logits_soft_cap=logits_soft_cap, + kv_len=self._max_kv_len, ) ) + if self._backend == "aiter" and _aiter_softcap_defect( + causal, logits_soft_cap, head_dim_qk, self._max_kv_len + ): + raise ValueError( + "AITER miscomputes logits_soft_cap for causal head_dim=128 prefill " + f"with kv_len >= 512 (through amd-aiter {_AITER_SOFTCAP_DEFECT_THROUGH}); " + "use backend='fa2' or backend='auto' instead." + ) if self._backend == "aiter" and pos_encoding_mode != "NONE": raise ValueError( f"AITER backend does not support pos_encoding_mode={pos_encoding_mode!r}; " diff --git a/include/flashinfer/attention/aiter/mha_fwd_args.h b/include/flashinfer/attention/aiter/mha_fwd_args.h index 9343b7ee42..7b6f1cde73 100644 --- a/include/flashinfer/attention/aiter/mha_fwd_args.h +++ b/include/flashinfer/attention/aiter/mha_fwd_args.h @@ -1,9 +1,15 @@ // SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 // -// Vendored aiter::mha_fwd_args from AITER amd-aiter>=0.1.10. +// Vendored aiter::mha_fwd_args from AITER amd-aiter>=0.1.16. // Extracted from aiter_meta/csrc/include/mha_fwd.h. // +// 0.1.16 is a hard minimum, not a recommendation: the block_scale_seqstart_* +// fields below were inserted before sink_ptr, so on 0.1.10 every field from +// sink_ptr on lands at the wrong offset. The struct goes by value through a +// dlsym'd pointer, so the mismatch corrupts silently rather than failing to +// load. flashinfer.aiter_utils enforces the floor at run time. +// // ABI note: aiter::mha_fwd() is called via dlsym. The struct layout here must // match the .so exactly. ck_tile::index_t = int32_t (ck_tile/core/numeric/integer.hpp). // Update this header and bump the amd-aiter version pin if AITER changes the struct. @@ -54,6 +60,9 @@ struct mha_fwd_args { const void* seqlen_k_ptr = nullptr; const void* cu_seqlen_q_ptr = nullptr; const void* cu_seqlen_k_ptr = nullptr; + // Added in 0.1.16 — inserted *before* sink_ptr, so every later field shifts. + const void* block_scale_seqstart_q_ptr = nullptr; + const void* block_scale_seqstart_k_ptr = nullptr; const void* sink_ptr = nullptr; // Dimensions (ck_tile::index_t = int32_t) @@ -84,6 +93,10 @@ struct mha_fwd_args { int32_t nhead_stride_randval = 0; int32_t nhead_stride_lse = 0; int32_t nhead_stride_o; + // Added in 0.1.16. + int32_t nhead_stride_q_descale = 0; + int32_t nhead_stride_k_descale = 0; + int32_t nhead_stride_v_descale = 0; int32_t batch_stride_q = 0; int32_t batch_stride_k = 0; @@ -92,6 +105,10 @@ struct mha_fwd_args { int32_t batch_stride_randval = 0; int32_t batch_stride_lse = 0; int32_t batch_stride_o = 0; + // Added in 0.1.16. + int32_t batch_stride_q_descale = 0; + int32_t batch_stride_k_descale = 0; + int32_t batch_stride_v_descale = 0; int32_t window_size_left = -1; int32_t window_size_right = -1; @@ -104,6 +121,11 @@ struct mha_fwd_args { // Dropout seed/offset (first variant = {0,0} when dropout disabled) std::variant, std::pair> drop_seed_offset; + + // Added in 0.1.16. Must match what AITER's own arg builders pass, which is + // an unconditional 128; the pipeline never sees 0 from AITER itself. + int32_t block_scale_size_q = 128; + int32_t block_scale_size_kv = 128; }; } // namespace aiter diff --git a/pyproject.toml b/pyproject.toml index 8104a886dd..0cad007588 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ authors = [ # Runtime dependencies # NOTE: torch is intentionally NOT listed here because ROCm users must install # torch from AMD's ROCm repository, not from PyPI. For example, -# pip install torch==2.7.1 --index-url https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/ +# pip install torch==2.9.1 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/ # See README.md for complete installation instructions. dependencies = [ "ninja", @@ -31,6 +31,7 @@ classifiers = [ "Intended Audience :: Science/Research", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.14", "Programming Language :: C++", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development :: Libraries :: Python Modules", diff --git a/tests/rocm_tests/test_aiter_version_gate_hip.py b/tests/rocm_tests/test_aiter_version_gate_hip.py new file mode 100644 index 0000000000..ab14c3185a --- /dev/null +++ b/tests/rocm_tests/test_aiter_version_gate_hip.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: Apache-2.0 +"""The amd-aiter ABI floor. + +The vendored structs under include/flashinfer/attention/aiter/ follow the 0.1.16 +layout and travel by value through dlsym'd pointers, so an older AITER shifts +field offsets instead of failing to load. Nothing downstream can detect that, +which is why the floor is enforced before routing rather than at the call. +""" + +import pathlib + +import pytest + +from flashinfer import aiter_utils + + +@pytest.mark.parametrize( + "version,supported", + [ + ("0.1.9", False), + ("0.1.10", False), # lexically ">= 0.1.16"; the pin this repo shipped before + ("0.1.15.post9", False), + ("0.1.16", True), + # The nightly wheel: PEP 440 sorts a .dev0 segment *below* its release, + # so a naive Version(...) >= Version(floor) rejects the only cp314 wheel. + ("0.1.16.post3.dev0+g620287969.d20260725", True), + ("0.1.21", True), + ("0.2.0", True), + ], +) +def test_version_floor(monkeypatch, version, supported): + monkeypatch.setattr(aiter_utils, "_aiter_installed_version", lambda: version) + assert aiter_utils._aiter_version_supported() is supported + + +def test_absent_aiter_is_not_a_version_failure(monkeypatch): + """A missing package must not be reported as an out-of-date one.""" + monkeypatch.setattr(aiter_utils, "_aiter_installed_version", lambda: None) + assert aiter_utils._aiter_version_supported() is False + + +def test_header_records_the_same_floor(): + """The vendored header's stated minimum must track AITER_MIN_VERSION.""" + header = ( + pathlib.Path(__file__).resolve().parents[2] + / "include/flashinfer/attention/aiter/mha_fwd_args.h" + ) + text = header.read_text() + assert f"amd-aiter>={aiter_utils.AITER_MIN_VERSION}" in text diff --git a/tests/rocm_tests/test_batch_prefill_kernels_hip.py b/tests/rocm_tests/test_batch_prefill_kernels_hip.py index 6f2ddad2ec..31ee9d29e7 100644 --- a/tests/rocm_tests/test_batch_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_batch_prefill_kernels_hip.py @@ -1037,3 +1037,100 @@ def _reject(*args, **kwargs): test_batch_prefill_with_ragged_kv_cache( 12, 54, 37, 8, 8, 128, True, "NONE", 0.0, False ) + + +@pytest.mark.parametrize("kv_len", [512, 2048]) +@pytest.mark.parametrize("qo_len", [37, 127]) +def test_ragged_softcap_avoids_broken_aiter_kernel(kv_len, qo_len): + """backend='auto' must stay numerically correct for causal soft-cap prefill. + + The ragged wrapper always dispatches through mha_varlen_fwd, which AITER + miscomputes when logits_soft_cap > 0; without the fallback this returns + plausible-looking values roughly 0.17 off an fp32 reference. + """ + device = torch.device("cuda:0") + if not is_aiter_supported(device) or not _aiter_ops_importable(): + pytest.skip("AITER requires a gfx942/gfx950 GPU and the aiter package") + + head_dim, num_heads, soft_cap = 128, 4, 8.0 + torch.manual_seed(0) + q = torch.randn(qo_len, num_heads, head_dim, dtype=torch.float16, device=device) + k = torch.randn(kv_len, num_heads, head_dim, dtype=torch.float16, device=device) + v = torch.randn(kv_len, num_heads, head_dim, dtype=torch.float16, device=device) + + qs, ks, vs = (t.transpose(0, 1).float() for t in (q, k, v)) + logits = soft_cap * torch.tanh( + (qs @ ks.transpose(-1, -2)) * head_dim**-0.5 / soft_cap + ) + mask = torch.ones(qo_len, kv_len, dtype=torch.bool, device=device).tril( + diagonal=kv_len - qo_len + ) + ref = ( + torch.softmax(logits.masked_fill(~mask, float("-inf")), dim=-1) @ vs + ).transpose(0, 1) + + workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device) + wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper( + workspace, "NHD", backend="auto" + ) + indptr_q = torch.tensor([0, qo_len], dtype=torch.int32, device=device) + indptr_kv = torch.tensor([0, kv_len], dtype=torch.int32, device=device) + wrapper.plan( + indptr_q, + indptr_kv, + num_heads, + num_heads, + head_dim, + causal=True, + logits_soft_cap=soft_cap, + q_data_type=torch.float16, + kv_data_type=torch.float16, + ) + torch.testing.assert_close(wrapper.run(q, k, v).float(), ref, rtol=1e-3, atol=1e-3) + + +def test_paged_softcap_guard_tracks_the_paging_route(): + """The explicit-aiter soft-cap guard must key on the route, not the shape. + + Native page sizes dispatch to mha_batch_prefill, which is exact; only + flat-gather carries the defect. A guard that ignores page_size rejects + calls that backend='auto' happily serves. + """ + device = torch.device("cuda:0") + if not is_aiter_supported(device) or not _aiter_ops_importable(): + pytest.skip("AITER requires a gfx942/gfx950 GPU and the aiter package") + native = sorted(_aiter_native_page_sizes()) + if not native: + pytest.skip("no native AITER page size on this aiter build") + + kv_len, qo_len, num_heads, head_dim, soft_cap = 512, 37, 4, 128, 8.0 + workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device) + + def plan(page_size): + num_pages = kv_len // page_size + wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + workspace, "NHD", backend="aiter" + ) + wrapper.plan( + torch.tensor([0, qo_len], dtype=torch.int32, device=device), + torch.tensor([0, num_pages], dtype=torch.int32, device=device), + torch.arange(num_pages, dtype=torch.int32, device=device), + torch.tensor([page_size], dtype=torch.int32, device=device), + num_heads, + num_heads, + head_dim, + page_size, + causal=True, + logits_soft_cap=soft_cap, + q_data_type=torch.float16, + kv_data_type=torch.float16, + ) + + plan(native[0]) + + non_native = next( + (p for p in (16, 32, 64, 8) if p not in _aiter_native_page_sizes()), None + ) + if non_native is not None: + with pytest.raises(ValueError, match="logits_soft_cap"): + plan(non_native) diff --git a/tests/rocm_tests/test_single_prefill_kernels_hip.py b/tests/rocm_tests/test_single_prefill_kernels_hip.py index 04c1838e09..52ba310cc2 100644 --- a/tests/rocm_tests/test_single_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_single_prefill_kernels_hip.py @@ -73,6 +73,18 @@ def test_single_prefill_with_kv_cache( if causal and qo_len > kv_len: pytest.skip("causal attention requires kv_len >= qo_len") + # A non-zero soft cap disables AITER's asm paths, leaving mha_varlen_fwd's + # CK kernel, which applies the cap wrongly. Non-causal is unaffected, and + # mha_batch_prefill is exact on the same inputs. + if ( + backend == "aiter" + and logits_soft_cap > 0 + and causal + and head_dim == 128 + and kv_len >= 512 + ): + pytest.skip("AITER mha_varlen_fwd soft-cap defect (aiter<=0.1.21)") + if kv_layout == "HND": k = torch.randn( num_kv_heads, kv_len, head_dim, device="cuda:0", dtype=torch.float16 @@ -316,3 +328,79 @@ def test_auto_backend_selects_aiter(head_dim, return_lse): q, k, v, causal=False, kv_layout="NHD", backend="aiter" ) torch.testing.assert_close(o_auto, o_aiter, rtol=0, atol=0) + + +# (causal, logits_soft_cap, head_dim, kv_len, expect_aiter) +_SOFTCAP_ROUTING = [ + (True, 8.0, 128, 512, False), # the defect region + (True, 8.0, 128, 2048, False), + (True, 0.0, 128, 512, True), # no cap: asm path, exact + (False, 8.0, 128, 512, True), # non-causal: exact + (True, 8.0, 64, 512, True), # other head dims unaffected + (True, 8.0, 256, 512, True), + (True, 8.0, 128, 128, True), # short kv unaffected +] + + +@pytest.mark.parametrize( + "causal,soft_cap,head_dim,kv_len,expect_aiter", _SOFTCAP_ROUTING +) +def test_auto_backend_avoids_aiter_softcap_defect( + causal, soft_cap, head_dim, kv_len, expect_aiter +): + """backend='auto' must not route the miscomputed soft-cap case to AITER. + + Guards the routing directly rather than the numerics, because the wrong + answer is silent: with AITER selected the call returns plausible values. + """ + device = torch.device("cuda:0") + if not is_aiter_supported(device) or not _aiter_ops_importable(): + pytest.skip("AITER requires a gfx942/gfx950 GPU and the aiter package") + + from flashinfer.prefill_rocm import _auto_select_prefill_backend + + chosen, reason = _auto_select_prefill_backend( + device, + dtype_q=torch.float16, + dtype_kv=torch.float16, + kv_layout="NHD", + has_custom_mask=False, + head_dim_qk=head_dim, + head_dim_vo=head_dim, + # Capability gates differ per op, so the row asserted here has to be the + # one single prefill actually consults. + op="single_prefill", + causal=causal, + logits_soft_cap=soft_cap, + kv_len=kv_len, + ) + assert chosen == ("aiter" if expect_aiter else "fa2"), reason + # Assert on the reason too, so a fallback for some unrelated cause cannot + # masquerade as the soft-cap guard working. + if not expect_aiter: + assert reason is not None and "logits_soft_cap" in reason, reason + + +def test_explicit_aiter_backend_rejects_softcap_defect(): + """An explicit backend='aiter' must fail loudly, not return wrong numbers. + + 'auto' silently falls back; asking for AITER by name is a deliberate choice, + so the defect region has to raise rather than degrade. + """ + device = torch.device("cuda:0") + if not is_aiter_supported(device) or not _aiter_ops_importable(): + pytest.skip("AITER requires a gfx942/gfx950 GPU and the aiter package") + + kv_len, qo_len, num_heads, head_dim = 512, 37, 4, 128 + q = torch.randn(qo_len, num_heads, head_dim, dtype=torch.float16, device=device) + k = torch.randn(kv_len, num_heads, head_dim, dtype=torch.float16, device=device) + v = torch.randn(kv_len, num_heads, head_dim, dtype=torch.float16, device=device) + + with pytest.raises(ValueError, match="logits_soft_cap"): + flashinfer.single_prefill_with_kv_cache( + q, k, v, causal=True, logits_soft_cap=8.0, backend="aiter" + ) + + # Same shape without the cap must still be served by AITER, so the guard is + # not quietly disabling the backend outright. + flashinfer.single_prefill_with_kv_cache(q, k, v, causal=True, backend="aiter")