build: update image cuBLAS to 13.8.1 (#6252)

Signed-off-by: Ajay Balasa <abalasa@nvidia.com>

Author: balasaajay

docker/Dockerfile.ci

@@ -62,6 +62,30 @@ ENV MBRIDGE_PRESERVE_BASE_RUNTIME=${PRESERVE_BASE_RUNTIME} \
     RDMA_CORE_HOME=/opt/rdma-core/build \
     LD_LIBRARY_PATH="/usr/local/cuda/lib64/:$LD_LIBRARY_PATH"
 
+ARG TARGETARCH
+ARG CUBLAS_VERSION=13.8.1.7
+RUN bash -euxo pipefail <<"EOF"
+    if [ "$MBRIDGE_PRESERVE_BASE_RUNTIME" = "True" ]; then exit 0; fi
+    # Match the cuBLAS update in NVIDIA/Megatron-LM#7685.
+    case "${TARGETARCH}" in
+        amd64) CUBLAS_ARCH=x86_64 ;;
+        arm64) CUBLAS_ARCH=sbsa ;;
+        *) echo "Unsupported TARGETARCH for cuBLAS: ${TARGETARCH}" >&2; exit 1 ;;
+    esac
+    CUBLAS_ARCHIVE="libcublas-linux-${CUBLAS_ARCH}-${CUBLAS_VERSION}_cuda13-archive"
+    CUBLAS_TMP=$(mktemp -d)
+    trap 'rm -rf "${CUBLAS_TMP}"' EXIT
+    curl -fLSs --retry 3 -o "${CUBLAS_TMP}/cublas.tar.xz" \
+        "https://developer.download.nvidia.com/compute/cublas/redist/libcublas/linux-${CUBLAS_ARCH}/${CUBLAS_ARCHIVE}.tar.xz"
+    tar -xJf "${CUBLAS_TMP}/cublas.tar.xz" -C "${CUBLAS_TMP}"
+    # Replace the toolkit copies so existing linker paths use the updated libraries.
+    CUDA_ROOT="${CUDA_HOME:-/usr/local/cuda}"
+    cp -a "${CUBLAS_TMP}/${CUBLAS_ARCHIVE}/include/." "${CUDA_ROOT}/include/"
+    cp -a "${CUBLAS_TMP}/${CUBLAS_ARCHIVE}/lib/." "${CUDA_ROOT}/lib64/"
+    install -Dm644 "${CUBLAS_TMP}/${CUBLAS_ARCHIVE}/LICENSE" "${CUDA_ROOT}/doc/cublas/LICENSE"
+    ldconfig
+EOF
+
 ##############################################################################
 ##
 ## Install DeepEP and nvshmem

docker/README.md

@@ -30,6 +30,7 @@ docker build \
 | Argument | Description |
 |---|---|
 | `BASE_IMAGE` | Base container |
+| `CUBLAS_VERSION` | CUDA 13 cuBLAS archive version; defaults to `13.8.1.7` for x86_64 and Arm SBSA |
 | `MCORE_TRIGGERED_TESTING` | When `true`, skips the uv lockfile check to allow testing against a different Megatron-LM version than the one pinned in the lockfile |
 | `INSTALL_DIFFUSION_DEPS` | When `true`, runs `scripts/install_diffusion_deps.sh` to add WAN codecs to a CI image; defaults to `false` so CVE-carrying codecs stay out of shipped Framework images |
 | `UV_CACHE_PRUNE_ARGS` | Extra arguments forwarded to `uv cache prune` after install |
@@ -58,7 +59,7 @@ For a base image with a validated, source-built Transformer Engine/PyTorch pairi
 to `Dockerfile.ci`. The default is `False`, preserving the normal installation path.
 
 When enabled, Bridge and FW-final dependency syncs skip installing Transformer
-Engine and cuDNN frontend, and Bridge skips its public CUTLASS DSL replacement.
+Engine and cuDNN frontend, and Bridge skips its public CUTLASS DSL and cuBLAS replacements.
 FW-final inherits the setting from the Bridge image; it does not need a second
 build argument. The build records the selected Torch, TE, cuDNN frontend, and
 CUTLASS distribution versions and installation locations in
@@ -199,14 +200,15 @@ docker build \
 | Argument | Description |
 |---|---|
 | `BASE_IMAGE` | Base container; set to the fw-base image when building the full stack |
+| `CUBLAS_VERSION` | CUDA 13 cuBLAS archive version; defaults to `13.8.1.7`. Installs cuBLAS/cuBLASLt and headers into the existing CUDA toolkit before native dependencies are built; skipped with `PRESERVE_BASE_RUNTIME=True` |
 | `APPLY_PYTORCH_LIBRARY_FINALIZER_PATCH` | Set to `True` to apply the torch.library finalizer patch; `False` for base images it does not apply against (e.g. Rubin) |
 | `INSTALL_DEEPEP` | Set to `True` to build and install DeepEP and nvshmem |
 | `DEEPEP_COMMIT` | DeepEP git commit SHA |
 | `REINSTALL_NVSHMEM` | Set to `True` to reinstall nvshmem (`nvidia-nvshmem-cu13`) over the base image version; only applied when `INSTALL_DEEPEP=True` |
 | `MCORE_TRIGGERED_TESTING` | Skip uv lockfile check for cross-version Megatron-LM testing |
 | `INSTALL_DIFFUSION_DEPS` | Install the test-only WAN diffusion dependencies; defaults to `false` for Framework/release images |
 | `UV_CACHE_PRUNE_ARGS` | Extra arguments for `uv cache prune` |
-| `PRESERVE_BASE_RUNTIME` | Keep validated base TE/cuDNN frontend/CUTLASS packages and verify package metadata; defaults to `False` |
+| `PRESERVE_BASE_RUNTIME` | Keep validated base TE/cuDNN frontend/CUTLASS/cuBLAS packages and verify package metadata; defaults to `False` |
 
 ### `Dockerfile.fw_final`