From ee7f2f411a09d09846f8819ccfee5988121cc8b0 Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Date: Sun, 2 Aug 2026 22:38:00 +0000 Subject: [PATCH 01/10] Upgrade CI to TheRock 7.14 --- .github/scripts/Dockerfile.ci.deps | 117 +++++++----------- .../attention/benchmark_attention_rocm.py | 24 +++- tests/pytorch/test_sanity_import.py | 21 +++- 3 files changed, 81 insertions(+), 81 deletions(-) diff --git a/.github/scripts/Dockerfile.ci.deps b/.github/scripts/Dockerfile.ci.deps index 722c8f6205..e369c6f476 100644 --- a/.github/scripts/Dockerfile.ci.deps +++ b/.github/scripts/Dockerfile.ci.deps @@ -3,120 +3,89 @@ # See LICENSE for license information. # # TE CI deps image -# -# ROCm installer: https://raw.githubusercontent.com/ROCm/TheRock/release/therock-7.13/dockerfiles/install_rocm_tarball.sh FROM ubuntu:24.04 ARG DEBIAN_FRONTEND=noninteractive SHELL ["/bin/bash", "-euo", "pipefail", "-c"] -ARG ROCK_RELEASE_VERSION=7.13.0 -ARG GPU_ARCH=gfx942 +ARG ROCK_RELEASE_VERSION=7.14.0 ARG PYTHON_VERSION=3.12 ARG PYTHON_ABI=cp312 -ARG TORCH_VERSION=2.10.0 -ARG TORCHVISION_VERSION=0.25.0 -ARG TORCHAUDIO_VERSION=2.10.0 -ARG TRITON_VERSION=3.6.0 -ARG JAX_VERSION=0.10.2 -ARG JAX_BRANCH=rocm-jax-v0.10.2 -ARG XLA_BRANCH=rocm-jaxlib-v0.10.2 -ARG FA_VERSION=v2.8.1 +ARG TORCH_VERSION=2.12.0 +ARG TORCHVISION_VERSION=0.27.0 +# triton 3.7.1 ships as a git snapshot: triton-3.7.1+.rocm +ARG TRITON_VERSION=3.7.1 +ARG TRITON_LOCAL=git0263a6a6 +ARG JAX_VERSION=0.11.0 +ARG FA_VERSION=v2.8.3 ARG AITER_COMMIT=77455e3ecf4f0d28756afc452e914940c45b944b -ARG INSTALL_ROCM_TARBALL_SH_URL=https://raw.githubusercontent.com/ROCm/TheRock/release/therock-7.13/dockerfiles/install_rocm_tarball.sh +ARG GPU_ARCH=gfx950;gfx942 # Base OS packages RUN apt-get update && apt-get install -y --no-install-recommends \ ca-certificates curl \ git vim \ build-essential cmake ninja-build pkg-config liblzma-dev \ - openjdk-17-jdk-headless \ + libnuma1 libnuma-dev \ python${PYTHON_VERSION} python${PYTHON_VERSION}-venv python${PYTHON_VERSION}-dev python3-pip \ && rm -rf /var/lib/apt/lists/* -# Native ROCm tarball → /opt/rocm -RUN case "${GPU_ARCH}" in \ - gfx942) _amd_gpu_family=gfx94X-dcgpu ;; \ - gfx950) _amd_gpu_family=gfx950-dcgpu ;; \ - esac \ - && curl -fsSL -o /tmp/install_rocm_tarball.sh "${INSTALL_ROCM_TARBALL_SH_URL}" \ - && chmod +x /tmp/install_rocm_tarball.sh \ - && /tmp/install_rocm_tarball.sh "${ROCK_RELEASE_VERSION}" "${_amd_gpu_family}" stable \ - && rm -f /tmp/install_rocm_tarball.sh - # Isolated Python env for pip packages RUN python${PYTHON_VERSION} -m venv /opt/venv -# Default container env: venv on PATH; /opt/rocm tarball for flash-attention / aiter. -ENV GPU_ARCH=${GPU_ARCH} \ - VIRTUAL_ENV=/opt/venv \ - PATH=/opt/venv/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +# Default container env: venv on PATH. ROCm is provided by the pip rocm-sdk and +# resolved on demand via `rocm-sdk path --root` (no /opt/rocm install). +ENV VIRTUAL_ENV=/opt/venv \ + PATH=/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin -RUN python -m pip install --upgrade pip setuptools wheel \ +RUN python -m pip install --upgrade pip setuptools wheel pyyaml \ && pip install --no-cache-dir ipython pytest fire pydantic pybind11 ninja pandas expecttest onnxscript -# Python ROCm SDK + torch from https://repo.amd.com/rocm/whl// -RUN case "${GPU_ARCH}" in \ - gfx942) _amd_gpu_family=gfx94X-dcgpu ;; \ - gfx950) _amd_gpu_family=gfx950-dcgpu ;; \ - esac \ - && W="https://repo.amd.com/rocm/whl/${_amd_gpu_family}" \ - && LIBS_PKG="rocm-sdk-libraries-$(echo "${_amd_gpu_family}" | tr '[:upper:]' '[:lower:]')" \ +# Python ROCm SDK + torch from https://repo.amd.com/rocm/whl-multi-arch/ +# Multi-arch layout: arch-agnostic wheels (rocm-sdk-libraries, torch, torchvision, +# triton) installed once, then a per-GPU device wheel (rocm-sdk-device-, +# amd_torch_device_, amd_torchvision_device_) for every arch in +# GPU_ARCH (semicolon-separated, overridable via --build-arg). +RUN W="https://repo.amd.com/rocm/whl-multi-arch" \ && ROCM_WHEEL_TAG="rocm${ROCK_RELEASE_VERSION}" \ && pip install --no-cache-dir \ --extra-index-url "${W}" \ "rocm-sdk-core==${ROCK_RELEASE_VERSION}" \ "rocm-sdk-devel==${ROCK_RELEASE_VERSION}" \ - "${LIBS_PKG}==${ROCK_RELEASE_VERSION}" \ + "rocm-sdk-libraries==${ROCK_RELEASE_VERSION}" \ "${W}/rocm-${ROCK_RELEASE_VERSION}.tar.gz" \ "${W}/torch-${TORCH_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" \ "${W}/torchvision-${TORCHVISION_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" \ - "${W}/torchaudio-${TORCHAUDIO_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" \ - "${W}/triton-${TRITON_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" + "${W}/triton-${TRITON_VERSION}%2B${TRITON_LOCAL}.${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" \ + && for arch in ${GPU_ARCH//;/ }; do \ + pip install --no-cache-dir --extra-index-url "${W}" \ + "rocm-sdk-device-${arch}==${ROCK_RELEASE_VERSION}" \ + "${W}/amd_torch_device_${arch}-${TORCH_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" \ + "${W}/amd_torchvision_device_${arch}-${TORCHVISION_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" ; \ + done -# Install Jax -RUN pip install --no-cache-dir patchelf auditwheel numpy scipy \ - && curl -fsSL -o /usr/local/bin/bazel \ - https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-amd64 \ - && chmod +x /usr/local/bin/bazel \ - && git clone --branch "${JAX_BRANCH}" --depth 1 https://github.com/ROCm/jax.git /tmp/jax-src \ - && git clone --branch "${XLA_BRANCH}" --depth 1 https://github.com/ROCm/xla.git /tmp/xla-src \ - && cd /tmp/jax-src \ - && JAX_DIST=/tmp/jax-dist \ - && mkdir -p "${JAX_DIST}" \ - && ROCM_SDK_ROOT="$(rocm-sdk path --root)" \ - && python build/build.py build \ - --wheels=jax-rocm-plugin,jax-rocm-pjrt \ - --bazel_startup_options="--bazelrc=build/rocm/rocm.bazelrc" \ - --bazel_startup_options="--output_user_root=/tmp/bazel-jax" \ - --bazel_options=--config=rocm_release_wheel \ - --bazel_options=--override_repository=xla=/tmp/xla-src \ - --bazel_options=--override_module=xla=/tmp/xla-src \ - --bazel_options=--repo_env=ML_WHEEL_TYPE=release \ - --bazel_options=--//jaxlib/tools:jaxlib_git_hash="$(git rev-parse HEAD)" \ - --bazel_options=--repo_env=ROCM_PATH="${ROCM_SDK_ROOT}" \ - --python_version="${PYTHON_VERSION}" \ - --rocm_path="${ROCM_SDK_ROOT}" \ - --rocm_version=7 \ - --rocm_amdgpu_targets="${GPU_ARCH}" \ - --verbose \ - --output_path="${JAX_DIST}" \ - && pip wheel . --no-deps -w "${JAX_DIST}" \ - && pip install --no-cache-dir "jaxlib==${JAX_VERSION}" "${JAX_DIST}"/*.whl \ - && pip install --no-cache-dir "jax==${JAX_VERSION}" \ - && rm -rf /tmp/jax-src /tmp/xla-src /tmp/jax-dist /tmp/bazel-jax +# Install Jax — prebuilt ROCm 7.14 (TheRock, multi-arch) wheels from PyPI. +# jax pulls jaxlib; jax-rocm7-plugin pulls jax-rocm7-pjrt. ROCm runtime is provided +# by the rocm-sdk-* wheels above, so the plugin's [with-rocm] extra (which would +# pull duplicate amd_*_instinct libs) is intentionally not requested. +RUN pip install --no-cache-dir \ + "jax==${JAX_VERSION}" \ + "jax-rocm7-pjrt==${JAX_VERSION}" \ + "jax-rocm7-plugin==${JAX_VERSION}" -RUN git clone --branch "${FA_VERSION}" --depth 1 https://github.com/Dao-AILab/flash-attention.git /tmp/flash-attention \ +RUN export ROCM_PATH="$(rocm-sdk path --root)" \ + && git clone --branch "${FA_VERSION}" --depth 1 https://github.com/Dao-AILab/flash-attention.git /tmp/flash-attention \ && cd /tmp/flash-attention \ - && ROCM_PATH=/opt/rocm GPU_ARCHS="${GPU_ARCH}" python setup.py install \ + && GPU_ARCHS="${GPU_ARCH}" python setup.py install \ && rm -rf /tmp/flash-attention -RUN git clone --no-checkout https://github.com/ROCm/aiter.git /tmp/aiter \ +RUN export ROCM_PATH="$(rocm-sdk path --root)" \ + && git clone --no-checkout https://github.com/ROCm/aiter.git /tmp/aiter \ && cd /tmp/aiter \ && git checkout "${AITER_COMMIT}" \ && git submodule update --init --recursive \ - && ROCM_PATH=/opt/rocm GPU_ARCHS="${GPU_ARCH}" pip install --no-build-isolation --no-cache-dir . \ + && pip install --no-build-isolation --no-cache-dir . \ && rm -rf /tmp/aiter WORKDIR /workspace/ diff --git a/benchmarks/attention/benchmark_attention_rocm.py b/benchmarks/attention/benchmark_attention_rocm.py index b234d374f5..87b57f94a7 100644 --- a/benchmarks/attention/benchmark_attention_rocm.py +++ b/benchmarks/attention/benchmark_attention_rocm.py @@ -139,6 +139,28 @@ def setup_backend_env(backend_name, use_ck_bwd_v3=True, use_ck_fwd_v3=True, use_ ROCPROF_STATS_CSV = "results.stats.csv" +def _resolve_rocprofv3(): + """Path to the rocprofv3 shipped in rocm-sdk-core. + + The `rocprofv3` on PATH is the rocm-sdk trampoline, which execs the copy in + rocm-sdk-devel. That copy loads librocprofiler-sdk.so.1 from _rocm_sdk_devel, + while torch/TE have already loaded the byte-identical copy from _rocm_sdk_core. + rocprofiler-register treats the two different paths as a conflict and aborts + (registration.cpp: "already set ... not overriding"), so rocprofv3 dies with + SIGABRT and no kernel stats are produced. Running the core copy makes the tool + and the profiled app share the same library, avoiding the conflict. + """ + try: + from rocm_sdk_core._cli import _get_core_module_path + + core_bin = _get_core_module_path() / "bin" / "rocprofv3" + if core_bin.is_file(): + return str(core_bin) + except Exception: + pass + return shutil.which("rocprofv3") or "rocprofv3" + + def _profiler_python_code(model, attention, column_name, benchmark_dir): return ( f"import sys; sys.path.insert(0, {benchmark_dir!r}); " @@ -167,7 +189,7 @@ def _run_attention_profiler(model, attention, column_name, dirname): benchmark_dir = os.path.dirname(os.path.abspath(__file__)) py_code = _profiler_python_code(model, attention, column_name, benchmark_dir) cmd = [ - "rocprofv3", + _resolve_rocprofv3(), "--hip-trace", "--kernel-trace", "--stats", diff --git a/tests/pytorch/test_sanity_import.py b/tests/pytorch/test_sanity_import.py index 84d2e8aa01..418a34e5f4 100644 --- a/tests/pytorch/test_sanity_import.py +++ b/tests/pytorch/test_sanity_import.py @@ -9,20 +9,29 @@ if __name__ == "__main__": print("OK") -AMDSMI_SRC = "/opt/rocm/share/amd_smi" - - def _amdsmi_pythonpath(): """Return ROCm amdsmi source dir if the package is not already installed. torch counts devices through amdsmi whenever it is importable, reading HIP_VISIBLE_DEVICES at call time rather than asking an already initialized HIP runtime -- which is what makes a visible-devices change after import observable. + + With TheRock pip wheels, bindings live under the rocm-sdk-core tree + (``/share/amd_smi``). """ - import importlib.util, os - if importlib.util.find_spec("amdsmi") is not None or not os.path.isdir(AMDSMI_SRC): + import importlib.util + + if importlib.util.find_spec("amdsmi") is not None: return None - return AMDSMI_SRC + try: + from rocm_sdk_core._cli import _get_core_module_path + + candidate = _get_core_module_path() / "share/amd_smi" + if candidate.is_dir(): + return str(candidate) + except (ImportError, ModuleNotFoundError, OSError): + pass + return None def _env_with_amdsmi_pythonpath(env, amdsmi_path): From 9bcbd0a8981c80473694405a22a996b7eff5205d Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Date: Sun, 2 Aug 2026 22:42:05 +0000 Subject: [PATCH 02/10] [ROCm] Work around torch 2.12 hipBLASLt graph-capture failure in test_gpt_cuda_graph test_gpt_cuda_graph fails on ROCm with torch>=2.12 with: Hip error: 'operation not permitted when stream is capturing'(900) at hipblaslt/library/src/amd_detail/hipblaslt.cpp:164 then hangs. Root cause is upstream PyTorch, not TE. pytorch/pytorch#179053 (in 2.12) changed the ROCm hipBLASLt handle cache from per-device to per-(device, stream). The graph-capture stream is fresh, so the first matmul on it does a lazy hipblasLtCreate, whose internal ~25 MB hipMalloc is illegal during capture. The capture_begin pre-init added in pytorch/pytorch#180692 only covers the calling thread, so torch's own bmm in the captured backward (fp32 -> unfused attention) still creates a handle on the autograd thread during capture. Not reproducible on torch 2.10 (per-device handle) and not on CUDA (the per-stream path is #ifdef USE_ROCM). Route torch's own matmul/bmm off hipBLASLt for this test via torch.backends.cuda.preferred_blas_library(cublas), gated on IS_HIP_EXTENSION and restored in a finally. TE's own GEMMs are unaffected. Remove once the upstream capture-time pre-init covers all threads/streams (or hipblasLtCreate becomes capture-safe). --- tests/pytorch/test_numerics.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 3a4405533a..b6ee0f4856 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -2126,8 +2126,24 @@ def test_gpt_cuda_graph(dtype, bs, model): for param1, param2 in zip(block.parameters(), graphed_block.parameters()): param2.copy_(param1) - out, grads = _test_gpt_e2e_cuda_graph(block, bs, dtype, config, False) - graphed_out, graphed_grads = _test_gpt_e2e_cuda_graph(graphed_block, bs, dtype, config, True) + # WAR (ROCm): torch>=2.12 (pytorch/pytorch#179053) makes hipBLASLt handles + # per-(device, stream). The graph-capture stream's handle is created lazily + # during capture, and hipblasLtCreate performs an internal hipMalloc that is + # illegal mid-capture, failing with HIP error 900 ("operation not permitted + # when stream is capturing"). The upstream capture_begin pre-init + # (pytorch/pytorch#180692) only covers the calling thread, not the autograd + # backward thread, so torch's own bmm in the captured backward still trips it. + # Route torch's matmul/bmm off hipBLASLt for this test until upstream fixes it. + _prev_blas_library = None + if IS_HIP_EXTENSION: + _prev_blas_library = torch.backends.cuda.preferred_blas_library() + torch.backends.cuda.preferred_blas_library("cublas") + try: + out, grads = _test_gpt_e2e_cuda_graph(block, bs, dtype, config, False) + graphed_out, graphed_grads = _test_gpt_e2e_cuda_graph(graphed_block, bs, dtype, config, True) + finally: + if _prev_blas_library is not None: + torch.backends.cuda.preferred_blas_library(_prev_blas_library) params = list(block.parameters()) graphed_params = list(graphed_block.parameters()) From dbcee5a3329717f8a8ce89562d0c965b1d252b53 Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Date: Sun, 2 Aug 2026 22:45:51 +0000 Subject: [PATCH 03/10] [torch.compile] Bunch of small changes needed for enabling torch.compile (#3130) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [PyTorch] torch.compile: wrap pybind11 UB methods as compile-time constants; fix SP memory leak; test suite hook-up Wrap CommOverlapCore pybind11 methods that return compile-time constants so torch.compile(fullgraph=True) can trace through them without graph breaks: - `is_fp8_ubuf()` → `ub_is_fp8()` / `get_ub_is_fp8()` in base.py; `_ub_is_fp8()` in gemm.py - `with_cublasmp()` → `ub_is_cublasmp()` in base.py All callers in linear.py, layernorm_linear.py, layernorm_mlp.py, base.py, gemm.py, userbuffers_backward_linear.py and userbuffers_forward_linear.py updated. Fix quantized grad_output not being freed early for column-parallel SP backward. Row-parallel SP already called clear_tensor_data(grad_output) to release the gathered tensor; column-parallel SP quantizes grad_output to Float8TensorStorage but never freed it before returning. Under torch.compile reduce-overhead this leaves 3 live pool tensors at recording end and triggers "Detected 3 tensor(s) in the cudagraph pool not tracked as outputs". Extend the existing clear_tensor_data guard to cover both parallel modes. Fix custom-recipe quantizer state being re-initialised on every forward call even when the recipe object has not changed. The existing early-exit for CustomRecipeState was missing an identity check on the recipe object, so any repeated call with the same recipe would bypass the early-return and rebuild quantizers unnecessarily. Add `if recipe_state.recipe is recipe: return` to restore the intended caching behaviour. Add test_torch_compile.py to L0_pytorch_unittest so the autocast and existing compile tests run in CI. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski * [PyTorch] Replace fp8_recipe in LinearBwdArgs with pre-resolved split-accumulator booleans LinearBwdArgs stored the entire FP8 recipe object so the backward could extract fp8_gemm_dgrad.use_split_accumulator and fp8_gemm_wgrad.use_split_accumulator at GEMM time. Recipe objects hold process-group references and are not serialisable as compile-time constants, making them incompatible with torch.compile custom-op paths. Replace fp8_recipe with two plain bool fields: - dgrad_use_split_accumulator (default _2X_ACC_DGRAD) - wgrad_use_split_accumulator (default _2X_ACC_WGRAD) These are resolved once in _linear_setup_ctx and passed into the args struct, so the backward consumes scalars instead of a live recipe object. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reset torch.compile state in destroy_ub to avoid stale assume_constant_result get_ub_is_fp8 bakes is_fp8_ubuf() as a compile-time constant; without a reset, destroy_ub + re-init with different FP8 settings would read stale values until recompile. Only affects in-memory caches, not disk. Signed-off-by: Pawel Gadzinski * Provide explicit QuantizerRoles in torch.compile custom-recipe test ToyLinear now overrides get_quantizer_roles so CustomRecipeState doesn't hit the no-roles warning, which graph-breaks under fullgraph=True. qfactory dispatches on role.tensor_type instead of a pre-baked string key. Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Co-authored-by: Claude Sonnet 4.6 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani (cherry picked from commit 43093f16d6dc3c0707c2b9ec3c3d07918fed598a) --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_torch_compile.py | 41 ++++++++++++---- transformer_engine/pytorch/module/base.py | 12 ++++- .../pytorch/module/layernorm_linear.py | 15 +++--- .../pytorch/module/layernorm_mlp.py | 3 +- transformer_engine/pytorch/module/linear.py | 47 ++++++++++--------- 6 files changed, 78 insertions(+), 41 deletions(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 92f73d5885..48b713ad66 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -35,6 +35,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_rope.xml $ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_nvfp4.xml $TE_PATH/tests/pytorch/nvfp4 || test_fail "test_nvfp4" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mxfp8.xml $TE_PATH/tests/pytorch/mxfp8 || test_fail "test_mxfp8" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_quantized_tensor.xml $TE_PATH/tests/pytorch/test_quantized_tensor.py || test_fail "test_quantized_tensor.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_torch_compile.xml $TE_PATH/tests/pytorch/test_torch_compile.py || test_fail "test_torch_compile.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8blockwisetensor.xml $TE_PATH/tests/pytorch/test_float8blockwisetensor.py || test_fail "test_float8blockwisetensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_scaling_exact.py || test_fail "test_float8_blockwise_scaling_exact.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_gemm_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_gemm_exact.py || test_fail "test_float8_blockwise_gemm_exact.py" diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 309a5d124e..1286492a6e 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -26,6 +26,7 @@ from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer +from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, @@ -123,21 +124,25 @@ def __fx_repr__(self): _Q = get_opaque_type_name(ToyQuantizer) def _make_qfactory(tag: str): - """Return a qfactory that produces ToyQuantizer instances tagged with *tag*.""" + """Return a qfactory that produces ToyQuantizer instances tagged with *tag*. + + The factory dispatches on ``QuantizerRole.tensor_type``; the roles are + supplied by :meth:`ToyLinear.get_quantizer_roles`. + """ quantizers = { - role: ToyQuantizer(tag=f"{tag}:{role}") - for role in ( - "linear_input", - "linear_weight", - "linear_output", - "linear_grad_output", - "linear_grad_input", + tensor_type: ToyQuantizer(tag=f"{tag}:{tensor_type}") + for tensor_type in ( + "input", + "weight", + "output", + "grad_output", + "grad_input", ) } - def qfactory(role: str): - return quantizers[role] + def qfactory(role: QuantizerRole): + return quantizers[role.tensor_type] return qfactory @@ -163,6 +168,22 @@ def __init__( ) torch.nn.init.normal_(self.weight) + def get_quantizer_roles(self, *, fwd: bool, num_quantizers: int): + # Supplying explicit roles keeps CustomRecipeState from emitting a + # warning (which would graph-break under fullgraph=True) and lets the + # qfactory dispatch per tensor slot. Order must match the module's + # quantizer array (FP8FwdTensorIdx / FP8BwdTensorIdx). + if fwd: + return [ + QuantizerRole(module_type="linear", tensor_type="input"), + QuantizerRole(module_type="linear", tensor_type="weight"), + QuantizerRole(module_type="linear", tensor_type="output"), + ] + return [ + QuantizerRole(module_type="linear", tensor_type="grad_output"), + QuantizerRole(module_type="linear", tensor_type="grad_input"), + ] + def _get_weight_tensors(self): return [self.weight] diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index f65052f9b4..7ca09f7015 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -595,6 +595,12 @@ def get_ub(name: str, use_fp8: bool): return _ub_communicators[key] +@torch.compiler.assume_constant_result +def get_ub_is_fp8(name: str, use_fp8: bool) -> bool: + """Query is_fp8_ubuf for a named UB communicator; treated as compile-time constant.""" + return get_ub(name, use_fp8).is_fp8_ubuf() + + def destroy_ub(): """Destroy all allocated userbuffer communicators.""" global _ub_communicators, _ub_with_cublasmp, _ub_initialized @@ -603,6 +609,9 @@ def destroy_ub(): _ub_initialized = False global layers_atomic_ring_exchange layers_atomic_ring_exchange = [] + # Compiled graphs may have baked is_fp8_ubuf() via assume_constant_result; + # reset so re-init with different settings doesn't read stale constants. + torch.compiler.reset() def fill_userbuffers_buffer_for_all_gather( @@ -1094,7 +1103,8 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState): return if recipe.custom() and isinstance(recipe_state, CustomRecipeState): - return + if recipe_state.recipe is recipe: + return # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and # 2 (grad_output and grad_input) for bwd diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index e9342c0646..988dab9fae 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -25,6 +25,7 @@ from .base import ( fill_userbuffers_buffer_for_all_gather, get_ub, + get_ub_is_fp8, is_ub_initialized, using_cublasmp_backend, quantize_weight, @@ -1094,8 +1095,10 @@ def wgrad_gemm( if ctx.ln_out_needs_gather: # Gathered input is internal clear_tensor_data(ln_out_total) - if ctx.parallel_mode == "row" and ctx.sequence_parallel: - # Gathered grad output tensor is internal + if ctx.sequence_parallel and ( + ctx.parallel_mode == "row" or (ctx.parallel_mode == "column" and ctx.fp8) + ): + # Gathered (row-SP) or quantized (column-SP FP8) grad_output is internal clear_tensor_data(grad_output) # Update grad input if overlapping reduce-scatter with wgrad GEMM @@ -1736,14 +1739,10 @@ def forward( is_first_microbatch = False if self.ub_overlap_rs_fprop: - if get_ub( - self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled() - ).is_fp8_ubuf(): + if get_ub_is_fp8(self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled()): fp8_output = True if self.ub_overlap_rs_dgrad: - if get_ub( - self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled() - ).is_fp8_ubuf(): + if get_ub_is_fp8(self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()): fp8_grad = True inp = self.prepare_forward( diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index cea110f675..e5a9675fcf 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -26,6 +26,7 @@ fill_userbuffers_buffer_for_all_gather, _ub_communicators, get_ub, + get_ub_is_fp8, is_ub_initialized, using_cublasmp_backend, quantize_weight, @@ -2367,7 +2368,7 @@ def forward( fp8_output = False if self.ub_overlap_rs: - if get_ub("fc2_fprop", FP8GlobalStateManager.is_fp8_enabled()).is_fp8_ubuf(): + if get_ub_is_fp8("fc2_fprop", FP8GlobalStateManager.is_fp8_enabled()): fp8_output = True inp = self.prepare_forward(inp, num_gemms=2) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index b91a3e8fd1..556348fde9 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -23,6 +23,7 @@ fill_userbuffers_buffer_for_all_gather, get_dummy_wgrad, get_ub, + get_ub_is_fp8, is_ub_initialized, using_cublasmp_backend, quantize_weight, @@ -122,6 +123,8 @@ class LinearFwdArgs: fp8_output: bool save_original_input: bool backward_override: Optional[str] + dgrad_use_split_accumulator: bool + wgrad_use_split_accumulator: bool custom: bool debug: bool @@ -192,7 +195,8 @@ class LinearBwdArgs: # --- Numerical / dtype config --- activation_dtype: Optional[torch.dtype] = None fp8: bool = False - fp8_recipe: Optional[Recipe] = None + dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD + wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD backward_override: Optional[str] = None is_weight_param_quantized: bool = False custom: bool = False @@ -695,7 +699,8 @@ def _linear_setup_ctx( # Numerical / dtype config bwd_args.activation_dtype = fwd_args.activation_dtype bwd_args.fp8 = fp8 - bwd_args.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator + bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator bwd_args.backward_override = backward_override bwd_args.is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) bwd_args.custom = fwd_args.custom @@ -1001,11 +1006,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. weight_fp8.update_usage(columnwise_usage=True) # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_DGRAD - if bwd_args.fp8: - recipe = bwd_args.fp8_recipe - if hasattr(recipe, "fp8_gemm_dgrad"): - use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator + use_split_accumulator = bwd_args.dgrad_use_split_accumulator # Update grad input quantizer if grad_input_quantizer is not None: @@ -1178,11 +1179,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. grad_output = grad_output_quantizer(grad_output) # Figure out whether to use split accumulator - use_split_accumulator = _2X_ACC_WGRAD - if bwd_args.fp8: - recipe = bwd_args.fp8_recipe - if hasattr(recipe, "fp8_gemm_wgrad"): - use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator + use_split_accumulator = bwd_args.wgrad_use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad if bwd_args.is_first_microbatch is not None: @@ -1273,8 +1270,11 @@ def wgrad_gemm( elif bwd_args.backward_input_needs_gather: # Gathered input tensor is internal clear_tensor_data(inputmat_total) - if bwd_args.parallel_mode == "row" and bwd_args.sequence_parallel: - # Gathered grad output tensor is internal + if bwd_args.sequence_parallel and ( + bwd_args.parallel_mode == "row" + or (bwd_args.parallel_mode == "column" and bwd_args.fp8) + ): + # Gathered (row-SP) or quantized (column-SP FP8) grad_output is internal clear_tensor_data(grad_output) # Update grad input if overlapping reduce-scatter with wgrad GEMM @@ -1888,14 +1888,10 @@ def forward( is_first_microbatch = False if self.ub_overlap_rs_fprop: - if get_ub( - self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled() - ).is_fp8_ubuf(): + if get_ub_is_fp8(self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled()): fp8_output = True if self.ub_overlap_rs_dgrad: - if get_ub( - self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled() - ).is_fp8_ubuf(): + if get_ub_is_fp8(self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()): fp8_grad = True inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor)) @@ -1933,8 +1929,15 @@ def forward( self._fp8_workspaces.get(cache_name) if cache_name is not None else None ) + dgrad_use_split_accumulator = _2X_ACC_DGRAD + wgrad_use_split_accumulator = _2X_ACC_WGRAD if self.fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + _recipe = FP8GlobalStateManager.get_fp8_recipe() + backward_override = _recipe.backward_override + if hasattr(_recipe, "fp8_gemm_dgrad"): + dgrad_use_split_accumulator = _recipe.fp8_gemm_dgrad.use_split_accumulator + if hasattr(_recipe, "fp8_gemm_wgrad"): + wgrad_use_split_accumulator = _recipe.fp8_gemm_wgrad.use_split_accumulator else: backward_override = None custom = is_custom(input_quantizer) or is_custom(weight_quantizer) @@ -1989,6 +1992,8 @@ def forward( fp8_output=fp8_output, save_original_input=self.save_original_input, backward_override=backward_override, + dgrad_use_split_accumulator=dgrad_use_split_accumulator, + wgrad_use_split_accumulator=wgrad_use_split_accumulator, custom=custom, debug=debug, # weight-workspace caching From e271bea10e62481906c048cb59f1522a35a71cff Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Date: Mon, 3 Aug 2026 21:09:25 +0000 Subject: [PATCH 04/10] [ROCm] Work around torch 2.12 HSA shutdown segfault in triton_kernels/test_norms --- tests/pytorch/triton_kernels/conftest.py | 48 ++++++++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 tests/pytorch/triton_kernels/conftest.py diff --git a/tests/pytorch/triton_kernels/conftest.py b/tests/pytorch/triton_kernels/conftest.py new file mode 100644 index 0000000000..9cbf0bffdd --- /dev/null +++ b/tests/pytorch/triton_kernels/conftest.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Work around a ROCm HSA-runtime teardown segfault after the triton_kernels tests. + +On ROCm 7.14 the HIP runtime's atexit handler calls ``hsa_shut_down()``, which +destroys the GPU agents and their AQL queues; ``AqlQueue::~AqlQueue()`` then +writes to an already-freed HSA doorbell signal and the process dies with SIGSEGV. +Observed backtrace (``libamdhip64.so.7`` / ``libhsa-runtime64.so.1``):: + + exit() -> -> rocr::HSA::hsa_shut_down() + -> Runtime::Release() -> Unload() -> DestroyAgents() + -> GpuAgent::~GpuAgent() -> AqlQueue::~AqlQueue() + -> hsa_signal_store_screlease <-- SIGSEGV + +The crash happens *after* every test has passed and pytest has already written +its JUnit XML report, so the run is functionally green, but the process exits +139 and ``ci/_utils.sh``'s exit-code gate records a suite error. ``test_norms.py`` +reliably trips it because it allocates the most Triton streams/queues of the +triton_kernels suites; the fault itself is purely in the ROCm libraries' shutdown +path, not in TE. + +Bypass the buggy C-level atexit handler by hard-exiting with pytest's real exit +status once the session (and its report writing) has finished. Real failures +still propagate through ``exitstatus``, and a crash *during* the run (before +``pytest_sessionfinish``) is unaffected and still surfaces as an error. Remove +once the ROCm HSA shutdown crash is fixed. +""" + +import os +import sys + +import pytest +import torch + + +@pytest.hookimpl(trylast=True) +def pytest_sessionfinish(session, exitstatus): + # Only ROCm hits the hsa_shut_down teardown segfault; leave CUDA/CPU exit + # semantics (and their normal atexit cleanup) untouched. + if getattr(torch.version, "hip", None) is None: + return + # trylast ensures the junitxml plugin and te_ci_result_sink have already + # written their reports in this same hook before we hard-exit. + sys.stdout.flush() + sys.stderr.flush() + os._exit(0 if exitstatus == 0 else int(exitstatus)) From cadc9443ef1e5bca65fa640d82562f73f689fc74 Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Date: Mon, 3 Aug 2026 21:17:52 +0000 Subject: [PATCH 05/10] Added comments --- .github/scripts/Dockerfile.ci.deps | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/.github/scripts/Dockerfile.ci.deps b/.github/scripts/Dockerfile.ci.deps index e369c6f476..14aa8f76e5 100644 --- a/.github/scripts/Dockerfile.ci.deps +++ b/.github/scripts/Dockerfile.ci.deps @@ -65,21 +65,20 @@ RUN W="https://repo.amd.com/rocm/whl-multi-arch" \ "${W}/amd_torchvision_device_${arch}-${TORCHVISION_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" ; \ done -# Install Jax — prebuilt ROCm 7.14 (TheRock, multi-arch) wheels from PyPI. -# jax pulls jaxlib; jax-rocm7-plugin pulls jax-rocm7-pjrt. ROCm runtime is provided -# by the rocm-sdk-* wheels above, so the plugin's [with-rocm] extra (which would -# pull duplicate amd_*_instinct libs) is intentionally not requested. +# Install Jax RUN pip install --no-cache-dir \ "jax==${JAX_VERSION}" \ "jax-rocm7-pjrt==${JAX_VERSION}" \ "jax-rocm7-plugin==${JAX_VERSION}" +# Install Flash Attention RUN export ROCM_PATH="$(rocm-sdk path --root)" \ && git clone --branch "${FA_VERSION}" --depth 1 https://github.com/Dao-AILab/flash-attention.git /tmp/flash-attention \ && cd /tmp/flash-attention \ && GPU_ARCHS="${GPU_ARCH}" python setup.py install \ && rm -rf /tmp/flash-attention +# Install Aiter RUN export ROCM_PATH="$(rocm-sdk path --root)" \ && git clone --no-checkout https://github.com/ROCm/aiter.git /tmp/aiter \ && cd /tmp/aiter \ From 9bf7e516c150f71f8afc3a126874500964ccab9c Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Date: Tue, 4 Aug 2026 17:24:24 +0000 Subject: [PATCH 06/10] Updated default docker image --- ci/ci_config.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/ci_config.json b/ci/ci_config.json index d2f01bfac5..db70c74555 100644 --- a/ci/ci_config.json +++ b/ci/ci_config.json @@ -1,5 +1,5 @@ { "docker_images": { - "default": "registry-sc-harbor.amd.com/framework/te-ci:therock_7.13.0_ubuntu24.04_py3.12_pytorch_2.10.0_triton_3.6.0_jax_0.10.2_fa_2.8.1_aiter_77455e3ecf" + "default": "registry-sc-harbor.amd.com/framework/te-ci:therock_7.14.0_ubuntu24.04_py3.12_pytorch_2.12.0_triton_3.7.1_git0263a6a6_jax_0.11.0_fa_2.8.3_aiter_77455e3ecf" } } From b99d341618c04801df9eb0f94fe63df697c78e2d Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Reddy Gopu Date: Tue, 4 Aug 2026 22:29:24 +0000 Subject: [PATCH 07/10] Added rocm_sdk support to core tests --- ci/_utils.sh | 4 ++-- transformer_engine/common/CMakeLists.txt | 2 ++ transformer_engine/common/util/cuda_runtime.cpp | 10 +++++++--- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/ci/_utils.sh b/ci/_utils.sh index f2478fc762..d09776b93a 100644 --- a/ci/_utils.sh +++ b/ci/_utils.sh @@ -313,8 +313,8 @@ check_test_filter() { start_message() { echo "Started with TEST_LEVEL=$TEST_LEVEL sGPU='$TEST_SGPU' mGPU='$TEST_MGPU' at `date`" - _rocm_path=$(resolve_rocm_path) - _rocm_path=`$REALPATH "$_rocm_path" 2>/dev/null || echo "$_rocm_path"` + export ROCM_PATH=$(resolve_rocm_path) + _rocm_path=`$REALPATH "$ROCM_PATH" 2>/dev/null || echo "$ROCM_PATH"` echo "ROCM PATH: $_rocm_path" python3 --version } diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index e646b72c38..43bd6b5a8c 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -757,6 +757,8 @@ make_string_header_from_file(transpose/rtc/swap_first_dims.cu make_string_header_from_file(utils.cuh string_code_utils_cuh) else() + make_string_header("${ROCM_PATH}" + string_path_rocm_root) make_string_header_from_file(utils_hip.cuh string_code_utils_cuh) make_string_header_from_file(transpose/rtc/cast_transpose_fusion.hip diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp index 2d29f7e06d..7f34d728d8 100644 --- a/transformer_engine/common/util/cuda_runtime.cpp +++ b/transformer_engine/common/util/cuda_runtime.cpp @@ -23,14 +23,17 @@ namespace transformer_engine { namespace cuda { -#ifndef __HIP_PLATFORM_AMD__ namespace { -// String with build-time CUDA include path +#ifndef __HIP_PLATFORM_AMD__ +// Build-time CUDA include path #include "string_path_cuda_include.h" +#else +// Build-time ROCm root path (from rocm_utils.cmake) +#include "string_path_rocm_root.h" +#endif } // namespace -#endif // #ifndef __HIP_PLATFORM_AMD__ int num_devices() { auto query_num_devices = []() -> int { @@ -197,6 +200,7 @@ const std::string &include_directory(bool required) { #ifdef __HIP_PLATFORM_AMD__ std::vector> search_paths = {{"ROCM_PATH", ""}, {"HIP_PATH", ""}, + {"", string_path_rocm_root}, {"", "/opt/rocm/core"}, {"", "/opt/rocm"}}; #else From a19283399509be3156faecba8c05c249763da237 Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Reddy Gopu Date: Wed, 5 Aug 2026 05:35:56 +0000 Subject: [PATCH 08/10] Updated docker image tag --- .github/workflows/rocm-ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/rocm-ci.yml b/.github/workflows/rocm-ci.yml index 5591f1f399..606b1e835c 100644 --- a/.github/workflows/rocm-ci.yml +++ b/.github/workflows/rocm-ci.yml @@ -109,7 +109,7 @@ jobs: timeout-minutes: 270 runs-on: ${{ matrix.arch_label == 'mi30x' && 'linux-te-mi30x-4' || matrix.arch_label == 'mi35x' && 'linux-te-mi35x-4' }} env: - DOCKER_IMAGE: ${{ inputs.docker_image_override || format('{0}{1}', needs.select_image.outputs.image-tag, matrix.arch_label == 'mi30x' && '_gfx942' || matrix.arch_label == 'mi35x' && '_gfx950') }} + DOCKER_IMAGE: ${{ inputs.docker_image_override || needs.select_image.outputs.image-tag }} strategy: fail-fast: false matrix: @@ -262,7 +262,7 @@ jobs: timeout-minutes: 210 runs-on: ${{ matrix.arch_label == 'mi30x' && 'linux-te-mi30x-8' || matrix.arch_label == 'mi35x' && 'linux-te-mi35x-8' }} env: - DOCKER_IMAGE: ${{ inputs.docker_image_override || format('{0}{1}', needs.select_image.outputs.image-tag, matrix.arch_label == 'mi30x' && '_gfx942' || matrix.arch_label == 'mi35x' && '_gfx950') }} + DOCKER_IMAGE: ${{ inputs.docker_image_override || needs.select_image.outputs.image-tag }} strategy: fail-fast: false matrix: From 32e2acad3161b0189d26dd7327f82514415d460e Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Reddy Gopu Date: Wed, 5 Aug 2026 05:58:30 +0000 Subject: [PATCH 09/10] Add runtime rocm_sdk resolution --- transformer_engine/common/CMakeLists.txt | 2 - .../common/util/cuda_runtime.cpp | 37 ++++++++++++++++--- 2 files changed, 31 insertions(+), 8 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 43bd6b5a8c..e646b72c38 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -757,8 +757,6 @@ make_string_header_from_file(transpose/rtc/swap_first_dims.cu make_string_header_from_file(utils.cuh string_code_utils_cuh) else() - make_string_header("${ROCM_PATH}" - string_path_rocm_root) make_string_header_from_file(utils_hip.cuh string_code_utils_cuh) make_string_header_from_file(transpose/rtc/cast_transpose_fusion.hip diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp index 7f34d728d8..802dda6e0c 100644 --- a/transformer_engine/common/util/cuda_runtime.cpp +++ b/transformer_engine/common/util/cuda_runtime.cpp @@ -10,6 +10,7 @@ #include +#include #include #include #include @@ -23,17 +24,14 @@ namespace transformer_engine { namespace cuda { +#ifndef __HIP_PLATFORM_AMD__ namespace { -#ifndef __HIP_PLATFORM_AMD__ // Build-time CUDA include path #include "string_path_cuda_include.h" -#else -// Build-time ROCm root path (from rocm_utils.cmake) -#include "string_path_rocm_root.h" -#endif } // namespace +#endif // __HIP_PLATFORM_AMD__ int num_devices() { auto query_num_devices = []() -> int { @@ -185,6 +183,33 @@ bool supports_multicast(int device_id) { } #endif // __HIP_PLATFORM_AMD__ +#ifdef __HIP_PLATFORM_AMD__ +namespace { + +using Path = std::filesystem::path; + +// Runtime fallback matching resolve_rocm_path() +Path rocm_sdk_root() { + FILE *pipe = popen("rocm-sdk path --root 2>/dev/null", "r"); + if (pipe == nullptr) { + return {}; + } + char buffer[4096]; + if (fgets(buffer, sizeof(buffer), pipe) == nullptr) { + pclose(pipe); + return {}; + } + pclose(pipe); + std::string root = buffer; + while (!root.empty() && (root.back() == '\n' || root.back() == '\r')) { + root.pop_back(); + } + return Path{root}; +} + +} // namespace +#endif // __HIP_PLATFORM_AMD__ + const std::string &include_directory(bool required) { static std::string path; @@ -200,7 +225,7 @@ const std::string &include_directory(bool required) { #ifdef __HIP_PLATFORM_AMD__ std::vector> search_paths = {{"ROCM_PATH", ""}, {"HIP_PATH", ""}, - {"", string_path_rocm_root}, + {"rocm-sdk", rocm_sdk_root()}, {"", "/opt/rocm/core"}, {"", "/opt/rocm"}}; #else From 026fa695847b6223e81277746c5d73595f63eaa3 Mon Sep 17 00:00:00 2001 From: Veera Rajasekhar Reddy Gopu Date: Wed, 5 Aug 2026 22:24:33 +0000 Subject: [PATCH 10/10] Addressed reviews --- .../attention/benchmark_attention_rocm.py | 29 +++---------------- tests/pytorch/test_numerics.py | 5 ++-- .../common/util/cuda_runtime.cpp | 29 ------------------- 3 files changed, 7 insertions(+), 56 deletions(-) diff --git a/benchmarks/attention/benchmark_attention_rocm.py b/benchmarks/attention/benchmark_attention_rocm.py index 87b57f94a7..60fcaa9aa3 100644 --- a/benchmarks/attention/benchmark_attention_rocm.py +++ b/benchmarks/attention/benchmark_attention_rocm.py @@ -138,29 +138,6 @@ def setup_backend_env(backend_name, use_ck_bwd_v3=True, use_ck_fwd_v3=True, use_ ROCPROF_STATS_CSV = "results.stats.csv" - -def _resolve_rocprofv3(): - """Path to the rocprofv3 shipped in rocm-sdk-core. - - The `rocprofv3` on PATH is the rocm-sdk trampoline, which execs the copy in - rocm-sdk-devel. That copy loads librocprofiler-sdk.so.1 from _rocm_sdk_devel, - while torch/TE have already loaded the byte-identical copy from _rocm_sdk_core. - rocprofiler-register treats the two different paths as a conflict and aborts - (registration.cpp: "already set ... not overriding"), so rocprofv3 dies with - SIGABRT and no kernel stats are produced. Running the core copy makes the tool - and the profiled app share the same library, avoiding the conflict. - """ - try: - from rocm_sdk_core._cli import _get_core_module_path - - core_bin = _get_core_module_path() / "bin" / "rocprofv3" - if core_bin.is_file(): - return str(core_bin) - except Exception: - pass - return shutil.which("rocprofv3") or "rocprofv3" - - def _profiler_python_code(model, attention, column_name, benchmark_dir): return ( f"import sys; sys.path.insert(0, {benchmark_dir!r}); " @@ -189,7 +166,7 @@ def _run_attention_profiler(model, attention, column_name, dirname): benchmark_dir = os.path.dirname(os.path.abspath(__file__)) py_code = _profiler_python_code(model, attention, column_name, benchmark_dir) cmd = [ - _resolve_rocprofv3(), + "rocprofv3", "--hip-trace", "--kernel-trace", "--stats", @@ -203,7 +180,9 @@ def _run_attention_profiler(model, attention, column_name, dirname): py_code, ] try: - result = subprocess.run(cmd, capture_output=True, text=True) + # Pass os.environ explicitly: C code may set ROCPROFILER_REGISTER_LIBRARY + # via setenv() without updating Python's environ cache, which breaks rocprofv3. + result = subprocess.run(cmd, capture_output=True, text=True, env=os.environ.copy()) except FileNotFoundError: print( "WARNING: rocprofv3 not found on PATH; kernel timing columns may be empty.", diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index b6ee0f4856..ec9cfcc22a 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -2130,10 +2130,11 @@ def test_gpt_cuda_graph(dtype, bs, model): # per-(device, stream). The graph-capture stream's handle is created lazily # during capture, and hipblasLtCreate performs an internal hipMalloc that is # illegal mid-capture, failing with HIP error 900 ("operation not permitted - # when stream is capturing"). The upstream capture_begin pre-init + # when stream is capturing"). The capture_begin pre-init # (pytorch/pytorch#180692) only covers the calling thread, not the autograd # backward thread, so torch's own bmm in the captured backward still trips it. - # Route torch's matmul/bmm off hipBLASLt for this test until upstream fixes it. + # Route torch's matmul/bmm off hipBLASLt for this test until PyTorch + # extends the pre-init to cover it. _prev_blas_library = None if IS_HIP_EXTENSION: _prev_blas_library = torch.backends.cuda.preferred_blas_library() diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp index 802dda6e0c..edb8650273 100644 --- a/transformer_engine/common/util/cuda_runtime.cpp +++ b/transformer_engine/common/util/cuda_runtime.cpp @@ -10,7 +10,6 @@ #include -#include #include #include #include @@ -183,33 +182,6 @@ bool supports_multicast(int device_id) { } #endif // __HIP_PLATFORM_AMD__ -#ifdef __HIP_PLATFORM_AMD__ -namespace { - -using Path = std::filesystem::path; - -// Runtime fallback matching resolve_rocm_path() -Path rocm_sdk_root() { - FILE *pipe = popen("rocm-sdk path --root 2>/dev/null", "r"); - if (pipe == nullptr) { - return {}; - } - char buffer[4096]; - if (fgets(buffer, sizeof(buffer), pipe) == nullptr) { - pclose(pipe); - return {}; - } - pclose(pipe); - std::string root = buffer; - while (!root.empty() && (root.back() == '\n' || root.back() == '\r')) { - root.pop_back(); - } - return Path{root}; -} - -} // namespace -#endif // __HIP_PLATFORM_AMD__ - const std::string &include_directory(bool required) { static std::string path; @@ -225,7 +197,6 @@ const std::string &include_directory(bool required) { #ifdef __HIP_PLATFORM_AMD__ std::vector> search_paths = {{"ROCM_PATH", ""}, {"HIP_PATH", ""}, - {"rocm-sdk", rocm_sdk_root()}, {"", "/opt/rocm/core"}, {"", "/opt/rocm"}}; #else