diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 679a5993d..872eb564c 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -582,6 +582,25 @@ RUN git clone ${AITER_REPO} \ git revert --no-edit --no-commit e708f6c15; \ fi +# The pinned AITER revision uses std::optional in topk_per_row_kernels.cu without +# including : ROCm/aiter#4702 removed the torch headers that previously +# supplied it transitively. ROCm 7.0 therefore fails module_top_k_per_row, while +# ROCm 7.2 toolchains still obtain from another header. +# GPU_ARCH keeps the full selected stage name; only unsuffixed gfx942/gfx950 +# denote ROCm 7.0, while newer flavors carry a -rocm... suffix. Remove this +# backport once AITER_COMMIT contains the upstream fix from ROCm/aiter#4853. +RUN python3 <<'PY' +import os +from pathlib import Path + +p = Path("/sgl-workspace/aiter/csrc/kernels/topk_per_row_kernels.cu") +anchor = "#include \n" +if os.environ["GPU_ARCH"] in {"gfx942", "gfx950"} and p.exists(): + s = p.read_text() + if "std::optional" in s and "#include " not in s and anchor in s: + p.write_text(s.replace(anchor, "#include \n" + anchor, 1)) +PY + RUN cd aiter \ && echo "[AITER] GPU_ARCH=${GPU_ARCH}" \ && echo "[AITER] AITER_USE_SYSTEM_TRITON=${AITER_USE_SYSTEM_TRITON}" \