Veergopu/upgrade ci rock 714

.github/scripts/Dockerfile.ci.deps

@@ -3,120 +3,88 @@
 # 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+<TRITON_LOCAL>.rocm<rel>
+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/<AMD_GPU_FAMILY>/
-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-<arch>,
+# amd_torch_device_<arch>, amd_torchvision_device_<arch>) 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
+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 \
+# 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 \
-    && 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 \
+# 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 \
     && 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/

.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:

benchmarks/attention/benchmark_attention_rocm.py

@@ -138,7 +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 _profiler_python_code(model, attention, column_name, benchmark_dir):
     return (
         f"import sys; sys.path.insert(0, {benchmark_dir!r}); "
@@ -181,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.",

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
 }

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"
   }
 }

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"

tests/pytorch/test_numerics.py

@@ -2126,8 +2126,25 @@ 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 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 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()
+        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())
 

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
+    (``<core-root>/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):

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]