From 00972521fab1405ca1032c467c9f7556a27e5eac Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sat, 22 Aug 2026 11:42:12 -0400 Subject: [PATCH 01/43] feat(rocm): support ROCm 7.14 / py3.14 / torch 2.12 with AITER 0.1.16 Moves the dev container to rocm/pytorch:rocm7.14_ubuntu26.04_py3.14_pytorch_release_2.12.0 and to amd-aiter 0.1.16.post3, which is the only AITER build published as a cp314 wheel. Four independent breakages had to be fixed to get there; the full "not slow" suite goes from 60 failures to 26 on gfx950. AITER's aiter::mha_fwd_args gained ten fields between 0.1.10 and 0.1.16, two of them inserted before sink_ptr so every later offset shifts. We pass that struct by value through a dlsym'd pointer, which does no type checking, so a stale mirror corrupted silently: 99.9% of elements wrong, no crash and no link error. The mirror is regenerated and checked against the installed aiter_meta headers (sizeof 424, offsets verified); mha_batch_prefill_args was checked too and had not drifted. torch 2.12 wraps the c10::hip compatibility namespace in `#ifdef USE_ROCM`, so COMMON_HIPCC_FLAGS now defines it. Without it every *_aiter.cu shim fails to compile on "no member named getCurrentHIPStream in namespace c10::hip", which reads like a missing include rather than a missing define. AITER moved its C++ API off torch types: 15 of 16 installed headers now take aiter_tensor_t, and silu_and_mul also gained a `limit` parameter. activation.h no longer pulls pybind11, so it is now included for real rather than forward-declared -- that is what turns the next signature change into a compile error instead of an undefined symbol at load. The at::Tensor -> aiter_tensor_t adapter deliberately includes AITER's own aiter_tensor.h for the same reason: layout drift should not be something we hand-maintain. AITER also split the rope modules by variant. "module_rope_pos_fwd" is simply not registered at 0.1.16; asking for it yields an empty source list and the JIT dies on `assert len(sources) > 0`. The entry point this shim calls now lives in module_rope_2c_cached_positions_fwd. rope.h itself still uses torch::Tensor, so that shim needed no port. Known remaining: 26 failures, all logits_soft_cap=8.0, on the mha_varlen_fwd group-mode path. cap=0 is exact (8e-4) while cap=8 is off by 0.166, so the cap is applied but wrongly. Not shipped as fixed. Co-Authored-By: Claude --- .devcontainer/rocm/Dockerfile | 132 +++++++++--------- flashinfer/compilation_context_hip.py | 10 ++ flashinfer/csrc_rocm/activation_aiter.cu | 26 ++-- flashinfer/csrc_rocm/aiter_tensor_compat.h | 70 ++++++++++ flashinfer/hip_utils.py | 2 +- flashinfer/jit/aiter_source.py | 39 ++++-- flashinfer/jit/rope.py | 10 +- .../flashinfer/attention/aiter/mha_fwd_args.h | 15 ++ 8 files changed, 214 insertions(+), 90 deletions(-) create mode 100644 flashinfer/csrc_rocm/aiter_tensor_compat.h diff --git a/.devcontainer/rocm/Dockerfile b/.devcontainer/rocm/Dockerfile index 34ac74ab58..7b762b45de 100644 --- a/.devcontainer/rocm/Dockerfile +++ b/.devcontainer/rocm/Dockerfile @@ -1,16 +1,37 @@ -ARG ROCM_VERSION=7.2 - -FROM mambaorg/micromamba:2.1.1 AS micromamba - -FROM rocm/dev-ubuntu-24.04:${ROCM_VERSION}-complete - -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 - -# Update package lists and install system dependencies +# 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 + +FROM rocm/pytorch:rocm${ROCM_VERSION}_ubuntu${UBUNTU_VERSION}_py${PY_VERSION}_pytorch_release_${TORCH_VERSION} + +# 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/ + +# 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)" + +# 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 +46,54 @@ 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. +# 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 + +# 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 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 +# Remove the default 'ubuntu' user (UID 1000) to prevent devcontainer permission +# conflicts. It still exists on Ubuntu 26.04. RUN rm -rf /home/ubuntu && if grep ubuntu:x:1000:1000 /etc/passwd >/dev/null; then userdel -f -r ubuntu; fi -# 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 /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" + +RUN pip install --no-cache-dir \ + ninja \ + "setuptools>=80" \ + "setuptools-scm>=9.2" \ + 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/flashinfer/compilation_context_hip.py b/flashinfer/compilation_context_hip.py index ee8cdd74f0..3be9e4462b 100644 --- a/flashinfer/compilation_context_hip.py +++ b/flashinfer/compilation_context_hip.py @@ -34,6 +34,16 @@ class CompilationContext: "-DFLASHINFER_ENABLE_FP8_E4M3", "-DFLASHINFER_ENABLE_FP8_E5M2", "-DHIP_ENABLE_WARP_SYNC_BUILTINS=1", + # Required from torch 2.12 on. The c10::hip / at::hip compatibility + # namespaces in c10/hip/HIPStream.h ("hipify v2 backward compat in + # external projects") are wrapped in `#ifdef USE_ROCM`. Without this, + # c10::hip::getCurrentHIPStream() -- used by every *_aiter.cu shim -- + # fails to resolve, while the namespace itself still exists via other + # headers, so the error reads "no member named ... in namespace + # 'c10::hip'" rather than a missing include. Earlier torch releases + # declared the block unconditionally, which is why this was not needed + # before. AITER's own builds already pass -DUSE_ROCM=1. + "-DUSE_ROCM=1", ] def __init__(self): diff --git a/flashinfer/csrc_rocm/activation_aiter.cu b/flashinfer/csrc_rocm/activation_aiter.cu index d4e6f28e53..4d56ab2fdd 100644 --- a/flashinfer/csrc_rocm/activation_aiter.cu +++ b/flashinfer/csrc_rocm/activation_aiter.cu @@ -9,15 +9,25 @@ #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 activation.h is included rather than forward-declared. That +// became possible at 0.1.16: the header now pulls only aiter_tensor.h, not +// , so it no longer drags in pybind11 and no longer clashes +// with FlashInfer's -DPy_LIMITED_API. +// +// Including it is what keeps this shim honest. AITER changed the signature to +// `silu_and_mul(const aiter_tensor_t&, const aiter_tensor_t&, float limit)`, and +// the old hand-written `at::Tensor&` declaration kept compiling happily and then +// failed at load with `undefined symbol`. A real declaration turns the next such +// change into a compile error instead. +#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); + const aiter_tensor_t out_a = flashinfer::aiter_compat::to_aiter(out); + const aiter_tensor_t in_a = flashinfer::aiter_compat::to_aiter(input); + // `limit` (new in 0.1.16) gates an optional clamp; 0.0f is AITER's own 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..ea5ca86447 --- /dev/null +++ b/flashinfer/csrc_rocm/aiter_tensor_compat.h @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +// SPDX-License-Identifier: Apache-2.0 +// +// at::Tensor -> aiter_tensor_t adapter. +// +// From 0.1.16 AITER's C++ API takes its own POD `aiter_tensor_t` instead of +// `at::Tensor` (15 of 16 public headers migrated). Shims therefore have to +// translate at the boundary. +// +// `aiter_tensor.h` is AITER's *real* header, and it is included deliberately +// rather than vendored: the struct layout is the part that fails silently when +// it drifts (see the mha_fwd_args incident), so taking it from AITER makes any +// future layout change a compile error instead of wrong numbers. The header is +// self-contained -- unlike rope.h and rmsnorm.h it pulls no pybind11, so it is +// safe under FlashInfer's -DPy_LIMITED_API build. +#pragma once + +#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; + 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; +} + +} // namespace flashinfer::aiter_compat diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index b74796e90f..19e633aa3f 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..f074894219 100644 --- a/flashinfer/jit/aiter_source.py +++ b/flashinfer/jit/aiter_source.py @@ -23,6 +23,7 @@ """ import functools +import inspect import os import re import shutil @@ -392,20 +393,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 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/include/flashinfer/attention/aiter/mha_fwd_args.h b/include/flashinfer/attention/aiter/mha_fwd_args.h index 9343b7ee42..6ce46e288c 100644 --- a/include/flashinfer/attention/aiter/mha_fwd_args.h +++ b/include/flashinfer/attention/aiter/mha_fwd_args.h @@ -54,6 +54,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 +87,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 +99,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 +115,10 @@ 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, at the tail of the struct. + int32_t block_scale_size_q = 0; + int32_t block_scale_size_kv = 0; }; } // namespace aiter From 193b115f1c8dff05fe664f29c54c04dca38ecd5d Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sat, 22 Aug 2026 12:50:28 -0400 Subject: [PATCH 02/43] fix(aiter): propagate the torch stream into AITER's POD entry point Two fixes from reviewing the 0.1.16 migration. silu_and_mul launched on the wrong stream. AITER's old torch-typed entry point read torch's current stream itself, but the POD API launches on aiter::getCurrentHIPStream() -- a thread_local in aiter_stream.h that defaults to nullptr and is otherwise only ever set by AITER's Python layer. The shim's OptionalHIPGuardMasqueradingAsCUDA restores the device, not the stream, so after the migration the kernel ran on the default stream while the surrounding torch ops ran on whatever stream the caller had current. That is correct on the default stream and an ordering hazard anywhere else -- CUDA graphs, multi-stream serving -- and no test on the default stream can catch it. Set the stream explicitly. block_scale_size_q/kv now default to 128 rather than 0. Every one of AITER's own arg builders (mha_fwd_kernels.cu, mha_varlen_fwd_kernels.cu, asm_mha_fwd.cu, asm_mha_varlen_fwd.cu) passes 128 unconditionally, so 0 is a value the pipeline never sees from AITER itself. This did not change any measured result, but matching the reference construction is right regardless. Co-Authored-By: Claude --- flashinfer/csrc_rocm/activation_aiter.cu | 12 ++++++++++++ include/flashinfer/attention/aiter/mha_fwd_args.h | 10 +++++++--- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/flashinfer/csrc_rocm/activation_aiter.cu b/flashinfer/csrc_rocm/activation_aiter.cu index 4d56ab2fdd..bdfcf1732b 100644 --- a/flashinfer/csrc_rocm/activation_aiter.cu +++ b/flashinfer/csrc_rocm/activation_aiter.cu @@ -20,11 +20,23 @@ // failed at load with `undefined symbol`. A real declaration turns the next such // change into a compile error instead. #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()); + + // The stream must be propagated explicitly. The old torch-typed entry point + // read torch's current stream itself; the POD API launches on + // aiter::getCurrentHIPStream(), a thread_local that defaults to nullptr and is + // otherwise only set by AITER's Python layer (aiter_stream.h). Without this the + // kernel silently runs on the default stream while the surrounding torch ops + // run on another — correct on the default stream, an ordering hazard anywhere + // else, which is exactly the case tests on the default stream cannot catch. + // OptionalHIPGuardMasqueradingAsCUDA above restores the device, not the stream. + aiter::setCurrentHIPStream(at::hip::getCurrentHIPStream()); + const aiter_tensor_t out_a = flashinfer::aiter_compat::to_aiter(out); const aiter_tensor_t in_a = flashinfer::aiter_compat::to_aiter(input); // `limit` (new in 0.1.16) gates an optional clamp; 0.0f is AITER's own default diff --git a/include/flashinfer/attention/aiter/mha_fwd_args.h b/include/flashinfer/attention/aiter/mha_fwd_args.h index 6ce46e288c..a4924a623b 100644 --- a/include/flashinfer/attention/aiter/mha_fwd_args.h +++ b/include/flashinfer/attention/aiter/mha_fwd_args.h @@ -116,9 +116,13 @@ 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, at the tail of the struct. - int32_t block_scale_size_q = 0; - int32_t block_scale_size_kv = 0; + // Added in 0.1.16, at the tail of the struct: per-block quantization block + // size. Defaulted to 128 because that is what every one of AITER's own arg + // builders passes unconditionally -- mha_fwd_kernels.cu, mha_varlen_fwd_kernels.cu, + // asm_mha_fwd.cu and asm_mha_varlen_fwd.cu all hard-code 128 -- so 0 is a + // value the pipeline never sees from AITER itself. + int32_t block_scale_size_q = 128; + int32_t block_scale_size_kv = 128; }; } // namespace aiter From ee0a0ce39c6f67e4ac49368a7ce153a0e952b816 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 24 Aug 2026 15:28:30 -0400 Subject: [PATCH 03/43] fix(rocm): restore build backend deps and harden the 0.1.16 shims Review findings against the two preceding commits on this branch. The devcontainer could not build the package it ships. The micromamba layer that was removed installed cmake and scikit-build-core>=0.4.3; the replacement pip block did not, so the documented `pip install --no-build-isolation -ve .` fails on a fresh container with ModuleNotFoundError. Dropping --no-build-isolation is not a workaround -- isolation resolves `torch >= 2.7` from PyPI and pulls a non-ROCm wheel, which is what the build-system comments and the Dockerfile's HIP guard exist to prevent. These become dead once the setuptools convergence (#291) lands. silu_and_mul gained a contiguity check. AITER's old torch-typed entry point validated its arguments itself; the POD aiter_tensor_t API has no torch wrapper left, and to_aiter passes strides through faithfully to a kernel that indexes linearly. A sliced or transposed input previously aborted on TORCH_CHECK and now returned silently wrong values. Tests only ever pass freshly-allocated contiguous tensors, so nothing covered it. Also: the aiter_source docstring still described the forward-declaration strategy that activation_aiter.cu just abandoned -- it is the one place a contributor reads to learn the shim convention, so it taught the pattern that caused the undefined-symbol failure. Comment blocks in the touched files are trimmed to the conclusion per the repo comment cap; the evidence they carried (the four AITER arg builders that hard-code block_scale_size 128, the torch 2.12 USE_ROCM gating) is here instead. Verified: silu_and_mul(..., float limit = 0.0f) is AITER's declared default in activation.h, so passing 0.0f preserves prior behaviour rather than clamping. --- .devcontainer/rocm/Dockerfile | 6 ++++ flashinfer/compilation_context_hip.py | 12 ++----- flashinfer/csrc_rocm/activation_aiter.cu | 35 ++++++++----------- flashinfer/csrc_rocm/aiter_tensor_compat.h | 18 ++++------ flashinfer/jit/aiter_source.py | 9 ++--- .../flashinfer/attention/aiter/mha_fwd_args.h | 7 ++-- 6 files changed, 37 insertions(+), 50 deletions(-) diff --git a/.devcontainer/rocm/Dockerfile b/.devcontainer/rocm/Dockerfile index 7b762b45de..4403485deb 100644 --- a/.devcontainer/rocm/Dockerfile +++ b/.devcontainer/rocm/Dockerfile @@ -77,8 +77,14 @@ RUN rm -rf /home/ubuntu && if grep ubuntu:x:1000:1000 /etc/passwd >/dev/null; th ENV VIRTUAL_ENV=/opt/venv ENV PATH="$VIRTUAL_ENV/bin:/usr/lib/llvm-19/bin:$PATH" +# cmake and scikit-build-core are build backend deps. They must be present in +# the image because the documented install uses --no-build-isolation, which is +# itself required: isolation would resolve `torch >= 2.7` from PyPI and pull a +# non-ROCm wheel. RUN pip install --no-cache-dir \ ninja \ + cmake \ + "scikit-build-core>=0.4.3" \ "setuptools>=80" \ "setuptools-scm>=9.2" \ pre-commit \ diff --git a/flashinfer/compilation_context_hip.py b/flashinfer/compilation_context_hip.py index 3be9e4462b..f6f461ccdf 100644 --- a/flashinfer/compilation_context_hip.py +++ b/flashinfer/compilation_context_hip.py @@ -34,15 +34,9 @@ class CompilationContext: "-DFLASHINFER_ENABLE_FP8_E4M3", "-DFLASHINFER_ENABLE_FP8_E5M2", "-DHIP_ENABLE_WARP_SYNC_BUILTINS=1", - # Required from torch 2.12 on. The c10::hip / at::hip compatibility - # namespaces in c10/hip/HIPStream.h ("hipify v2 backward compat in - # external projects") are wrapped in `#ifdef USE_ROCM`. Without this, - # c10::hip::getCurrentHIPStream() -- used by every *_aiter.cu shim -- - # fails to resolve, while the namespace itself still exists via other - # headers, so the error reads "no member named ... in namespace - # 'c10::hip'" rather than a missing include. Earlier torch releases - # declared the block unconditionally, which is why this was not needed - # before. AITER's own builds already pass -DUSE_ROCM=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", ] diff --git a/flashinfer/csrc_rocm/activation_aiter.cu b/flashinfer/csrc_rocm/activation_aiter.cu index bdfcf1732b..73ec53eeaf 100644 --- a/flashinfer/csrc_rocm/activation_aiter.cu +++ b/flashinfer/csrc_rocm/activation_aiter.cu @@ -9,16 +9,10 @@ #include #include -// AITER's real activation.h is included rather than forward-declared. That -// became possible at 0.1.16: the header now pulls only aiter_tensor.h, not -// , so it no longer drags in pybind11 and no longer clashes -// with FlashInfer's -DPy_LIMITED_API. -// -// Including it is what keeps this shim honest. AITER changed the signature to -// `silu_and_mul(const aiter_tensor_t&, const aiter_tensor_t&, float limit)`, and -// the old hand-written `at::Tensor&` declaration kept compiling happily and then -// failed at load with `undefined symbol`. A real declaration turns the next such -// change into a compile error instead. +// 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 @@ -27,19 +21,20 @@ void silu_and_mul_aiter(at::Tensor out, at::Tensor input) { const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(input.device()); - // The stream must be propagated explicitly. The old torch-typed entry point - // read torch's current stream itself; the POD API launches on - // aiter::getCurrentHIPStream(), a thread_local that defaults to nullptr and is - // otherwise only set by AITER's Python layer (aiter_stream.h). Without this the - // kernel silently runs on the default stream while the surrounding torch ops - // run on another — correct on the default stream, an ordering hazard anywhere - // else, which is exactly the case tests on the default stream cannot catch. - // OptionalHIPGuardMasqueradingAsCUDA above restores the device, not the stream. + // The POD API launches on aiter::getCurrentHIPStream(), a thread_local that + // defaults to nullptr and is otherwise set only by AITER's Python layer; the + // old torch-typed entry point read torch's stream itself. The device guard + // above restores the device, not the stream. aiter::setCurrentHIPStream(at::hip::getCurrentHIPStream()); + // 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); - // `limit` (new in 0.1.16) gates an optional clamp; 0.0f is AITER's own default - // and preserves the previous behaviour. + // `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 index ea5ca86447..639873f9d5 100644 --- a/flashinfer/csrc_rocm/aiter_tensor_compat.h +++ b/flashinfer/csrc_rocm/aiter_tensor_compat.h @@ -1,22 +1,14 @@ // SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 // -// at::Tensor -> aiter_tensor_t adapter. +// at::Tensor -> aiter_tensor_t adapter, for AITER's POD C++ API (0.1.16+). // -// From 0.1.16 AITER's C++ API takes its own POD `aiter_tensor_t` instead of -// `at::Tensor` (15 of 16 public headers migrated). Shims therefore have to -// translate at the boundary. -// -// `aiter_tensor.h` is AITER's *real* header, and it is included deliberately -// rather than vendored: the struct layout is the part that fails silently when -// it drifts (see the mha_fwd_args incident), so taking it from AITER makes any -// future layout change a compile error instead of wrong numbers. The header is -// self-contained -- unlike rope.h and rmsnorm.h it pulls no pybind11, so it is -// safe under FlashInfer's -DPy_LIMITED_API build. +// 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 @@ -30,6 +22,8 @@ inline AiterDtype to_aiter_dtype(at::ScalarType t) { 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; diff --git a/flashinfer/jit/aiter_source.py b/flashinfer/jit/aiter_source.py index f074894219..4a5479e129 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 diff --git a/include/flashinfer/attention/aiter/mha_fwd_args.h b/include/flashinfer/attention/aiter/mha_fwd_args.h index a4924a623b..a0b74a9ca6 100644 --- a/include/flashinfer/attention/aiter/mha_fwd_args.h +++ b/include/flashinfer/attention/aiter/mha_fwd_args.h @@ -116,11 +116,8 @@ 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, at the tail of the struct: per-block quantization block - // size. Defaulted to 128 because that is what every one of AITER's own arg - // builders passes unconditionally -- mha_fwd_kernels.cu, mha_varlen_fwd_kernels.cu, - // asm_mha_fwd.cu and asm_mha_varlen_fwd.cu all hard-code 128 -- so 0 is a - // value the pipeline never sees from AITER itself. + // 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; }; From 516b866d8f71f23c0343a67448d5945ab27df2c0 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 24 Aug 2026 15:28:48 -0400 Subject: [PATCH 04/43] fix(aiter): route causal soft-cap away from AITER's miscomputing kernel AITER's mha_varlen_fwd applies logits_soft_cap incorrectly. A non-zero cap disables every AITER assembly path (aiter/ops/mha.py gates can_impl_fmha_v3_fwd and friends on logits_soft_cap == 0.0), leaving the CK _logits kernel, which is off by ~0.17 against an fp32 reference while cap=0 is exact to 8e-4. The defect is AITER's, not ours: a probe calling aiter.ops.mha.mha_varlen_fwd directly, with no FlashInfer in the call path, reproduces it on all 10 failing shapes, and mha_batch_prefill -- the entry point AITER's own op_tests cover at logits_soft_cap=[0.0, 30.0] -- is tanh-exact on the same inputs. Scored against both CK soft-cap formulas (tanh and the softsign variant selected by CK_TILE_ATTENTION_LOGITS_SOFT_CAP_DEFAULT) it matches neither, so it is not a build-flag mismatch. Our mha_fwd_args is field-equivalent to AITER's own builder, including mask_type=2 (bottom-right), so there is nothing to fix on the caller side. No fix commits exist through v0.1.21. Scope is narrow and measured, not assumed: causal only (test_logits_cap_hip.py exercises non-causal soft-cap at 30.0/50.0 across head_dim 128/256 and passes), head_dim=128, kv_len >= 512. Guarding only the test would have left the default path silently wrong -- backend="auto" is the default at five public entry points, and Gemma-2/Grok shapes (causal, hd128, cap 30/50, long context) sit squarely in the defect region. _auto_select_prefill_backend now falls back to fa2 with the same one-time warning it uses for its other AITER constraints, and explicit backend="aiter" warns rather than returning plausible-looking wrong numbers. The routing is tested rather than the numerics, since the failure is silent. A/B: 7/7 cases pass with the guard; neutering _aiter_softcap_defect fails exactly the two defect-region cases and nothing else. _AITER_LAST_VALIDATED also corrected to the full pinned string. Under PEP 440 "0.1.16.post3.dev0+g..." sorts BELOW "0.1.16.post3", so naming the release claimed validation for builds never exercised -- the exact silent widening of the support boundary the adjacent comment warns against. --- flashinfer/prefill_rocm.py | 44 ++++++++++++++- .../test_single_prefill_kernels_hip.py | 56 +++++++++++++++++++ 2 files changed, 99 insertions(+), 1 deletion(-) diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index 618074b500..aabf0b2b3e 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -20,6 +20,7 @@ import math import os import threading +import warnings from importlib.metadata import PackageNotFoundError from types import SimpleNamespace from typing import Any, Dict, List, Literal, Optional, Tuple, Union, overload @@ -62,7 +63,7 @@ # 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" @functools.cache @@ -329,6 +330,24 @@ 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: 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. Present through at least aiter 0.1.21. + """ + if not (causal and logits_soft_cap and logits_soft_cap > 0): + return False + if head_dim != 128: + return False + # kv_len is unknown at plan() time for some wrappers; assume the worst, + # since a wrong answer costs more than the fa2 slowdown. + return kv_len is None or kv_len >= 512 + + def _auto_select_prefill_backend( device: torch.device, *, @@ -340,6 +359,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. @@ -377,6 +399,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) @@ -1516,12 +1543,27 @@ 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=k.shape[0] if kv_layout == "NHD" else k.shape[1], ) if backend == "aiter": # Outside the probe on purpose: this raises ArchCapabilityError, which # gates known-bad toolchains and must never be demoted to a silent fa2. _require_aiter_runtime(q.device, "single_prefill") + if _aiter_softcap_defect( + causal, + logits_soft_cap, + q.shape[-1], + k.shape[0] if kv_layout == "NHD" else k.shape[1], + ): + warnings.warn( + "AITER computes logits_soft_cap incorrectly for causal " + "head_dim=128 prefill with kv_len >= 512; results will be wrong. " + "Use backend='fa2' or backend='auto'.", + stacklevel=2, + ) if pos_encoding_mode != "NONE": raise ValueError( f"AITER backend does not support pos_encoding_mode={pos_encoding_mode!r}; " diff --git a/tests/rocm_tests/test_single_prefill_kernels_hip.py b/tests/rocm_tests/test_single_prefill_kernels_hip.py index 04c1838e09..403b986ac3 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,47 @@ 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 = _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, + causal=causal, + logits_soft_cap=soft_cap, + kv_len=kv_len, + ) + assert chosen == ("aiter" if expect_aiter else "fa2") From 730a7e83b6966f14677e5dc8b059d9074266c25c Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 24 Aug 2026 15:29:03 -0400 Subject: [PATCH 05/43] docs(rocm): describe the ROCm 7.14 / py3.14 / torch 2.12 stack The Dockerfile moved to ROCm 7.14 / Ubuntu 26.04 / Python 3.14 / torch 2.12 and AITER 0.1.16.post3; every doc still described 7.2 / py3.12 / torch 2.9.1 / AITER 0.1.10. The build-arg block was the most acutely wrong -- it documents `docker build` against .devcontainer/rocm/Dockerfile and every default in it had changed underneath, including two args (AITER_VERSION, AITER_INDEX) that were undocumented entirely. ROCm 7.14 needs its own wording rather than a version bump: repo.radeon.com publishes no rocm-rel-7.14/ directory, so the `pip install torch -f ...` recipe has no valid form for it and torch must come from the base image. CLAUDE.md also claimed torch installs via --index-url while README said the opposite; README was right (the radeon repo is a flat listing, not a PEP 503 index). The AITER section now carries both pins and says why they differ -- the CI image stays on 0.1.10 from pypi.amd.com, and the devcontainer takes the only cp314 wheel that exists, from the vllm-cdna nightlies, spelled out in full because pip will not select a local version from a loose specifier. The prior claim that both images install the same version is no longer true. Known Limitations gains the AITER soft-cap defect, scoped to causal since non-causal soft-capped prefill is exercised and correct. Left alone deliberately: the published rocm/flashinfer image table and the micromamba note under it (that image is still 7.2 and still micromamba-based), docker/Dockerfile.rocm_ci (still 7.1.1, and its parameterised -f install cannot reach 7.14), and the "since amd-aiter >= 0.1.10" feature floors, which document history rather than the current stack. --- .claude/skills/debug-rocm-crash/SKILL.md | 2 +- CLAUDE.md | 43 ++++++++++++++++-------- README.md | 11 ++++-- amd-flashinfer-jit-cache/pyproject.toml | 5 +-- pyproject.toml | 3 +- 5 files changed, 43 insertions(+), 21 deletions(-) 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/CLAUDE.md b/CLAUDE.md index 9e8eae875e..b43be34971 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -77,6 +77,10 @@ ranks higher than a same-version PyPI wheel. Supported versions are in the [Supported hardware and toolchain](README.md#supported-hardware-and-toolchain) table in `README.md`. +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 @@ -101,26 +105,37 @@ 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 +`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`, and [`docs/rocm/backends.md`](docs/rocm/backends.md) 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. -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. +`aiter_utils.AITER_MIN_VERSION` is the hard floor below which the vendored +struct layouts stop matching. 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 fedf30ac3a..9a91eccd61 100644 --- a/README.md +++ b/README.md @@ -98,14 +98,19 @@ python examples/single_prefill_example.py | | Supported | | :--- | :--- | | GPUs | gfx942 (CDNA3 — MI300X, MI325X), gfx950 (CDNA4 — MI350X, MI355X) | -| ROCm | 7.0.2, 7.1.1, 7.2 | -| PyTorch+ROCm | 2.8.0, 2.9.1 | -| Python | 3.10+ (the published images and devcontainer use 3.12) | +| ROCm | 7.0.2, 7.1.1, 7.2, 7.14 | +| PyTorch+ROCm | 2.8.0, 2.9.1, 2.12.0 | +| Python | 3.10+ (the published images use 3.12; the devcontainer uses 3.14) | Other versions may work but are untested. Replace `7.2` in the torch install command with the ROCm version you need; see for what is available. +ROCm 7.14 is the exception: `repo.radeon.com` publishes no `rocm-rel-7.14/` +directory, so there is no pip recipe for it. Take torch from the +`rocm/pytorch:rocm7.14_ubuntu26.04_py3.14_pytorch_release_2.12.0` image +instead, as the devcontainer does. + ## Support matrix Every op has an in-tree HIP kernel unless noted; a subset also has an 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/pyproject.toml b/pyproject.toml index e141a65857..bdff3d2a29 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", @@ -36,6 +36,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", From 9b7fe67ca8dff3e749b8a4921c532b840376ef3c Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 24 Aug 2026 15:46:16 -0400 Subject: [PATCH 06/43] docs(rocm): pass the host's numeric render GID, not the group name Found by building the devcontainer image and running it as the devuser: the README run command fails with "No CUDA GPUs are available". The old base image shipped a `render` group at the host's conventional GID, so `--group-add render` worked. The rewritten Dockerfile does `groupadd -r render` when the base has none, which takes an arbitrary free system GID -- 995 here against a host device owned by GID 109 -- and docker resolves the *name* against the image, so the container user never gets access to /dev/dri/renderD*. The Dockerfile comment already noted callers must pass the numeric GID; the docs were never updated to match. Missed until now because every test so far ran as root, which bypasses the group check entirely. Verified after the change: pip install --no-build-isolation -ve . succeeds in the image, torch reports MI350X, and import flashinfer loads. --- CONTRIBUTING.md | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 37c2958f37..c86ca65353 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -16,12 +16,15 @@ building from source and everything else specific to contributing code. Build the development image with the repository's Dockerfile: ```bash -docker build -t flashinfer-dev:rocm7.2 -f .devcontainer/rocm/Dockerfile . +docker build -t flashinfer-dev:rocm7.14 -f .devcontainer/rocm/Dockerfile . ``` -`ROCM_VERSION`, `PY_VERSION`, and `TORCH_VERSION` default to 7.2, 3.12, and -2.9.1; override with `--build-arg` if you need a different combination. Pass -`--build-arg USERNAME=$USER --build-arg USER_UID=$(id -u) --build-arg +`ROCM_VERSION`, `UBUNTU_VERSION`, `PY_VERSION`, and `TORCH_VERSION` default to +7.14, 26.04, 3.14, and 2.12.0. They select the `rocm/pytorch` base image tag, +so they are not independent knobs — any override has to name a tag that exists +on Docker Hub. `AITER_VERSION` and `AITER_INDEX` pin the AITER wheel; the +default is the vllm-cdna nightly, which is the only cp314 build published. +Pass `--build-arg USERNAME=$USER --build-arg USER_UID=$(id -u) --build-arg USER_GID=$(id -g)` to match container file ownership to your host user — without them, build artifacts come out root-owned. @@ -30,12 +33,17 @@ docker run -it \ --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ --privileged --ipc=host --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-dev:rocm7.2 + flashinfer-dev:rocm7.14 ``` +`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, which surfaces as "No CUDA GPUs are available". + # Building and Installing **Editable install** — what you want for day-to-day work: From 4ff8bcc883ff8a8336d06b11bf81ff47668c334f Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 24 Aug 2026 17:56:09 -0400 Subject: [PATCH 07/43] fix(aiter): extend the soft-cap fallback to the batch prefill wrappers The previous commit guarded only single prefill, which left the production path exposed: BatchPrefillWithRaggedKVCacheWrapper always dispatches through mha_varlen_fwd, so it carries the same defect, and it is what a serving stack actually calls. Measured on gfx950, ragged at logits_soft_cap=8.0 is off by 0.0595-0.1915 against an fp32 reference -- the same magnitudes single prefill produced, from the same kernel -- while cap=0 and fa2 are exact. That also corrects an earlier claim of mine that batch prefill was merely untested at cap>0. It is untested (the suite parametrizes logits_soft_cap=[0.0] only, which is why this went unnoticed) and it is broken. Paged is guarded only where the defect can actually occur. Its native-paging route calls mha_batch_prefill, which is exact at cap=8 on all 10 shapes measured, so guarding it unconditionally would trade correct fast kernels for slower ones. A page size outside the native set forces flat-gather, which does go through mha_varlen_fwd, and that is the case ruled out at plan() time. One residual gap: when native paging is claimed but the runtime probe later falls back to flat-gather, the call escapes the guard. Closing that means moving the probe ahead of backend selection, which is a larger change than this fix. The new test drives the ragged wrapper end to end rather than asserting on routing, because routing is not the property that matters. A/B: 4/4 pass with the guard, 4/4 fail without it. Two review suggestions declined. Version-gating the guard on _AITER_SOFTCAP_DEFECT_THROUGH would auto-expire it on a nightly AITER bump and silently re-enable a wrong-answer path; deoptimizing until someone re-measures is the safer failure. And the test skip is not over-broad: it drops 28 cases, 26 of which are the measured failures, not 1. --- flashinfer/csrc_rocm/activation_aiter.cu | 11 ++-- flashinfer/prefill_rocm.py | 60 ++++++++++++------- .../test_batch_prefill_kernels_hip.py | 50 ++++++++++++++++ 3 files changed, 96 insertions(+), 25 deletions(-) diff --git a/flashinfer/csrc_rocm/activation_aiter.cu b/flashinfer/csrc_rocm/activation_aiter.cu index 73ec53eeaf..97b7422d19 100644 --- a/flashinfer/csrc_rocm/activation_aiter.cu +++ b/flashinfer/csrc_rocm/activation_aiter.cu @@ -21,17 +21,18 @@ void silu_and_mul_aiter(at::Tensor out, at::Tensor input) { const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(input.device()); + // 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. Checked + // before the stream is set, so a rejected call leaves no thread_local behind. + TORCH_CHECK(input.is_contiguous(), "silu_and_mul: input must be contiguous"); + TORCH_CHECK(out.is_contiguous(), "silu_and_mul: out must be contiguous"); + // The POD API launches on aiter::getCurrentHIPStream(), a thread_local that // defaults to nullptr and is otherwise set only by AITER's Python layer; the // old torch-typed entry point read torch's stream itself. The device guard // above restores the device, not the stream. aiter::setCurrentHIPStream(at::hip::getCurrentHIPStream()); - // 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); // `limit` (new in 0.1.16) gates an optional clamp; 0.0f is AITER's declared diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index aabf0b2b3e..d7cdb767e7 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -20,7 +20,6 @@ import math import os import threading -import warnings from importlib.metadata import PackageNotFoundError from types import SimpleNamespace from typing import Any, Dict, List, Literal, Optional, Tuple, Union, overload @@ -64,6 +63,9 @@ # bumping it must not silently move the support boundary. _AITER_NATIVE_PAGING_SINCE = "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 @@ -331,21 +333,25 @@ def _require_aiter_runtime(device: torch.device, op: str = "batch_prefill") -> N def _aiter_softcap_defect( - causal: bool, logits_soft_cap: float, head_dim: int, kv_len: Optional[int] + 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. Present through at least aiter 0.1.21. + 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 - if head_dim != 128: - return False - # kv_len is unknown at plan() time for some wrappers; assume the worst, - # since a wrong answer costs more than the fa2 slowdown. - return kv_len is None or kv_len >= 512 + return head_dim == 128 and kv_len is not None and kv_len >= 512 def _auto_select_prefill_backend( @@ -1532,6 +1538,8 @@ def single_prefill_with_kv_cache( scale_v = torch.ones(v.shape[1], dtype=torch.float32, device=q.device) resolved_from_auto = backend == "auto" + kv_len = k.shape[0] if kv_layout == "NHD" else k.shape[1] + if backend == "auto": backend, _ = _auto_select_prefill_backend( q.device, @@ -1545,24 +1553,18 @@ def single_prefill_with_kv_cache( op="single_prefill", causal=causal, logits_soft_cap=logits_soft_cap, - kv_len=k.shape[0] if kv_layout == "NHD" else k.shape[1], + kv_len=kv_len, ) if backend == "aiter": # Outside the probe on purpose: this raises ArchCapabilityError, which # gates known-bad toolchains and must never be demoted to a silent fa2. _require_aiter_runtime(q.device, "single_prefill") - if _aiter_softcap_defect( - causal, - logits_soft_cap, - q.shape[-1], - k.shape[0] if kv_layout == "NHD" else k.shape[1], - ): - warnings.warn( - "AITER computes logits_soft_cap incorrectly for causal " - "head_dim=128 prefill with kv_len >= 512; results will be wrong. " - "Use backend='fa2' or backend='auto'.", - stacklevel=2, + 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( @@ -2295,6 +2297,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, + # 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. + kv_len=( + None + if page_size in _aiter_native_page_sizes() + else self._max_kv_len + ), ) ) if self._backend == "aiter" and pos_encoding_mode != "NONE": @@ -3295,6 +3310,11 @@ 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 pos_encoding_mode != "NONE": diff --git a/tests/rocm_tests/test_batch_prefill_kernels_hip.py b/tests/rocm_tests/test_batch_prefill_kernels_hip.py index 6f2ddad2ec..51b834d4ef 100644 --- a/tests/rocm_tests/test_batch_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_batch_prefill_kernels_hip.py @@ -1037,3 +1037,53 @@ 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) From 89f9835f4431be20fdfbd739b1446f4375a6dfd0 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 24 Aug 2026 20:16:41 -0400 Subject: [PATCH 08/43] fix(rocm): drop the scikit-build-core deps the setuptools switch made dead Rebasing onto amd-integration picked up #291, which replaced scikit-build-core + CMake with the in-tree backend. The image no longer needs cmake or scikit-build-core; it needs packaging>=24, which the new build-system requires and the pip block did not install. The commit that added them was written against a base where #291 had not landed, and was correct there. --- .devcontainer/rocm/Dockerfile | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/.devcontainer/rocm/Dockerfile b/.devcontainer/rocm/Dockerfile index 4403485deb..eb62f14827 100644 --- a/.devcontainer/rocm/Dockerfile +++ b/.devcontainer/rocm/Dockerfile @@ -77,16 +77,15 @@ RUN rm -rf /home/ubuntu && if grep ubuntu:x:1000:1000 /etc/passwd >/dev/null; th ENV VIRTUAL_ENV=/opt/venv ENV PATH="$VIRTUAL_ENV/bin:/usr/lib/llvm-19/bin:$PATH" -# cmake and scikit-build-core are build backend deps. They must be present in -# the image because the documented install uses --no-build-isolation, which is -# itself required: isolation would resolve `torch >= 2.7` from PyPI and pull a -# non-ROCm wheel. +# 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 \ - cmake \ - "scikit-build-core>=0.4.3" \ "setuptools>=80" \ "setuptools-scm>=9.2" \ + "packaging>=24" \ pre-commit \ numpy \ pytest pytest-cov pytest-xdist pytest-rerunfailures \ From 572284400f38d8376ba52a9b2d081f0421c205a8 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 24 Aug 2026 21:11:44 -0400 Subject: [PATCH 09/43] fix(rocm): remove the default ubuntu user before creating devuser Copilot review, reproduced: with USER_UID=1000 the build dies at step 4 with `groupadd: GID '1000' already exists`. Ubuntu 26.04's base image ships `ubuntu:x:1000:1000`, and the userdel that clears it ran twenty lines *after* the groupadd that needs the ID free. This is the common case, not a corner: README tells users to pass $(id -u), and a single-user host is almost always 1000. It went unnoticed because every build I ran passed a six-digit domain UID. Also from the same review: use the canonical `amd-aiter` spelling everywhere rather than mixing in `amd_aiter`, matching the importlib.metadata.version("amd-aiter") lookups; and cover the explicit backend="aiter" raise, which had no test -- only the auto-fallback did. --- .devcontainer/rocm/Dockerfile | 13 ++++++---- CLAUDE.md | 2 +- .../test_single_prefill_kernels_hip.py | 25 +++++++++++++++++++ 3 files changed, 34 insertions(+), 6 deletions(-) diff --git a/.devcontainer/rocm/Dockerfile b/.devcontainer/rocm/Dockerfile index eb62f14827..0c82681977 100644 --- a/.devcontainer/rocm/Dockerfile +++ b/.devcontainer/rocm/Dockerfile @@ -50,6 +50,13 @@ ARG USERNAME=devuser ARG USER_UID=1003 ARG USER_GID=$USER_UID +# 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 && \ @@ -67,10 +74,6 @@ RUN (getent group render >/dev/null || groupadd -r render) && \ RUN echo "set-option -g default-command \"/bin/bash -i\"" >> /home/$USERNAME/.tmux.conf -# Remove the default 'ubuntu' user (UID 1000) to prevent devcontainer permission -# conflicts. It still exists on Ubuntu 26.04. -RUN rm -rf /home/ubuntu && if grep ubuntu:x:1000:1000 /etc/passwd >/dev/null; then userdel -f -r ubuntu; fi - # 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. @@ -92,7 +95,7 @@ RUN pip install --no-cache-dir \ pybind11 \ ruff \ filelock && \ - pip install --no-cache-dir "amd_aiter==${AITER_VERSION}" --extra-index-url "${AITER_INDEX}" && \ + 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. diff --git a/CLAUDE.md b/CLAUDE.md index b43be34971..f126c420cd 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -126,7 +126,7 @@ 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 \ +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/ ``` diff --git a/tests/rocm_tests/test_single_prefill_kernels_hip.py b/tests/rocm_tests/test_single_prefill_kernels_hip.py index 403b986ac3..85e5182963 100644 --- a/tests/rocm_tests/test_single_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_single_prefill_kernels_hip.py @@ -372,3 +372,28 @@ def test_auto_backend_avoids_aiter_softcap_defect( kv_len=kv_len, ) assert chosen == ("aiter" if expect_aiter else "fa2") + + +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") From 9e2d1dd3275025c12ab23ef8738a9a0d07aa7def Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 25 Aug 2026 07:45:46 -0400 Subject: [PATCH 10/43] fix(rocm): graft ROCm headers and dev symlinks into the devcontainer The devcontainer could not compile a single kernel. Every JIT build died on `/opt/rocm/bin/amdclang++: not found`, which took the whole rocm_tests suite with it -- 13079 failures and 15000+ errors, none of them code. This is a regression from the base image swap, not a pre-existing gap. The old `rocm/dev-ubuntu-*:complete` base shipped a full /opt/rocm; `rocm/pytorch:*` carries the pip ROCm SDK instead, which is runtime-only. Three things are missing from it, and all three are needed: - math/utility headers (no thrust, rocPRIM, rocBLAS, hipBLASLt), while torch 2.12's complex.h includes thrust/complex.h -- grafted from a rocm/dev-ubuntu donor stage; - unversioned .so names, so -lamdhip64 cannot resolve against libamdhip64.so.7 alone; - an `amdclang++`, which flashinfer invokes as $ROCM_HOME/bin/amdclang++; the SDK ships bin/hipcc and lib/llvm/bin/clang++, and puts device bitcode at lib/llvm/amdgcn rather than $ROCM_PATH/amdgcn. ROCM_PATH points at the SDK root and not at the grafted header tree, because hipcc resolves its own clang through $ROCM_PATH/lib/llvm/bin. I had diagnosed all of this while gating the base image and closed it in a throwaway build layer, then never folded it into the Dockerfile this PR ships -- so every green test run since was in an image with capabilities the deliverable lacked. Caught by running the full suite in the real image. Verified in the built image at USER_UID=1000: amdclang++ resolves, and a fa2 single-prefill kernel JIT-compiles from a cold cache and runs on gfx950. --- .devcontainer/rocm/Dockerfile | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/.devcontainer/rocm/Dockerfile b/.devcontainer/rocm/Dockerfile index 0c82681977..3699efebbf 100644 --- a/.devcontainer/rocm/Dockerfile +++ b/.devcontainer/rocm/Dockerfile @@ -19,7 +19,15 @@ 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 # 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 @@ -80,6 +88,33 @@ RUN echo "set-option -g default-command \"/bin/bash -i\"" >> /home/$USERNAME/.tm 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" +# 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" +ENV LIBRARY_PATH="$ROCM_SDK/lib" +ENV 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 From 495378ee253d99c84597deafe371fd30be4235b6 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 25 Aug 2026 09:09:41 -0400 Subject: [PATCH 11/43] test(rocm): assert single-prefill routing against the single-prefill row Copilot review: the routing test called _auto_select_prefill_backend without op=, taking the default op="batch_prefill" while asserting single-prefill behaviour. The capability gates are not uniform across ops -- that is the reason the parameter exists -- so the test was passing against the wrong row. --- tests/rocm_tests/test_single_prefill_kernels_hip.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/rocm_tests/test_single_prefill_kernels_hip.py b/tests/rocm_tests/test_single_prefill_kernels_hip.py index 85e5182963..650ca82de9 100644 --- a/tests/rocm_tests/test_single_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_single_prefill_kernels_hip.py @@ -367,6 +367,9 @@ def test_auto_backend_avoids_aiter_softcap_defect( 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, From 348e019f3750b5dff527ffbc486c5966d5531082 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 25 Aug 2026 09:47:03 -0400 Subject: [PATCH 12/43] fix(aiter): scope the AITER stream override to the call Copilot review. aiter::setCurrentHIPStream writes a thread_local that AITER never clears, so the value outlived silu_and_mul: a caller inside a temporary torch.cuda.Stream would strand a freed handle for the next AITER POD call on that thread, and any such call that does not set the stream itself would silently inherit ours. Low risk today -- this is flashinfer's only POD caller and AITER's Python layer always sets its own stream -- but the fix is a save/restore guard, so there is no reason to leave it. aiter_stream.h exposes getCurrentHIPStream(), so the previous value is readable. The tensor conversions now happen before the guard is taken: to_aiter can throw on an unsupported dtype, and there is no point entering the scope for a call that will not launch. --- flashinfer/csrc_rocm/activation_aiter.cu | 35 ++++++++++++++++++------ 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/flashinfer/csrc_rocm/activation_aiter.cu b/flashinfer/csrc_rocm/activation_aiter.cu index 97b7422d19..bd531ad093 100644 --- a/flashinfer/csrc_rocm/activation_aiter.cu +++ b/flashinfer/csrc_rocm/activation_aiter.cu @@ -18,23 +18,42 @@ #include "aiter_tensor_compat.h" +namespace { + +// The POD API launches on aiter::getCurrentHIPStream(), a thread_local that +// defaults to nullptr and is otherwise set only by AITER's Python layer; the +// old torch-typed entry point read torch's stream itself. 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 AiterStreamGuard { + public: + explicit AiterStreamGuard(hipStream_t stream) : prev_(aiter::getCurrentHIPStream()) { + aiter::setCurrentHIPStream(stream); + } + ~AiterStreamGuard() { aiter::setCurrentHIPStream(prev_); } + + AiterStreamGuard(const AiterStreamGuard&) = delete; + AiterStreamGuard& operator=(const AiterStreamGuard&) = delete; + + private: + hipStream_t prev_; +}; + +} // namespace + void silu_and_mul_aiter(at::Tensor out, at::Tensor input) { const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(input.device()); // 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. Checked - // before the stream is set, so a rejected call leaves no thread_local behind. + // 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"); - // The POD API launches on aiter::getCurrentHIPStream(), a thread_local that - // defaults to nullptr and is otherwise set only by AITER's Python layer; the - // old torch-typed entry point read torch's stream itself. The device guard - // above restores the device, not the stream. - aiter::setCurrentHIPStream(at::hip::getCurrentHIPStream()); - 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 AiterStreamGuard 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); From 02dd016a94b6b45650a38f7669a679e1bc85b364 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 25 Aug 2026 12:11:42 -0400 Subject: [PATCH 13/43] fix(aiter): set GPU_ARCHS for AITER's own JIT, not just our shim build AITER 0.1.16 asserts on GPU_ARCHS in its JIT: unset reaches its validator as [''] and every build dies with "One of GPU archs of [''] is invalid or not supported". _build_aiter_lib sets it for the shim build and restores it after, so anything using AITER's *Python* ops -- batch decode, paged-append, fused MoE -- builds with it unset. Measured on gfx950: test_batch_decode_aiter_hip.py alone goes 108 failed -> 0, and passing GPU_ARCHS=gfx950 into the container has the identical effect, which is what identified the variable. Across the suite this accounted for 303 of the failures, spread over decode, POD, paged-append and fused MoE -- every AITER surface except the three shims this branch already ported. Fixed in the library rather than the devcontainer: the requirement comes from the AITER version, so a Dockerfile ENV would fix this image and leave every other install broken. Only fills a missing value, so an operator-set GPU_ARCHS still wins. Found only because the full suite was run against the real devcontainer image. The earlier runs used a throwaway gate image and stopped at the prefill failures, so the rest of the AITER surface was never exercised under 0.1.16. --- flashinfer/aiter_utils.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/flashinfer/aiter_utils.py b/flashinfer/aiter_utils.py index 4d5a4b83b5..6b2052f461 100644 --- a/flashinfer/aiter_utils.py +++ b/flashinfer/aiter_utils.py @@ -30,6 +30,25 @@ def is_aiter_supported(device: torch.device) -> bool: @functools.lru_cache(maxsize=1) +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 + + def _aiter_importable() -> bool: """True when the AITER packages needed for the C++ backends actually import. @@ -39,6 +58,7 @@ def _aiter_importable() -> bool: that raises at build/load time. """ try: + _ensure_aiter_gpu_archs() import aiter # noqa: F401 import aiter_meta # noqa: F401 from aiter.jit import core as _core # noqa: F401 From 542be2d65417872fc2441b5d0113b79c2b920cf6 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 25 Aug 2026 20:56:39 -0400 Subject: [PATCH 14/43] fix(aiter): port the paged-append shim to AITER's POD API page_aiter.cu still forward-declared reshape_and_cache_flash with the old at::Tensor& signature. AITER 0.1.16 moved cache.h to aiter_tensor_t, so the shim compiled cleanly and then failed at dlopen with undefined symbol: _ZN5aiter23reshape_and_cache_flashERN2at6TensorE... taking all 48 tests in test_append_paged_kv_cache_aiter_hip.py with it. 48 -> 0 on gfx950 after this change. Exactly the failure mode the activation shim's own comment warns about, in a file the earlier port missed: three shims were moved to the real headers and this one was not, because the investigation stopped at the prefill failures. cache.h no longer pulls , so the reason for the forward declaration is gone and the header can be included like activation.h. The stream guard moves into aiter_tensor_compat.h now that a second shim needs it -- it was written as a file-local class in activation_aiter.cu, which a review had already flagged as a single-call-site abstraction. Still outstanding on this branch: fused MoE fails the same way (28 tests), and six torch_compile tests fail on an unrelated gfx942/gfx950 arch mismatch. --- flashinfer/csrc_rocm/activation_aiter.cu | 26 +---------------- flashinfer/csrc_rocm/aiter_tensor_compat.h | 22 ++++++++++++++ flashinfer/csrc_rocm/page_aiter.cu | 34 +++++++++++++++------- 3 files changed, 46 insertions(+), 36 deletions(-) diff --git a/flashinfer/csrc_rocm/activation_aiter.cu b/flashinfer/csrc_rocm/activation_aiter.cu index bd531ad093..e106e47dc0 100644 --- a/flashinfer/csrc_rocm/activation_aiter.cu +++ b/flashinfer/csrc_rocm/activation_aiter.cu @@ -18,30 +18,6 @@ #include "aiter_tensor_compat.h" -namespace { - -// The POD API launches on aiter::getCurrentHIPStream(), a thread_local that -// defaults to nullptr and is otherwise set only by AITER's Python layer; the -// old torch-typed entry point read torch's stream itself. 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 AiterStreamGuard { - public: - explicit AiterStreamGuard(hipStream_t stream) : prev_(aiter::getCurrentHIPStream()) { - aiter::setCurrentHIPStream(stream); - } - ~AiterStreamGuard() { aiter::setCurrentHIPStream(prev_); } - - AiterStreamGuard(const AiterStreamGuard&) = delete; - AiterStreamGuard& operator=(const AiterStreamGuard&) = delete; - - private: - hipStream_t prev_; -}; - -} // namespace - void silu_and_mul_aiter(at::Tensor out, at::Tensor input) { const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(input.device()); @@ -53,7 +29,7 @@ void silu_and_mul_aiter(at::Tensor out, at::Tensor input) { 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 AiterStreamGuard stream_guard(at::hip::getCurrentHIPStream()); + 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 index 639873f9d5..2741853291 100644 --- a/flashinfer/csrc_rocm/aiter_tensor_compat.h +++ b/flashinfer/csrc_rocm/aiter_tensor_compat.h @@ -10,6 +10,7 @@ #include #include +#include #include namespace flashinfer::aiter_compat { @@ -61,4 +62,25 @@ inline aiter_tensor_t to_aiter(const at::Tensor& t) { 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/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); } From bc279ad3924db6c139ca069af9b79eb206654264 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 25 Aug 2026 21:12:08 -0400 Subject: [PATCH 15/43] fix(aiter): track moe_ck.h's new is_shuffled parameter AITER 0.1.16 appended `bool is_shuffled` to ck_moe_stage1 and ck_moe_stage2. The forward declarations here still described the 0.1.10 signature, so the mangled names no longer matched and every fused-MoE test died at dlopen with undefined symbol: _Z13ck_moe_stage1RN2at6TensorE... 28 -> 0 on gfx950 after adding the parameter. Passed as true: this API's contract is that the caller supplies pre-shuffled weights, which flashinfer.fused_moe documents and shuffle_moe_weight exists to do -- test_unshuffled_weights_are_wrong pins the ~1.3 relative error you get without it. Unlike cache.h and activation.h, moe_ck.h still includes , so the forward declaration cannot be replaced by the real header and has to track upstream by hand. That is the fragile half of the shim strategy and it is what broke here; the comment now says so at the declaration. --- flashinfer/csrc_rocm/fused_moe_aiter.cu | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/flashinfer/csrc_rocm/fused_moe_aiter.cu b/flashinfer/csrc_rocm/fused_moe_aiter.cu index 1503b181af..b07e7e7a30 100644 --- a/flashinfer/csrc_rocm/fused_moe_aiter.cu +++ b/flashinfer/csrc_rocm/fused_moe_aiter.cu @@ -28,6 +28,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. These three are at global namespace. +// +// They 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 CK 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, @@ -40,7 +44,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, @@ -48,7 +52,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); #ifdef FLASHINFER_MOE_AITER_PER_TOKEN // Unlike the three above, this one AITER declares inside `namespace aiter`. @@ -287,7 +291,7 @@ void fused_moe_aiter(at::Tensor out, at::Tensor hidden_states, at::Tensor w1, at topk_i32, kernel_name, w1_scale, a1_scale, block_m_i32, /*sorted_weights=*/std::nullopt, quant_type, activation_i32, /*splitk=*/1, /*nt=*/false, - /*dst_type=*/std::nullopt); + /*dst_type=*/std::nullopt, /*is_shuffled=*/true); at::Tensor stage2_in = inter_states; std::optional a2_scale; @@ -304,5 +308,5 @@ void fused_moe_aiter(at::Tensor out, at::Tensor hidden_states, at::Tensor w1, at ck_moe_stage2(stage2_in, w1, w2, sorted_token_ids, sorted_expert_ids, num_valid_ids, out, topk_i32, kernel_name, w2_scale, a2_scale, block_m_i32, sorted_weights, quant_type, activation_i32, /*splitk=*/1, - /*nt=*/false, /*dst_type=*/std::nullopt); + /*nt=*/false, /*dst_type=*/std::nullopt, /*is_shuffled=*/true); } From 1de02432c806e40ecff26e6b603d6bec717498ae Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 25 Aug 2026 23:10:44 -0400 Subject: [PATCH 16/43] fix(aiter): close the review gaps in the soft-cap guard and the graft Four findings from reviewing the unpushed commits, two of them self-inflicted. _aiter_importable lost its @functools.lru_cache: the new _ensure_aiter_gpu_archs was inserted directly beneath the decorator, so the decorator bound to the new function instead. Restored. That is the hot half of is_aiter_available, which every backend="auto" dispatch calls. The routing test compared a 2-tuple to a string and so could never pass. The latest rebase changed _auto_select_prefill_backend to return (backend, reason); I merged the signature and did not re-run the test that reads it. It now unpacks, and asserts the reason mentions logits_soft_cap so a fallback for an unrelated cause cannot masquerade as this guard working. GPU_ARCHS was only set on one of the two import paths. prefill_rocm's _aiter_ops_importable is the one prefill and decode reach first, and it imported aiter without it. _ensure_aiter_gpu_archs is also no longer cached: it early -returns when GPU_ARCHS is already set, and _build_aiter_lib sets that variable for its own scope and pops it again, so a cached "already done" taken during that window would leave it unset for the process lifetime. The batch wrappers now raise on an explicit backend="aiter" in the defect region, matching single prefill. Previously only the auto path was guarded, so asking for AITER by name on a ragged wrapper still ran the miscomputing kernel. Dockerfile: CPATH/LIBRARY_PATH/LD_LIBRARY_PATH append rather than assign. The base image sets these for torch's own bundled ROCm libs; the runtime paths keep the base first so torch wins, while CPATH puts the grafted headers first, which is what the graft is for. Checked and refuted: rmsnorm2d was reported as a fourth signature drift on the grounds that rmsnorm.h declares four parameters and the shim declares five. flashinfer.rmsnorm(backend="aiter") builds and runs, so the installed library exports the overload the shim expects. Left alone. --- .devcontainer/rocm/Dockerfile | 8 ++++--- flashinfer/aiter_utils.py | 6 ++++- flashinfer/prefill_rocm.py | 23 +++++++++++++++++++ .../test_single_prefill_kernels_hip.py | 8 +++++-- 4 files changed, 39 insertions(+), 6 deletions(-) diff --git a/.devcontainer/rocm/Dockerfile b/.devcontainer/rocm/Dockerfile index 3699efebbf..5513b2e999 100644 --- a/.devcontainer/rocm/Dockerfile +++ b/.devcontainer/rocm/Dockerfile @@ -94,7 +94,7 @@ 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" +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 @@ -112,8 +112,10 @@ RUN set -eu; \ # clang via $ROCM_PATH/lib/llvm/bin, which only exists under the SDK. ENV ROCM_PATH="$ROCM_SDK" ENV ROCM_HOME="$ROCM_SDK" -ENV LIBRARY_PATH="$ROCM_SDK/lib" -ENV LD_LIBRARY_PATH="$ROCM_SDK/lib" +# Appended, not assigned: the base image sets these for torch's own bundled +# ROCm libs (RCCL, MIOpen), and clobbering them surfaces far from here. +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 diff --git a/flashinfer/aiter_utils.py b/flashinfer/aiter_utils.py index 6b2052f461..0bf0e67e1e 100644 --- a/flashinfer/aiter_utils.py +++ b/flashinfer/aiter_utils.py @@ -29,7 +29,6 @@ def is_aiter_supported(device: torch.device) -> bool: return arch in FLASHINFER_SUPPORTED_ROCM_ARCHS -@functools.lru_cache(maxsize=1) def _ensure_aiter_gpu_archs() -> None: """Give AITER's JIT a GPU_ARCHS, since from 0.1.16 it requires one. @@ -37,6 +36,10 @@ def _ensure_aiter_gpu_archs() -> None: 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. + + Deliberately uncached: ``_build_aiter_lib`` sets GPU_ARCHS for its own scope + and pops it again, so a cached "already done" taken while that value was + live would leave the variable unset for the rest of the process. """ if os.environ.get("GPU_ARCHS"): return @@ -49,6 +52,7 @@ def _ensure_aiter_gpu_archs() -> None: os.environ["GPU_ARCHS"] = arch +@functools.lru_cache(maxsize=1) def _aiter_importable() -> bool: """True when the AITER packages needed for the C++ backends actually import. diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index d7cdb767e7..6abb011a69 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -304,6 +304,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 @@ -2312,6 +2319,14 @@ def plan( ), ) ) + 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}; " @@ -3317,6 +3332,14 @@ def plan( 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/tests/rocm_tests/test_single_prefill_kernels_hip.py b/tests/rocm_tests/test_single_prefill_kernels_hip.py index 650ca82de9..52ba310cc2 100644 --- a/tests/rocm_tests/test_single_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_single_prefill_kernels_hip.py @@ -359,7 +359,7 @@ def test_auto_backend_avoids_aiter_softcap_defect( from flashinfer.prefill_rocm import _auto_select_prefill_backend - chosen = _auto_select_prefill_backend( + chosen, reason = _auto_select_prefill_backend( device, dtype_q=torch.float16, dtype_kv=torch.float16, @@ -374,7 +374,11 @@ def test_auto_backend_avoids_aiter_softcap_defect( logits_soft_cap=soft_cap, kv_len=kv_len, ) - assert chosen == ("aiter" if expect_aiter else "fa2") + 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(): From 6b866bf0295afa6eda4b67a39c7e385de8861e3e Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 26 Aug 2026 15:14:17 -0400 Subject: [PATCH 17/43] fix(aiter): key the paged soft-cap guard on the paging route The explicit backend="aiter" guard added for the paged wrapper passed self._max_kv_len unconditionally, while the auto path seven lines above passes kv_len=None whenever page_size is a native AITER page size -- native paging goes through mha_batch_prefill, which does not carry the soft-cap defect. The two therefore disagreed about the same call: BatchPrefillWithPagedKVCacheWrapper(backend="aiter") with page_size=128, causal=True, head_dim_qk=128, logits_soft_cap=8.0 and max kv_len >= 512 raised ValueError from plan(), while backend="auto" on the identical call selected aiter and ran the exact kernel. That configuration worked before the guard landed. _aiter_native_page_sizes() is {128, 256, 1024} on amd-aiter 0.1.16.post3, so page sizes 128/256/1024 were all affected. Derive the value once as softcap_kv_len and feed both, so the routing decision and the guard cannot drift apart again. The new test A/B's: with the fix it passes; with prefill_rocm.py reverted it fails at prefill_rocm.py:2218 with the ValueError above. Measured on gfx942 / MI300X, ROCm 7.14.60850, torch 2.12.0, amd-aiter 0.1.16.post3.dev0+g620287969.d20260725. Not closed here: when page_size is native but plan()'s run-time probe sets use_native_paging=False, the call degrades to flat-gather and the defect does apply, yet neither the auto path nor this guard raises. That gap predates the guard and is noted in the comment at the auto call. --- flashinfer/prefill_rocm.py | 24 +++++----- .../test_batch_prefill_kernels_hip.py | 47 +++++++++++++++++++ 2 files changed, 59 insertions(+), 12 deletions(-) diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index 6abb011a69..6b779aa024 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -2292,6 +2292,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( @@ -2306,21 +2316,11 @@ def plan( op="batch_prefill", causal=causal, logits_soft_cap=logits_soft_cap, - # 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. - kv_len=( - None - if page_size in _aiter_native_page_sizes() - else self._max_kv_len - ), + kv_len=softcap_kv_len, ) ) if self._backend == "aiter" and _aiter_softcap_defect( - causal, logits_soft_cap, head_dim_qk, self._max_kv_len + causal, logits_soft_cap, head_dim_qk, softcap_kv_len ): raise ValueError( "AITER miscomputes logits_soft_cap for causal head_dim=128 prefill " diff --git a/tests/rocm_tests/test_batch_prefill_kernels_hip.py b/tests/rocm_tests/test_batch_prefill_kernels_hip.py index 51b834d4ef..31ee9d29e7 100644 --- a/tests/rocm_tests/test_batch_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_batch_prefill_kernels_hip.py @@ -1087,3 +1087,50 @@ def test_ragged_softcap_avoids_broken_aiter_kernel(kv_len, qo_len): 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) From 534b36011819e1b739a72c0d7ed3b578b94ca79f Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 26 Aug 2026 15:14:30 -0400 Subject: [PATCH 18/43] fix(aiter): let the shim build's GPU_ARCHS outlive its own scope _ensure_aiter_gpu_archs was left uncached to survive _build_aiter_lib setting GPU_ARCHS and popping it again. That does not work: both of its callers are themselves cached (_aiter_importable is lru_cache(maxsize=1), _aiter_ops_importable is functools.cache), and they are the only two in the repo. The window still loses -- thread A holds _BUILD_LOCK with GPU_ARCHS live, thread B makes the first backend="auto" prefill call, the probe early-returns because the value is set and caches True, thread A pops the variable, and it stays unset for the rest of the process. A later AITER Python-op build then asserts on ['']. Fix it at the source instead: when _build_aiter_lib set GPU_ARCHS from unset, leave it set. It is the same architecture _ensure_aiter_gpu_archs would resolve, so nothing downstream sees a value it would not have chosen, and the interleaving disappears. Drop the now-false rationale from the docstring; the function stays uncached because an environ lookup is cheaper than reasoning about when it would be safe to cache. --- flashinfer/aiter_utils.py | 4 ---- flashinfer/jit/aiter_source.py | 6 +++++- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/flashinfer/aiter_utils.py b/flashinfer/aiter_utils.py index 0bf0e67e1e..38680c66e8 100644 --- a/flashinfer/aiter_utils.py +++ b/flashinfer/aiter_utils.py @@ -36,10 +36,6 @@ def _ensure_aiter_gpu_archs() -> None: 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. - - Deliberately uncached: ``_build_aiter_lib`` sets GPU_ARCHS for its own scope - and pops it again, so a cached "already done" taken while that value was - live would leave the variable unset for the rest of the process. """ if os.environ.get("GPU_ARCHS"): return diff --git a/flashinfer/jit/aiter_source.py b/flashinfer/jit/aiter_source.py index 4a5479e129..09a61ee376 100644 --- a/flashinfer/jit/aiter_source.py +++ b/flashinfer/jit/aiter_source.py @@ -430,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 From 99a6068f2b16ed781ef12dc850346d38cc54764e Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 26 Aug 2026 15:14:30 -0400 Subject: [PATCH 19/43] docs(rocm): stop claiming the base image sets the ROCm env vars The comment justified appending to CPATH/LIBRARY_PATH/LD_LIBRARY_PATH by saying the base sets them for torch's bundled RCCL and MIOpen. It does not: the only ENV in rocm/pytorch:rocm7.14_ubuntu26.04_py3.14's history is VIRTUAL_ENV/PATH, and nothing under /etc/profile.d, /etc/environment or /etc/ld.so.conf.d sets them either, so the append resolves to the same string the assignment did. Keep the append -- it is the right form if a future base does set them -- and say that instead. --- .devcontainer/rocm/Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.devcontainer/rocm/Dockerfile b/.devcontainer/rocm/Dockerfile index 5513b2e999..6cc456ed1b 100644 --- a/.devcontainer/rocm/Dockerfile +++ b/.devcontainer/rocm/Dockerfile @@ -112,8 +112,8 @@ RUN set -eu; \ # clang via $ROCM_PATH/lib/llvm/bin, which only exists under the SDK. ENV ROCM_PATH="$ROCM_SDK" ENV ROCM_HOME="$ROCM_SDK" -# Appended, not assigned: the base image sets these for torch's own bundled -# ROCm libs (RCCL, MIOpen), and clobbering them surfaces far from here. +# 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" From e8f1a0af564b4d8b0a50bb9ff006fab4bba5b3a6 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 26 Aug 2026 15:27:29 -0400 Subject: [PATCH 20/43] fix(aiter): reject non-NHD kv_layout on explicit single-prefill aiter Both batch prefill wrappers raise when backend="aiter" is asked for with kv_layout != "NHD", and the auto-selector falls back to fa2 for the same reason (prefill_rocm.py:400). single_prefill_with_kv_cache checked the soft-cap defect and pos_encoding_mode but not the layout, so an explicit backend="aiter" passed TensorLayout[kv_layout] straight through to an NHD-only kernel. The HND cases are skipped in the tests, so nothing covered it. Raise the same error the batch wrappers do. --- flashinfer/prefill_rocm.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index 6b779aa024..57d2afdcd7 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -1578,6 +1578,11 @@ def single_prefill_with_kv_cache( 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. From 63a7b64382c028a4d11d703b7f10b10f775331a8 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 26 Aug 2026 15:27:29 -0400 Subject: [PATCH 21/43] fix(aiter): enforce an amd-aiter 0.1.16 ABI floor before routing The vendored mha_fwd_args gained block_scale_seqstart_q_ptr and block_scale_seqstart_k_ptr in 0.1.16, inserted before sink_ptr, so every field from sink_ptr on sits at a different offset than in 0.1.10. The struct goes by value through a dlsym'd pointer: an older AITER links and runs, and corrupts arguments silently. Nothing downstream can detect it, so the floor has to sit before routing. _aiter_importable() now also requires the version, which makes both paths safe at once -- auto stops selecting aiter, and require_aiter raises naming the installed version instead of the generic "not installed" message. Compare on base_version, and re-wrap it in Version: base_version is a str, so the obvious `Version(a).base_version >= Version(b).base_version` compares lexically and rates 0.1.10 and 0.1.9 as >= 0.1.16. Reverting to that form fails test_version_floor[0.1.9-False], which is what the parametrization is for. base_version is also what lets the only cp314 wheel through: 0.1.16.post3.dev0+g620287969.d20260725 carries a .dev0 segment that PEP 440 sorts below plain 0.1.16. Verified on gfx942 / MI300X, ROCm 7.14.60850, torch 2.12.0. --- flashinfer/aiter_utils.py | 48 ++++++++++++++++- .../flashinfer/attention/aiter/mha_fwd_args.h | 8 ++- .../rocm_tests/test_aiter_version_gate_hip.py | 51 +++++++++++++++++++ 3 files changed, 104 insertions(+), 3 deletions(-) create mode 100644 tests/rocm_tests/test_aiter_version_gate_hip.py diff --git a/flashinfer/aiter_utils.py b/flashinfer/aiter_utils.py index 38680c66e8..a25820cc74 100644 --- a/flashinfer/aiter_utils.py +++ b/flashinfer/aiter_utils.py @@ -4,6 +4,7 @@ import functools import os +from typing import Optional import torch @@ -48,6 +49,40 @@ def _ensure_aiter_gpu_archs() -> None: 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. @@ -55,7 +90,8 @@ 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() @@ -64,7 +100,7 @@ def _aiter_importable() -> bool: 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: @@ -102,6 +138,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/include/flashinfer/attention/aiter/mha_fwd_args.h b/include/flashinfer/attention/aiter/mha_fwd_args.h index a0b74a9ca6..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. 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 From cbe8f0d7befe2da9bf2ec87f9369ce7291d5f971 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 26 Aug 2026 16:47:36 -0400 Subject: [PATCH 22/43] docs(skills): drop line-number anchors that no longer resolve prefill_rocm.py:1978 was the kv_layout ValueError when the debug skill was written; it is docstring prose now, and prefill_rocm.py:59 and :311 have drifted the same way. Point at the error text and the function name, both of which survive edits, instead of a line. Two claims were stale as well: the kv_layout error is raised in three places now, not two, and the {16, 1024} arm of _aiter_native_page_sizes() became unreachable when aiter_utils.AITER_MIN_VERSION floored AITER at 0.1.16. Note the probe too -- the native set is a hint, not the route. --- .claude/skills/benchmark-kernel/SKILL.md | 6 +++--- .claude/skills/debug-rocm-crash/SKILL.md | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.claude/skills/benchmark-kernel/SKILL.md b/.claude/skills/benchmark-kernel/SKILL.md index f1df84dce4..d9f1413c52 100644 --- a/.claude/skills/benchmark-kernel/SKILL.md +++ b/.claude/skills/benchmark-kernel/SKILL.md @@ -22,11 +22,11 @@ For the in-repo profiler wrapper, see [`rocm_profiler/rocm_profiler.py`](../../. - **CUPTI is NVIDIA-only — `enable_cupti=True` on ROCm warns and falls back.** [`flashinfer/testing/utils.py:1010`](../../../flashinfer/testing/utils.py) routes through `bench_gpu_time_with_cupti`, which `try/except`s the `cupti` import, emits a `UserWarning`, and reverts to CUDA/HIP event timing. No functional benefit on ROCm; just leave `enable_cupti=False` (the default) so `bench_gpu_time` uses `torch.cuda.Event` (HIP events) directly without the warning. - **AITER backend constraints, accurately:** - - Explicit `backend="aiter"` + `kv_layout != "NHD"` → `ValueError` at `plan()` time. Raised in the prefill wrapper, e.g. [`prefill_rocm.py:1978`](../../../flashinfer/prefill_rocm.py) (single/paged) and the batch-paged wrapper around line 2920. Not raised by auto-selection — that path silently falls back to `fa2`. + - Explicit `backend="aiter"` + `kv_layout != "NHD"` → `ValueError`. Grep [`prefill_rocm.py`](../../../flashinfer/prefill_rocm.py) for `only supports kv_layout`; single prefill raises at the call, both batch wrappers at `plan()`. Not raised by auto-selection — that path silently falls back to `fa2`. - Explicit `backend="aiter"` on non-gfx942/gfx950 → `RuntimeError`. - `amd-aiter` not importable → `ImportError`. - - **"Native" page sizes** (no flat-gather): `{128, 256, 1024}` for `amd-aiter >= 0.1.10`, else `{16, 1024}` — see `_aiter_native_page_sizes()` in [`prefill_rocm.py:59`](../../../flashinfer/prefill_rocm.py). **Non-native page sizes are NOT rejected** — they go through a flat-gather code path. So the "{1, 16, 1024}" guidance from older docs is wrong. - - Auto-selection (no explicit `backend=`) silently falls back to `fa2` for any of: `kv_layout != "NHD"`, custom mask, dtype not in `{fp16, bf16}`, `dtype_q != dtype_kv`, `head_dim_qk != head_dim_vo`, `pos_encoding_mode != "NONE"`, or `amd-aiter` not importable. See `_auto_select_prefill_backend()` in [`prefill_rocm.py:311`](../../../flashinfer/prefill_rocm.py) for the authoritative list. + - **"Native" page sizes** (no flat-gather): `{128, 256, 1024}` — see `_aiter_native_page_sizes()` in [`prefill_rocm.py`](../../../flashinfer/prefill_rocm.py). The `{16, 1024}` arm there is unreachable now that `aiter_utils.AITER_MIN_VERSION` floors AITER at 0.1.16. It is only a hint: `plan()` probes with `_aiter_native_paging_available()` and degrades to flat-gather if the installed AITER cannot serve the config. **Non-native page sizes are NOT rejected** — they flat-gather. So the "{1, 16, 1024}" guidance from older docs is wrong. + - Auto-selection (no explicit `backend=`) silently falls back to `fa2` for any of: `kv_layout != "NHD"`, custom mask, dtype not in `{fp16, bf16}`, `dtype_q != dtype_kv`, `head_dim_qk != head_dim_vo`, `pos_encoding_mode != "NONE"`, or `amd-aiter` not importable. See `_auto_select_prefill_backend()` in [`prefill_rocm.py`](../../../flashinfer/prefill_rocm.py) for the authoritative list; it returns `(backend, reason)`, and the reason names the constraint that forced the fallback. - **Always verify numerical parity before trusting perf numbers.** Compare default-HIP vs AITER outputs with `torch.testing.assert_close(rtol=1e-2, atol=1e-2)` for BF16/FP16 first. - **`gcnArchName` is the unambiguous arch marker.** Device strings show `cuda:0` on AMD too. Record `torch.cuda.get_device_properties(0).gcnArchName` and `torch.version.hip` alongside every number — a `gfx942` / ROCm 7.2 result is not comparable to a `gfx950` / ROCm 7.0.2 result. diff --git a/.claude/skills/debug-rocm-crash/SKILL.md b/.claude/skills/debug-rocm-crash/SKILL.md index 7645aafda1..8b9b06c5fe 100644 --- a/.claude/skills/debug-rocm-crash/SKILL.md +++ b/.claude/skills/debug-rocm-crash/SKILL.md @@ -29,7 +29,7 @@ For an in-script view of what's being passed, wrap the suspect call with `print( | Symptom | First check | | --- | --- | | `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"` `ValueError` before launch | `kv_layout != "NHD"` (only NHD is allowed). Grep [`prefill_rocm.py`](../../../flashinfer/prefill_rocm.py) for `only supports kv_layout` — single prefill and both batch wrappers each raise it. | | `backend="aiter"` `RuntimeError` | Non-gfx942/gfx950 GPU. | | `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. | From 1e4a5d79bd891d9be7b968d22277e3ef4accc0ee Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 11:24:57 -0400 Subject: [PATCH 23/43] fix(aiter): port the fp8 per-token quant call to AITER's POD API #309 declared aiter::dynamic_per_token_scaled_quant with at::Tensor parameters, which was right for amd-aiter 0.1.10. In 0.1.16 quant.h takes aiter_tensor_t, so the declaration compiled cleanly and then failed at dlopen with undefined symbol: _ZN5aiter30dynamic_per_token_scaled_quantERN2at6\ TensorERKS1_S2_St8optionalIS1_EbS6_i taking down every fp8 MoE test -- including the pure argument-validation ones, since the module never loaded. 21 failures in test_fused_moe_aiter_hip.py, all from this one symbol; 0 after. Same migration as page_aiter.cu, so the existing aiter_tensor_compat.h adapter covers it: convert to POD locals first, since the out and scales parameters are non-const references. Only surfaces on this branch because it is what moves the devcontainer to 0.1.16; #309 is correct against the 0.1.10 stack it was written on. Measured on gfx942 / MI300X, ROCm 7.14.60850, torch 2.12.0, amd-aiter 0.1.16.post3.dev0+g620287969.d20260725. --- flashinfer/csrc_rocm/fused_moe_aiter.cu | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/flashinfer/csrc_rocm/fused_moe_aiter.cu b/flashinfer/csrc_rocm/fused_moe_aiter.cu index b07e7e7a30..efe5cdfca3 100644 --- a/flashinfer/csrc_rocm/fused_moe_aiter.cu +++ b/flashinfer/csrc_rocm/fused_moe_aiter.cu @@ -24,6 +24,8 @@ #include #include +#include "aiter_tensor_compat.h" + // AITER's public headers (moe_sorting.h, moe_ck.h) pull in → // 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 @@ -55,11 +57,14 @@ void ck_moe_stage2(at::Tensor& inter_states, at::Tensor& w1, at::Tensor& w2, std::optional dst_type, bool is_shuffled); #ifdef FLASHINFER_MOE_AITER_PER_TOKEN -// Unlike the three above, this one AITER declares inside `namespace aiter`. +// Unlike the three above, this one AITER declares inside `namespace aiter`, and +// on the POD API rather than at::Tensor -- see quant.h. Declaring it with +// at::Tensor compiles and then fails at dlopen on the mangled name. namespace aiter { -void dynamic_per_token_scaled_quant(at::Tensor& out, at::Tensor const& input, at::Tensor& scales, - std::optional scale_ub, bool shuffle_scale, - std::optional num_rows, int num_rows_factor); +void dynamic_per_token_scaled_quant(aiter_tensor_t& out, const aiter_tensor_t& input, + aiter_tensor_t& scales, std::optional scale_ub, + bool shuffle_scale, std::optional num_rows, + int num_rows_factor); } // namespace aiter #endif @@ -83,7 +88,12 @@ std::pair quantize_per_token(const at::Tensor& x, at::Sc scale_sizes.back() = 1; at::Tensor q = at::empty(x.sizes(), x.options().dtype(fp8)); at::Tensor scale = at::empty(scale_sizes, x.options().dtype(at::kFloat)); - aiter::dynamic_per_token_scaled_quant(q, x, scale, /*scale_ub=*/std::nullopt, + + namespace compat = flashinfer::aiter_compat; + aiter_tensor_t q_a = compat::to_aiter(q); + aiter_tensor_t x_a = compat::to_aiter(x); + aiter_tensor_t scale_a = compat::to_aiter(scale); + aiter::dynamic_per_token_scaled_quant(q_a, x_a, scale_a, /*scale_ub=*/std::nullopt, /*shuffle_scale=*/false, /*num_rows=*/std::nullopt, static_cast(num_rows_factor)); return {q, scale}; From d4d03a92afc8435d42093488b27a35a3247a71b5 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 11:48:52 -0400 Subject: [PATCH 24/43] fix(aiter): run the fp8 quant kernel on the caller's stream dynamic_per_token_scaled_quant takes its stream from AITER's thread-local (quant_kernels.cu calls aiter::getCurrentHIPStream()), not from torch, so the POD call needs the same StreamGuard the other two POD call sites in this tree already hold -- page_aiter.cu:186 and activation_aiter.cu:32. Without it, a caller inside torch.cuda.stream(s) gets the quant kernel on the legacy null stream while at::empty for q/scale and both ck_moe stages run on s. Torch's pool streams are created hipStreamNonBlocking, so the null stream does not implicitly synchronise with them: stage 1 can read q/scale while quant is still writing, and the caching allocator can recycle them on s with the null-stream kernel pending. Under HIP graph capture the stray launch invalidates the capture outright. It passes today only because the default stream is nullptr on both sides, which is the same value AITER starts with. Not covered by a test: every existing test runs on the default stream, where the bug is invisible, and a non-default-stream test would detect it only through a race that is not reliably reproducible. The three POD call sites now all hold the guard, which is the invariant worth keeping. --- flashinfer/csrc_rocm/fused_moe_aiter.cu | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/flashinfer/csrc_rocm/fused_moe_aiter.cu b/flashinfer/csrc_rocm/fused_moe_aiter.cu index efe5cdfca3..73211cdcd3 100644 --- a/flashinfer/csrc_rocm/fused_moe_aiter.cu +++ b/flashinfer/csrc_rocm/fused_moe_aiter.cu @@ -93,6 +93,11 @@ std::pair quantize_per_token(const at::Tensor& x, at::Sc aiter_tensor_t q_a = compat::to_aiter(q); aiter_tensor_t x_a = compat::to_aiter(x); aiter_tensor_t scale_a = compat::to_aiter(scale); + // The POD entry point reads AITER's thread-local stream, not torch's. Without + // this the quant kernel lands on the null stream while the CK stages consume + // q/scale on the caller's -- and torch's pool streams are non-blocking, so + // nothing synchronises them. + const compat::StreamGuard stream_guard(at::hip::getCurrentHIPStream()); aiter::dynamic_per_token_scaled_quant(q_a, x_a, scale_a, /*scale_ub=*/std::nullopt, /*shuffle_scale=*/false, /*num_rows=*/std::nullopt, static_cast(num_rows_factor)); From f7bae0dcd6d021e2c4f3cb9a323ad6dd9db1aaa3 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 13:40:29 -0400 Subject: [PATCH 25/43] fix(rocm): only check PyTorch's arch list when PYTORCH_ROCM_ARCH pins it torch 2.12's _get_rocm_arch_flags() does not report what PyTorch was built for. With PYTORCH_ROCM_ARCH unset it enumerates the *visible cards*: with a GPU: ['--offload-arch=gfx942', '-fno-gpu-rdc'] no GPU: ['--offload-arch=', '-fno-gpu-rdc'] Validating against that is wrong in both directions. On a host with no visible device the list is empty, so every requested arch reads as unsupported and `import flashinfer` raises -- undoing #316. And on a gfx942 box it rejects a gfx950 request, which is exactly what an ahead-of-time build is for. Both showed up as real failures once the devcontainer moved to ROCm 7.14 / torch 2.12: test_jit_env_hip.py::test_import_succeeds_with_no_visible_device and test_aot_hip.py::test_publishing_preserves_the_requested_order. The check is only meaningful when PYTORCH_ROCM_ARCH is set, since that is what actually constrains PyTorch's build; gate it on that and say so in the error. The two existing tests now set the variable, matching the condition they were always describing. The new parametrized test A/B's: with the gate removed both cases fail with "PyTorch does not support the following architectures: --offload-arch=gfx950". Measured on gfx942 / MI300X, ROCm 7.14.60850, torch 2.12.0. --- flashinfer/hip_utils.py | 16 ++++++++-- .../rocm_tests/test_aiter_version_gate_hip.py | 4 ++- tests/rocm_tests/test_hip_utils.py | 31 +++++++++++++++++-- 3 files changed, 45 insertions(+), 6 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 19e633aa3f..82a7643256 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -4,6 +4,7 @@ import functools import logging +import os # arch_caps imports nothing (in particular, not torch), so importing it here # keeps this module importable without torch -- which the hardware-less @@ -487,9 +488,16 @@ def validate_flashinfer_rocm_arch( ) requested_archs = supported_in_request - # Step 3: Validate against PyTorch's available architectures (if module provided) + # Step 3: Validate against PyTorch's architectures -- but only when PyTorch was + # actually told which ones to build. With PYTORCH_ROCM_ARCH unset, + # _get_rocm_arch_flags() enumerates the *visible cards* rather than anything + # about the build: it yields ["--offload-arch=", ...] when no device is + # visible, and lists only the local card otherwise. Enforcing against that + # turns an import on a GPU-free host into a hard error, and rejects + # cross-compiling for gfx950 from a gfx942 box -- which is precisely what an + # ahead-of-time build is for. arch_flags = [f"--offload-arch={arch}" for arch in requested_archs] - if torch_cpp_ext_module is not None: + if torch_cpp_ext_module is not None and os.environ.get("PYTORCH_ROCM_ARCH"): pytorch_arch_flags = torch_cpp_ext_module._get_rocm_arch_flags() missing_in_pytorch = [ flag for flag in arch_flags if flag not in pytorch_arch_flags @@ -497,7 +505,9 @@ def validate_flashinfer_rocm_arch( if missing_in_pytorch: raise RuntimeError( f"PyTorch does not support the following architectures: {', '.join(missing_in_pytorch)}.\n" - f"PyTorch was compiled with: {', '.join(pytorch_arch_flags)}" + f"PyTorch was built for: {', '.join(pytorch_arch_flags)}\n" + "This is checked because PYTORCH_ROCM_ARCH is set; unset it to build " + "for whatever FLASHINFER_ROCM_ARCH_LIST requests." ) if verbose: diff --git a/tests/rocm_tests/test_aiter_version_gate_hip.py b/tests/rocm_tests/test_aiter_version_gate_hip.py index ab14c3185a..b59fb0822c 100644 --- a/tests/rocm_tests/test_aiter_version_gate_hip.py +++ b/tests/rocm_tests/test_aiter_version_gate_hip.py @@ -19,8 +19,10 @@ @pytest.mark.parametrize( "version,supported", [ + # "0.1.9" > "0.1.16" as strings, so this row is the one that catches a + # comparison done on base_version without re-wrapping it in Version. ("0.1.9", False), - ("0.1.10", False), # lexically ">= 0.1.16"; the pin this repo shipped before + ("0.1.10", False), # the pin this repo shipped before the 0.1.16 floor ("0.1.15.post9", False), ("0.1.16", True), # The nightly wheel: PEP 440 sorts a .dev0 segment *below* its release, diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index 25e9cb3031..d450179c7f 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -458,7 +458,8 @@ def test_warns_and_filters_when_some_archs_unsupported_by_flashinfer(self): assert flags == ["--offload-arch=gfx942"] assert arch_set == {"gfx942"} - def test_pytorch_validation_passes_when_all_flags_present(self): + def test_pytorch_validation_passes_when_all_flags_present(self, monkeypatch): + monkeypatch.setenv("PYTORCH_ROCM_ARCH", "gfx942;gfx950") torch_cpp_ext = MagicMock() torch_cpp_ext._get_rocm_arch_flags.return_value = [ "--offload-arch=gfx942", @@ -470,7 +471,8 @@ def test_pytorch_validation_passes_when_all_flags_present(self): ) assert flags == ["--offload-arch=gfx942"] - def test_pytorch_validation_raises_when_flag_missing(self): + def test_pytorch_validation_raises_when_flag_missing(self, monkeypatch): + monkeypatch.setenv("PYTORCH_ROCM_ARCH", "gfx950") torch_cpp_ext = MagicMock() torch_cpp_ext._get_rocm_arch_flags.return_value = ["--offload-arch=gfx950"] with ( @@ -481,6 +483,31 @@ def test_pytorch_validation_raises_when_flag_missing(self): arch_list="gfx942", torch_cpp_ext_module=torch_cpp_ext ) + @pytest.mark.parametrize( + "torch_flags", + [ + pytest.param(["--offload-arch=", "-fno-gpu-rdc"], id="no-visible-device"), + pytest.param(["--offload-arch=gfx942"], id="cross-compile-from-gfx942"), + ], + ) + def test_pytorch_validation_skipped_without_pytorch_rocm_arch( + self, monkeypatch, torch_flags + ): + """Unset PYTORCH_ROCM_ARCH makes torch report visible cards, not its build. + + Enforcing against that breaks a GPU-free import and any cross-compile, + so the check only applies when the variable pins the set explicitly. + """ + monkeypatch.delenv("PYTORCH_ROCM_ARCH", raising=False) + torch_cpp_ext = MagicMock() + torch_cpp_ext._get_rocm_arch_flags.return_value = torch_flags + with self._patch_validate_rocm_arch("gfx950"): + flags, arch_set = validate_flashinfer_rocm_arch( + arch_list="gfx950", torch_cpp_ext_module=torch_cpp_ext + ) + assert flags == ["--offload-arch=gfx950"] + assert arch_set == {"gfx950"} + def test_reads_arch_from_env_when_none_given(self, monkeypatch): monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx942") with self._patch_validate_rocm_arch("gfx942"): From 84c0e3cc16cbd052698c04a5f05ab7759aa57616 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 16:37:05 -0400 Subject: [PATCH 26/43] docs(rocm): describe the arch check as the PYTORCH_ROCM_ARCH gate it is The gate landed in the previous commit but three docstrings and the error message still described the old always-on check against PyTorch's build. The message was the worst of them: in the only branch that can print it PYTORCH_ROCM_ARCH is set, so _get_rocm_arch_flags() is echoing that variable back -- "PyTorch was built for: --offload-arch=gfx1030" is a claim about the wheel that the variable alone cannot support, and it sends the reader off to reinstall torch when the fix is to unset an env var. Name the variable in all four places, and move the test's match string onto the new text. --- flashinfer/compilation_context_hip.py | 2 +- flashinfer/hip_utils.py | 10 +++++----- flashinfer/jit/rocm/core.py | 2 +- tests/rocm_tests/test_hip_utils.py | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/flashinfer/compilation_context_hip.py b/flashinfer/compilation_context_hip.py index f6f461ccdf..d56a3c520b 100644 --- a/flashinfer/compilation_context_hip.py +++ b/flashinfer/compilation_context_hip.py @@ -47,7 +47,7 @@ def __init__(self): Performs comprehensive validation: 1. System ROCm version compatibility 2. FlashInfer AMD port availability - 3. PyTorch ROCm compilation support + 3. PYTORCH_ROCM_ARCH, when it is set """ import torch.utils.cpp_extension as torch_cpp_ext diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 82a7643256..21a3063204 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -436,7 +436,8 @@ def validate_flashinfer_rocm_arch( Validates in order: 1. System ROCm version supports the architectures (ROCM_COMPAT_MATRIX) 2. FlashInfer has AMD ports for the architectures (FLASHINFER_SUPPORTED_ROCM_ARCHS) - 3. PyTorch was compiled with the architectures (torch.utils.cpp_extension) + 3. PYTORCH_ROCM_ARCH, if set, permits the architectures (skipped when unset, + where torch reports the visible cards rather than its build) Args: arch_list: Comma-separated list (e.g., "gfx942,gfx90a") or None for default @@ -504,10 +505,9 @@ def validate_flashinfer_rocm_arch( ] if missing_in_pytorch: raise RuntimeError( - f"PyTorch does not support the following architectures: {', '.join(missing_in_pytorch)}.\n" - f"PyTorch was built for: {', '.join(pytorch_arch_flags)}\n" - "This is checked because PYTORCH_ROCM_ARCH is set; unset it to build " - "for whatever FLASHINFER_ROCM_ARCH_LIST requests." + f"PYTORCH_ROCM_ARCH excludes the following architectures: {', '.join(missing_in_pytorch)}.\n" + f"It restricts extension builds to: {', '.join(pytorch_arch_flags)}\n" + "Unset it to build for whatever FLASHINFER_ROCM_ARCH_LIST requests." ) if verbose: diff --git a/flashinfer/jit/rocm/core.py b/flashinfer/jit/rocm/core.py index a656ca5470..1e2cb7fd72 100644 --- a/flashinfer/jit/rocm/core.py +++ b/flashinfer/jit/rocm/core.py @@ -12,7 +12,7 @@ def check_rocm_arch() -> None: """Validate that this GPU, ROCm and torch can build FlashInfer's kernels. Delegates to hip_utils so the ROCm version, the ported-architecture list and - torch's build architectures are checked in one place. + PYTORCH_ROCM_ARCH (when set) are checked in one place. """ import torch.utils.cpp_extension as torch_cpp_ext diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index d450179c7f..08df47e62b 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -477,7 +477,7 @@ def test_pytorch_validation_raises_when_flag_missing(self, monkeypatch): torch_cpp_ext._get_rocm_arch_flags.return_value = ["--offload-arch=gfx950"] with ( self._patch_validate_rocm_arch("gfx942"), - pytest.raises(RuntimeError, match="PyTorch does not support"), + pytest.raises(RuntimeError, match="PYTORCH_ROCM_ARCH excludes"), ): validate_flashinfer_rocm_arch( arch_list="gfx942", torch_cpp_ext_module=torch_cpp_ext From 2ce803f4b324bfe236d8c1eca7f8dbedf5e84152 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 17:32:08 -0400 Subject: [PATCH 27/43] fix(aiter): compare AITER versions without stripping the pre-release _aiter_version_supported compared on base_version, which strips .dev and .pre segments. That admits 0.1.16.dev0 -- a pre-release *of* the floor, built before the ABI it guards -- as though it were 0.1.16. A straight PEP 440 comparison is right for every case, including the one base_version was reached for: the nightly is 0.1.16.post3.dev0+g, a dev build of post3, which sorts above 0.1.16 rather than below it. The comment claiming otherwise was wrong. The new 0.1.16.dev0 row A/B's: restoring base_version fails it. Also stop hard-coding backend='native' in require_aiter's two errors. Prefill and decode do not accept 'native' -- their fallback is 'fa2' -- so read it from the capability table, which already records one per op. Measured on gfx942 / MI300X, ROCm 7.14.60850, torch 2.12.0. --- flashinfer/aiter_utils.py | 25 ++++++++++++------- flashinfer/arch_caps.py | 11 ++++++++ .../rocm_tests/test_aiter_version_gate_hip.py | 7 ++++-- 3 files changed, 32 insertions(+), 11 deletions(-) diff --git a/flashinfer/aiter_utils.py b/flashinfer/aiter_utils.py index a25820cc74..f06482399d 100644 --- a/flashinfer/aiter_utils.py +++ b/flashinfer/aiter_utils.py @@ -10,6 +10,7 @@ from .arch_caps import ( ArchCapabilityError, + aiter_fallback_backend, capability_available, normalize_arch, require_capability, @@ -74,11 +75,12 @@ def _aiter_version_supported() -> bool: 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) + # Straight PEP 440 comparison, which already sorts the cases that matter: + # 0.1.16.post3.dev0+g (the nightly) is a dev build *of post3* and + # sorts above 0.1.16, while 0.1.16.dev0 is a pre-release of 0.1.16 and + # sorts below it. base_version would strip both segments and wrongly + # admit the latter. + return Version(installed) >= Version(AITER_MIN_VERSION) except Exception: return False @@ -138,18 +140,23 @@ def require_aiter(device: torch.device, op: str) -> None: """ require_capability(device, op, "aiter") if not _aiter_importable(): + # The usable fallback is per-op -- "fa2" for attention, "native" for the + # rest -- so take it from the table rather than naming one that this op + # does not accept. + alt = aiter_fallback_backend(op) + advice = f"use backend={alt!r}" if alt else "use a non-AITER backend" 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'." + f"match older releases and would corrupt arguments silently. Upgrade " + f"amd-aiter or {advice}." ) 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 " - "section in the README) or use backend='native'." + f"installed or failed to import. Install it (see the AITER Support " + f"section in the README) or {advice}." ) diff --git a/flashinfer/arch_caps.py b/flashinfer/arch_caps.py index 9f9db6aba1..4db19d0eb9 100644 --- a/flashinfer/arch_caps.py +++ b/flashinfer/arch_caps.py @@ -34,6 +34,7 @@ "Capability", "KnownBad", "Support", + "aiter_fallback_backend", "capability_available", "capability_reason", "normalize_arch", @@ -597,6 +598,16 @@ def _blocking_reason(op: str, backend: str, arch: str) -> Optional[str]: return None +def aiter_fallback_backend(op: str) -> Optional[str]: + """The backend to suggest when AITER cannot serve ``op``, or None if unknown. + + "fa2" for attention, "native" elsewhere -- naming the wrong one sends the + reader to a value the op rejects. + """ + cap = _BY_KEY.get((op, "aiter")) + return cap.fallback if cap is not None and cap.fallback else None + + def capability_reason(device, op: str, backend: str) -> Optional[str]: """Why ``backend`` cannot serve ``op`` on ``device``, or ``None`` if it can. diff --git a/tests/rocm_tests/test_aiter_version_gate_hip.py b/tests/rocm_tests/test_aiter_version_gate_hip.py index b59fb0822c..fd9dba970e 100644 --- a/tests/rocm_tests/test_aiter_version_gate_hip.py +++ b/tests/rocm_tests/test_aiter_version_gate_hip.py @@ -24,9 +24,12 @@ ("0.1.9", False), ("0.1.10", False), # the pin this repo shipped before the 0.1.16 floor ("0.1.15.post9", False), + # A pre-release of the floor is below it. Comparing on base_version + # strips the segment and wrongly admits this one. + ("0.1.16.dev0", 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. + # The nightly wheel, and the only cp314 build: .dev0 *of post3*, which + # sorts above 0.1.16 rather than below it. ("0.1.16.post3.dev0+g620287969.d20260725", True), ("0.1.21", True), ("0.2.0", True), From f7317335dbb9a3ec735cbd3e5d2afc1dba614430 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 18:50:34 -0400 Subject: [PATCH 28/43] fix(rocm): re-guard the soft cap when a native page size degrades softcap_kv_len is None whenever page_size is in _aiter_native_page_sizes(), which disarms the guard -- correct when native paging happens, since mha_batch_prefill is exact. But that predicate is only a hint: the runtime probe _aiter_native_paging_available() runs later, and when it degrades the call to flat-gather it lands on the defective mha_varlen_fwd having already passed a guard that assumed otherwise. This is the residual gap the PR description names. Re-check against the real _max_kv_len once the probe has settled the route. auto demotes to fa2 with a reason and a warning, matching every other fallback in this file; explicit aiter raises. Both honour the existing rule that a demotion cannot happen once graph-captured flat-gather buffers exist. The gap is not a corner case on this stack. The probe returns False for page_size=128 -- a size the predicate advertises as native -- on BOTH gfx942 and gfx950 with amd-aiter 0.1.16.post3, so flat-gather is the route soft-cap prefill actually takes here. That also invalidated test_paged_softcap_guard_tracks_the_paging_route, whose premise is that a native page size dispatches to mha_batch_prefill. It failed identically on both arches before this commit; it now skips when the probe says the page size cannot be served natively. Version boundary measured while scoping this, driving mha_varlen_fwd directly (max abs err vs fp32, fp16, causal, head_dim=128, cap=8.0): aiter 0.1.10 gfx942 0.0011 gfx950 0.0010 aiter 0.1.16 gfx950 0.309 So the defect is 0.1.16-specific, which is consistent with the guard being ungated here given this branch floors AITER at 0.1.16. Verified on gfx942 (MI300X) and gfx950 (MI350X), rocm 7.14 / torch 2.12 / aiter 0.1.16.post3. New test A/B'd: both parametrizations fail with the re-check neutered. --- flashinfer/prefill_rocm.py | 61 +++++++++---- .../test_batch_prefill_kernels_hip.py | 90 +++++++++++++++++++ 2 files changed, 134 insertions(+), 17 deletions(-) diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index 57d2afdcd7..2b84501bdc 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -308,12 +308,15 @@ def _aiter_ops_importable() -> bool: # 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 + from .aiter_utils import _aiter_version_supported, _ensure_aiter_gpu_archs _ensure_aiter_gpu_archs() import aiter.ops # noqa: F401 - return True + # Same ABI floor aiter_utils._aiter_importable applies. This probe is the + # one prefill and decode reach first, so leaving it out would let auto + # route into the layouts the vendored structs no longer match. + return _aiter_version_supported() except Exception: return False @@ -1567,12 +1570,9 @@ def single_prefill_with_kv_cache( # Outside the probe on purpose: this raises ArchCapabilityError, which # gates known-bad toolchains and must never be demoted to a silent fa2. _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." - ) + # Hard constraints first: a caller in the defect region *and* on an + # unsupported layout should hear about the layout, which is the thing + # they control, not about a soft-cap defect they may not have hit. if pos_encoding_mode != "NONE": raise ValueError( f"AITER backend does not support pos_encoding_mode={pos_encoding_mode!r}; " @@ -1583,6 +1583,12 @@ def single_prefill_with_kv_cache( f"AITER backend only supports kv_layout='NHD'; got {kv_layout!r}. " "use backend='fa2' or backend='auto' instead." ) + 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." + ) # 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. @@ -2371,24 +2377,45 @@ def plan( dev_idx, ) if not use_native_paging: - # The flat-gather route dispatches through - # get_aiter_mha_varlen_fwd_handle, whose .so variant is keyed on - # (dtype, causal, has_lse, has_logits_cap). AITER pre-ships only - # a subset of that family, so bootstrap the rest here rather than - # let the C++ dlopen fail inside run(). The helper compiles both - # has_lse variants because plan() can't know which one run() will - # request; causal is fixed per plan() call. + # The guard above disarmed itself because the page size looked + # native; the probe just proved otherwise, so this call takes + # flat-gather after all. Re-check against the real kv_len. + softcap_now = softcap_kv_len is None and _aiter_softcap_defect( + causal, logits_soft_cap, head_dim_qk, self._max_kv_len + ) # Demoting after a graph capture would null the flat-gather # buffers the captured graph still points at, so once they # exist under capture the failure has to stay an exception. - if resolved_from_auto and not ( + demotable = resolved_from_auto and not ( self.is_cuda_graph_enabled and self._aiter_flat_gather_idx is not None - ): + ) + if softcap_now and not demotable: + raise ValueError( + "AITER miscomputes logits_soft_cap for causal head_dim=128 " + f"prefill with kv_len >= 512 (through amd-aiter " + f"{_AITER_SOFTCAP_DEFECT_THROUGH}); this page size fell back " + "to the flat-gather kernel. Use backend='fa2'." + ) + if softcap_now: + reason = ( + "aiter native paging was unavailable for page_size=" + f"{page_size}, and the flat-gather kernel miscomputes " + "logits_soft_cap for causal head_dim=128 with kv_len >= 512 " + f"(through amd-aiter {_AITER_SOFTCAP_DEFECT_THROUGH})" + ) + logger.warning("auto backend falling back to fa2: %s", reason) + elif demotable: reason = _aiter_batch_ragged_available( q_data_type, has_logits, causal, head_dim_qk, dev_idx ) else: + # The flat-gather route dispatches through + # get_aiter_mha_varlen_fwd_handle, whose .so variant is keyed + # on (dtype, causal, has_lse, has_logits_cap). AITER + # pre-ships only a subset, so bootstrap the rest here rather + # than let the C++ dlopen fail inside run(); both has_lse + # variants, since plan() cannot know which run() will want. _aiter_bootstrap_batch_ragged_prefill( q_data_type, has_logits, diff --git a/tests/rocm_tests/test_batch_prefill_kernels_hip.py b/tests/rocm_tests/test_batch_prefill_kernels_hip.py index 31ee9d29e7..ddc3c33061 100644 --- a/tests/rocm_tests/test_batch_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_batch_prefill_kernels_hip.py @@ -957,6 +957,89 @@ def _reject(*args, **kwargs): torch.testing.assert_close(o, wrapper_ref.run(q, kv_data), rtol=2e-2, atol=2e-2) +def _plan_softcap_flat_gather(wrapper, device, page_size, kv_len, dtype): + """plan() a defect-shape call (causal, cap>0, head_dim=128, kv_len>=512).""" + batch_size, qo_len = 1, 16 + num_qo_heads = num_kv_heads = 8 + num_pages = (kv_len + page_size - 1) // page_size + qo_indptr = ( + torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * qo_len + ) + kv_indptr = ( + torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * num_pages + ) + kv_indices = torch.arange( + 0, num_pages * batch_size, dtype=torch.int32, device=device + ) + kv_last_page_len = torch.full( + (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32, device=device + ) + wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + kv_last_page_len, + num_qo_heads, + num_kv_heads, + 128, + page_size, + causal=True, + logits_soft_cap=30.0, + q_data_type=dtype, + kv_data_type=dtype, + ) + + +@pytest.mark.parametrize("backend", ["auto", "aiter"]) +def test_softcap_guard_survives_a_native_page_size_degrading(backend, monkeypatch): + """The soft-cap guard disarms on a native page size; the probe can undo that. + + softcap_kv_len is None whenever page_size looks native, because native paging + uses mha_batch_prefill and is exact. When the runtime probe then degrades the + call to flat-gather, it lands on the defective mha_varlen_fwd and has to be + re-guarded against the real kv_len. + """ + # Plumbing only -- no numbers compared, so the gfx950 causal gate is irrelevant. + monkeypatch.setenv("FLASHINFER_ARCH_ALLOW_KNOWN_BAD", "1") + # Strict turns the probe into a raise, so there would be nothing to demote. + monkeypatch.delenv("FLASHINFER_AITER_STRICT", raising=False) + 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") + page_size = 128 + if page_size not in _aiter_native_page_sizes(): + pytest.skip(f"page_size={page_size} is not native on this amd-aiter build") + + def _reject(*args, **kwargs): + raise RuntimeError( + "invalid argument for batch_prefill: no matching kernel found. " + f"page_size={page_size}, num_pages=1, dtype=bf16" + ) + + _aiter_native_paging_available.cache_clear() + monkeypatch.setattr( + flashinfer.prefill_rocm, "_aiter_bootstrap_batch_prefill", _reject + ) + workspace = torch.empty(256 * 1024 * 1024, dtype=torch.int8, device=device) + try: + wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + workspace, "NHD", backend=backend + ) + if backend == "aiter": + with pytest.raises(ValueError, match="logits_soft_cap"): + _plan_softcap_flat_gather( + wrapper, device, page_size, 1024, torch.bfloat16 + ) + else: + _plan_softcap_flat_gather(wrapper, device, page_size, 1024, torch.bfloat16) + assert wrapper._backend == "fa2", ( + "auto stayed on the defective flat-gather kernel" + ) + assert "logits_soft_cap" in (wrapper.backend_fallback_reason or "") + finally: + _aiter_native_paging_available.cache_clear() + + def test_batch_prefill_aiter_strict_mode_raises(monkeypatch): """FLASHINFER_AITER_STRICT=1 must surface the AITER failure instead of degrading.""" # Asserts plumbing and compares no numbers, so the ROCm 7.2 gfx950 causal @@ -1126,6 +1209,13 @@ def plan(page_size): kv_data_type=torch.float16, ) + # "Native page size" is a hint; only the probe settles the route. Where it + # degrades to flat-gather the guard is meant to fire, so the premise is gone. + if not _aiter_native_paging_available( + torch.float16, True, True, native[0], head_dim, device.index or 0 + ): + pytest.skip(f"aiter cannot serve page_size={native[0]} natively on this build") + plan(native[0]) non_native = next( From 0b01a4fc648ae85bf39aecef1a65a32ca8ea9db4 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 21:07:46 -0400 Subject: [PATCH 29/43] fix(aiter): reject a cross-device silu_and_mul instead of guessing The device guard was taken from input.device() while the kernel writes to out, so a caller passing tensors on two devices got a guard for the wrong one and AITER launched against mismatched pointers -- undefined behaviour rather than an error. activation.cu already checks this; match it. --- flashinfer/csrc_rocm/activation_aiter.cu | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flashinfer/csrc_rocm/activation_aiter.cu b/flashinfer/csrc_rocm/activation_aiter.cu index e106e47dc0..c58e7187d3 100644 --- a/flashinfer/csrc_rocm/activation_aiter.cu +++ b/flashinfer/csrc_rocm/activation_aiter.cu @@ -19,6 +19,8 @@ #include "aiter_tensor_compat.h" void silu_and_mul_aiter(at::Tensor out, at::Tensor input) { + TORCH_CHECK(out.device() == input.device(), "silu_and_mul: out is on ", out.device(), + " but input is on ", input.device()); const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(input.device()); // The kernel indexes linearly, so strides in aiter_tensor_t are not honoured. From 824014a268df31c85cb01002c0c3ca6d5fcc8121 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 21:07:47 -0400 Subject: [PATCH 30/43] test(rocm): only page the soft-cap guard test at sizes that divide kv_len num_pages was kv_len // page_size, which silently yields zero pages for a native page size larger than kv_len and drops a partial trailing page otherwise -- neither of which the test models, since it passes a single kv_last_page_len of page_size. Today _aiter_native_page_sizes() starts at 128 and kv_len is 512, so nothing is affected; a future AITER whose smallest native size exceeds kv_len would have built invalid paging input and failed for a reason unrelated to the guard. Filter to the divisors and skip when none remain. --- tests/rocm_tests/test_batch_prefill_kernels_hip.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/rocm_tests/test_batch_prefill_kernels_hip.py b/tests/rocm_tests/test_batch_prefill_kernels_hip.py index ddc3c33061..3d61b33542 100644 --- a/tests/rocm_tests/test_batch_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_batch_prefill_kernels_hip.py @@ -1182,11 +1182,14 @@ def test_paged_softcap_guard_tracks_the_paging_route(): 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()) + kv_len, qo_len, num_heads, head_dim, soft_cap = 512, 37, 4, 128, 8.0 + # Only page sizes that divide kv_len: a partial trailing page would need a + # kv_last_page_len this test does not model, and one larger than kv_len + # floor-divides to zero pages. + native = sorted(p for p in _aiter_native_page_sizes() if p <= kv_len and kv_len % p == 0) if not native: - pytest.skip("no native AITER page size on this aiter build") + pytest.skip(f"no native AITER page size divides kv_len={kv_len}") - 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): From cf8662e43b74f506fd53d97645dc7c93de378739 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 21:07:47 -0400 Subject: [PATCH 31/43] docs: say that causal soft-capped prefill does not use AITER The support matrix documents the gfx950 ROCm 7.2 miscompile but not this one, so a user whose causal soft-cap call quietly ran on fa2 had nothing to read. The routing is deliberate and not version-gated, which makes it a standing property of the backend rather than a transient workaround. --- README.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/README.md b/README.md index 9a91eccd61..ccd98de610 100644 --- a/README.md +++ b/README.md @@ -188,6 +188,13 @@ matching `tests/rocm_tests/test_*_hip.py`; `single_decode` is exercised from the batch-decode, sliding-window, and logits-cap files, and `quantization` by `tests/utils/test_quantization.py`. +**Soft-capped causal prefill does not use AITER.** AITER miscomputes +`logits_soft_cap` for causal prefill with `head_dim=128` and `kv_len >= 512` +(through amd-aiter 0.1.21), so `backend="auto"` serves those calls with `fa2` +and `backend="aiter"` raises rather than returning wrong numbers. Every other +soft-cap shape — non-causal, other head dims, shorter contexts — still goes to +AITER. + ## `torch.compile` Set `FLASHINFER_USE_TORCH_CUSTOM_OPS=1` **before** importing `flashinfer` From 73c9037a698624f91aa90a96e3f98a6cd369dee7 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 21:29:59 -0400 Subject: [PATCH 32/43] fix(rocm): skip the symlink loop when the versioned-lib glob matches nothing An unmatched glob stays literal in POSIX sh, so the loop would run once with so='.../lib*.so.[0-9]*' and create a symlink literally named lib*.so. Reproduced in an empty directory: $ sh -c 'set -eu; for so in "$PWD"/lib/lib*.so.[0-9]*; do base="${so%%.so.*}.so"; [ -e "$base" ] || ln -s "$(basename "$so")" "$base"; done' $ ls lib*.so Harmless today because the pip ROCm SDK does ship versioned libs, but it would turn a base-image change into a confusing link-time failure rather than a clean no-op. Same command with the guard leaves the directory untouched. --- .devcontainer/rocm/Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/.devcontainer/rocm/Dockerfile b/.devcontainer/rocm/Dockerfile index 6cc456ed1b..ba02c57515 100644 --- a/.devcontainer/rocm/Dockerfile +++ b/.devcontainer/rocm/Dockerfile @@ -102,6 +102,7 @@ ENV CPATH="/opt/rocm-headers/include${CPATH:+:$CPATH}" # 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 \ + [ -e "$so" ] || continue; \ base="${so%%.so.*}.so"; [ -e "$base" ] || ln -s "$(basename "$so")" "$base"; \ done; \ ln -sf ../lib/llvm/bin/clang++ "$ROCM_SDK/bin/amdclang++"; \ From ca134ca0657dc7e147ece921774971abe92449a7 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 27 Aug 2026 23:56:41 -0400 Subject: [PATCH 33/43] style: wrap the native page-size filter to ruff's line length The CI pre-commit job failed on 824014a26: ruff format wraps the sorted() generator, and my local run passed only because it executed in the main checkout rather than this worktree. --- tests/rocm_tests/test_batch_prefill_kernels_hip.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/rocm_tests/test_batch_prefill_kernels_hip.py b/tests/rocm_tests/test_batch_prefill_kernels_hip.py index 3d61b33542..f7206574f5 100644 --- a/tests/rocm_tests/test_batch_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_batch_prefill_kernels_hip.py @@ -1186,7 +1186,9 @@ def test_paged_softcap_guard_tracks_the_paging_route(): # Only page sizes that divide kv_len: a partial trailing page would need a # kv_last_page_len this test does not model, and one larger than kv_len # floor-divides to zero pages. - native = sorted(p for p in _aiter_native_page_sizes() if p <= kv_len and kv_len % p == 0) + native = sorted( + p for p in _aiter_native_page_sizes() if p <= kv_len and kv_len % p == 0 + ) if not native: pytest.skip(f"no native AITER page size divides kv_len={kv_len}") From eacfced8686870a22b49274637feba7611043f74 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Fri, 28 Aug 2026 00:25:21 -0400 Subject: [PATCH 34/43] fix(aiter): say "too old" rather than "not installed" for a below-floor AITER _aiter_ops_importable() now returns False for an installed-but-pre-0.1.16 AITER as well as a missing one, so _require_aiter_runtime told an explicit backend='aiter' caller to git-clone a package they already have. Split the two conditions, matching what aiter_utils.require_aiter already does. --- flashinfer/prefill_rocm.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index 2b84501bdc..ad89034bb5 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -331,6 +331,19 @@ def _require_aiter_runtime(device: torch.device, op: str = "batch_prefill") -> N """ require_capability(device, op, "aiter") if not _aiter_ops_importable(): + from .aiter_utils import ( + AITER_MIN_VERSION, + _aiter_installed_version, + _aiter_version_supported, + ) + + installed = _aiter_installed_version() + if installed is not None and not _aiter_version_supported(): + raise ImportError( + f"The AITER backend requires amd-aiter >= {AITER_MIN_VERSION}, but " + f"{installed} is installed. The vendored struct layouts do not match " + "older releases and would corrupt arguments silently." + ) raise ImportError( "The 'aiter' package is required for the AITER backend. " "Install it via:\n" From 85e7868a258ecbd432650abe518d087f8f2b2621 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Fri, 28 Aug 2026 00:38:15 -0400 Subject: [PATCH 35/43] docs(aiter): say what kv_len=None actually signals to the defect check The docstring read as though None only meant "the caller does not know the length". The paged wrapper also passes None deliberately, to disarm the check when a native page size means the call will not reach mha_varlen_fwd -- and re-checks after the runtime probe in case that prediction does not hold. --- flashinfer/prefill_rocm.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index ad89034bb5..8d0c328434 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -365,8 +365,12 @@ def _aiter_softcap_defect( 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. + kv_len >= 512. Non-causal is exact. + + kv_len=None disarms the check. Callers pass it either because they do not + know the length or because they know the call will not reach + mha_varlen_fwd -- the paged wrapper does the latter for a natively-paged + page size, and re-checks after the runtime probe in case it degrades. Deliberately not version-gated: auto-expiring on an AITER newer than _AITER_SOFTCAP_DEFECT_THROUGH would silently re-enable a wrong-answer path From 4583a59379589a69ce31b4567690e0e6e6ca5ba0 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Fri, 28 Aug 2026 00:38:15 -0400 Subject: [PATCH 36/43] test(rocm): assert paged soft-cap numerics on both paging routes Covers the guard by its output rather than its decision: page_size=128 is advertised native but the runtime probe can still demote it to flat-gather, and 16 never is, so both reach the defective mha_varlen_fwd unless the guard redirects. Comparing against an fp32 reference catches a wrong answer whichever route plan() takes, where asserting on the chosen backend would not. warmup_jit gains the soft-cap variants so the modules these cases need are built with the rest. 8 passed, 1 skipped on gfx942 / MI300X, ROCm 7.14.60850, torch 2.12.0, amd-aiter 0.1.16.post3. --- .../test_batch_prefill_kernels_hip.py | 60 ++++++++++++++++++- 1 file changed, 59 insertions(+), 1 deletion(-) diff --git a/tests/rocm_tests/test_batch_prefill_kernels_hip.py b/tests/rocm_tests/test_batch_prefill_kernels_hip.py index f7206574f5..86a27f89ed 100644 --- a/tests/rocm_tests/test_batch_prefill_kernels_hip.py +++ b/tests/rocm_tests/test_batch_prefill_kernels_hip.py @@ -48,7 +48,7 @@ def warmup_jit(): [128, 256], # head_dims [0], # pos_encoding_modes [False], # use_sliding_windows - [False], # use_logits_soft_caps + [False, True], # use_logits_soft_caps [False], # use_fp16_qk_reductions ), verbose=False, @@ -1172,6 +1172,64 @@ def test_ragged_softcap_avoids_broken_aiter_kernel(kv_len, qo_len): torch.testing.assert_close(wrapper.run(q, k, v).float(), ref, rtol=1e-3, atol=1e-3) +@pytest.mark.parametrize("kv_len", [512, 2048]) +@pytest.mark.parametrize("page_size", [128, 16]) +def test_paged_softcap_is_numerically_correct(kv_len, page_size): + """Causal soft-cap paged prefill must be correct whichever route plan() picks. + + page_size=128 is advertised as native but the probe can degrade it to + flat-gather, and 16 never is; both land on the defective mha_varlen_fwd + unless the guard redirects. Asserts the numbers, not the route. + """ + 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, qo_len = 128, 4, 8.0, 37 + torch.manual_seed(0) + q = torch.randn(qo_len, num_heads, head_dim, dtype=torch.float16, device=device) + num_pages = (kv_len + page_size - 1) // page_size + kv_data = torch.randn( + num_pages, 2, page_size, num_heads, head_dim, dtype=torch.float16, device=device + ) + # Flatten the pages back into the [kv_len, heads, dim] view the reference wants. + k = kv_data[:, 0].reshape(-1, num_heads, head_dim)[:kv_len] + v = kv_data[:, 1].reshape(-1, num_heads, head_dim)[:kv_len] + + 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.prefill.BatchPrefillWithPagedKVCacheWrapper( + workspace, "NHD", backend="auto" + ) + 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([(kv_len - 1) % page_size + 1], 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, + ) + torch.testing.assert_close( + wrapper.run(q, kv_data).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. From 53df88c4673c85181caf9f9cd85dfbf3f82b6500 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Fri, 28 Aug 2026 01:22:49 -0400 Subject: [PATCH 37/43] fix(aiter): reject a host tensor in to_aiter instead of encoding device_id=-1 aiter_tensor_t::is_gpu() keys off device_id >= 0, so a CPU tensor was converted to a well-formed POD carrying a host pointer and -1. Every POD entry point here is GPU-only, so that can only end as a fault inside AITER, further from the mistake than it needs to be. --- flashinfer/csrc_rocm/aiter_tensor_compat.h | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/flashinfer/csrc_rocm/aiter_tensor_compat.h b/flashinfer/csrc_rocm/aiter_tensor_compat.h index 2741853291..23351c1389 100644 --- a/flashinfer/csrc_rocm/aiter_tensor_compat.h +++ b/flashinfer/csrc_rocm/aiter_tensor_compat.h @@ -57,8 +57,10 @@ inline aiter_tensor_t to_aiter(const at::Tensor& t) { 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()); + // Every POD entry point is GPU-only, so reject a host tensor here rather than + // hand AITER a device_id of -1 and let it fault on a host pointer. + TORCH_CHECK(t.is_cuda(), "aiter_tensor_t requires a GPU tensor, got ", t.device()); + out.device_id = static_cast(t.device().index()); return out; } From 4f042bb05e8ddcf62ac9205eca066eadd655be50 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Fri, 28 Aug 2026 01:22:49 -0400 Subject: [PATCH 38/43] fix(aiter): name the version floor in the auto-fallback reason Third site of the same distinction: _aiter_ops_importable() is False for an installed-but-too-old AITER as well as a missing one, so the auto fallback logged "aiter package not installed" for what is really an ABI mismatch -- the failure mode most worth naming, since it is the one that would otherwise corrupt arguments silently. The reason string also reaches backend_fallback_reason and the benchmark harness's CSV, so it is what a user sees when asking why AITER was not used. --- flashinfer/prefill_rocm.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index 8d0c328434..eead99ebc7 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -446,10 +446,23 @@ def _auto_select_prefill_backend( return "fa2", reason if not _aiter_ops_importable(): - reason = ( - "aiter package not installed " - "(see https://github.com/ROCm/aiter for install instructions)" + from .aiter_utils import ( + AITER_MIN_VERSION, + _aiter_installed_version, + _aiter_version_supported, ) + + installed = _aiter_installed_version() + if installed is not None and not _aiter_version_supported(): + reason = ( + f"amd-aiter {installed} is below the {AITER_MIN_VERSION} ABI floor " + "(the vendored struct layouts do not match older releases)" + ) + else: + reason = ( + "aiter package not installed " + "(see https://github.com/ROCm/aiter for install instructions)" + ) key = (device, "import_failed") if key not in _aiter_auto_warned: _aiter_auto_warned.add(key) From 20e5308daf1518a8276aa7c9665301a86a6ed8a5 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Fri, 28 Aug 2026 01:34:28 -0400 Subject: [PATCH 39/43] docs(aiter): point the install error at the page that has the instructions The message named an "AITER Support" section of README.md, which #307 removed when it split backend and install detail into docs/rocm/backends.md. --- flashinfer/aiter_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flashinfer/aiter_utils.py b/flashinfer/aiter_utils.py index f06482399d..7025b0e6df 100644 --- a/flashinfer/aiter_utils.py +++ b/flashinfer/aiter_utils.py @@ -155,8 +155,8 @@ def require_aiter(device: torch.device, op: str) -> None: ) raise ValueError( f"backend='aiter' for {op} requires the aiter package, which is not " - f"installed or failed to import. Install it (see the AITER Support " - f"section in the README) or {advice}." + f"installed or failed to import. Install it (see docs/rocm/backends.md) " + f"or {advice}." ) From d5a3550db112602633482c99f4f9a7049f55e555 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Fri, 28 Aug 2026 01:42:12 -0400 Subject: [PATCH 40/43] docs(aiter): stop the install error recommending a source build The message told users to git-clone AITER and `setup.py develop`, which is the one install this PR documents as untested: master is many releases past the pin with a different C ABI, so following it lands exactly the silent by-value corruption AITER_MIN_VERSION exists to prevent. Point at the wheel and the page that carries the index. --- flashinfer/prefill_rocm.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index eead99ebc7..12dda096ef 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -345,10 +345,11 @@ def _require_aiter_runtime(device: torch.device, op: str = "batch_prefill") -> N "older releases and would corrupt arguments silently." ) raise ImportError( - "The 'aiter' package is required for the AITER backend. " - "Install it via:\n" - " git clone --recursive https://github.com/ROCm/aiter.git\n" - " cd aiter && python3 setup.py develop" + "The 'aiter' package is required for the AITER backend and is not " + f"installed. Install a wheel >= {AITER_MIN_VERSION}; see " + "docs/rocm/backends.md for the index and the pinned version. A source " + "build tracks master, whose C ABI does not match the structs vendored " + "here." ) From 5e838b3c95bc33dc871efc427a22f72b6e8abdbf Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Fri, 28 Aug 2026 02:00:48 -0400 Subject: [PATCH 41/43] fix(aiter): let both import-failure reasons warn, and drop the last source build Two problems in the auto-fallback import branch. The warn-once key was the constant "import_failed" while the branch now produces two distinct reasons, so whichever fired first hid the other for the rest of the process -- and "installed but below the ABI floor" and "not installed" call for different actions. Key on the reason, matching the branch above. Its message also still pointed at the AITER repo, which reads as an invitation to build from source. mla_rocm.py spelled the source build out in full; nobody reported that one, but it is the same defect, and `grep -rn "ROCm/aiter.git\|setup.py develop" flashinfer/` is now empty. --- flashinfer/mla_rocm.py | 10 ++++++---- flashinfer/prefill_rocm.py | 9 ++++++--- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/flashinfer/mla_rocm.py b/flashinfer/mla_rocm.py index b68739749f..173ee57440 100644 --- a/flashinfer/mla_rocm.py +++ b/flashinfer/mla_rocm.py @@ -130,11 +130,13 @@ def _require_aiter_mla(device: torch.device) -> None: try: _aiter_mla() except ImportError as exc: + from .aiter_utils import AITER_MIN_VERSION + raise ImportError( - "The 'aiter' package is required for MLA on ROCm. " - "Install it via:\n" - " git clone --recursive https://github.com/ROCm/aiter.git\n" - " cd aiter && python3 setup.py develop" + "The 'aiter' package is required for MLA on ROCm. Install a wheel >= " + f"{AITER_MIN_VERSION}; see docs/rocm/backends.md for the index and the " + "pinned version. A source build tracks master, whose C ABI does not " + "match the structs vendored here." ) from exc diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index 12dda096ef..61d2896866 100755 --- a/flashinfer/prefill_rocm.py +++ b/flashinfer/prefill_rocm.py @@ -461,10 +461,13 @@ def _auto_select_prefill_backend( ) else: reason = ( - "aiter package not installed " - "(see https://github.com/ROCm/aiter for install instructions)" + "aiter package not installed (see docs/rocm/backends.md for the " + "wheel index and pinned version)" ) - key = (device, "import_failed") + # Keyed on the reason, like the branch above: a constant key would let + # whichever condition fired first hide the other for the rest of the + # process, and "too old" and "not installed" want different actions. + key = (device, reason) if key not in _aiter_auto_warned: _aiter_auto_warned.add(key) logger.warning("auto backend falling back to fa2: %s", reason) From fd468ab4419edf13f0ddf63289b151dad9570cff Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Fri, 28 Aug 2026 02:19:34 -0400 Subject: [PATCH 42/43] docs: stop the README overstating the soft-cap fallback It read as though causal soft-capped prefill never uses AITER, which contradicts the routing in this same PR: paged prefill with a native page size dispatches to mha_batch_prefill, which is exact, and only falls back when the run-time probe demotes the call to a flat gather. Single and ragged prefill do always fall back, since both go through mha_varlen_fwd -- name the kernel rather than the API. --- README.md | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index ccd98de610..21f935979b 100644 --- a/README.md +++ b/README.md @@ -188,12 +188,15 @@ matching `tests/rocm_tests/test_*_hip.py`; `single_decode` is exercised from the batch-decode, sliding-window, and logits-cap files, and `quantization` by `tests/utils/test_quantization.py`. -**Soft-capped causal prefill does not use AITER.** AITER miscomputes -`logits_soft_cap` for causal prefill with `head_dim=128` and `kv_len >= 512` -(through amd-aiter 0.1.21), so `backend="auto"` serves those calls with `fa2` -and `backend="aiter"` raises rather than returning wrong numbers. Every other -soft-cap shape — non-causal, other head dims, shorter contexts — still goes to -AITER. +**Soft-capped causal prefill avoids one AITER kernel.** AITER's +`mha_varlen_fwd` miscomputes `logits_soft_cap` for causal prefill with +`head_dim=128` and `kv_len >= 512` (through amd-aiter 0.1.21). Single and +ragged prefill always dispatch through it, so `backend="auto"` serves those +calls with `fa2` and `backend="aiter"` raises rather than returning wrong +numbers. Paged prefill keeps using AITER when the page size is native, since +that route takes `mha_batch_prefill`, which is exact — it falls back only if +the run-time probe demotes the call to a flat gather. Every other soft-cap +shape — non-causal, other head dims, shorter contexts — is unaffected. ## `torch.compile` From 5c6070491846f88c437583087a3bc9ba300a316a Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Fri, 28 Aug 2026 11:35:07 -0400 Subject: [PATCH 43/43] docs(rocm): stop check_rocm_arch claiming an unconditional torch check The summary said it validates that "this GPU, ROCm and torch" can build the kernels, but the torch half only runs when PYTORCH_ROCM_ARCH is set -- unset, _get_rocm_arch_flags() reports the visible cards rather than the build, which is why the check is gated. Say which part is conditional and why. --- flashinfer/jit/rocm/core.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/flashinfer/jit/rocm/core.py b/flashinfer/jit/rocm/core.py index 1e2cb7fd72..f3ad774b7c 100644 --- a/flashinfer/jit/rocm/core.py +++ b/flashinfer/jit/rocm/core.py @@ -9,10 +9,12 @@ def check_rocm_arch() -> None: - """Validate that this GPU, ROCm and torch can build FlashInfer's kernels. + """Validate that this GPU and ROCm can build FlashInfer's kernels. Delegates to hip_utils so the ROCm version, the ported-architecture list and - PYTORCH_ROCM_ARCH (when set) are checked in one place. + PYTORCH_ROCM_ARCH are checked in one place. Torch's own architecture list is + only consulted when PYTORCH_ROCM_ARCH is set; unset, it reports the visible + cards rather than anything about the build. """ import torch.utils.cpp_extension as torch_cpp_ext