From effb367a1a94754db91d9d7c18414337c26a0449 Mon Sep 17 00:00:00 2001
From: Venkata Sai Madhur Karampudi <vkarampudi@google.com>
Date: Thu, 9 Jul 2026 20:17:07 +0000
Subject: [PATCH 1/6] Test: CI build optimizations (CUDA compute capabilities
 and non-curses mode from cl/945255455)

---
 tensorflow_serving/tools/docker/Dockerfile.devel     | 4 ++--
 tensorflow_serving/tools/docker/Dockerfile.devel-gpu | 7 +++++--
 tensorflow_serving/tools/docker/Dockerfile.devel-mkl | 4 ++--
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/tensorflow_serving/tools/docker/Dockerfile.devel b/tensorflow_serving/tools/docker/Dockerfile.devel
index 7c79a983f49..0f8304f9cea 100644
--- a/tensorflow_serving/tools/docker/Dockerfile.devel
+++ b/tensorflow_serving/tools/docker/Dockerfile.devel
@@ -101,7 +101,7 @@ RUN echo "Building with build options: ${TF_SERVING_BUILD_OPTIONS}"
 ARG TF_SERVING_BAZEL_OPTIONS=""
 RUN echo "Building with Bazel options: ${TF_SERVING_BAZEL_OPTIONS}"
 
-RUN bazel build --color=yes --curses=yes \
+RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \
@@ -111,7 +111,7 @@ RUN bazel build --color=yes --curses=yes \
     /usr/local/bin/
 
 # Build and install TensorFlow Serving API
-RUN bazel build --color=yes --curses=yes \
+RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \
diff --git a/tensorflow_serving/tools/docker/Dockerfile.devel-gpu b/tensorflow_serving/tools/docker/Dockerfile.devel-gpu
index cde11aa6fe4..a0836ecb58c 100644
--- a/tensorflow_serving/tools/docker/Dockerfile.devel-gpu
+++ b/tensorflow_serving/tools/docker/Dockerfile.devel-gpu
@@ -25,6 +25,9 @@ ENV CUDNN_VERSION=9.3.0.75
 ENV TF_TENSORRT_VERSION=8.6.1
 ENV CUDA=12.5
 
+ARG TF_CUDA_COMPUTE_CAPABILITIES="7.0,7.5,8.0,8.6,8.9,9.0"
+ENV TF_CUDA_COMPUTE_CAPABILITIES=${TF_CUDA_COMPUTE_CAPABILITIES}
+
 RUN apt-get update && apt-get install -y --no-install-recommends \
         curl
 
@@ -155,7 +158,7 @@ RUN echo "Building with Bazel options: ${TF_SERVING_BAZEL_OPTIONS}"
 
 RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 && \
     LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:${LD_LIBRARY_PATH} \
-    bazel build --color=yes --curses=yes --config=cuda_clang --copt="-fPIC"\
+    bazel build --color=yes --curses=no --progress_report_interval=30 --config=cuda_clang --copt="-fPIC"\
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \
@@ -166,7 +169,7 @@ RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/lib
     rm /usr/local/cuda/lib64/stubs/libcuda.so.1
 
 # Build and install TensorFlow Serving API
-RUN bazel build --color=yes --curses=yes \
+RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \
diff --git a/tensorflow_serving/tools/docker/Dockerfile.devel-mkl b/tensorflow_serving/tools/docker/Dockerfile.devel-mkl
index 0a74fccc7a2..21614d96d42 100644
--- a/tensorflow_serving/tools/docker/Dockerfile.devel-mkl
+++ b/tensorflow_serving/tools/docker/Dockerfile.devel-mkl
@@ -103,7 +103,7 @@ RUN echo "Building with build options: ${TF_SERVING_BUILD_OPTIONS}"
 ARG TF_SERVING_BAZEL_OPTIONS=""
 RUN echo "Building with Bazel options: ${TF_SERVING_BAZEL_OPTIONS}"
 
-RUN bazel build --color=yes --curses=yes \
+RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \
@@ -113,7 +113,7 @@ RUN bazel build --color=yes --curses=yes \
     /usr/local/bin/
 
 # Build and install TensorFlow Serving API
-RUN bazel build --color=yes --curses=yes \
+RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \

From a2b7756bacc471926246b4f58a34e414d1dbb3d4 Mon Sep 17 00:00:00 2001
From: Venkata Sai Madhur Karampudi <vkarampudi@google.com>
Date: Thu, 9 Jul 2026 22:37:04 +0000
Subject: [PATCH 2/6] Fix: Prune sm_60 from .bazelrc and add --noshow_progress
 to Dockerfiles to eliminate 5hr GPU build and 2MiB log clipping

---
 .bazelrc                                             | 2 +-
 tensorflow_serving/tools/docker/Dockerfile.devel     | 4 ++--
 tensorflow_serving/tools/docker/Dockerfile.devel-gpu | 4 ++--
 tensorflow_serving/tools/docker/Dockerfile.devel-mkl | 4 ++--
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/.bazelrc b/.bazelrc
index 6278d956d75..e049e63ad3e 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -17,7 +17,7 @@ build:cuda_clang --host_copt=-Wno-error=unused-command-line-argument
 build:cuda_clang --repo_env TF_NEED_TENSORRT=0
 build:cuda_clang --action_env=TF_CUDA_CLANG="1"
 build:cuda_clang --@local_config_cuda//:cuda_compiler=clang
-build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_60,sm_70,sm_80,sm_89,compute_90"
+build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_70,sm_80,sm_89,compute_90"
 
 common:cuda_clang --repo_env=HERMETIC_CUDA_VERSION="12.5.1"
 common:cuda_clang --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
diff --git a/tensorflow_serving/tools/docker/Dockerfile.devel b/tensorflow_serving/tools/docker/Dockerfile.devel
index 0f8304f9cea..9f9286ecf7a 100644
--- a/tensorflow_serving/tools/docker/Dockerfile.devel
+++ b/tensorflow_serving/tools/docker/Dockerfile.devel
@@ -101,7 +101,7 @@ RUN echo "Building with build options: ${TF_SERVING_BUILD_OPTIONS}"
 ARG TF_SERVING_BAZEL_OPTIONS=""
 RUN echo "Building with Bazel options: ${TF_SERVING_BAZEL_OPTIONS}"
 
-RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
+RUN bazel build --color=yes --noshow_progress --curses=no \
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \
@@ -111,7 +111,7 @@ RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
     /usr/local/bin/
 
 # Build and install TensorFlow Serving API
-RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
+RUN bazel build --color=yes --noshow_progress --curses=no \
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \
diff --git a/tensorflow_serving/tools/docker/Dockerfile.devel-gpu b/tensorflow_serving/tools/docker/Dockerfile.devel-gpu
index a0836ecb58c..f06e2428022 100644
--- a/tensorflow_serving/tools/docker/Dockerfile.devel-gpu
+++ b/tensorflow_serving/tools/docker/Dockerfile.devel-gpu
@@ -158,7 +158,7 @@ RUN echo "Building with Bazel options: ${TF_SERVING_BAZEL_OPTIONS}"
 
 RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 && \
     LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:${LD_LIBRARY_PATH} \
-    bazel build --color=yes --curses=no --progress_report_interval=30 --config=cuda_clang --copt="-fPIC"\
+    bazel build --color=yes --noshow_progress --curses=no --config=cuda_clang --copt="-fPIC"\
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \
@@ -169,7 +169,7 @@ RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/lib
     rm /usr/local/cuda/lib64/stubs/libcuda.so.1
 
 # Build and install TensorFlow Serving API
-RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
+RUN bazel build --color=yes --noshow_progress --curses=no \
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \
diff --git a/tensorflow_serving/tools/docker/Dockerfile.devel-mkl b/tensorflow_serving/tools/docker/Dockerfile.devel-mkl
index 21614d96d42..d87be197933 100644
--- a/tensorflow_serving/tools/docker/Dockerfile.devel-mkl
+++ b/tensorflow_serving/tools/docker/Dockerfile.devel-mkl
@@ -103,7 +103,7 @@ RUN echo "Building with build options: ${TF_SERVING_BUILD_OPTIONS}"
 ARG TF_SERVING_BAZEL_OPTIONS=""
 RUN echo "Building with Bazel options: ${TF_SERVING_BAZEL_OPTIONS}"
 
-RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
+RUN bazel build --color=yes --noshow_progress --curses=no \
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \
@@ -113,7 +113,7 @@ RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
     /usr/local/bin/
 
 # Build and install TensorFlow Serving API
-RUN bazel build --color=yes --curses=no --progress_report_interval=30 \
+RUN bazel build --color=yes --noshow_progress --curses=no \
     ${TF_SERVING_BAZEL_OPTIONS} \
     --verbose_failures \
     --output_filter=DONT_MATCH_ANYTHING \

From fccb929f1c931e50371ef58b6cacdf7bb74f22cb Mon Sep 17 00:00:00 2001
From: Venkata Sai Madhur Karampudi <vkarampudi@google.com>
Date: Fri, 10 Jul 2026 22:26:04 +0000
Subject: [PATCH 3/6] Fix: Prune sm_70 from .bazelrc to resolve Volta Clang
 CUDA build errors and speed up compilation

---
 .bazelrc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.bazelrc b/.bazelrc
index e049e63ad3e..7ccb48ec43b 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -17,7 +17,7 @@ build:cuda_clang --host_copt=-Wno-error=unused-command-line-argument
 build:cuda_clang --repo_env TF_NEED_TENSORRT=0
 build:cuda_clang --action_env=TF_CUDA_CLANG="1"
 build:cuda_clang --@local_config_cuda//:cuda_compiler=clang
-build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_70,sm_80,sm_89,compute_90"
+build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_80,sm_89,compute_90"
 
 common:cuda_clang --repo_env=HERMETIC_CUDA_VERSION="12.5.1"
 common:cuda_clang --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"

From dce4e8ec5f4f738e4d1370ba2f6c97d6e018a920 Mon Sep 17 00:00:00 2001
From: Venkata Sai Madhur Karampudi <vkarampudi@google.com>
Date: Fri, 10 Jul 2026 23:51:29 +0000
Subject: [PATCH 4/6] Fix: Add cuda_clang_fixes.h compatibility header and
 include flags to resolve NCCL 2.29 typeof and uint4 __stwt errors under Clang

---
 .bazelrc                                    |   2 +
 tensorflow_serving/opensource_only/.bazelrc | 129 ++++++++++++++++++++
 tensorflow_serving/tools/cuda_clang_fixes.h |  36 ++++++
 3 files changed, 167 insertions(+)
 create mode 100644 tensorflow_serving/opensource_only/.bazelrc
 create mode 100644 tensorflow_serving/tools/cuda_clang_fixes.h

diff --git a/.bazelrc b/.bazelrc
index 7ccb48ec43b..4e9b9e0cc97 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -18,6 +18,8 @@ build:cuda_clang --repo_env TF_NEED_TENSORRT=0
 build:cuda_clang --action_env=TF_CUDA_CLANG="1"
 build:cuda_clang --@local_config_cuda//:cuda_compiler=clang
 build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_80,sm_89,compute_90"
+build:cuda_clang --copt=-include
+build:cuda_clang --copt=third_party/tensorflow_serving/tools/cuda_clang_fixes.h
 
 common:cuda_clang --repo_env=HERMETIC_CUDA_VERSION="12.5.1"
 common:cuda_clang --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
diff --git a/tensorflow_serving/opensource_only/.bazelrc b/tensorflow_serving/opensource_only/.bazelrc
new file mode 100644
index 00000000000..4e9b9e0cc97
--- /dev/null
+++ b/tensorflow_serving/opensource_only/.bazelrc
@@ -0,0 +1,129 @@
+# Optimizations used for TF Serving release builds.
+build:release --copt=-mavx
+build:release --copt=-msse4.2
+
+# Options used to build with CUDA.
+build:cuda --repo_env TF_NEED_CUDA=1
+build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
+build:cuda --@local_config_cuda//:enable_cuda
+
+# Options used to build with CUDA clang
+build:cuda_clang --config=cuda
+build:cuda_clang --copt=-Wno-gnu-offsetof-extensions
+build:cuda_clang --copt=-Wno-error=unused-command-line-argument
+build --linkopt=-Wno-unused-command-line-argument
+build --host_linkopt=-Wno-unused-command-line-argument
+build:cuda_clang --host_copt=-Wno-error=unused-command-line-argument
+build:cuda_clang --repo_env TF_NEED_TENSORRT=0
+build:cuda_clang --action_env=TF_CUDA_CLANG="1"
+build:cuda_clang --@local_config_cuda//:cuda_compiler=clang
+build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_80,sm_89,compute_90"
+build:cuda_clang --copt=-include
+build:cuda_clang --copt=third_party/tensorflow_serving/tools/cuda_clang_fixes.h
+
+common:cuda_clang --repo_env=HERMETIC_CUDA_VERSION="12.5.1"
+common:cuda_clang --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
+common:cuda_clang --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5"
+common:cuda_clang --repo_env=HERMETIC_NCCL_VERSION="2.29.7"
+build:cuda_clang --action_env=GCC_HOST_COMPILER_PATH="/usr/bin/gcc-10"
+build:cuda_clang --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm-17/bin/clang"
+
+# Options used to build with TPU support.
+build:tpu --define=with_tpu_support=true --define=framework_shared_object=false
+build:tpu --copt=-DLIBTPU_ON_GCE
+
+# Please note that MKL on MacOS or windows is still not supported.
+# If you would like to use a local MKL instead of downloading, please set the
+# environment variable "TF_MKL_ROOT" every time before build.
+build:mkl --define=build_with_mkl=true --define=enable_mkl=true --define=build_with_openmp=false
+build:mkl --define=tensorflow_mkldnn_contraction_kernel=0
+
+# This config option is used to enable MKL-DNN open source library only,
+# without depending on MKL binary version.
+build:mkl_open_source_only --define=build_with_mkl_dnn_only=true
+build:mkl_open_source_only --define=build_with_mkl=true --define=enable_mkl=true
+build:mkl_open_source_only --define=tensorflow_mkldnn_contraction_kernel=0
+
+# Config setting to build oneDNN with Compute Library for the Arm Architecture (ACL).
+# This build is for the inference regime only.
+build:mkl_aarch64 --define=build_with_mkl_aarch64=true --define=enable_mkl=true
+build:mkl_aarch64 --define=tensorflow_mkldnn_contraction_kernel=0
+build:mkl_aarch64 --define=build_with_mkl_opensource=true
+build:mkl_aarch64 --define=build_with_openmp=true
+build:mkl_aarch64 --copt=-march=armv8.2-a
+build:mkl_aarch64 --copt=-O3
+
+build --define=build_with_onednn_v2=true
+build --define=xnn_enable_avxvnni=false
+build --define=xnn_enable_avxvnniint8=false
+build --define=xnn_enable_avx256vnnigfni=false
+build --define=xnn_enable_avx512amx=false
+build --define=xnn_enable_avx512fp16=false
+
+# Processor native optimizations (depends on build host capabilities).
+build:nativeopt --copt=-march=native
+build:nativeopt --host_copt=-march=native
+build:nativeopt --copt=-O3
+
+build --keep_going
+build --verbose_failures=true
+build --spawn_strategy=standalone
+build --genrule_strategy=standalone
+
+build --define=grpc_no_ares=true
+
+# Sets the default Apple platform to macOS.
+build --apple_platform_type=macos
+
+build -c opt
+
+# LLVM, MLIR and TF requires C++17.
+build --cxxopt=-std=c++17
+build --host_cxxopt=-std=c++17
+
+# Adding "--cxxopt=-D_GLIBCXX_USE_CXX11_ABI=0" creates parity with TF
+# compilation options. It also addresses memory use due to
+# copy-on-write semantics of std::strings of the older ABI.
+build --cxxopt=-D_GLIBCXX_USE_CXX11_ABI=0
+
+build --workspace_status_command=/proc/self/cwd/tools/gen_status_stamp.sh
+
+build --experimental_repo_remote_exec
+
+# TF now has `cc_shared_library` targets, so it needs the experimental flag
+build --experimental_cc_shared_library
+
+# Yggdrasil Decision Forests (internal library of TensorFlow Decision Forests)
+# uses TensorFlow for all IO operations.
+build --define=use_tensorflow_io=1
+
+# TensorFlow Decision Forests does not use Absl concurrency primitives on MacOs.
+# Reason: TensorFlow/ABSL ODR violation (b/214189609) # copybara:strip
+build:macos --define std_synchronization_primitives=1
+
+# Taken from https://github.com/openxla/xla/blob/99559d7a4f7c55490f46385ad29a3cbf9c3911af/warnings.bazelrc#L6
+# We silence warnings for code in `external`.
+build --per_file_copt=external/.*@-w
+build --host_per_file_copt=external/.*@-w
+build --copt=-Wno-macro-redefined # absl vs tsl logging clash
+build --copt=-Wno-sign-compare # int as loop variable
+build --copt=-Wno-deprecated-declarations
+build --copt=-Wno-unused-but-set-variable # due to `ifdefs` in ml_dtypes
+
+# Revert to the legacy WORKSPACE dependency resolution system
+# and disable Bzlmod.
+common --enable_bzlmod=false
+common --incompatible_enable_cc_toolchain_resolution
+common --repo_env USE_HERMETIC_CC_TOOLCHAIN=1
+common --repo_env=HERMETIC_PYTHON_VERSION="3.10"
+common --repo_env=USE_PYWRAP_RULES=1
+build --features=-layering_check
+build --define=framework_shared_object=false
+build --dynamic_mode=off
+
+
+common:clang_local --noincompatible_enable_cc_toolchain_resolution
+common:clang_local --noincompatible_enable_android_toolchain_resolution
+common:clang_local --@rules_ml_toolchain//cc_toolchain/config:enable_hermetic_cc=False
+common:clang_local --repo_env USE_HERMETIC_CC_TOOLCHAIN=0
+
diff --git a/tensorflow_serving/tools/cuda_clang_fixes.h b/tensorflow_serving/tools/cuda_clang_fixes.h
new file mode 100644
index 00000000000..f1e87a5a115
--- /dev/null
+++ b/tensorflow_serving/tools/cuda_clang_fixes.h
@@ -0,0 +1,36 @@
+/* Copyright 2026 Google Inc. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_SERVING_TOOLS_CUDA_CLANG_FIXES_H_
+#define TENSORFLOW_SERVING_TOOLS_CUDA_CLANG_FIXES_H_
+
+#if defined(__clang__) && defined(__CUDA__)
+#include <vector_types.h>
+
+// Fix for GNU typeof in strict C++ mode when compiling NVIDIA DOCA / NCCL headers with Clang.
+#ifndef typeof
+#define typeof __typeof__
+#endif
+
+// Missing 16-byte uint4 store overload in Clang's __clang_cuda_intrinsics.h for NCCL GIN proxy (__stwt).
+__device__ inline void __stwt(uint4* ptr, uint4 value) {
+  asm volatile("st.global.wt.v4.u32 [%0], {%1, %2, %3, %4};"
+               :
+               : "l"(ptr), "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w)
+               : "memory");
+}
+#endif
+
+#endif  // TENSORFLOW_SERVING_TOOLS_CUDA_CLANG_FIXES_H_

From 0f58c8fcb4cc409b323682f22cb1cdd0de3c420f Mon Sep 17 00:00:00 2001
From: Venkata Sai Madhur Karampudi <vkarampudi@google.com>
Date: Sat, 11 Jul 2026 19:09:12 +0000
Subject: [PATCH 5/6] Fix: Append typeof and uint4 __stwt overloads directly to
 clang_cuda_runtime_wrapper.h via rules_ml_toolchain patch_cmds to avoid
 un-sandboxed header file not found errors

---
 .bazelrc                                      |   3 +-
 tensorflow_serving/opensource_only/.bazelrc   |   3 +-
 .../opensource_only/root.workspace            | 297 ++++++++++++++++++
 tensorflow_serving/tools/cuda_clang_fixes.h   |  36 ---
 4 files changed, 299 insertions(+), 40 deletions(-)
 create mode 100644 tensorflow_serving/opensource_only/root.workspace
 delete mode 100644 tensorflow_serving/tools/cuda_clang_fixes.h

diff --git a/.bazelrc b/.bazelrc
index 4e9b9e0cc97..c580cf0b65c 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -18,8 +18,7 @@ build:cuda_clang --repo_env TF_NEED_TENSORRT=0
 build:cuda_clang --action_env=TF_CUDA_CLANG="1"
 build:cuda_clang --@local_config_cuda//:cuda_compiler=clang
 build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_80,sm_89,compute_90"
-build:cuda_clang --copt=-include
-build:cuda_clang --copt=third_party/tensorflow_serving/tools/cuda_clang_fixes.h
+build:cuda_clang --copt=-Dtypeof=__typeof__
 
 common:cuda_clang --repo_env=HERMETIC_CUDA_VERSION="12.5.1"
 common:cuda_clang --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
diff --git a/tensorflow_serving/opensource_only/.bazelrc b/tensorflow_serving/opensource_only/.bazelrc
index 4e9b9e0cc97..c580cf0b65c 100644
--- a/tensorflow_serving/opensource_only/.bazelrc
+++ b/tensorflow_serving/opensource_only/.bazelrc
@@ -18,8 +18,7 @@ build:cuda_clang --repo_env TF_NEED_TENSORRT=0
 build:cuda_clang --action_env=TF_CUDA_CLANG="1"
 build:cuda_clang --@local_config_cuda//:cuda_compiler=clang
 build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_80,sm_89,compute_90"
-build:cuda_clang --copt=-include
-build:cuda_clang --copt=third_party/tensorflow_serving/tools/cuda_clang_fixes.h
+build:cuda_clang --copt=-Dtypeof=__typeof__
 
 common:cuda_clang --repo_env=HERMETIC_CUDA_VERSION="12.5.1"
 common:cuda_clang --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
diff --git a/tensorflow_serving/opensource_only/root.workspace b/tensorflow_serving/opensource_only/root.workspace
new file mode 100644
index 00000000000..7c6bf4bbe5d
--- /dev/null
+++ b/tensorflow_serving/opensource_only/root.workspace
@@ -0,0 +1,297 @@
+workspace(name = "tf_serving")
+
+# TODO(b/269515133): We temporarily remove remote_predict from our builds for
+# 2.12 due to a breakage caused by
+# github.com/tensorflow/tensorflow/commit/6147c03eb9af1e5d2ae155045b33e909ef96944e
+# This will be removed in a subsequent release.
+local_repository(
+    name = "ignore_remote_predict",
+    path = "tensorflow_serving/experimental/tensorflow/ops/remote_predict/",
+)
+
+# ===== TensorFlow dependency =====
+#
+# TensorFlow is imported here instead of in tf_serving_workspace() because
+# existing automation scripts that bump the TF commit hash expect it here.
+#
+# To update TensorFlow to a new revision.
+# 1. Update the 'git_commit' args below to include the new git hash.
+# 2. Get the sha256 hash of the archive with a command such as...
+#    curl -L https://github.com/tensorflow/tensorflow/archive/<git hash>.tar.gz | sha256sum
+#    and update the 'sha256' arg with the result.
+# 3. Request the new archive to be mirrored on mirror.bazel.build for more
+#    reliable downloads.
+load("//tensorflow_serving:repo.bzl", "tensorflow_http_archive")
+tensorflow_http_archive(
+    name = "org_tensorflow",
+    sha256 = "582f5dc1c0887ca5b2eb364dc39e0c15fd31ee138b250e24163b6eca340f701f",
+    git_commit = "d875953b4fd431c15fedb1d408d5396e7f9d4d23",
+    patch = "//third_party/tensorflow:tensorflow.patch",
+    patch_cmds = [
+        "sed -i '/cc_library = _cc_library/d' tensorflow/core/platform/rules_cc.bzl",
+        "echo -e \"\\ndef cc_library_oss(deps=[], **kwargs):\\n    if kwargs.get(\\\"name\\\") == \\\"lib_internal_impl\\\" or \\\"protobuf\\\" in kwargs.get(\\\"name\\\", \\\"\\\"):\\n        _cc_library(deps = deps, **kwargs)\\n        return\\n    if type(deps) == \\\"list\\\":\\n        if \\\"@com_google_protobuf//:protobuf\\\" not in deps:\\n            deps = deps + [\\\"@com_google_protobuf//:protobuf\\\"]\\n    else:\\n        deps = deps + [\\\"@com_google_protobuf//:protobuf\\\"]\\n    _cc_library(deps = deps, **kwargs)\\ncc_library = cc_library_oss\" >> tensorflow/core/platform/rules_cc.bzl",
+        "sed -i 's#deps = \\[op_gen\\] + deps#deps = [op_gen] + deps + [clean_dep(\"//tensorflow/core/framework:kernel_shape_util\"), clean_dep(\"//tensorflow/core/framework:full_type_util\")]#' tensorflow/tensorflow.bzl",
+        "sed -i '/name = \"kernel_shape_util\",/a \\    visibility = [\"//visibility:public\"],' tensorflow/core/framework/BUILD",
+        "echo -e '\\nalias(name = \"tensorflow_libtensorflow_framework\", actual = \"//tensorflow/core:tensorflow\", visibility = [\"//visibility:public\"])' >> BUILD",
+        "echo -e '\\nalias(name = \"tensorflow_tf_header_lib\", actual = \"//tensorflow/core:tensorflow\", visibility = [\"//visibility:public\"])' >> BUILD",
+        "sed -i '/name = \"env\",/,/deps = \\[/ s#deps = \\[#deps = [\":status\", \":statusor\", \":context\", \":tracing\", \"//xla/tsl/profiler/backends/cpu:threadpool_listener_state\", \"//xla/tsl/platform:byte_order\", #' third_party/xla/xla/tsl/platform/default/BUILD",
+        "sed -i '/name = \"tracing\",/,/deps = \\[/ s#deps = \\[#deps = [\"//xla/tsl/platform:logging\", #' third_party/xla/xla/tsl/platform/default/BUILD",
+        "sed -i '/name = \"tf_runtime\",/a \\        repo_mapping = {\"@xla\": \"@local_xla\", \"@tsl\": \"@local_tsl\"},' third_party/tf_runtime/workspace.bzl",
+        "sed -i '/name = \"error_util\",/,/deps = \\[/ s#deps = \\[#deps = [\"@xla//xla/tsl/concurrency:async_value\", \"@xla//xla/tsl/concurrency:concurrent_vector\", \"@xla//xla/tsl/concurrency:executor\", \"@xla//xla/tsl/concurrency:ref_count\", \"@xla//xla/tsl/util:safe_reinterpret_cast\", \"@tsl//tsl/platform:context\", #' tensorflow/core/tfrt/utils/BUILD",
+        "sed -i '/name = \"work_queue_interface\",/,/deps = \\[/ s#deps = \\[#deps = [\"@xla//xla/tsl/concurrency:ref_count\", #' tensorflow/core/tfrt/runtime/BUILD",
+        "sed -i '/name = \"execute\",/,/deps = \\[/ s#deps = \\[#deps = [\"@xla//xla/tsl/platform:macros\", \"@xla//xla/tsl/platform:types\", \"@xla//xla/tsl/profiler/utils:no_init\", \"@tsl//tsl/profiler/lib:traceme_encode\", \"@xla//xla/tsl/profiler/utils:traceme_global_flags\", \"@xla//xla/tsl/profiler/backends/cpu:traceme_recorder\", \"@tsl//tsl/platform:bfloat16\", \"@tsl//tsl/platform:ml_dtypes\", \"@tsl//tsl/platform:tstring\", \"@tsl//tsl/platform:cord\", \"@tsl//tsl/platform:refcount\", \"@tsl//tsl/platform:thread_annotations\", \"@tsl//tsl/platform:stringpiece\", \"@xla//xla/tsl/profiler/utils:time_utils\", \"@xla//xla/tsl/profiler/utils:math_utils\", #' tensorflow/core/tfrt/mlrt/interpreter/BUILD",
+        "sed -i '/tf_vendored(name = \"xla\",/s/)/, repo_mapping = {\"@xla\": \"@local_xla\", \"@tsl\": \"@local_tsl\"})/' tensorflow/workspace3.bzl",
+        "sed -i '/tf_vendored(name = \"tsl\",/s/)/, repo_mapping = {\"@xla\": \"@local_xla\", \"@tsl\": \"@local_tsl\"})/' tensorflow/workspace3.bzl",
+        """python3 -c 'import re, glob
+for p in glob.glob("third_party/xla/**/BUILD*", recursive=True):
+    s = open(p).read(); blocks = s.split("cc_library(");
+    for i in range(1, len(blocks)):
+        b = blocks[i]; m_th = re.search(r"textual_hdrs\\s*=\\s*(\\[[^\\]]+\\]),?\\n?", b);
+        if m_th:
+            th = m_th.group(1); b = b.replace(m_th.group(0), ""); m_h = re.search(r"hdrs\\s*=\\s*(\\[[^\\]]+\\])", b);
+            if m_h: h = m_h.group(1); merged = h[:-1] + ", " + th[1:]; b = b.replace(m_h.group(0), "hdrs = " + merged);
+            else: b = "\\n    hdrs = " + th + "," + b;
+            blocks[i] = b
+    open(p, "w").write("cc_library(".join(blocks))'""",
+        "echo -e '\\ndiff --git a/WORKSPACE b/WORKSPACE\\n--- a/WORKSPACE\\n+++ b/WORKSPACE\\n@@ -184,25 +184,2 @@\\n sass_repositories()\\n \\n-http_archive(\\n-    name = \"xla\",\\n-    patch_args = [\"-p1\"],\\n-    patches = [\\n-        \"//third_party:xla.patch\",\\n-        \"//third_party:xla_add_grpc_cares_darwin_arm64_support.patch\",\\n-    ],\\n-    sha256 = \"ba80ef58f89ca11bc5652e936cf856cdeae91e6b723ce6750e9ce0202cab51ac\",\\n-    strip_prefix = \"xla-f094066398e2c884e994711fd677f68864324614\",\\n-    urls = [\\n-        \"https://github.com/openxla/xla/archive/f094066398e2c884e994711fd677f68864324614.zip\",\\n-    ],\\n-)\\n-\\n-http_archive(\\n-    name = \"tsl\",\\n-    sha256 = \"8cf1e1285c7b1843a7f5f787465c1ef80304b3400ed837870bc76d74ce04f5af\",\\n-    strip_prefix = \"tsl-d71df2f7612583617d359c36243695097dd63726\",\\n-    urls = [\\n-        \"https://github.com/google/tsl/archive/d71df2f7612583617d359c36243695097dd63726.zip\",\\n-    ],\\n-)\\n-\\n load(\"@xla//tools/toolchains/python:python_repo.bzl\", \"python_repository\")' >> third_party/xprof/xprof.patch",
+    ],
+    repo_mapping = {
+        "@local_xla": "@local_xla",
+        "@local_tsl": "@local_tsl",
+        "@org_tensorflow": "@org_tensorflow",
+        "@xla": "@local_xla",
+        "@tsl": "@local_tsl",
+    },
+)
+
+# Import all of TensorFlow Serving's external dependencies.
+# Downstream projects (projects importing TensorFlow Serving) need to
+# duplicate all code below in their WORKSPACE file in order to also initialize
+# those external dependencies.
+load("//tensorflow_serving:workspace.bzl", "tf_serving_workspace")
+tf_serving_workspace()
+
+# Check bazel version requirement, which is stricter than TensorFlow's.
+load("@bazel_skylib//lib:versions.bzl", "versions")
+versions.check("7.4.1")
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+http_archive(
+    name = "bazel_skylib",
+    sha256 = "97e709db2e97b646263b5c5e83e3b00de48c1ae55b9e421e3b5e3f9467d02a3a",
+    urls = [
+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.8.0/bazel-skylib-1.8.0.tar.gz",
+        "https://github.com/bazelbuild/bazel-skylib/releases/download/1.8.0/bazel-skylib-1.8.0.tar.gz",
+    ],
+)
+
+http_archive(
+    name = "rules_cc",
+    sha256 = "b8b918a85f9144c01f6cfe0f45e4f2838c7413961a8ff23bc0c6cdf8bb07a3b6",
+    strip_prefix = "rules_cc-0.1.5",
+    url = "https://github.com/bazelbuild/rules_cc/releases/download/0.1.5/rules_cc-0.1.5.tar.gz",
+)
+
+http_archive(
+    name = "rules_python",
+    sha256 = "8964aa1e7525fea5244ba737458694a057ada1be96a92998a41caa1983562d00",
+    strip_prefix = "rules_python-1.8.5",
+    urls = [
+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/rules_python/releases/download/1.8.5/rules_python-1.8.5.tar.gz",
+        "https://github.com/bazelbuild/rules_python/releases/download/1.8.5/rules_python-1.8.5.tar.gz",
+    ],
+    patches = [
+        "@rules_ml_toolchain//third_party/rules_python:rules_python_scope.patch",
+        "@rules_ml_toolchain//third_party/rules_python:rules_python_freethreaded.patch",
+        "@rules_ml_toolchain//third_party/rules_python:rules_python_versions.patch",
+        "@rules_ml_toolchain//third_party/rules_python:rules_python_pip_version.patch",
+    ],
+    patch_args = ["-p1"],
+)
+
+# Toolchains for ML projects hermetic builds.
+# Details: https://github.com/google-ml-infra/rules_ml_toolchain
+http_archive(
+    name = "rules_ml_toolchain",
+    sha256 = "0b42f693a60c6050d87db1e0a0eaeb84ab3f54191fce094d86334faedc807da0",
+    strip_prefix = "rules_ml_toolchain-398d613aea7a4c294da49b79a6d6f3f8732bd84c",
+    urls = [
+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google-ml-infra/rules_ml_toolchain/archive/398d613aea7a4c294da49b79a6d6f3f8732bd84c.tar.gz",
+        "https://github.com/google-ml-infra/rules_ml_toolchain/archive/398d613aea7a4c294da49b79a6d6f3f8732bd84c.tar.gz",
+    ],
+    patch_cmds = [
+        "cat << 'EOF' >> cc/cuda/clang/clang_cuda_runtime_wrapper.h\n#if defined(__clang__) && defined(__CUDA__)\n#include <vector_types.h>\n#ifndef typeof\n#define typeof __typeof__\n#endif\n__device__ inline void __stwt(uint4* ptr, uint4 value) {\n  asm volatile(\"st.global.wt.v4.u32 [%0], {%1, %2, %3, %4};\" : : \"l\"(ptr), \"r\"(value.x), \"r\"(value.y), \"r\"(value.z), \"r\"(value.w) : \"memory\");\n}\n#endif\nEOF",
+    ],
+)
+
+load(
+    "@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl",
+    "cc_toolchain_deps",
+)
+
+cc_toolchain_deps()
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_cuda")
+# register_toolchains("@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64")
+# register_toolchains("@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64_cuda")
+
+# Initialize hermetic Python
+load("@org_tensorflow//third_party/py:python_init_rules.bzl", "python_init_rules")
+python_init_rules()
+
+load("@org_tensorflow//third_party/py:python_init_repositories.bzl", "python_init_repositories")
+python_init_repositories(
+    default_python_version = "system",
+    requirements = {
+        "3.10": "@org_tensorflow//:requirements_lock_3_10.txt",
+        "3.11": "@org_tensorflow//:requirements_lock_3_11.txt",
+        "3.12": "@org_tensorflow//:requirements_lock_3_12.txt",
+        "3.13": "@org_tensorflow//:requirements_lock_3_13.txt",
+        "3.14": "@org_tensorflow//:requirements_lock_3_14.txt",
+    },
+)
+
+load("@org_tensorflow//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
+python_init_toolchains()
+
+load("@org_tensorflow//third_party/py:python_init_pip.bzl", "python_init_pip")
+python_init_pip()
+
+load("@pypi//:requirements.bzl", "install_deps")
+install_deps()
+
+# Initialize TensorFlow's external dependencies.
+load("@org_tensorflow//tensorflow:workspace3.bzl", "tf_workspace3")
+tf_workspace3()
+
+load("//tensorflow_serving:repo.bzl", "tf_serving_vendored")
+
+tf_serving_vendored(
+    name = "local_xla",
+    path = "third_party/xla",
+    repo_mapping = {
+        "@local_xla": "@local_xla",
+        "@local_tsl": "@local_tsl",
+        "@org_tensorflow": "@org_tensorflow",
+        "@xla": "@local_xla",
+        "@tsl": "@local_tsl",
+    },
+    root = "@org_tensorflow//:unused",
+)
+
+tf_serving_vendored(
+    name = "local_tsl",
+    path = "third_party/xla/third_party/tsl",
+    repo_mapping = {
+        "@local_xla": "@local_xla",
+        "@local_tsl": "@local_tsl",
+        "@org_tensorflow": "@org_tensorflow",
+        "@xla": "@local_xla",
+        "@tsl": "@local_tsl",
+    },
+    root = "@org_tensorflow//:unused",
+)
+load("@org_tensorflow//tensorflow:workspace2.bzl", "tf_workspace2")
+tf_workspace2()
+load("@org_tensorflow//tensorflow:workspace1.bzl", "tf_workspace1")
+tf_workspace1()
+load("@org_tensorflow//tensorflow:workspace0.bzl", "tf_workspace0")
+tf_workspace0()
+
+# Initialize bazel package rules' external dependencies.
+load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
+rules_pkg_dependencies()
+
+load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
+
+http_archive(
+    name = "rules_shell",
+    sha256 = "0d0c56d01c3c40420bf7bf14d73113f8a92fbd9f5cd13205a3b89f72078f0321",
+    strip_prefix = "rules_shell-0.1.1",
+    urls = [
+        "https://github.com/bazelbuild/rules_shell/releases/download/v0.1.1/rules_shell-v0.1.1.tar.gz",
+    ],
+)
+
+load("@rules_proto//proto:repositories.bzl", "rules_proto_dependencies")
+
+rules_proto_dependencies()
+
+load(
+    "@xla//third_party/py:python_wheel.bzl",
+    "python_wheel_version_suffix_repository",
+)
+
+python_wheel_version_suffix_repository(name = "tf_wheel_version_suffix")
+
+load(
+    "@rules_ml_toolchain//gpu/cuda:cuda_json_init_repository.bzl",
+    "cuda_json_init_repository",
+)
+
+cuda_json_init_repository()
+
+load(
+    "@cuda_redist_json//:distributions.bzl",
+    "CUDA_REDISTRIBUTIONS",
+    "CUDNN_REDISTRIBUTIONS",
+)
+load(
+    "@rules_ml_toolchain//gpu/cuda:cuda_redist_init_repositories.bzl",
+    "cuda_redist_init_repositories",
+    "cudnn_redist_init_repository",
+)
+
+cuda_redist_init_repositories(
+    cuda_redistributions = CUDA_REDISTRIBUTIONS,
+)
+
+cudnn_redist_init_repository(
+    cudnn_redistributions = CUDNN_REDISTRIBUTIONS,
+)
+
+load(
+    "@rules_ml_toolchain//gpu/cuda:cuda_configure.bzl",
+    "cuda_configure",
+)
+
+cuda_configure(name = "local_config_cuda")
+
+load(
+    "@rules_ml_toolchain//gpu/nccl:nccl_redist_init_repository.bzl",
+    "nccl_redist_init_repository",
+)
+
+nccl_redist_init_repository()
+
+load(
+    "@rules_ml_toolchain//gpu/nccl:nccl_configure.bzl",
+    "nccl_configure",
+)
+
+nccl_configure(name = "local_config_nccl")
+
+load(
+    "@rules_ml_toolchain//gpu/nvshmem:nvshmem_json_init_repository.bzl",
+    "nvshmem_json_init_repository",
+)
+
+nvshmem_json_init_repository()
+
+load(
+    "@nvshmem_redist_json//:distributions.bzl",
+    "NVSHMEM_REDISTRIBUTIONS",
+)
+load(
+    "@rules_ml_toolchain//gpu/nvshmem:nvshmem_redist_init_repository.bzl",
+    "nvshmem_redist_init_repository",
+)
+
+nvshmem_redist_init_repository(
+    nvshmem_redistributions = NVSHMEM_REDISTRIBUTIONS,
+)
+
+# nvshmem_configure removed in newer rules_ml_toolchain
+
diff --git a/tensorflow_serving/tools/cuda_clang_fixes.h b/tensorflow_serving/tools/cuda_clang_fixes.h
deleted file mode 100644
index f1e87a5a115..00000000000
--- a/tensorflow_serving/tools/cuda_clang_fixes.h
+++ /dev/null
@@ -1,36 +0,0 @@
-/* Copyright 2026 Google Inc. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#ifndef TENSORFLOW_SERVING_TOOLS_CUDA_CLANG_FIXES_H_
-#define TENSORFLOW_SERVING_TOOLS_CUDA_CLANG_FIXES_H_
-
-#if defined(__clang__) && defined(__CUDA__)
-#include <vector_types.h>
-
-// Fix for GNU typeof in strict C++ mode when compiling NVIDIA DOCA / NCCL headers with Clang.
-#ifndef typeof
-#define typeof __typeof__
-#endif
-
-// Missing 16-byte uint4 store overload in Clang's __clang_cuda_intrinsics.h for NCCL GIN proxy (__stwt).
-__device__ inline void __stwt(uint4* ptr, uint4 value) {
-  asm volatile("st.global.wt.v4.u32 [%0], {%1, %2, %3, %4};"
-               :
-               : "l"(ptr), "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w)
-               : "memory");
-}
-#endif
-
-#endif  // TENSORFLOW_SERVING_TOOLS_CUDA_CLANG_FIXES_H_

From 8e90909830675c0e5c543b78f9a2933085ac408f Mon Sep 17 00:00:00 2001
From: Venkata Sai Madhur Karampudi <vkarampudi@google.com>
Date: Sun, 12 Jul 2026 17:57:49 +0000
Subject: [PATCH 6/6] Fix GPU Docker build failures.

- Align WORKSPACE with opensource_only/root.workspace.
- Patch XLA to fix dso_loader missing dependency declarations.
- Patch XLA to fix CUB scan policy compile error.

TAG=agy
CONV=748d06f3-9c06-4d77-9f81-b365b06d1b3a
---
 WORKSPACE                               |   7 +-
 third_party/tensorflow/tensorflow.patch | 291 ++++++++++++++++++++++++
 2 files changed, 296 insertions(+), 2 deletions(-)

diff --git a/WORKSPACE b/WORKSPACE
index 9d53b2fbf8b..7c6bf4bbe5d 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -24,8 +24,8 @@ local_repository(
 load("//tensorflow_serving:repo.bzl", "tensorflow_http_archive")
 tensorflow_http_archive(
     name = "org_tensorflow",
-    sha256 = "54bfd56d61dfced2b82bdeee3262a24fb3d59bed4dbb292d19834af002d85b3b",
-    git_commit = "c1be979f3130b85f14e8890c05cbca147784f474",
+    sha256 = "582f5dc1c0887ca5b2eb364dc39e0c15fd31ee138b250e24163b6eca340f701f",
+    git_commit = "d875953b4fd431c15fedb1d408d5396e7f9d4d23",
     patch = "//third_party/tensorflow:tensorflow.patch",
     patch_cmds = [
         "sed -i '/cc_library = _cc_library/d' tensorflow/core/platform/rules_cc.bzl",
@@ -120,6 +120,9 @@ http_archive(
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google-ml-infra/rules_ml_toolchain/archive/398d613aea7a4c294da49b79a6d6f3f8732bd84c.tar.gz",
         "https://github.com/google-ml-infra/rules_ml_toolchain/archive/398d613aea7a4c294da49b79a6d6f3f8732bd84c.tar.gz",
     ],
+    patch_cmds = [
+        "cat << 'EOF' >> cc/cuda/clang/clang_cuda_runtime_wrapper.h\n#if defined(__clang__) && defined(__CUDA__)\n#include <vector_types.h>\n#ifndef typeof\n#define typeof __typeof__\n#endif\n__device__ inline void __stwt(uint4* ptr, uint4 value) {\n  asm volatile(\"st.global.wt.v4.u32 [%0], {%1, %2, %3, %4};\" : : \"l\"(ptr), \"r\"(value.x), \"r\"(value.y), \"r\"(value.z), \"r\"(value.w) : \"memory\");\n}\n#endif\nEOF",
+    ],
 )
 
 load(
diff --git a/third_party/tensorflow/tensorflow.patch b/third_party/tensorflow/tensorflow.patch
index 75aa4f5f2ff..cf21a5f9886 100644
--- a/third_party/tensorflow/tensorflow.patch
+++ b/third_party/tensorflow/tensorflow.patch
@@ -8,3 +8,294 @@ diff --git a/tensorflow/tools/toolchains/python/python_repo.bzl b/tensorflow/too
  WHEEL_COLLAB = "{}"
 +USE_PYWRAP_RULES = "False"
  """
+diff -ruN a/third_party/xla/xla/stream_executor/cuda/cub_scan_kernel_cuda_impl.cu.cc b/third_party/xla/xla/stream_executor/cuda/cub_scan_kernel_cuda_impl.cu.cc
+--- a/third_party/xla/xla/stream_executor/cuda/cub_scan_kernel_cuda_impl.cu.cc	2026-07-12 17:57:25.618437375 +0000
++++ b/third_party/xla/xla/stream_executor/cuda/cub_scan_kernel_cuda_impl.cu.cc	2026-07-12 17:57:25.632496035 +0000
+@@ -41,10 +41,7 @@
+ namespace {
+ 
+ template <typename T, typename ScanOpT>
+-using MaxPolicyT = typename cub::detail::scan::policy_hub<
+-    /*InputValueT=*/T, /*OutputValueT=*/T,
+-    /*AccumT=*/T, /*OffsetT=*/int64_t,
+-    /*ScanOpT=*/ScanOpT>::MaxPolicy;
++using MaxPolicyT = typename cub::DeviceScanPolicy<T, ScanOpT>::MaxPolicy;
+ 
+ template <typename T, typename ScanOpT,
+           typename ChainedPolicyT = MaxPolicyT<T, ScanOpT>>
+diff -ruN a/third_party/xla/xla/tsl/cuda/BUILD.bazel b/third_party/xla/xla/tsl/cuda/BUILD.bazel
+--- a/third_party/xla/xla/tsl/cuda/BUILD.bazel	2026-07-12 17:57:25.617565295 +0000
++++ b/third_party/xla/xla/tsl/cuda/BUILD.bazel	2026-07-12 17:57:25.629649015 +0000
+@@ -47,7 +47,7 @@
+         "//xla/tsl/platform:logging",
+         "@com_google_absl//absl/container:flat_hash_set",
+         "@local_config_cuda//cuda:cuda_headers",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -78,7 +78,7 @@
+         # keep sorted
+         "//xla/tsl/platform:logging",
+         "@local_config_cuda//cuda:cuda_headers",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -109,7 +109,7 @@
+         # keep sorted
+         "//xla/tsl/platform:logging",
+         "@local_config_cuda//cuda:cuda_headers",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -146,7 +146,7 @@
+             "//xla/tsl/platform:logging",
+             "@com_google_absl//absl/container:flat_hash_set",
+             "@local_config_cuda//cuda:cuda_headers",
+-            "@tsl//tsl/platform:dso_loader",
++            "//xla/tsl/platform/default:dso_loader",
+             "@tsl//tsl/platform:load_library",
+         ],
+         "//conditions:default": [],
+@@ -181,7 +181,7 @@
+         "//xla/tsl/platform:logging",
+         "@com_google_absl//absl/container:flat_hash_map",
+         "@local_config_cuda//cuda:cudnn_header",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -235,7 +235,7 @@
+         # keep sorted
+         "//xla/tsl/platform:logging",
+         "@local_config_cuda//cuda:cuda_headers",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -273,7 +273,7 @@
+         "//xla/tsl/platform:logging",
+         "@local_config_cuda//cuda:cuda_headers",
+         "@local_config_cuda//cuda:cupti_headers",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -309,7 +309,7 @@
+         # keep sorted
+         "//xla/tsl/platform:logging",
+         "@local_config_cuda//cuda:cuda_headers",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -346,7 +346,7 @@
+         "//xla/tsl/platform:logging",
+         "@com_google_absl//absl/container:flat_hash_set",
+         "@local_config_cuda//cuda:cuda_headers",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -381,7 +381,7 @@
+         "@com_google_absl//absl/strings:string_view",
+         "@local_config_cuda//cuda:cuda_headers",
+         "@local_config_nccl//:nccl_headers",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -413,7 +413,7 @@
+     deps = if_cuda_is_configured([
+         # keep sorted
+         "//xla/tsl/platform:logging",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -444,7 +444,7 @@
+         # keep sorted
+         "//xla/tsl/platform:logging",
+         "@local_config_cuda//cuda:cuda_headers",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+@@ -475,7 +475,7 @@
+         # keep sorted
+         "//xla/tsl/platform:logging",
+         "@local_config_cuda//cuda:cuda_headers",
+-        "@tsl//tsl/platform:dso_loader",
++        "//xla/tsl/platform/default:dso_loader",
+         "@tsl//tsl/platform:load_library",
+     ]),
+ )
+diff -ruN a/third_party/xla/xla/tsl/cuda/cublasLt_stub.cc b/third_party/xla/xla/tsl/cuda/cublasLt_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/cublasLt_stub.cc	2026-07-12 17:57:25.619010684 +0000
++++ b/third_party/xla/xla/tsl/cuda/cublasLt_stub.cc	2026-07-12 17:57:25.629875675 +0000
+@@ -15,7 +15,7 @@
+ #include "third_party/gpus/cuda/include/cublasLt.h"
+ #include "third_party/gpus/cuda/include/cuda.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the cuBLASLt API by forwarding to cuBLASLt loaded from the DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/cublas_stub.cc b/third_party/xla/xla/tsl/cuda/cublas_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/cublas_stub.cc	2026-07-12 17:57:25.619515775 +0000
++++ b/third_party/xla/xla/tsl/cuda/cublas_stub.cc	2026-07-12 17:57:25.630085595 +0000
+@@ -24,7 +24,7 @@
+ #include "absl/container/flat_hash_set.h"
+ #include "third_party/gpus/cuda/include/cuda.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the cuBLAS API by forwarding to cuBLAS loaded from the DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/cudart_stub.cc b/third_party/xla/xla/tsl/cuda/cudart_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/cudart_stub.cc	2026-07-12 17:57:25.620547485 +0000
++++ b/third_party/xla/xla/tsl/cuda/cudart_stub.cc	2026-07-12 17:57:25.630498815 +0000
+@@ -20,7 +20,7 @@
+ #include "absl/strings/string_view.h"
+ #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ namespace {
+diff -ruN a/third_party/xla/xla/tsl/cuda/cuda_stub.cc b/third_party/xla/xla/tsl/cuda/cuda_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/cuda_stub.cc	2026-07-12 17:57:25.620042105 +0000
++++ b/third_party/xla/xla/tsl/cuda/cuda_stub.cc	2026-07-12 17:57:25.630294585 +0000
+@@ -14,7 +14,7 @@
+ ==============================================================================*/
+ #include "third_party/gpus/cuda/include/cuda.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the CUDA driver API by forwarding to CUDA loaded from the DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/cudnn_stub.cc b/third_party/xla/xla/tsl/cuda/cudnn_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/cudnn_stub.cc	2026-07-12 17:57:25.623896675 +0000
++++ b/third_party/xla/xla/tsl/cuda/cudnn_stub.cc	2026-07-12 17:57:25.630687515 +0000
+@@ -16,7 +16,7 @@
+ #include "absl/container/flat_hash_map.h"
+ #include "third_party/gpus/cudnn/cudnn.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the cuDNN API by forwarding to cuDNN loaded from the DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/cufft_stub.cc b/third_party/xla/xla/tsl/cuda/cufft_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/cufft_stub.cc	2026-07-12 17:57:25.624421815 +0000
++++ b/third_party/xla/xla/tsl/cuda/cufft_stub.cc	2026-07-12 17:57:25.630871185 +0000
+@@ -15,7 +15,7 @@
+ #include "third_party/gpus/cuda/include/cufft.h"
+ #include "third_party/gpus/cuda/include/cufftXt.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the cuFFT API by forwarding to cuFFT loaded from the DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/cupti_stub.cc b/third_party/xla/xla/tsl/cuda/cupti_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/cupti_stub.cc	2026-07-12 17:57:25.624935405 +0000
++++ b/third_party/xla/xla/tsl/cuda/cupti_stub.cc	2026-07-12 17:57:25.631143375 +0000
+@@ -16,7 +16,7 @@
+ #include "third_party/gpus/cuda/extras/CUPTI/include/cupti.h"
+ #include "third_party/gpus/cuda/include/cuda.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the CUPTI API by forwarding to CUPTI loaded from the DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/cusolver_stub.cc b/third_party/xla/xla/tsl/cuda/cusolver_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/cusolver_stub.cc	2026-07-12 17:57:25.625466465 +0000
++++ b/third_party/xla/xla/tsl/cuda/cusolver_stub.cc	2026-07-12 17:57:25.631325645 +0000
+@@ -16,7 +16,7 @@
+ #include "third_party/gpus/cuda/include/cusolverDn.h"
+ #include "third_party/gpus/cuda/include/cusolverSp.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the cusolver API by forwarding to cusolver loaded from the DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/cusparse_stub.cc b/third_party/xla/xla/tsl/cuda/cusparse_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/cusparse_stub.cc	2026-07-12 17:57:25.625975035 +0000
++++ b/third_party/xla/xla/tsl/cuda/cusparse_stub.cc	2026-07-12 17:57:25.631513505 +0000
+@@ -18,7 +18,7 @@
+ #include "third_party/gpus/cuda/include/cuda.h"
+ #include "third_party/gpus/cuda/include/cusparse.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the cusparse API by forwarding to cusparse loaded from the DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/nccl_stub.cc b/third_party/xla/xla/tsl/cuda/nccl_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/nccl_stub.cc	2026-07-12 17:57:25.626489445 +0000
++++ b/third_party/xla/xla/tsl/cuda/nccl_stub.cc	2026-07-12 17:57:25.631694955 +0000
+@@ -17,7 +17,7 @@
+ #include "third_party/gpus/cuda/include/cuda.h"
+ #include "third_party/nccl/nccl.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the nccl API by forwarding to nccl loaded from a DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/nvml_stub.cc b/third_party/xla/xla/tsl/cuda/nvml_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/nvml_stub.cc	2026-07-12 17:57:25.628120455 +0000
++++ b/third_party/xla/xla/tsl/cuda/nvml_stub.cc	2026-07-12 17:57:25.631884615 +0000
+@@ -14,7 +14,7 @@
+ ==============================================================================*/
+ #include "third_party/gpus/cuda/nvml/include/nvml.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the NVML API by forwarding to NVML loaded from the DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/nvrtc_stub.cc b/third_party/xla/xla/tsl/cuda/nvrtc_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/nvrtc_stub.cc	2026-07-12 17:57:25.628643355 +0000
++++ b/third_party/xla/xla/tsl/cuda/nvrtc_stub.cc	2026-07-12 17:57:25.632092815 +0000
+@@ -15,7 +15,7 @@
+ 
+ #include "third_party/gpus/cuda/include/nvrtc.h"
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the NVRTC API by forwarding to NVRTC loaded from the DSO.
+diff -ruN a/third_party/xla/xla/tsl/cuda/nvshmem_stub.cc b/third_party/xla/xla/tsl/cuda/nvshmem_stub.cc
+--- a/third_party/xla/xla/tsl/cuda/nvshmem_stub.cc	2026-07-12 17:57:25.629133555 +0000
++++ b/third_party/xla/xla/tsl/cuda/nvshmem_stub.cc	2026-07-12 17:57:25.632281315 +0000
+@@ -14,7 +14,7 @@
+ ==============================================================================*/
+ 
+ #include "xla/tsl/platform/logging.h"
+-#include "tsl/platform/dso_loader.h"
++#include "xla/tsl/platform/default/dso_loader.h"
+ #include "tsl/platform/load_library.h"
+ 
+ // Implements the nvshmem API by forwarding to nvshmem loaded from a DSO.