[ROCm][Perf][Bugfix] Multi-stream perf improvements; rocprofiler fixes (#55099)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Matt committed
Sep 8, 2026 at 17:42 UTC
bc8587f82944b0f1836ed3f9c1ae6e99421d9e6b
4 files changed
+161
-104
docker/Dockerfile.rock_base
+80
@@ -10,6 +10,8 @@ ARG AITER_BRANCH="v0.1.21.post1"
10
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
11
ARG MORI_BRANCH="f95cd5e0ca7b730f92f9ffdfb358484d6965d9f7" # TheRock compatible v1.1.0
12
ARG MORI_REPO="https://github.com/ROCm/mori.git"
13
+ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git"
14
+ARG ROCM_RUNTIME_COMMIT="25e14349f2606a6b95c4ccd45eb92a1835646cf2"
15
16
17
# TheRock configuration
@@ -246,6 +248,61 @@ RUN pip install pyyaml && cd aiter \
248
RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
249
250
251
+###
252
+### ROCr + CLR build
253
+###
254
+# Brings in multi-stream perf improvements; drop after these arrive in an official release
255
+FROM base AS build_rocm_runtime
256
+ARG ROCM_SYSTEMS_REPO
257
+ARG ROCM_RUNTIME_COMMIT
258
+# libgl-dev: rocclr does find_package(OpenGL REQUIRED) unconditionally and the wheel image has
259
+# no GL headers. The link must not pick up a libGL DT_NEEDED, which is asserted after the build.
260
+RUN apt-get update \
261
+ && apt-get install -y --no-install-recommends g++ libdw-dev xxd libgl-dev libglx-dev \
262
+ && rm -rf /var/lib/apt/lists/*
263
+RUN pip install --no-cache-dir CppHeaderParser
264
+RUN git init -q /src && cd /src \
265
+ && git remote add origin ${ROCM_SYSTEMS_REPO} \
266
+ && git config core.sparseCheckout true \
267
+ && git config remote.origin.promisor true \
268
+ && git config remote.origin.partialclonefilter blob:none \
269
+ && git sparse-checkout init --cone \
270
+ && git sparse-checkout set projects/rocr-runtime projects/clr projects/hip shared cmake \
271
+ && git fetch --filter=blob:none --depth 1 origin ${ROCM_RUNTIME_COMMIT} \
272
+ && git checkout -q FETCH_HEAD
273
+# libelf is taken from the SDK's own vendored sysdeps rather than apt -- the wheel image has no
274
+# /usr/include/libelf.h -- so ROCr links against the copy the wheel itself was built against.
275
+ENV SYSDEPS="$SDK_DEV/lib/rocm_sysdeps"
276
+# ROCr first, and installed into $SDK_DEV, because the clr build below needs its headers.
277
+RUN cd /src/projects/rocr-runtime \
278
+ && cmake -B build -G Ninja \
279
+ -DCMAKE_BUILD_TYPE=Release \
280
+ -DCMAKE_PREFIX_PATH="$SDK_DEV;$SYSDEPS" \
281
+ -DCMAKE_INSTALL_PREFIX="$SDK_DEV" \
282
+ -DLIBELF_INCLUDE_DIRS="$SYSDEPS/include" \
283
+ -DLIBELF_LIBRARIES="$SYSDEPS/lib/libelf.so" \
284
+ && cmake --build build --parallel 16 \
285
+ && cmake --install build \
286
+ && cmake --install build --prefix /rocr-install --strip
287
+RUN cd /src/projects/clr \
288
+ && cmake -B build -G Ninja \
289
+ -DCMAKE_BUILD_TYPE=Release \
290
+ -DCLR_BUILD_HIP=ON \
291
+ -DCLR_BUILD_OCL=OFF \
292
+ -DROCM_KPACK_ENABLED=ON \
293
+ -DHIP_COMMON_DIR=/src/projects/hip \
294
+ -DROCM_PATH="$SDK_DEV" \
295
+ -DCMAKE_PREFIX_PATH="$SDK_DEV;$SYSDEPS" \
296
+ -DCMAKE_INSTALL_PREFIX="$SDK_DEV" \
297
+ -DHIP_LLVM_ROOT="$SDK_DEV/lib/llvm" \
298
+ && cmake --build build --parallel 16 \
299
+ && cmake --install build --prefix /clr-install --strip
300
+RUN mkdir -p /staging/lib \
301
+ && cp -P /rocr-install/lib/libhsa-runtime64.so* /staging/lib/ \
302
+ && cp -P /clr-install/lib/libamdhip64.so* /staging/lib/ \
303
+ && ls -l /staging/lib
304
+
305
+
306
###
307
### Final Build
308
###
@@ -274,6 +331,25 @@ FROM base AS final
331
# teardown, so disable it. TODO: drop once ROCm's stock rocprofiler carries the fix.
332
ENV ROCPROFILER_QUEUE_INTERPOSITION=0
333
334
+COPY --from=build_rocm_runtime /staging/lib/ ${SDK_DEV}/lib/
335
+# Hard link the patched library onto every stock file name, in both SDK trees, so that the
336
+# SONAME lookup and the hard coded paths all resolve to it. Hard links rather than symlinks
337
+# because ldconfig rewrites SONAME symlinks but cannot re-point an inode. $SDK_CORE keeps its
338
+# own libhsa-runtime64 copy, which is on LD_LIBRARY_PATH and in ld.so.conf; relinking it is
339
+# belt and braces rather than strictly required, but it leaves no stock ROCr in the image.
340
+RUN set -eux; \
341
+ cd "$SDK_DEV/lib"; \
342
+ ours_hsa="$(readlink -f libhsa-runtime64.so.1)"; \
343
+ ours_hip="$(readlink -f libamdhip64.so.7)"; \
344
+ for f in "$SDK_DEV"/lib/libhsa-runtime64.so.1* "$SDK_CORE"/lib/libhsa-runtime64.so.1* \
345
+ "$SDK_DEV"/lib/libamdhip64.so.7* "$SDK_CORE"/lib/libamdhip64.so.7*; do \
346
+ [ -f "$f" ] && [ ! -L "$f" ] || continue; \
347
+ case "${f##*/}" in libhsa*) ours="$ours_hsa";; *) ours="$ours_hip";; esac; \
348
+ [ "$f" -ef "$ours" ] || ln -f "$ours" "$f"; \
349
+ done
350
+
351
+RUN ldconfig
352
+
353
ARG TRITON_VERSION
354
RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
355
pip install /install/*.whl \
@@ -304,6 +380,8 @@ ARG AITER_BRANCH
380
ARG AITER_REPO
381
ARG MORI_BRANCH
382
ARG MORI_REPO
383
+ARG ROCM_SYSTEMS_REPO
384
+ARG ROCM_RUNTIME_COMMIT
385
# Record hashes for Triton and pybind11 so we can verify later.
386
RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
387
&& echo "PYTHON_VERSION: ${PYTHON_VERSION}" >> /app/versions.txt \
@@ -319,6 +397,8 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
397
&& echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \
398
&& echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \
399
&& echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt \
400
+ && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \
401
+ && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \
402
# Fingerprint our triton so Dockerfile.rocm can prove it survived the build.
403
&& echo "TRITON_VERSION: $(pip show triton | awk '/^Version:/ {print $2}')" >> /app/versions.txt \
404
&& echo "TRITON_SHA256: $(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \
docker/Dockerfile.rocm_base
+81
-85
@@ -13,10 +13,9 @@ ARG AITER_BRANCH="v0.1.21.post1"
13
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14
ARG MORI_BRANCH="v1.1.0"
15
ARG MORI_REPO="https://github.com/ROCm/mori.git"
16
-ARG ROCPROFILER_SDK_REPO="https://github.com/ROCm/rocm-systems.git"
16
+ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git"
17
ARG ROCPROFILER_SDK_COMMIT="2b22ab0195cc1461cd9abf3b969e9dd7c10af350"
18
-ARG ROCM_RUNTIME_REPO="https://github.com/mawong-amd/rocm-systems.git"
19
-ARG ROCM_RUNTIME_COMMIT="c7e40e16bca10311af72aee09cfa0cd1726f70a7"
18
+ARG ROCM_RUNTIME_COMMIT="b539bf7eebfd99ad0a69668caa1f4037034d501f"
19
20
# Sccache configuration (only used in release pipeline)
21
ARG USE_SCCACHE
@@ -27,85 +26,12 @@ ARG SCCACHE_REGION_NAME=us-west-2
26
ARG SCCACHE_S3_NO_CREDENTIALS=0
27
28
30
-###
31
-### ROCr + CLR build
32
-###
33
-# Pull in updates to ROCr + clr from ROCm 7.2.4 and beyond. Happens before base so they
34
-# are present in each downstream stage
35
-FROM ${BASE_IMAGE} AS build_rocm_runtime
36
-ARG ROCM_RUNTIME_REPO
37
-ARG ROCM_RUNTIME_COMMIT
38
-RUN apt-get update \
39
- && apt-get install -y --no-install-recommends git cmake ninja-build g++ pkg-config \
40
- python3-pip libelf-dev libdrm-dev libnuma-dev libdw-dev xxd \
41
- && rm -rf /var/lib/apt/lists/*
42
-RUN pip install --no-cache-dir CppHeaderParser
43
-RUN git init -q /src && cd /src \
44
- && git remote add origin ${ROCM_RUNTIME_REPO} \
45
- && git config core.sparseCheckout true \
46
- && git config remote.origin.promisor true \
47
- && git config remote.origin.partialclonefilter blob:none \
48
- && git sparse-checkout init --cone \
49
- && git sparse-checkout set projects/rocr-runtime projects/clr projects/hip shared cmake \
50
- && git fetch --filter=blob:none --depth 1 origin ${ROCM_RUNTIME_COMMIT} \
51
- && git checkout -q FETCH_HEAD
52
-# ROCr's trap_handler and blit_shaders call find_package(Clang/LLVM REQUIRED) only to get
53
-# IMPORTED executables. The ROCm dev image ships the binaries but no CMake package config, so
54
-# supply one pointing at them.
55
-RUN mkdir -p /opt/rocm-cmake-shim \
56
- && printf 'if(NOT TARGET clang)\n add_executable(clang IMPORTED GLOBAL)\n set_target_properties(clang PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/clang")\nendif()\nset(Clang_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/ClangConfig.cmake \
57
- && printf 'if(NOT TARGET llvm-objcopy)\n add_executable(llvm-objcopy IMPORTED GLOBAL)\n set_target_properties(llvm-objcopy PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/llvm-objcopy")\nendif()\nset(LLVM_FOUND TRUE)\nset(LLVM_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/LLVMConfig.cmake
58
-# ROCr first, and installed into /opt/rocm, because the clr build below needs its headers
59
-RUN cd /src/projects/rocr-runtime \
60
- && cmake -B build -G Ninja \
61
- -DCMAKE_BUILD_TYPE=Release \
62
- -DCMAKE_PREFIX_PATH=/opt/rocm \
63
- -DCMAKE_INSTALL_PREFIX=/opt/rocm \
64
- -DClang_DIR=/opt/rocm-cmake-shim \
65
- -DLLVM_DIR=/opt/rocm-cmake-shim \
66
- && cmake --build build --parallel 16 \
67
- && cmake --install build \
68
- && cmake --install build --prefix /rocr-install --strip
69
-RUN cd /src/projects/clr \
70
- && cmake -B build -G Ninja \
71
- -DCMAKE_BUILD_TYPE=Release \
72
- -DCLR_BUILD_HIP=ON \
73
- -DCLR_BUILD_OCL=OFF \
74
- -DHIP_COMMON_DIR=/src/projects/hip \
75
- -DROCM_PATH=/opt/rocm \
76
- -DCMAKE_PREFIX_PATH=/opt/rocm \
77
- -DCMAKE_INSTALL_PREFIX=/opt/rocm \
78
- -DLLVM_DIR=/opt/rocm-cmake-shim \
79
- -DHIP_LLVM_ROOT=/opt/rocm/llvm \
80
- && cmake --build build --parallel 16 \
81
- && cmake --install build --prefix /clr-install --strip
82
-RUN mkdir -p /staging/lib \
83
- && cp -P /rocr-install/lib/libhsa-runtime64.so* /staging/lib/ \
84
- && cp -P /clr-install/lib/libamdhip64.so* /staging/lib/ \
85
- && ls -l /staging/lib
86
-
87
-
29
FROM ${BASE_IMAGE} AS base
30
31
ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
32
ENV ROCM_PATH=/opt/rocm
33
ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
34
94
-COPY --from=build_rocm_runtime /staging/ /opt/rocm/
95
-# Hard link the patched library onto every stock file name so that both the
96
-# SONAME lookup and the hard coded paths resolve to it.
97
-RUN set -eux; \
98
- cd /opt/rocm/lib; \
99
- ours_hsa="$(readlink -f libhsa-runtime64.so.1)"; \
100
- ours_hip="$(readlink -f libamdhip64.so.7)"; \
101
- for f in libhsa-runtime64.so.1.*; do \
102
- [ "$f" = "${ours_hsa##*/}" ] || ln -f "$ours_hsa" "$f"; \
103
- done; \
104
- for f in libamdhip64.so.7.*; do \
105
- [ "$f" = "${ours_hip##*/}" ] || ln -f "$ours_hip" "$f"; \
106
- done; \
107
- ldconfig
108
-
35
ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151
36
ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
37
ENV AITER_ROCM_ARCH=gfx942;gfx950
@@ -195,7 +121,7 @@ RUN cd triton \
121
&& git checkout ${TRITON_BRANCH} \
122
&& git config --global user.email "you@example.com" && git config --global user.name "Your Name" \
123
&& if [ ! -f setup.py ]; then cd python; fi \
198
- && python3 setup.py bdist_wheel --dist-dir=dist \
124
+ && MAX_JOBS=16 python3 setup.py bdist_wheel --dist-dir=dist \
125
&& mkdir -p /app/install && cp dist/*.whl /app/install
126
127
@@ -392,16 +318,72 @@ RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
318
RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
319
cp /install/*.whl /app/debs
320
321
+###
322
+### ROCr + CLR build
323
+###
324
+# Pull in updates to ROCr + clr from ROCm 7.2.4 and beyond.
325
+FROM base AS build_rocm_runtime
326
+ARG ROCM_SYSTEMS_REPO
327
+ARG ROCM_RUNTIME_COMMIT
328
+RUN apt-get update \
329
+ && apt-get install -y --no-install-recommends g++ \
330
+ libelf-dev libdrm-dev libnuma-dev libdw-dev xxd \
331
+ && rm -rf /var/lib/apt/lists/*
332
+RUN pip install --no-cache-dir CppHeaderParser
333
+RUN git init -q /src && cd /src \
334
+ && git remote add origin ${ROCM_SYSTEMS_REPO} \
335
+ && git config core.sparseCheckout true \
336
+ && git config remote.origin.promisor true \
337
+ && git config remote.origin.partialclonefilter blob:none \
338
+ && git sparse-checkout init --cone \
339
+ && git sparse-checkout set projects/rocr-runtime projects/clr projects/hip shared cmake \
340
+ && git fetch --filter=blob:none --depth 1 origin ${ROCM_RUNTIME_COMMIT} \
341
+ && git checkout -q FETCH_HEAD
342
+# ROCr's trap_handler and blit_shaders call find_package(Clang/LLVM REQUIRED) only to get
343
+# IMPORTED executables. The ROCm dev image ships the binaries but no CMake package config, so
344
+# supply one pointing at them.
345
+RUN mkdir -p /opt/rocm-cmake-shim \
346
+ && printf 'if(NOT TARGET clang)\n add_executable(clang IMPORTED GLOBAL)\n set_target_properties(clang PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/clang")\nendif()\nset(Clang_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/ClangConfig.cmake \
347
+ && printf 'if(NOT TARGET llvm-objcopy)\n add_executable(llvm-objcopy IMPORTED GLOBAL)\n set_target_properties(llvm-objcopy PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/llvm-objcopy")\nendif()\nset(LLVM_FOUND TRUE)\nset(LLVM_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/LLVMConfig.cmake
348
+# ROCr first, and installed into /opt/rocm, because the clr build below needs its headers
349
+RUN cd /src/projects/rocr-runtime \
350
+ && cmake -B build -G Ninja \
351
+ -DCMAKE_BUILD_TYPE=Release \
352
+ -DCMAKE_PREFIX_PATH=/opt/rocm \
353
+ -DCMAKE_INSTALL_PREFIX=/opt/rocm \
354
+ -DClang_DIR=/opt/rocm-cmake-shim \
355
+ -DLLVM_DIR=/opt/rocm-cmake-shim \
356
+ && cmake --build build --parallel 16 \
357
+ && cmake --install build \
358
+ && cmake --install build --prefix /rocr-install --strip
359
+RUN cd /src/projects/clr \
360
+ && cmake -B build -G Ninja \
361
+ -DCMAKE_BUILD_TYPE=Release \
362
+ -DCLR_BUILD_HIP=ON \
363
+ -DCLR_BUILD_OCL=OFF \
364
+ -DHIP_COMMON_DIR=/src/projects/hip \
365
+ -DROCM_PATH=/opt/rocm \
366
+ -DCMAKE_PREFIX_PATH=/opt/rocm \
367
+ -DCMAKE_INSTALL_PREFIX=/opt/rocm \
368
+ -DLLVM_DIR=/opt/rocm-cmake-shim \
369
+ -DHIP_LLVM_ROOT=/opt/rocm/llvm \
370
+ && cmake --build build --parallel 16 \
371
+ && cmake --install build --prefix /clr-install --strip
372
+RUN mkdir -p /staging/lib \
373
+ && cp -P /rocr-install/lib/libhsa-runtime64.so* /staging/lib/ \
374
+ && cp -P /clr-install/lib/libamdhip64.so* /staging/lib/ \
375
+ && ls -l /staging/lib
376
+
377
# rocprofiler-sdk 1.3.2 fixes torch 2.12 profiler on ROCm 7.2.3
378
FROM base AS build_rocprofiler_sdk
397
-ARG ROCPROFILER_SDK_REPO
379
+ARG ROCM_SYSTEMS_REPO
380
ARG ROCPROFILER_SDK_COMMIT
381
ARG USE_SCCACHE
382
RUN apt-get update \
383
&& apt-get install -y --no-install-recommends libdw-dev libelf-dev libsqlite3-dev libzstd-dev \
384
&& rm -rf /var/lib/apt/lists/*
385
RUN git init -q /src && cd /src \
404
- && git remote add origin ${ROCPROFILER_SDK_REPO} \
386
+ && git remote add origin ${ROCM_SYSTEMS_REPO} \
387
&& git config core.sparseCheckout true \
388
&& git config remote.origin.promisor true \
389
&& git config remote.origin.partialclonefilter blob:none \
@@ -417,8 +399,8 @@ ARG ROCPROFILER_SDK_PR_7796=fae7fd57ffb8a76e811d6410b8a0cee1e94d710a
399
ARG ROCPROFILER_SDK_PR_7924=f8db22adbe17686a9b05ae03367dd1ec73011653
400
RUN cd /src \
401
&& HSA_PATHS='projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/*' \
420
- && curl -fsSL -o /tmp/qi-7796.patch "${ROCPROFILER_SDK_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7796}.patch" \
421
- && curl -fsSL -o /tmp/qi-7924.patch "${ROCPROFILER_SDK_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7924}.patch" \
402
+ && curl -fsSL -o /tmp/qi-7796.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7796}.patch" \
403
+ && curl -fsSL -o /tmp/qi-7924.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7924}.patch" \
404
&& git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7796.patch \
405
&& git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7924.patch
406
RUN cd /src/projects/rocprofiler-sdk \
@@ -441,6 +423,22 @@ RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
423
pip install /install/*.whl
424
425
COPY --from=build_rocprofiler_sdk /staging/ /opt/rocm/
426
+ENV ROCPROFILER_QUEUE_INTERPOSITION=0
427
+
428
+COPY --from=build_rocm_runtime /staging/ /opt/rocm/
429
+# Hard link the patched library onto every stock file name so that both the
430
+# SONAME lookup and the hard coded paths resolve to it.
431
+RUN set -eux; \
432
+ cd /opt/rocm/lib; \
433
+ ours_hsa="$(readlink -f libhsa-runtime64.so.1)"; \
434
+ ours_hip="$(readlink -f libamdhip64.so.7)"; \
435
+ for f in libhsa-runtime64.so.1.*; do \
436
+ [ "$f" = "${ours_hsa##*/}" ] || ln -f "$ours_hsa" "$f"; \
437
+ done; \
438
+ for f in libamdhip64.so.7.*; do \
439
+ [ "$f" = "${ours_hip##*/}" ] || ln -f "$ours_hip" "$f"; \
440
+ done
441
+
442
RUN ldconfig
443
444
ARG BASE_IMAGE
@@ -458,14 +456,12 @@ ARG AITER_BRANCH
456
ARG AITER_REPO
457
ARG MORI_BRANCH
458
ARG MORI_REPO
461
-ARG ROCPROFILER_SDK_REPO
459
+ARG ROCM_SYSTEMS_REPO
460
ARG ROCPROFILER_SDK_COMMIT
463
-ARG ROCM_RUNTIME_REPO
461
ARG ROCM_RUNTIME_COMMIT
462
RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
466
- && echo "ROCPROFILER_SDK_REPO: ${ROCPROFILER_SDK_REPO}" >> /app/versions.txt \
463
+ && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \
464
&& echo "ROCPROFILER_SDK_COMMIT: ${ROCPROFILER_SDK_COMMIT}" >> /app/versions.txt \
468
- && echo "ROCM_RUNTIME_REPO: ${ROCM_RUNTIME_REPO}" >> /app/versions.txt \
465
&& echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \
466
&& echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
467
&& echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
tests/kernels/moe/parallel_utils.py
-12
@@ -6,7 +6,6 @@ DeepEP test utilities
6
7
import dataclasses
8
import os
9
-import sys
9
import traceback
10
from collections.abc import Callable
11
from typing import Concatenate
@@ -16,7 +15,6 @@ from torch.distributed import ProcessGroup
15
from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage]
16
from typing_extensions import ParamSpec
17
19
-from vllm.platforms import current_platform
18
from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2
19
from vllm.utils.network_utils import get_open_port
20
@@ -74,7 +72,6 @@ def _worker_parallel_launch(
72
barrier = torch.tensor([rank], device=device)
73
torch.distributed.all_reduce(barrier)
74
77
- exit_code = 0
75
try:
76
worker(
77
ProcessGroupInfo(
@@ -91,18 +88,9 @@ def _worker_parallel_launch(
88
except Exception as ex:
89
print(ex)
90
traceback.print_exc()
94
- exit_code = 1
91
raise
92
finally:
93
torch.distributed.destroy_process_group()
98
- if current_platform.is_rocm():
99
- # Bypass a ROCm teardown use-after-free (SIGSEGV in HIP's atexit
100
- # handler; fixed by https://github.com/ROCm/rocm-systems/pull/6942).
101
- # Flush first since os._exit skips it, to keep a failing traceback.
102
- # TODO (Rohan138): Remove after ROCm 10.0 lands
103
- sys.stdout.flush()
104
- sys.stderr.flush()
105
- os._exit(exit_code)
94
95
96
def parallel_launch(
tests/kernels/moe/test_deepep_moe.py
-7
@@ -488,17 +488,10 @@ def test_deep_ep_moe(
488
world_dp_size: tuple[int, int],
489
per_act_token_quant: bool,
490
workspace_init,
491
- monkeypatch,
491
):
492
low_latency_mode = False
493
use_fp8_dispatch = False
494
496
- if current_platform.is_rocm():
497
- # The cooperative kernel launches on this path segfault inside the ROCm
498
- # HSA runtime at process exit while rocprofiler-sdk is attached, and
499
- # torch attaches it implicitly. Workers are spawned, so they inherit it.
500
- monkeypatch.setenv("ROCPROFILER_REGISTER_ENABLED", "0")
501
-
495
set_random_seed(7)
496
world_size, dp_size = world_dp_size
497
config = TestConfig(dtype=dtype, topk=topk, m=m, k=k, n=n, num_experts=num_experts)