From ce7e79b32c5dc3e2c7536dbe7f3b4e041195c1a6 Mon Sep 17 00:00:00 2001 From: Michael <13900043+michaelzhang-ai@users.noreply.github.com> Date: Tue, 1 Sep 2026 04:06:32 -0700 Subject: [PATCH] [AMD] Fix nightly ROCm 7.0 image build: patch missing include in AITER topk kernel (#36216) Co-authored-by: Cursor Agent Co-authored-by: quitenode Co-authored-by: Chen Co-authored-by: YC Yen-Ching Tseng --- docker/rocm.Dockerfile | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) 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}" \